From 81830e6848ceefd055aa4d08f933d1154455a0f6 Mon Sep 17 00:00:00 2001
From: Chris Smith <chrisdsmith@google.com>
Date: Thu, 4 Apr 2024 09:27:27 -0600
Subject: [PATCH] feat(auth/credentials/externalaccount): add default TokenURL
 (#9700)

---
 .../externalaccount/externalaccount.go        | 19 +++++++-
 .../externalaccount/externalaccount_test.go   | 44 +++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/auth/credentials/internal/externalaccount/externalaccount.go b/auth/credentials/internal/externalaccount/externalaccount.go
index b33279387fe3..b19c6edeae5a 100644
--- a/auth/credentials/internal/externalaccount/externalaccount.go
+++ b/auth/credentials/internal/externalaccount/externalaccount.go
@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"regexp"
 	"strconv"
+	"strings"
 	"time"
 
 	"cloud.google.com/go/auth"
@@ -32,6 +33,10 @@ import (
 const (
 	timeoutMinimum = 5 * time.Second
 	timeoutMaximum = 120 * time.Second
+
+	universeDomainPlaceholder = "UNIVERSE_DOMAIN"
+	defaultTokenURL           = "https://sts.UNIVERSE_DOMAIN/v1/token"
+	defaultUniverseDomain     = "googleapis.com"
 )
 
 var (
@@ -176,12 +181,25 @@ func (o *Options) validate() error {
 	return nil
 }
 
+// resolveTokenURL sets the default STS token endpoint with the configured
+// universe domain.
+func (o *Options) resolveTokenURL() {
+	if o.TokenURL != "" {
+		return
+	} else if o.UniverseDomain != "" {
+		o.TokenURL = strings.Replace(defaultTokenURL, universeDomainPlaceholder, o.UniverseDomain, 1)
+	} else {
+		o.TokenURL = strings.Replace(defaultTokenURL, universeDomainPlaceholder, defaultUniverseDomain, 1)
+	}
+}
+
 // NewTokenProvider returns a [cloud.google.com/go/auth.TokenProvider]
 // configured with the provided options.
 func NewTokenProvider(opts *Options) (auth.TokenProvider, error) {
 	if err := opts.validate(); err != nil {
 		return nil, err
 	}
+	opts.resolveTokenURL()
 	stp, err := newSubjectTokenProvider(opts)
 	if err != nil {
 		return nil, err
@@ -282,7 +300,6 @@ func (tp *tokenProvider) Token(ctx context.Context) (*auth.Token, error) {
 // subjectTokenProvider
 func newSubjectTokenProvider(o *Options) (subjectTokenProvider, error) {
 	reqOpts := &RequestOptions{Audience: o.Audience, SubjectTokenType: o.SubjectTokenType}
-
 	if o.AwsSecurityCredentialsProvider != nil {
 		return &awsSubjectProvider{
 			securityCredentialsProvider: o.AwsSecurityCredentialsProvider,
diff --git a/auth/credentials/internal/externalaccount/externalaccount_test.go b/auth/credentials/internal/externalaccount/externalaccount_test.go
index 50fa912be264..6e32c98a910f 100644
--- a/auth/credentials/internal/externalaccount/externalaccount_test.go
+++ b/auth/credentials/internal/externalaccount/externalaccount_test.go
@@ -462,3 +462,47 @@ func TestOptionsValidate(t *testing.T) {
 		})
 	}
 }
+
+func TestOptionsResolveTokenURL(t *testing.T) {
+	tests := []struct {
+		name string
+		o    *Options
+		want string
+	}{
+		{
+			name: "default",
+			o:    &Options{},
+			want: "https://sts.googleapis.com/v1/token",
+		},
+		{
+			name: "Options TokenURL",
+			o: &Options{
+				TokenURL: "http://localhost:8080/v1/token",
+			},
+			want: "http://localhost:8080/v1/token",
+		},
+		{
+			name: "Options UniverseDomain",
+			o: &Options{
+				UniverseDomain: "example.com",
+			},
+			want: "https://sts.example.com/v1/token",
+		},
+		{
+			name: "Options TokenURL overrides UniverseDomain",
+			o: &Options{
+				TokenURL:       "http://localhost:8080/v1/token",
+				UniverseDomain: "example.com",
+			},
+			want: "http://localhost:8080/v1/token",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.o.resolveTokenURL()
+			if tc.o.TokenURL != tc.want {
+				t.Errorf("got %s, want %s", tc.o.TokenURL, tc.want)
+			}
+		})
+	}
+}