Fix Astra timeouts (#96)

This adds Context and timeouts to several paths where calls could block when contacting Astra services. The timeout can be changed by providing the --astra-timeout flag with a duration.

Author: mpenick

README.md

@@ -45,6 +45,7 @@ Flags:
                                                           or contact points option ($ASTRA_TOKEN).
   -i, --astra-database-id=STRING                          Database ID of the Astra database. Requires '--astra-token' ($ASTRA_DATABASE_ID)
       --astra-api-url="https://api.astra.datastax.com"    URL for the Astra API ($ASTRA_API_URL)
+      --astra-timeout=10s                                 Timeout for contacting Astra when retrieving the bundle and metadata ($ASTRA_TIMEOUT)
   -c, --contact-points=CONTACT-POINTS,...                 Contact points for cluster. Ignored if using the bundle path or token option ($CONTACT_POINTS).
   -u, --username=STRING                                   Username to use for authentication ($USERNAME)
   -p, --password=STRING                                   Password to use for authentication ($PASSWORD)

astra/bundle.go

@@ -33,9 +33,9 @@ import (
 )
 
 type Bundle struct {
-	tlsConfig *tls.Config
-	host      string
-	port      int
+	TLSConfig *tls.Config
+	Host      string
+	Port      int
 }
 
 func LoadBundleZip(reader *zip.Reader) (*Bundle, error) {
@@ -69,13 +69,13 @@ func LoadBundleZip(reader *zip.Reader) (*Bundle, error) {
 	}
 
 	return &Bundle{
-		tlsConfig: &tls.Config{
+		TLSConfig: &tls.Config{
 			RootCAs:      rootCAs,
 			Certificates: []tls.Certificate{cert},
 			ServerName:   config.Host,
 		},
-		host: config.Host,
-		port: config.Port,
+		Host: config.Host,
+		Port: config.Port,
 	}, nil
 }
 
@@ -93,7 +93,7 @@ func LoadBundleZipFromPath(path string) (*Bundle, error) {
 }
 
 func LoadBundleZipFromURL(url, databaseID, token string, timeout time.Duration) (*Bundle, error) {
-	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(timeout))
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
 	credsURL, err := generateSecureBundleURLWithResponse(url, databaseID, token, ctx)
@@ -154,18 +154,6 @@ func generateSecureBundleURLWithResponse(url, databaseID, token string, ctx cont
 	return res.JSON200, nil
 }
 
-func (b *Bundle) Host() string {
-	return b.host
-}
-
-func (b *Bundle) Port() int {
-	return b.port
-}
-
-func (b *Bundle) TLSConfig() *tls.Config {
-	return b.tlsConfig.Clone()
-}
-
 func extract(reader *zip.Reader) (map[string][]byte, error) {
 	contents := make(map[string][]byte)
 

astra/bundle_test.go

@@ -40,22 +40,22 @@ func TestLoadBundleZip(t *testing.T) {
 	b, err := LoadBundleZipFromPath(path)
 	require.NoError(t, err)
 
-	assert.Equal(t, hostname, b.Host())
-	assert.Equal(t, port, b.Port())
+	assert.Equal(t, hostname, b.Host)
+	assert.Equal(t, port, b.Port)
 
 	block, _ := pem.Decode(testCAPEM)
 	ca, _ := x509.ParseCertificate(block.Bytes)
 
 	// Verify CA added to cert pool
 	caSub, err := asn1.Marshal(ca.Subject.ToRDNSequence())
 	found := false
-	for _, sub := range b.TLSConfig().RootCAs.Subjects() {
+	for _, sub := range b.TLSConfig.RootCAs.Subjects() {
 		if bytes.Compare(caSub, sub) == 0 {
 			found = true
 		}
 	}
 	assert.True(t, found)
-	require.Equal(t, 1, len(b.TLSConfig().Certificates))
+	require.Equal(t, 1, len(b.TLSConfig.Certificates))
 }
 
 func TestLoadBundleZip_InvalidJson(t *testing.T) {

astra/endpoint.go

@@ -15,12 +15,12 @@
 package astra
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"net/http"
 	"sync"
 	"time"
@@ -33,6 +33,7 @@ type astraResolver struct {
 	region          string
 	bundle          *Bundle
 	mu              *sync.Mutex
+	timeout         time.Duration
 }
 
 type astraEndpoint struct {
@@ -41,28 +42,38 @@ type astraEndpoint struct {
 	tlsConfig *tls.Config
 }
 
-func NewResolver(bundle *Bundle) proxycore.EndpointResolver {
+func NewResolver(bundle *Bundle, timeout time.Duration) proxycore.EndpointResolver {
 	return &astraResolver{
-		bundle: bundle,
-		mu:     &sync.Mutex{},
+		bundle:  bundle,
+		mu:      &sync.Mutex{},
+		timeout: timeout,
 	}
 }
 
-func (r *astraResolver) Resolve() ([]proxycore.Endpoint, error) {
+func (r *astraResolver) Resolve(ctx context.Context) ([]proxycore.Endpoint, error) {
 	var metadata *astraMetadata
 
-	url := fmt.Sprintf("https://%s:%d/metadata", r.bundle.Host(), r.bundle.Port())
+	ctx, cancel := context.WithTimeout(ctx, r.timeout)
+	defer cancel()
+
 	httpsClient := &http.Client{
 		Transport: &http.Transport{
-			TLSClientConfig: r.bundle.TLSConfig(),
+			TLSClientConfig: r.bundle.TLSConfig.Clone(),
 		},
 	}
-	response, err := httpsClient.Get(url)
+
+	url := fmt.Sprintf("https://%s:%d/metadata", r.bundle.Host, r.bundle.Port)
+	req, err := http.NewRequestWithContext(ctx, "GET", url, http.NoBody)
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := httpsClient.Do(req)
 	if err != nil {
-		return nil, fmt.Errorf("unable to get metadata from %s: %v", url, err)
+		return nil, fmt.Errorf("unable to get metadata from %s: %w", url, err)
 	}
 
-	body, err := ioutil.ReadAll(response.Body)
+	body, err := readAllWithTimeout(response.Body, ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -145,7 +156,7 @@ func (a astraEndpoint) TLSConfig() *tls.Config {
 }
 
 func copyTLSConfig(bundle *Bundle, serverName string) *tls.Config {
-	tlsConfig := bundle.TLSConfig()
+	tlsConfig := bundle.TLSConfig.Clone()
 	tlsConfig.ServerName = serverName
 	tlsConfig.InsecureSkipVerify = true
 	tlsConfig.VerifyPeerCertificate = func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
@@ -161,7 +172,7 @@ func copyTLSConfig(bundle *Bundle, serverName string) *tls.Config {
 		opts := x509.VerifyOptions{
 			Roots:         tlsConfig.RootCAs,
 			CurrentTime:   time.Now(),
-			DNSName:       bundle.Host(),
+			DNSName:       bundle.Host,
 			Intermediates: x509.NewCertPool(),
 		}
 		for _, cert := range certs[1:] {

astra/endpoint_test.go

@@ -15,13 +15,15 @@
 package astra
 
 import (
+	"context"
 	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"net"
 	"net/http"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/datastax/cql-proxy/proxycore"
 	"github.com/datastax/go-cassandra-native-protocol/datatype"
@@ -51,7 +53,7 @@ func TestMain(m *testing.M) {
 
 func TestAstraResolver_Resolve(t *testing.T) {
 	resolver := createResolver(t)
-	endpoints, err := resolver.Resolve()
+	endpoints, err := resolver.Resolve(context.Background())
 	require.NoError(t, err)
 
 	for _, endpoint := range endpoints {
@@ -63,7 +65,7 @@ func TestAstraResolver_Resolve(t *testing.T) {
 
 func TestAstraResolver_NewEndpoint(t *testing.T) {
 	resolver := createResolver(t)
-	_, err := resolver.Resolve()
+	_, err := resolver.Resolve(context.Background())
 	require.NoError(t, err)
 
 	const hostId = "a2e24181-d732-402a-ab06-894a8b2f6094"
@@ -101,7 +103,7 @@ func TestAstraResolver_NewEndpoint(t *testing.T) {
 
 func TestAstraResolver_NewEndpoint_Ignored(t *testing.T) {
 	resolver := createResolver(t)
-	_, err := resolver.Resolve()
+	_, err := resolver.Resolve(context.Background())
 	require.NoError(t, err)
 
 	const hostId = "a2e24181-d732-402a-ab06-894a8b2f6094"
@@ -138,7 +140,7 @@ func TestAstraResolver_NewEndpoint_Ignored(t *testing.T) {
 
 func TestAstraResolver_NewEndpointInvalidHostID(t *testing.T) {
 	resolver := createResolver(t)
-	_, err := resolver.Resolve()
+	_, err := resolver.Resolve(context.Background())
 	require.NoError(t, err)
 
 	rs := proxycore.NewResultSet(&message.RowsResult{
@@ -164,14 +166,23 @@ func TestAstraResolver_NewEndpointInvalidHostID(t *testing.T) {
 	assert.Error(t, err, "ignoring host because its `host_id` is not set or is invalid")
 }
 
+func TestAstraResolver_Timeout(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 1) // Very short timeout
+	defer cancel()
+
+	resolver := createResolver(t)
+	_, err := resolver.Resolve(ctx)
+	assert.ErrorIs(t, err, context.DeadlineExceeded) // Expect a timeout
+}
+
 func createResolver(t *testing.T) proxycore.EndpointResolver {
 	path, err := writeBundle("127.0.0.1", 8080)
 	require.NoError(t, err)
 
 	bundle, err := LoadBundleZipFromPath(path)
 	require.NoError(t, err)
 
-	return NewResolver(bundle)
+	return NewResolver(bundle, 10*time.Second)
 }
 
 func runTestMetaSvcAsync(sniProxyAddr string, contactPoints []string) (*http.Server, error) {

proxy/run.go

@@ -44,6 +44,7 @@ type runConfig struct {
 	AstraToken         string        `yaml:"astra-token" help:"Token used to authenticate to an Astra database. Requires '--astra-database-id'. Ignored if using the bundle path or contact points option." short:"t" env:"ASTRA_TOKEN"`
 	AstraDatabaseID    string        `yaml:"astra-database-id" help:"Database ID of the Astra database. Requires '--astra-token'" short:"i" env:"ASTRA_DATABASE_ID"`
 	AstraApiURL        string        `yaml:"astra-api-url" help:"URL for the Astra API" default:"https://api.astra.datastax.com" env:"ASTRA_API_URL"`
+	AstraTimeout       time.Duration `yaml:"astra-timeout" help:"Timeout for contacting Astra when retrieving the bundle and metadata" default:"10s" env:"ASTRA_TIMEOUT"`
 	ContactPoints      []string      `yaml:"contact-points" help:"Contact points for cluster. Ignored if using the bundle path or token option." short:"c" env:"CONTACT_POINTS"`
 	Username           string        `yaml:"username" help:"Username to use for authentication" short:"u" env:"USERNAME"`
 	Password           string        `yaml:"password" help:"Password to use for authentication" short:"p" env:"PASSWORD"`
@@ -101,19 +102,18 @@ func Run(ctx context.Context, args []string) int {
 	if len(cfg.AstraBundle) > 0 {
 		if bundle, err := astra.LoadBundleZipFromPath(cfg.AstraBundle); err != nil {
 			cliCtx.Errorf("unable to open bundle %s from file: %v", cfg.AstraBundle, err)
-			return 1
 		} else {
-			resolver = astra.NewResolver(bundle)
+			resolver = astra.NewResolver(bundle, cfg.AstraTimeout)
 		}
 	} else if len(cfg.AstraToken) > 0 {
 		if len(cfg.AstraDatabaseID) == 0 {
 			cliCtx.Fatalf("database ID is required when using a token")
 		}
-		bundle, err := astra.LoadBundleZipFromURL(cfg.AstraApiURL, cfg.AstraDatabaseID, cfg.AstraToken, 10*time.Second)
-		if err != nil {
+		if bundle, err := astra.LoadBundleZipFromURL(cfg.AstraApiURL, cfg.AstraDatabaseID, cfg.AstraToken, cfg.AstraTimeout); err != nil {
 			cliCtx.Fatalf("unable to load bundle %s from astra: %v", cfg.AstraBundle, err)
+		} else {
+			resolver = astra.NewResolver(bundle, cfg.AstraTimeout)
 		}
-		resolver = astra.NewResolver(bundle)
 		cfg.Username = "token"
 		cfg.Password = cfg.AstraToken
 	} else if len(cfg.ContactPoints) > 0 {

proxycore/cluster.go

@@ -150,7 +150,7 @@ func ConnectCluster(ctx context.Context, config ClusterConfig) (*Cluster, error)
 		listeners:        make([]ClusterListener, 0),
 	}
 
-	endpoints, err := config.Resolver.Resolve()
+	endpoints, err := config.Resolver.Resolve(ctx)
 	if err != nil {
 		return nil, err
 	}

proxycore/endpoint.go

@@ -15,6 +15,7 @@
 package proxycore
 
 import (
+	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -59,7 +60,7 @@ func (e defaultEndpoint) TLSConfig() *tls.Config {
 }
 
 type EndpointResolver interface {
-	Resolve() ([]Endpoint, error)
+	Resolve(ctx context.Context) ([]Endpoint, error)
 	NewEndpoint(row Row) (Endpoint, error)
 }
 
@@ -87,14 +88,15 @@ func NewResolverWithDefaultPort(contactPoints []string, defaultPort int) Endpoin
 	}
 }
 
-func (r *defaultEndpointResolver) Resolve() ([]Endpoint, error) {
+func (r *defaultEndpointResolver) Resolve(ctx context.Context) ([]Endpoint, error) {
 	var endpoints []Endpoint
+	var resolver net.Resolver
 	for _, cp := range r.contactPoints {
 		host, port, err := net.SplitHostPort(cp)
 		if err != nil {
 			host = cp
 		}
-		addrs, err := net.LookupHost(host)
+		addrs, err := resolver.LookupHost(ctx, host)
 		if err != nil {
 			return nil, fmt.Errorf("unable to resolve contact point %s: %v", cp, err)
 		}