feat: Support --universe-domain CLI flag for TPC environments

cmd/root.go

@@ -641,8 +641,10 @@ to close after receiving a TERM signal. The proxy will shut
 down when the number of open connections reaches 0 or when
 the maximum time has passed. Defaults to 0s.`)
 	localFlags.StringVar(&c.conf.APIEndpointURL, "alloydbadmin-api-endpoint",
-		"https://alloydb.googleapis.com",
+		"",
 		"When set, the proxy uses this host as the base API path.")
+	localFlags.StringVar(&c.conf.UniverseDomain, "universe-domain", "",
+		"Universe Domain for non-GDU environments. (default: googleapis.com)")
 	localFlags.StringVar(&c.conf.FUSEDir, "fuse", "",
 		"Mount a directory at the path using FUSE to access AlloyDB instances.")
 	localFlags.StringVar(&c.conf.FUSETempDir, "fuse-tmp-dir",
@@ -906,17 +908,18 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 		return newBadCommandError("cannot specify --json-credentials and --gcloud-auth flags at the same time")
 	}
 
-	if userHasSetLocal(cmd, "alloydbadmin-api-endpoint") {
+	if userHasSetLocal(cmd, "alloydbadmin-api-endpoint") && userHasSetLocal(cmd, "universe-domain") {
+		return newBadCommandError("cannot specify --alloydbadmin-api-endpoint and --universe-domain flags at the same time")
+	}
+	if conf.APIEndpointURL != "" {
+		conf.APIEndpointURL = strings.TrimSuffix(conf.APIEndpointURL, "/")
 		_, err := url.Parse(conf.APIEndpointURL)
 		if err != nil {
 			return newBadCommandError(fmt.Sprintf(
-				"provided value for --alloydbadmin-api-endpoint is not a valid url, %v",
+				"value %q is not a valid URL",
 				conf.APIEndpointURL,
 			))
 		}
-
-		// Remove trailing '/' if included
-		conf.APIEndpointURL = strings.TrimSuffix(conf.APIEndpointURL, "/")
 		cmd.logger.Infof("Using API Endpoint %v", conf.APIEndpointURL)
 	}
 

cmd/root_test.go

@@ -101,9 +101,6 @@ func withDefaults(c *proxy.Config) *proxy.Config {
 	if c.TelemetryTracingSampleRate == 0 {
 		c.TelemetryTracingSampleRate = 10_000
 	}
-	if c.APIEndpointURL == "" {
-		c.APIEndpointURL = "https://alloydb.googleapis.com"
-	}
 	return c
 }
 
@@ -1347,3 +1344,46 @@ func TestQuitQuitQuitWithErrors(t *testing.T) {
 		t.Fatalf("want = %v, got = %v", errCloseFailed, got)
 	}
 }
+
+func TestUniverseDomainFlag(t *testing.T) {
+	t.Parallel()
+	tcs := []struct {
+		desc string
+		args []string
+		want string
+	}{
+		{
+			desc: "with universe domain specified",
+			args: []string{"--universe-domain", "my-universe.cloud", "projects/p/locations/r/clusters/c/instances/i"},
+			want: "my-universe.cloud",
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			c, err := invokeProxyCommand(tc.args)
+			if err != nil {
+				t.Fatalf("invokeProxyCommand failed: %v", err)
+			}
+			if c.conf.UniverseDomain != tc.want {
+				t.Errorf("got %q, want %q", c.conf.UniverseDomain, tc.want)
+			}
+		})
+	}
+}
+
+func TestUniverseDomainMutualExclusion(t *testing.T) {
+	t.Parallel()
+	args := []string{
+		"--alloydbadmin-api-endpoint", "https://alloydb.googleapis.com",
+		"--universe-domain", "my-universe.cloud",
+		"projects/p/locations/r/clusters/c/instances/i",
+	}
+	_, err := invokeProxyCommand(args)
+	if err == nil {
+		t.Fatal("expected error when both --alloydbadmin-api-endpoint and --universe-domain are provided, got nil")
+	}
+	want := "cannot specify --alloydbadmin-api-endpoint and --universe-domain flags at the same time"
+	if !strings.Contains(err.Error(), want) {
+		t.Errorf("got %q, want error containing %q", err, want)
+	}
+}

docs/cmd/alloydb-auth-proxy.md

@@ -297,7 +297,7 @@ alloydb-auth-proxy instance_uri... [flags]
 ```
   -a, --address string                       (*) Address on which to bind AlloyDB instance listeners. (default "127.0.0.1")
       --admin-port string                    Port for localhost-only admin server (default "9091")
-      --alloydbadmin-api-endpoint string     When set, the proxy uses this host as the base API path. (default "https://alloydb.googleapis.com")
+      --alloydbadmin-api-endpoint string     When set, the proxy uses this host as the base API path.
   -i, --auto-iam-authn                       (*) Enables Automatic IAM Authentication for all instances
       --config-file string                   Path to a TOML file containing configuration options.
   -c, --credentials-file string              Path to a service account key to use for authentication.
@@ -354,6 +354,7 @@ alloydb-auth-proxy instance_uri... [flags]
       --telemetry-sample-rate int            Configure the denominator of the probabilistic sample rate of traces sent to Cloud Trace
                                              (e.g., 10,000 traces 1/10,000 calls). (default 10000)
   -t, --token string                         Bearer token used for authorization.
+      --universe-domain string               Universe Domain for non-GDU environments. (default: googleapis.com)
   -u, --unix-socket string                   (*) Enables Unix sockets for all listeners using the provided directory.
       --user-agent string                    Space separated list of additional user agents, e.g. custom-agent/0.0.1
   -v, --version                              Print the alloydb-auth-proxy version

internal/proxy/proxy.go

@@ -17,6 +17,7 @@ package proxy
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -137,6 +138,10 @@ type Config struct {
 	// APIEndpointURL is the URL of the AlloyDB Admin API.
 	APIEndpointURL string
 
+	// UniverseDomain is the universe domain for the TPC environment. When left
+	// blank, the proxy will use the Google Default Universe (GDU): googleapis.com
+	UniverseDomain string
+
 	// Instances are configuration for individual instances. Instance
 	// configuration takes precedence over global configuration.
 	Instances []InstanceConnConfig
@@ -272,6 +277,9 @@ func credentialsOpt(c Config, l alloydb.Logger) (alloydbconn.Option, error) {
 	// credentials token source.
 	if c.ImpersonationChain != "" {
 		var iopts []option.ClientOption
+		if c.UniverseDomain != "" {
+			iopts = append(iopts, option.WithUniverseDomain(c.UniverseDomain))
+		}
 		switch {
 		case c.Token != "":
 			l.Infof("Impersonating service account with OAuth2 token")
@@ -366,6 +374,10 @@ func (c *Config) DialerOptions(l alloydb.Logger) ([]alloydbconn.Option, error) {
 		opts = append(opts, alloydbconn.WithAdminAPIEndpoint(c.APIEndpointURL))
 	}
 
+	if c.UniverseDomain != "" {
+		opts = append(opts, alloydbconn.WithUniverseDomain(c.UniverseDomain))
+	}
+
 	if c.AutoIAMAuthNEnabled() {
 		opts = append(opts, alloydbconn.WithIAMAuthN())
 		switch {
@@ -745,6 +757,9 @@ func (c *Client) serveSocketMount(ctx context.Context, s *socketMount) error {
 		}
 		cConn, err := s.Accept()
 		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				return nil
+			}
 			if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
 				c.logger.Errorf("[%s] Error accepting connection: %v", s.instShort, err)
 				// For transient errors, wait a small amount of time to see if it resolves itself

tests/alloydb_test.go

@@ -57,6 +57,11 @@ projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE>`,
 		os.Getenv("ALLOYDB_DB"),
 		"Name of the database to connect to.",
 	)
+	alloydbUniverseDomain = flag.String(
+		"alloydb_universe_domain",
+		os.Getenv("ALLOYDB_UNIVERSE_DOMAIN"),
+		"Universe domain for the AlloyDB instances.",
+	)
 )
 
 func requirePostgresVars(t *testing.T) {

tests/connection_test.go

@@ -81,6 +81,9 @@ func proxyConnTestWithReady(t *testing.T, args []string, driver, dsn string, rea
 	ctx, cancel := context.WithTimeout(context.Background(), connTestTimeout)
 	defer cancel()
 	// Start the proxy
+	if *alloydbUniverseDomain != "" {
+		args = append(args, "--universe-domain", *alloydbUniverseDomain)
+	}
 	p, err := StartProxy(ctx, args...)
 	if err != nil {
 		t.Fatalf("unable to start proxy: %v", err)