feat: support static connection info (#572)

In development environments it can be useful to seed the Dialer with
static connection info properties. This commit makes it possible to
replace the internal lazy or refresh ahead cache with a static cache
that is populated from JSON data.

The JSON format looks like this:

{
    "publicKey": "<PEM Encoded public RSA key>",
    "privateKey": "<PEM Encoded private RSA key>",
    "projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE>": {
        "ipAddress": "<PSA-based private IP address>",
        "publicIpAddress": "<public IP address>",
        "pscInstanceConfig": {
            "pscDnsName": "<PSC DNS name>"
        },
        "pemCertificateChain": [
            "<client cert>", "<intermediate cert>", "<CA cert>"
        ],
        "caCert": "<CA cert>"
    }
}

Author: enocom

dialer.go

@@ -108,6 +108,8 @@ type Dialer struct {
 	// ahead cache assumes a background goroutine may run consistently.
 	lazyRefresh bool
 
+	staticConnInfo io.Reader
+
 	client *alloydbadmin.AlloyDBAdminClient
 	logger debug.Logger
 
@@ -194,6 +196,7 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 		closed:         make(chan struct{}),
 		cache:          make(map[alloydb.InstanceURI]monitoredCache),
 		lazyRefresh:    cfg.lazyRefresh,
+		staticConnInfo: cfg.staticConnInfo,
 		key:            cfg.rsaKey,
 		refreshTimeout: cfg.refreshTimeout,
 		client:         client,
@@ -563,14 +566,25 @@ func (d *Dialer) connectionInfoCache(
 				uri.String(),
 			)
 			var cache connectionInfoCache
-			if d.lazyRefresh {
+			switch {
+			case d.lazyRefresh:
 				cache = alloydb.NewLazyRefreshCache(
 					uri,
 					d.logger,
 					d.client, d.key,
 					d.refreshTimeout, d.dialerID,
 				)
-			} else {
+			case d.staticConnInfo != nil:
+				var err error
+				cache, err = alloydb.NewStaticConnectionInfoCache(
+					uri,
+					d.logger,
+					d.staticConnInfo,
+				)
+				if err != nil {
+					return monitoredCache{}, err
+				}
+			default:
 				cache = alloydb.NewRefreshAheadCache(
 					uri,
 					d.logger,

dialer_test.go

@@ -15,7 +15,13 @@
 package alloydbconn
 
 import (
+	"bytes"
 	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/json"
+	"encoding/pem"
 	"errors"
 	"fmt"
 	"io"
@@ -89,7 +95,81 @@ func TestDialerCanConnectToInstance(t *testing.T) {
 			}
 		})
 	}
+}
+
+func writeStaticInfo(t *testing.T, i mock.FakeAlloyDBInstance) io.Reader {
+	t.Helper()
+	key, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pub := x509.MarshalPKCS1PublicKey(&key.PublicKey)
+	pubPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PUBLIC KEY", Bytes: pub})
+	if pubPEM == nil {
+		t.Fatal("public key encoding failed")
+	}
+	priv := x509.MarshalPKCS1PrivateKey(key)
+	privPEM := pem.EncodeToMemory(
+		&pem.Block{Type: "OPENSSH PRIVATE KEY", Bytes: priv},
+	)
+	if privPEM == nil {
+		t.Fatal("private key encoding failed")
+	}
+
+	static := map[string]interface{}{}
+	static["publicKey"] = string(pubPEM)
+	static["privateKey"] = string(privPEM)
+	info := make(map[string]interface{})
+	info["ipAddress"] = "127.0.0.1" // "private" IP is localhost in testing
+	chain, err := i.GeneratePEMCertificateChain(&key.PublicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+	info["pemCertificateChain"] = chain
+	info["caCert"] = chain[len(chain)-1] // CA cert is last in chain
+	static[i.String()] = info
+
+	data, err := json.Marshal(static)
+	if err != nil {
+		t.Fatal(err)
+	}
 
+	return bytes.NewReader(data)
+}
+
+func TestDialerWorksWithStaticConnectionInfo(t *testing.T) {
+	ctx := context.Background()
+	inst := mock.NewFakeInstance(
+		"my-project", "my-region", "my-cluster", "my-instance",
+	)
+	stop := mock.StartServerProxy(t, inst)
+	t.Cleanup(stop)
+
+	staticPath := writeStaticInfo(t, inst)
+
+	d, err := NewDialer(
+		ctx,
+		WithTokenSource(stubTokenSource{}),
+		WithStaticConnectionInfo(staticPath),
+	)
+	if err != nil {
+		t.Fatalf("expected NewDialer to succeed, but got error: %v", err)
+	}
+
+	conn, err := d.Dial(ctx, testInstanceURI)
+	if err != nil {
+		t.Fatalf("expected Dial to succeed, but got error: %v", err)
+	}
+	defer conn.Close()
+
+	data, err := io.ReadAll(conn)
+	if err != nil {
+		t.Fatalf("expected ReadAll to succeed, got error %v", err)
+	}
+	if string(data) != "my-instance" {
+		t.Fatalf("expected known response from the server, but got %v", string(data))
+	}
 }
 
 func TestDialWithAdminAPIErrors(t *testing.T) {

internal/alloydb/instance.go

@@ -62,6 +62,15 @@ type InstanceURI struct {
 	name    string
 }
 
+// URI returns the full URI specifying an instance.
+func (i *InstanceURI) URI() string {
+	return fmt.Sprintf(
+		"projects/%s/locations/%s/clusters/%s/instances/%s",
+		i.project, i.region, i.cluster, i.name,
+	)
+}
+
+// String returns a short-hand representation of an instance URI.
 func (i *InstanceURI) String() string {
 	return fmt.Sprintf("%s/%s/%s/%s", i.project, i.region, i.cluster, i.name)
 }

internal/alloydb/refresh.go

@@ -149,13 +149,25 @@ func fetchClientCertificate(
 		)
 	}
 
-	certPEMBlock := []byte(strings.Join(resp.PemCertificateChain, "\n"))
 	keyPEMBlock := &pem.Block{
 		Type:  "RSA PRIVATE KEY",
 		Bytes: x509.MarshalPKCS1PrivateKey(key),
 	}
+	keyPEM := pem.EncodeToMemory(keyPEMBlock)
 
-	cert, err := tls.X509KeyPair(certPEMBlock, pem.EncodeToMemory(keyPEMBlock))
+	return newClientCertificate(
+		inst, keyPEM, resp.PemCertificateChain, resp.CaCert,
+	)
+}
+
+func newClientCertificate(
+	inst InstanceURI,
+	keyPEM []byte,
+	chain []string,
+	caCertRaw string,
+) (cc *clientCertificate, err error) {
+	certPEMBlock := []byte(strings.Join(chain, "\n"))
+	cert, err := tls.X509KeyPair(certPEMBlock, keyPEM)
 	if err != nil {
 		return nil, errtype.NewRefreshError(
 			"create ephemeral cert failed",
@@ -164,7 +176,7 @@ func fetchClientCertificate(
 		)
 	}
 
-	caCertPEMBlock, _ := pem.Decode([]byte(resp.CaCert))
+	caCertPEMBlock, _ := pem.Decode([]byte(caCertRaw))
 	if caCertPEMBlock == nil {
 		return nil, errtype.NewRefreshError(
 			"create ephemeral cert failed",
@@ -182,7 +194,7 @@ func fetchClientCertificate(
 	}
 
 	// Extract expiry from client certificate.
-	clientCertPEMBlock, _ := pem.Decode([]byte(resp.PemCertificateChain[0]))
+	clientCertPEMBlock, _ := pem.Decode([]byte(chain[0]))
 	if clientCertPEMBlock == nil {
 		return nil, errtype.NewRefreshError(
 			"create ephemeral cert failed",

internal/alloydb/static.go

@@ -0,0 +1,142 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package alloydb
+
+import (
+	"context"
+	"crypto/x509"
+	"encoding/json"
+	"io"
+
+	"cloud.google.com/go/alloydbconn/debug"
+	"cloud.google.com/go/alloydbconn/errtype"
+)
+
+type staticPSCConfig struct {
+	PSCDNSName string `json:"pscDnsName"`
+}
+
+// staticConnectionInfo is an amalgamation of the generate ephemeral
+// certificate and instance metadata endpoints. Its structure concatenates the
+// IP address information with certificate information. As such it provides all
+// the necessary properties needed for the Dialer to connect to an instance's
+// Auth Proxy server.
+type staticConnectionInfo struct {
+	IPAddress           string          `json:"ipAddress"`
+	PublicIPAddress     string          `json:"publicIPAddress"`
+	PSCInstanceConfig   staticPSCConfig `json:"pscInstanceConfig"`
+	PEMCertificateChain []string        `json:"pemCertificateChain"`
+	CACert              string          `json:"caCert"`
+}
+
+// staticInstanceInfo correlates instance URIs with static connection info.
+type staticInstanceInfo map[string]staticConnectionInfo
+
+// staticData represent a collection of static connection info.
+type staticData struct {
+	PublicKey    string
+	PrivateKey   string
+	InstanceInfo staticInstanceInfo
+}
+
+func (s *staticData) UnmarshalJSON(data []byte) error {
+	inner := map[string]json.RawMessage{}
+	if err := json.Unmarshal(data, &inner); err != nil {
+		return err
+	}
+	if err := json.Unmarshal(inner["privateKey"], &s.PrivateKey); err != nil {
+		return err
+	}
+	delete(inner, "privateKey")
+	if err := json.Unmarshal(inner["publicKey"], &s.PublicKey); err != nil {
+		return err
+	}
+	delete(inner, "publicKey")
+
+	s.InstanceInfo = staticInstanceInfo{}
+	for k, v := range inner {
+		var sci staticConnectionInfo
+		if err := json.Unmarshal(v, &sci); err != nil {
+			return err
+		}
+		s.InstanceInfo[k] = sci
+	}
+	return nil
+}
+
+// StaticConnectionInfoCache provides connection info that is never refreshed.
+type StaticConnectionInfoCache struct {
+	logger debug.Logger
+	info   ConnectionInfo
+}
+
+// NewStaticConnectionInfoCache creates a connection info cache that will
+// always return the predefined connection info within the provided io.Reader
+func NewStaticConnectionInfoCache(
+	inst InstanceURI,
+	l debug.Logger,
+	r io.Reader,
+) (*StaticConnectionInfoCache, error) {
+	data, err := io.ReadAll(r)
+	if err != nil {
+		return nil, err
+	}
+	var d staticData
+	if err := json.Unmarshal(data, &d); err != nil {
+		return nil, err
+	}
+	static, ok := d.InstanceInfo[inst.URI()]
+	if !ok {
+		return nil, errtype.NewConfigError("unknown instance", inst.String())
+	}
+	cc, err := newClientCertificate(
+		inst, []byte(d.PrivateKey), static.PEMCertificateChain, static.CACert,
+	)
+	if err != nil {
+		return nil, err
+	}
+	pool := x509.NewCertPool()
+	pool.AddCert(cc.caCert)
+	info := ConnectionInfo{
+		Instance: inst,
+		IPAddrs: map[string]string{
+			PublicIP:  static.PublicIPAddress,
+			PrivateIP: static.IPAddress,
+			PSC:       static.PSCInstanceConfig.PSCDNSName,
+		},
+		ClientCert: cc.certChain,
+		RootCAs:    pool,
+		Expiration: cc.expiry,
+	}
+	return &StaticConnectionInfoCache{
+		logger: l,
+		info:   info,
+	}, nil
+}
+
+// ConnectionInfo returns the connection info for the specified instance URI as
+// loaded from the provided io.Reader.
+func (c *StaticConnectionInfoCache) ConnectionInfo(
+	_ context.Context,
+) (ConnectionInfo, error) {
+	return c.info, nil
+}
+
+// ForceRefresh is a no-op as the cache holds only static connection
+// information and does no refresh.
+func (*StaticConnectionInfoCache) ForceRefresh() {}
+
+// Close is a no-op.
+func (*StaticConnectionInfoCache) Close() error { return nil }