From db67c6a89175cbccbe84ae2de59c895dbb1f37cf Mon Sep 17 00:00:00 2001
From: ale <ale@incal.net>
Date: Sun, 17 Dec 2017 22:05:27 +0000
Subject: [PATCH] Support partitioned keystore service

Uses the clientutil.Backend abstraction for the keystore client API.
---
 saml/saml.go                                  |   2 +-
 server/config.go                              |  14 +--
 server/http.go                                |  14 ++-
 server/service_test.go                        |   2 +-
 .../ai3/go-common/clientutil/backend.go       | 119 ++++++++++++++++++
 .../id/keystore/client/client.go              |  54 +++-----
 vendor/vendor.json                            |  12 +-
 7 files changed, 162 insertions(+), 55 deletions(-)
 create mode 100644 vendor/git.autistici.org/ai3/go-common/clientutil/backend.go

diff --git a/saml/saml.go b/saml/saml.go
index f130454..a75dffd 100644
--- a/saml/saml.go
+++ b/saml/saml.go
@@ -19,7 +19,7 @@ import (
 	"github.com/crewjam/saml"
 	"github.com/crewjam/saml/logger"
 	"github.com/gorilla/mux"
-	yaml "gopkg.in/yaml.v2"
+	"gopkg.in/yaml.v2"
 
 	"git.autistici.org/id/go-sso/httpsso"
 )
diff --git a/server/config.go b/server/config.go
index 1a4748e..4750c8c 100644
--- a/server/config.go
+++ b/server/config.go
@@ -6,9 +6,9 @@ import (
 	"regexp"
 	"time"
 
-	ksclient "git.autistici.org/id/keystore/client"
 	"github.com/gorilla/securecookie"
 
+	"git.autistici.org/ai3/go-common/clientutil"
 	"git.autistici.org/id/go-sso/server/device"
 )
 
@@ -29,12 +29,12 @@ type Config struct {
 		TTLSeconds int    `yaml:"ttl"`
 		rx         *regexp.Regexp
 	} `yaml:"service_ttls"`
-	AuthSessionLifetimeSeconds int              `yaml:"auth_session_lifetime"`
-	SessionSecrets             []string         `yaml:"session_secrets"`
-	CSRFSecret                 string           `yaml:"csrf_secret"`
-	AuthService                string           `yaml:"auth_service"`
-	DeviceManager              *device.Config   `yaml:"device_manager"`
-	KeyStore                   *ksclient.Config `yaml:"keystore"`
+	AuthSessionLifetimeSeconds int                       `yaml:"auth_session_lifetime"`
+	SessionSecrets             []string                  `yaml:"session_secrets"`
+	CSRFSecret                 string                    `yaml:"csrf_secret"`
+	AuthService                string                    `yaml:"auth_service"`
+	DeviceManager              *device.Config            `yaml:"device_manager"`
+	KeyStore                   *clientutil.BackendConfig `yaml:"keystore"`
 
 	allowedServicesRx []*regexp.Regexp
 }
diff --git a/server/http.go b/server/http.go
index 39397a7a..a277bf9 100644
--- a/server/http.go
+++ b/server/http.go
@@ -86,7 +86,7 @@ type Server struct {
 	authSessionLifetime time.Duration
 	loginHandler        *loginHandler
 	loginService        *LoginService
-	keystore            *ksclient.Client
+	keystore            ksclient.Client
 	csrfSecret          []byte
 	tpl                 *template.Template
 }
@@ -147,7 +147,11 @@ func (h *Server) loginCallback(w http.ResponseWriter, req *http.Request, usernam
 	// authenticate. Set the TTL to the duration of the
 	// authenticated session.
 	if h.keystore != nil {
-		if err := h.keystore.Open(req.Context(), username, password, int(h.authSessionLifetime.Seconds())); err != nil {
+		var shard string
+		if userinfo != nil {
+			shard = userinfo.Shard
+		}
+		if err := h.keystore.Open(req.Context(), shard, username, password, int(h.authSessionLifetime.Seconds())); err != nil {
 			log.Printf("failed to unlock keystore for user %s: %v", username, err)
 			return err
 		}
@@ -257,7 +261,11 @@ func (h *Server) handleLogout(w http.ResponseWriter, req *http.Request, session
 
 		// Close the keystore.
 		if h.keystore != nil {
-			if err := h.keystore.Close(req.Context(), session.Username); err != nil {
+			var shard string
+			if session.UserInfo != nil {
+				shard = session.UserInfo.Shard
+			}
+			if err := h.keystore.Close(req.Context(), shard, session.Username); err != nil {
 				log.Printf("failed to wipe keystore for user %s: %v", session.Username, err)
 			}
 		}
diff --git a/server/service_test.go b/server/service_test.go
index 51aeb7d..b61f697 100644
--- a/server/service_test.go
+++ b/server/service_test.go
@@ -46,7 +46,7 @@ auth_service: login
 	if keystoreURL != "" {
 		cfgstr += fmt.Sprintf(`
 keystore:
-  backend_url: "%s"
+  url: "%s"
 `, keystoreURL)
 	}
 	ioutil.WriteFile(filepath.Join(tmpdir, "config"), []byte(cfgstr), 0600)
diff --git a/vendor/git.autistici.org/ai3/go-common/clientutil/backend.go b/vendor/git.autistici.org/ai3/go-common/clientutil/backend.go
new file mode 100644
index 0000000..1f49a7e
--- /dev/null
+++ b/vendor/git.autistici.org/ai3/go-common/clientutil/backend.go
@@ -0,0 +1,119 @@
+package clientutil
+
+import (
+	"crypto/tls"
+	"fmt"
+	"net/http"
+	"net/url"
+	"sync"
+	"time"
+)
+
+// BackendConfig specifies the configuration to access a service.
+//
+// Services with multiple backends can be replicated or partitioned,
+// depending on a configuration switch, making it a deployment-time
+// decision. Clients are expected to compute their own sharding
+// function (either by database lookup or other methods), and expose a
+// 'shard' parameter on their APIs.
+type BackendConfig struct {
+	URL       string           `yaml:"url"`
+	Sharded   bool             `yaml:"sharded"`
+	TLSConfig *TLSClientConfig `yaml:"tls_config"`
+}
+
+// Backend is a runtime class that provides http Clients for use with
+// a specific service backend. If the service can't be partitioned,
+// pass an empty string to the Client method.
+type Backend interface {
+	// URL for the service for a specific shard.
+	URL(string) string
+
+	// Client that can be used to make a request to the service.
+	Client(string) *http.Client
+}
+
+// NewBackend returns a new Backend with the given config.
+func NewBackend(config *BackendConfig) (Backend, error) {
+	u, err := url.Parse(config.URL)
+	if err != nil {
+		return nil, err
+	}
+
+	var tlsConfig *tls.Config
+	if config.TLSConfig != nil {
+		tlsConfig, err = config.TLSConfig.TLSConfig()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if config.Sharded {
+		return &replicatedClient{
+			u: u,
+			c: newHTTPClient(u, tlsConfig),
+		}, nil
+	}
+	return &shardedClient{
+		baseURL:   u,
+		tlsConfig: tlsConfig,
+		urls:      make(map[string]*url.URL),
+		shards:    make(map[string]*http.Client),
+	}, nil
+}
+
+type replicatedClient struct {
+	c *http.Client
+	u *url.URL
+}
+
+func (r *replicatedClient) Client(_ string) *http.Client { return r.c }
+func (r *replicatedClient) URL(_ string) string          { return r.u.String() }
+
+type shardedClient struct {
+	baseURL   *url.URL
+	tlsConfig *tls.Config
+	mx        sync.Mutex
+	urls      map[string]*url.URL
+	shards    map[string]*http.Client
+}
+
+func (s *shardedClient) getShardURL(shard string) *url.URL {
+	if shard == "" {
+		return s.baseURL
+	}
+	u, ok := s.urls[shard]
+	if !ok {
+		var tmp = *s.baseURL
+		tmp.Host = fmt.Sprintf("%s.%s", shard, tmp.Host)
+		u = &tmp
+		s.urls[shard] = u
+	}
+	return u
+}
+
+func (s *shardedClient) URL(shard string) string {
+	s.mx.Lock()
+	defer s.mx.Unlock()
+	return s.getShardURL(shard).String()
+}
+
+func (s *shardedClient) Client(shard string) *http.Client {
+	s.mx.Lock()
+	defer s.mx.Unlock()
+
+	client, ok := s.shards[shard]
+	if !ok {
+		u := s.getShardURL(shard)
+		client = newHTTPClient(u, s.tlsConfig)
+		s.shards[shard] = client
+	}
+	return client
+}
+
+func newHTTPClient(u *url.URL, tlsConfig *tls.Config) *http.Client {
+	return &http.Client{
+		Transport: NewTransport([]string{u.Host}, tlsConfig, nil),
+		Timeout:   30 * time.Second,
+	}
+}
diff --git a/vendor/git.autistici.org/id/keystore/client/client.go b/vendor/git.autistici.org/id/keystore/client/client.go
index cbceea1..e0278fd 100644
--- a/vendor/git.autistici.org/id/keystore/client/client.go
+++ b/vendor/git.autistici.org/id/keystore/client/client.go
@@ -2,10 +2,6 @@ package client
 
 import (
 	"context"
-	"crypto/tls"
-	"net/http"
-	"net/url"
-	"time"
 
 	"git.autistici.org/ai3/go-common/clientutil"
 
@@ -13,66 +9,50 @@ import (
 )
 
 // Client for the keystore API.
-type Client struct {
-	*http.Client
-	backendURL string
+type Client interface {
+	Open(context.Context, string, string, string, int) error
+	Get(context.Context, string, string, string) ([]byte, error)
+	Close(context.Context, string, string) error
 }
 
-// Config for a keystore client.
-type Config struct {
-	BackendURL string                      `yaml:"backend_url"`
-	TLSConfig  *clientutil.TLSClientConfig `yaml:"tls_config"`
+type ksClient struct {
+	be clientutil.Backend
 }
 
-// New returns a new Client with the given Config.
-func New(config *Config) (*Client, error) {
-	u, err := url.Parse(config.BackendURL)
+// New returns a new Client with the given Config. Use this when the
+// keystore service runs on a single global instance.
+func New(config *clientutil.BackendConfig) (*ksClient, error) {
+	be, err := clientutil.NewBackend(config)
 	if err != nil {
 		return nil, err
 	}
-
-	var tlsConfig *tls.Config
-	if config.TLSConfig != nil {
-		tlsConfig, err = config.TLSConfig.TLSConfig()
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	c := &http.Client{
-		Transport: clientutil.NewTransport([]string{u.Host}, tlsConfig, nil),
-		Timeout:   20 * time.Second,
-	}
-	return &Client{
-		Client:     c,
-		backendURL: config.BackendURL,
-	}, nil
+	return &ksClient{be}, nil
 }
 
-func (c *Client) Open(ctx context.Context, username, password string, ttl int) error {
+func (c *ksClient) Open(ctx context.Context, shard, username, password string, ttl int) error {
 	req := keystore.OpenRequest{
 		Username: username,
 		Password: password,
 		TTL:      ttl,
 	}
 	var resp keystore.OpenResponse
-	return clientutil.DoJSONHTTPRequest(ctx, c.Client, c.backendURL+"/api/open", &req, &resp)
+	return clientutil.DoJSONHTTPRequest(ctx, c.be.Client(shard), c.be.URL(shard)+"/api/open", &req, &resp)
 }
 
-func (c *Client) Get(ctx context.Context, username, ssoTicket string) ([]byte, error) {
+func (c *ksClient) Get(ctx context.Context, shard, username, ssoTicket string) ([]byte, error) {
 	req := keystore.GetRequest{
 		Username:  username,
 		SSOTicket: ssoTicket,
 	}
 	var resp keystore.GetResponse
-	err := clientutil.DoJSONHTTPRequest(ctx, c.Client, c.backendURL+"/api/get", &req, &resp)
+	err := clientutil.DoJSONHTTPRequest(ctx, c.be.Client(shard), c.be.URL(shard)+"/api/get", &req, &resp)
 	return resp.Key, err
 }
 
-func (c *Client) Close(ctx context.Context, username string) error {
+func (c *ksClient) Close(ctx context.Context, shard, username string) error {
 	req := keystore.CloseRequest{
 		Username: username,
 	}
 	var resp keystore.CloseResponse
-	return clientutil.DoJSONHTTPRequest(ctx, c.Client, c.backendURL+"/api/close", &req, &resp)
+	return clientutil.DoJSONHTTPRequest(ctx, c.be.Client(shard), c.be.URL(shard)+"/api/close", &req, &resp)
 }
diff --git a/vendor/vendor.json b/vendor/vendor.json
index 87be133..67136f2 100644
--- a/vendor/vendor.json
+++ b/vendor/vendor.json
@@ -9,10 +9,10 @@
 			"revisionTime": "2017-12-16T15:39:23Z"
 		},
 		{
-			"checksumSHA1": "2X2UMundICtpGTb8pTdBk7PCKss=",
+			"checksumSHA1": "o+rWKVQIDy79ZwrItwa5/whAL6g=",
 			"path": "git.autistici.org/ai3/go-common/clientutil",
-			"revision": "8cedcb1d73128f5566216cb3e39ad1ccea318213",
-			"revisionTime": "2017-12-16T15:39:23Z"
+			"revision": "9b20acad90c411c48f7ddc837a35ef3d0d6f98d4",
+			"revisionTime": "2017-12-17T20:32:41Z"
 		},
 		{
 			"checksumSHA1": "wY0SM35qAhX3P2IZzDnYa068cPw=",
@@ -39,10 +39,10 @@
 			"revisionTime": "2017-12-15T13:50:57Z"
 		},
 		{
-			"checksumSHA1": "HGK52MX+2CEKVzb9I5y1BfgDkWQ=",
+			"checksumSHA1": "MgtHklQMI/3fNcZZzkg+fmQUrCQ=",
 			"path": "git.autistici.org/id/keystore/client",
-			"revision": "b09f1210471f6a60402e8ced4783be3889a4074f",
-			"revisionTime": "2017-12-15T13:50:57Z"
+			"revision": "ae260514708a9eb3e81b554e1b7b2c65ef802584",
+			"revisionTime": "2017-12-17T21:51:41Z"
 		},
 		{
 			"checksumSHA1": "usT4LCSQItkFvFOQT7cBlkCuGaE=",
-- 
GitLab