From 915799102ba8f553bd683cf370f1e2b42e030d10 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Wed, 7 Jun 2023 13:08:00 -0600
Subject: [PATCH] feat: add support for connection test (#367)

This is a port of
https://github.com/GoogleCloudPlatform/cloud-sql-proxy/pull/1832.
---
 cmd/root.go                  |  7 ++++++
 cmd/root_test.go             | 16 +++++++++++-
 internal/proxy/proxy.go      | 13 ++++++++++
 internal/proxy/proxy_test.go | 49 ++++++++++++++++++++++++++++++++++++
 4 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/cmd/root.go b/cmd/root.go
index c4e8e119..08e3a3dd 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -425,6 +425,9 @@ the maximum time has passed. Defaults to 0s.`)
 		`Enables HTTP endpoints /startup, /liveness, and /readiness
 that report on the proxy's health. Endpoints are available on localhost
 only. Uses the port specified by the http-port flag.`)
+	pflags.BoolVar(&c.conf.RunConnectionTest, "run-connection-test", false, `Runs a connection test
+against all specified instances. If an instance is unreachable, the Proxy exits with a failure
+status code.`)
 
 	// Global and per instance flags
 	pflags.StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
@@ -461,6 +464,10 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 	}
 
 	if conf.FUSEDir != "" {
+		if conf.RunConnectionTest {
+			return newBadCommandError("cannot run connection tests in FUSE mode")
+		}
+
 		if err := proxy.SupportsFUSE(); err != nil {
 			return newBadCommandError(
 				fmt.Sprintf("--fuse is not supported: %v", err),
diff --git a/cmd/root_test.go b/cmd/root_test.go
index 3bbed517..a45bcc50 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -333,6 +333,14 @@ func TestNewCommandArguments(t *testing.T) {
 				QuitQuitQuit: true,
 			}),
 		},
+		{
+			desc: "using the run-connection-test flag",
+			args: []string{"--run-connection-test",
+				"projects/proj/locations/region/clusters/clust/instances/inst"},
+			want: withDefaults(&proxy.Config{
+				RunConnectionTest: true,
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -828,7 +836,13 @@ func TestNewCommandWithErrors(t *testing.T) {
 			desc: "using fuse-tmp-dir without fuse",
 			args: []string{"--fuse-tmp-dir", "/mydir"},
 		},
-	}
+		{
+			desc: "run-connection-test with fuse",
+			args: []string{
+				"--run-connection-test",
+				"--fuse", "myfusedir",
+			},
+		}}
 
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 61268c4d..f2d5cded 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -166,6 +166,10 @@ type Config struct {
 	// OtherUserAgents is a list of space separate user agents that will be
 	// appended to the default user agent.
 	OtherUserAgents string
+
+	// RunConnectionTest determines whether the Proxy should attempt a connection
+	// to all specified instances to verify the network path is valid.
+	RunConnectionTest bool
 }
 
 func parseImpersonationChain(chain string) (string, []string) {
@@ -491,6 +495,15 @@ func (c *Client) Serve(ctx context.Context, notify func()) error {
 		return c.serveFuse(ctx, notify)
 	}
 
+	if c.conf.RunConnectionTest {
+		c.logger.Infof("Connection test started")
+		if _, err := c.CheckConnections(ctx); err != nil {
+			c.logger.Errorf("Connection test failed")
+			return err
+		}
+		c.logger.Infof("Connection test passed")
+	}
+
 	exitCh := make(chan error)
 	for _, m := range c.mnts {
 		go func(mnt *socketMount) {
diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go
index 5ed6604c..050026a1 100644
--- a/internal/proxy/proxy_test.go
+++ b/internal/proxy/proxy_test.go
@@ -677,3 +677,52 @@ func TestCheckConnections(t *testing.T) {
 		t.Fatalf("CheckConnections number of connections: want = %v, got = %v", want, got)
 	}
 }
+
+func TestRunConnectionCheck(t *testing.T) {
+	in := &proxy.Config{
+		Addr: "127.0.0.1",
+		Port: 50002,
+		Instances: []proxy.InstanceConnConfig{
+			{Name: "projects/proj/locations/region/clusters/clust/instances/inst1"},
+		},
+		RunConnectionTest: true,
+	}
+	d := &fakeDialer{}
+	c, err := proxy.NewClient(context.Background(), d, testLogger, in)
+	if err != nil {
+		t.Fatalf("proxy.NewClient error: %v", err)
+	}
+	defer func(c *proxy.Client) {
+		err := c.Close()
+		if err != nil {
+			t.Log(err)
+		}
+	}(c)
+	go func() {
+		// Serve alone without any connections will still verify that the
+		// provided instances are reachable.
+		err := c.Serve(context.Background(), func() {})
+		if err != nil {
+			t.Log(err)
+		}
+	}()
+
+	verifyDialAttempts := func() error {
+		var tries int
+		for {
+			tries++
+			if tries == 10 {
+				return errors.New("failed to verify dial tries after 10 tries")
+			}
+			if got := d.dialAttempts(); got > 0 {
+				return nil
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+
+	if err := verifyDialAttempts(); err != nil {
+		t.Fatal(err)
+	}
+
+}