google: add UniverseDomain to CredentialsParams

Change-Id: I7925b8341e1f047d0115acd7a01a34679a489ee0
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/552716
Reviewed-by: Cody Oss <codyoss@google.com>
Run-TryBot: Cody Oss <codyoss@google.com>
Reviewed-by: Viacheslav Rostovtsev <virost@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

2 files changed, 92 insertions, 4 deletions

diff --git a/google/default.go b/google/default.go
index 12b12a3..04ebdc0 100644
--- a/google/default.go
+++ b/google/default.go
@@ -91,6 +91,12 @@ type CredentialsParams struct {
 	// Note: This option is currently only respected when using credentials
 	// fetched from the GCE metadata server.
 	EarlyTokenRefresh time.Duration
+
+	// UniverseDomain is the default service domain for a given Cloud universe.
+	// Only supported in authentication flows that support universe domains.
+	// This value takes precedence over a universe domain explicitly specified
+	// in a credentials config file or by the GCE metadata server. Optional.
+	UniverseDomain string
 }
 
 func (params CredentialsParams) deepCopy() CredentialsParams {
@@ -175,8 +181,9 @@ func FindDefaultCredentialsWithParams(ctx context.Context, params CredentialsPar
 	if metadata.OnGCE() {
 		id, _ := metadata.ProjectID()
 		return &Credentials{
-			ProjectID:   id,
-			TokenSource: computeTokenSource("", params.EarlyTokenRefresh, params.Scopes...),
+			ProjectID:      id,
+			TokenSource:    computeTokenSource("", params.EarlyTokenRefresh, params.Scopes...),
+			universeDomain: params.UniverseDomain,
 		}, nil
 	}
 
@@ -217,6 +224,9 @@ func CredentialsFromJSONWithParams(ctx context.Context, jsonData []byte, params
 	}
 
 	universeDomain := f.UniverseDomain
+	if params.UniverseDomain != "" {
+		universeDomain = params.UniverseDomain
+	}
 	// Authorized user credentials are only supported in the googleapis.com universe.
 	if f.Type == userCredentialsKey {
 		universeDomain = universeDomainDefault
diff --git a/google/default_test.go b/google/default_test.go
index 1f76bae..439887a 100644
--- a/google/default_test.go
+++ b/google/default_test.go
@@ -53,6 +53,10 @@ var userJSONUniverseDomain = []byte(`{
   "universe_domain": "example.com"
 }`)
 
+var universeDomain = "example.com"
+
+var universeDomain2 = "apis-tpclp.goog"
+
 func TestCredentialsFromJSONWithParams_SA(t *testing.T) {
 	ctx := context.Background()
 	scope := "https://www.googleapis.com/auth/cloud-platform"
@@ -72,6 +76,26 @@ func TestCredentialsFromJSONWithParams_SA(t *testing.T) {
 	}
 }
 
+func TestCredentialsFromJSONWithParams_SA_Params_UniverseDomain(t *testing.T) {
+	ctx := context.Background()
+	scope := "https://www.googleapis.com/auth/cloud-platform"
+	params := CredentialsParams{
+		Scopes:         []string{scope},
+		UniverseDomain: universeDomain2,
+	}
+	creds, err := CredentialsFromJSONWithParams(ctx, saJSONJWT, params)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want := "fake_project"; creds.ProjectID != want {
+		t.Fatalf("got %q, want %q", creds.ProjectID, want)
+	}
+	if creds.UniverseDomain() != universeDomain2 {
+		t.Fatalf("got %q, want %q", creds.UniverseDomain(), universeDomain2)
+	}
+}
+
 func TestCredentialsFromJSONWithParams_SA_UniverseDomain(t *testing.T) {
 	ctx := context.Background()
 	scope := "https://www.googleapis.com/auth/cloud-platform"
@@ -86,8 +110,28 @@ func TestCredentialsFromJSONWithParams_SA_UniverseDomain(t *testing.T) {
 	if want := "fake_project"; creds.ProjectID != want {
 		t.Fatalf("got %q, want %q", creds.ProjectID, want)
 	}
-	if want := "example.com"; creds.UniverseDomain() != want {
-		t.Fatalf("got %q, want %q", creds.UniverseDomain(), want)
+	if creds.UniverseDomain() != universeDomain {
+		t.Fatalf("got %q, want %q", creds.UniverseDomain(), universeDomain)
+	}
+}
+
+func TestCredentialsFromJSONWithParams_SA_UniverseDomain_Params_UniverseDomain(t *testing.T) {
+	ctx := context.Background()
+	scope := "https://www.googleapis.com/auth/cloud-platform"
+	params := CredentialsParams{
+		Scopes:         []string{scope},
+		UniverseDomain: universeDomain2,
+	}
+	creds, err := CredentialsFromJSONWithParams(ctx, saJSONJWTUniverseDomain, params)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want := "fake_project"; creds.ProjectID != want {
+		t.Fatalf("got %q, want %q", creds.ProjectID, want)
+	}
+	if creds.UniverseDomain() != universeDomain2 {
+		t.Fatalf("got %q, want %q", creds.UniverseDomain(), universeDomain2)
 	}
 }
 
@@ -107,6 +151,23 @@ func TestCredentialsFromJSONWithParams_User(t *testing.T) {
 	}
 }
 
+func TestCredentialsFromJSONWithParams_User_Params_UniverseDomain(t *testing.T) {
+	ctx := context.Background()
+	scope := "https://www.googleapis.com/auth/cloud-platform"
+	params := CredentialsParams{
+		Scopes:         []string{scope},
+		UniverseDomain: universeDomain2,
+	}
+	creds, err := CredentialsFromJSONWithParams(ctx, userJSON, params)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want := "googleapis.com"; creds.UniverseDomain() != want {
+		t.Fatalf("got %q, want %q", creds.UniverseDomain(), want)
+	}
+}
+
 func TestCredentialsFromJSONWithParams_User_UniverseDomain(t *testing.T) {
 	ctx := context.Background()
 	scope := "https://www.googleapis.com/auth/cloud-platform"
@@ -122,3 +183,20 @@ func TestCredentialsFromJSONWithParams_User_UniverseDomain(t *testing.T) {
 		t.Fatalf("got %q, want %q", creds.UniverseDomain(), want)
 	}
 }
+
+func TestCredentialsFromJSONWithParams_User_UniverseDomain_Params_UniverseDomain(t *testing.T) {
+	ctx := context.Background()
+	scope := "https://www.googleapis.com/auth/cloud-platform"
+	params := CredentialsParams{
+		Scopes:         []string{scope},
+		UniverseDomain: universeDomain2,
+	}
+	creds, err := CredentialsFromJSONWithParams(ctx, userJSONUniverseDomain, params)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want := "googleapis.com"; creds.UniverseDomain() != want {
+		t.Fatalf("got %q, want %q", creds.UniverseDomain(), want)
+	}
+}