From ce20bea74c1d9f88301c35e672ad27b7c807daa0 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Mon, 13 May 2024 22:18:04 -0600
Subject: [PATCH] feat: support static connection info

Add support for seeding the connection info cache with static
certificates and IP addresses. This is useful in development contexts
and should otherwise be considered a non-production feature.

With this commit, callers may now provide the path to a JSON file which
contains an RSA key pair, IP addresses, and certificate chains for any
number of AlloyDB instances.

NOTE: the file format is subject to breaking changes. As such this
feature should be considered for development only.
---
 cmd/root.go             | 26 ++++++++++++++++++++++++++
 cmd/root_test.go        | 11 +++++++++++
 internal/proxy/proxy.go | 19 +++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/cmd/root.go b/cmd/root.go
index b0c0f381..c9205656 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -370,6 +370,30 @@ Third Party Licenses
   distribution please see:
 
   https://storage.googleapis.com/alloydb-auth-proxy/v1.9.0/third_party/licenses.tar.gz {x-release-please-version}
+
+Static Connection Info
+
+  In development contexts, it can be helpful to populate the Proxy with static
+  connection info. This is a *dev-only* feature and NOT for use in production.
+  The file format is subject so breaking changes.
+
+  The format is:
+
+  {
+    "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>"
+    }
+  }
 `
 
 var waitHelp = `
@@ -597,6 +621,8 @@ the cached copy has expired. Use this setting in environments where the
 CPU may be throttled and a background refresh cannot run reliably
 (e.g., Cloud Run)`,
 	)
+	localFlags.StringVar(&c.conf.StaticConnectionInfo, "static-connection-info",
+		"", "JSON file with static connection info. See --help for format.")
 
 	// Global and per instance flags
 	localFlags.StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
diff --git a/cmd/root_test.go b/cmd/root_test.go
index 202862c7..0a64dba4 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -510,6 +510,17 @@ func TestNewCommandArguments(t *testing.T) {
 				DebugLogs: true,
 			}),
 		},
+		{
+			desc: "using the static connection info flag",
+			args: []string{
+				"--static-connection-info",
+				"myfile.json",
+				"projects/proj/locations/region/clusters/clust/instances/inst",
+			},
+			want: withDefaults(&proxy.Config{
+				StaticConnectionInfo: "myfile.json",
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index f3bffffd..81c4f82f 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -15,6 +15,7 @@
 package proxy
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -206,6 +207,10 @@ type Config struct {
 	// RunConnectionTest determines whether the Proxy should attempt a connection
 	// to all specified instances to verify the network path is valid.
 	RunConnectionTest bool
+
+	// StaticConnectionInfo is the file path for a static connection info JSON
+	// file. See the proxy help message for details on its format.
+	StaticConnectionInfo string
 }
 
 // dialOptions interprets appropriate dial options for a particular instance
@@ -335,6 +340,20 @@ func (c *Config) DialerOptions(l alloydb.Logger) ([]alloydbconn.Option, error) {
 	if c.LazyRefresh {
 		opts = append(opts, alloydbconn.WithLazyRefresh())
 	}
+	if c.StaticConnectionInfo != "" {
+		f, err := os.Open(c.StaticConnectionInfo)
+		if err != nil {
+			return nil, err
+		}
+		defer f.Close()
+		data, err := io.ReadAll(f)
+		if err != nil {
+			return nil, err
+		}
+		opts = append(opts, alloydbconn.WithStaticConnectionInfo(
+			bytes.NewReader(data),
+		))
+	}
 
 	return opts, nil
 }