From deefa7e836e5aa043b2b8102b5b098da53a43903 Mon Sep 17 00:00:00 2001
From: Chris Smith <chrisdsmith@google.com>
Date: Fri, 19 Jan 2024 11:51:13 -0700
Subject: [PATCH] google/downscope: add DownscopingConfig.UniverseDomain to
 support TPC

Change-Id: I3669352b382414ea640ca176afa4071995fc5ff1
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/557135
Reviewed-by: Cody Oss <codyoss@google.com>
TryBot-Bypass: Cody Oss <codyoss@google.com>
Auto-Submit: Cody Oss <codyoss@google.com>
---
 google/downscope/downscoping.go      | 30 ++++++++++++++++++++----
 google/downscope/downscoping_test.go | 35 ++++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/google/downscope/downscoping.go b/google/downscope/downscoping.go
index 3d4b5532d..ca1f35462 100644
--- a/google/downscope/downscoping.go
+++ b/google/downscope/downscoping.go
@@ -42,13 +42,16 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	"golang.org/x/oauth2"
 )
 
-var (
-	identityBindingEndpoint = "https://sts.googleapis.com/v1/token"
+const (
+	universeDomainPlaceholder       = "UNIVERSE_DOMAIN"
+	identityBindingEndpointTemplate = "https://sts.UNIVERSE_DOMAIN/v1/token"
+	universeDomainDefault           = "googleapis.com"
 )
 
 type accessBoundary struct {
@@ -105,6 +108,18 @@ type DownscopingConfig struct {
 	// access (or set of accesses) that the new token has to a given resource.
 	// There can be a maximum of 10 AccessBoundaryRules.
 	Rules []AccessBoundaryRule
+	// UniverseDomain is the default service domain for a given Cloud universe.
+	// The default value is "googleapis.com". Optional.
+	UniverseDomain string
+}
+
+// identityBindingEndpoint returns the identity binding endpoint with the
+// configured universe domain.
+func (dc *DownscopingConfig) identityBindingEndpoint() string {
+	if dc.UniverseDomain == "" {
+		return strings.Replace(identityBindingEndpointTemplate, universeDomainPlaceholder, universeDomainDefault, 1)
+	}
+	return strings.Replace(identityBindingEndpointTemplate, universeDomainPlaceholder, dc.UniverseDomain, 1)
 }
 
 // A downscopingTokenSource is used to retrieve a downscoped token with restricted
@@ -114,6 +129,9 @@ type downscopingTokenSource struct {
 	ctx context.Context
 	// config holds the information necessary to generate a downscoped Token.
 	config DownscopingConfig
+	// identityBindingEndpoint is the identity binding endpoint with the
+	// configured universe domain.
+	identityBindingEndpoint string
 }
 
 // NewTokenSource returns a configured downscopingTokenSource.
@@ -135,7 +153,11 @@ func NewTokenSource(ctx context.Context, conf DownscopingConfig) (oauth2.TokenSo
 			return nil, fmt.Errorf("downscope: all rules must provide at least one permission: %+v", val)
 		}
 	}
-	return downscopingTokenSource{ctx: ctx, config: conf}, nil
+	return downscopingTokenSource{
+		ctx:                     ctx,
+		config:                  conf,
+		identityBindingEndpoint: conf.identityBindingEndpoint(),
+	}, nil
 }
 
 // Token() uses a downscopingTokenSource to generate an oauth2 Token.
@@ -171,7 +193,7 @@ func (dts downscopingTokenSource) Token() (*oauth2.Token, error) {
 	form.Add("options", string(b))
 
 	myClient := oauth2.NewClient(dts.ctx, nil)
-	resp, err := myClient.PostForm(identityBindingEndpoint, form)
+	resp, err := myClient.PostForm(dts.identityBindingEndpoint, form)
 	if err != nil {
 		return nil, fmt.Errorf("unable to generate POST Request %v", err)
 	}
diff --git a/google/downscope/downscoping_test.go b/google/downscope/downscoping_test.go
index d5adda19c..ecdd98691 100644
--- a/google/downscope/downscoping_test.go
+++ b/google/downscope/downscoping_test.go
@@ -38,18 +38,43 @@ func Test_DownscopedTokenSource(t *testing.T) {
 		w.Write([]byte(standardRespBody))
 
 	}))
-	new := []AccessBoundaryRule{
+	myTok := oauth2.Token{AccessToken: "Mellon"}
+	tmpSrc := oauth2.StaticTokenSource(&myTok)
+	rules := []AccessBoundaryRule{
 		{
 			AvailableResource:    "test1",
 			AvailablePermissions: []string{"Perm1", "Perm2"},
 		},
 	}
-	myTok := oauth2.Token{AccessToken: "Mellon"}
-	tmpSrc := oauth2.StaticTokenSource(&myTok)
-	dts := downscopingTokenSource{context.Background(), DownscopingConfig{tmpSrc, new}}
-	identityBindingEndpoint = ts.URL
+	dts := downscopingTokenSource{
+		ctx: context.Background(),
+		config: DownscopingConfig{
+			RootSource: tmpSrc,
+			Rules:      rules,
+		},
+		identityBindingEndpoint: ts.URL,
+	}
 	_, err := dts.Token()
 	if err != nil {
 		t.Fatalf("NewDownscopedTokenSource failed with error: %v", err)
 	}
 }
+
+func Test_DownscopingConfig(t *testing.T) {
+	tests := []struct {
+		universeDomain string
+		want           string
+	}{
+		{"", "https://sts.googleapis.com/v1/token"},
+		{"googleapis.com", "https://sts.googleapis.com/v1/token"},
+		{"example.com", "https://sts.example.com/v1/token"},
+	}
+	for _, tt := range tests {
+		c := DownscopingConfig{
+			UniverseDomain: tt.universeDomain,
+		}
+		if got := c.identityBindingEndpoint(); got != tt.want {
+			t.Errorf("got %q, want %q", got, tt.want)
+		}
+	}
+}