diff --git a/cmd/root.go b/cmd/root.go index d6d6cb6d..5ddaadce 100644 --- a/cmd/root.go +++ b/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) } diff --git a/cmd/root_test.go b/cmd/root_test.go index 23c1384b..bf92314b 100644 --- a/cmd/root_test.go +++ b/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) + } +} diff --git a/docs/cmd/alloydb-auth-proxy.md b/docs/cmd/alloydb-auth-proxy.md index bff89a70..47b4b218 100644 --- a/docs/cmd/alloydb-auth-proxy.md +++ b/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 diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 72ab24db..88ced0a4 100644 --- a/internal/proxy/proxy.go +++ b/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 diff --git a/tests/alloydb_test.go b/tests/alloydb_test.go index f884953d..1b4b83fe 100644 --- a/tests/alloydb_test.go +++ b/tests/alloydb_test.go @@ -57,6 +57,11 @@ projects//locations//clusters//instances/`, 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) { diff --git a/tests/connection_test.go b/tests/connection_test.go index a215f799..3d3c7b9f 100644 --- a/tests/connection_test.go +++ b/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)