diff --git a/client/client.go b/client/client.go index cbceea16cd9fd25e2b0e925101115df83c546c6b..3988c1f2a6dd36be37831a0820bcdefe0c84572d 100644 --- a/client/client.go +++ b/client/client.go @@ -3,8 +3,10 @@ package client import ( "context" "crypto/tls" + "fmt" "net/http" "net/url" + "sync" "time" "git.autistici.org/ai3/go-common/clientutil" @@ -12,7 +14,7 @@ import ( "git.autistici.org/id/keystore" ) -// Client for the keystore API. +// Client for the keystore API (for a specific backend). type Client struct { *http.Client backendURL string @@ -24,7 +26,8 @@ type Config struct { TLSConfig *clientutil.TLSClientConfig `yaml:"tls_config"` } -// New returns a new Client with the given Config. +// New returns a new Client with the given Config. Use this when the +// keystore service runs on a single global instance. func New(config *Config) (*Client, error) { u, err := url.Parse(config.BackendURL) if err != nil { @@ -76,3 +79,60 @@ func (c *Client) Close(ctx context.Context, username string) error { var resp keystore.CloseResponse return clientutil.DoJSONHTTPRequest(ctx, c.Client, c.backendURL+"/api/close", &req, &resp) } + +// ShardedClient for the keystore API (sharded service). +type ShardedClient struct { + baseURL *url.URL + tlsConfig *tls.Config + mx sync.Mutex + shards map[string]*Client +} + +// NewSharded creates a ShardedClient for the keystore service. Use it +// when the service is partitioned (sharded) across multiple backends. +func NewSharded(config *Config) (*ShardedClient, error) { + u, err := url.Parse(config.BackendURL) + 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 + } + } + + return &ShardedClient{ + baseURL: u, + tlsConfig: tlsConfig, + shards: make(map[string]*Client), + }, nil +} + +func (c *ShardedClient) getShardURL(shard string) *url.URL { + u := *c.baseURL + u.Host = fmt.Sprintf("%s.%s", shard, u.Host) + return &u +} + +// Shard returns the Client for a specific service shard. +func (c *ShardedClient) Shard(shard string) *Client { + c.mx.Lock() + defer c.mx.Unlock() + + client, ok := c.shards[shard] + if !ok { + u := c.getShardURL(shard) + client = &Client{ + Client: &http.Client{ + Transport: clientutil.NewTransport([]string{u.Host}, c.tlsConfig, nil), + Timeout: 20 * time.Second, + }, + backendURL: u.String(), + } + c.shards[shard] = client + } + return client +}