From e3caf05831011dc05d3a8b01ebf79305eda70183 Mon Sep 17 00:00:00 2001
From: aeitzman <12433791+aeitzman@users.noreply.github.com>
Date: Thu, 11 Apr 2024 13:13:08 -0700
Subject: [PATCH] fix: makes default token url universe aware (#1383)

* fix: makes default token url universe aware

* lint and add test

* Update oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java

Co-authored-by: Leo <39062083+lsirac@users.noreply.github.com>

* add back else

* move code into override

---------

Co-authored-by: Leo <39062083+lsirac@users.noreply.github.com>
---
 .../oauth2/ExternalAccountCredentials.java    | 21 +++++++++++++++++--
 .../ExternalAccountCredentialsTest.java       | 18 ++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java
index ad9633da8..a9b3ef9eb 100644
--- a/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java
+++ b/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java
@@ -73,7 +73,7 @@ public abstract class ExternalAccountCredentials extends GoogleCredentials {
   static final String EXTERNAL_ACCOUNT_FILE_TYPE = "external_account";
   static final String EXECUTABLE_SOURCE_KEY = "executable";
 
-  static final String DEFAULT_TOKEN_URL = "https://sts.googleapis.com/v1/token";
+  static final String DEFAULT_TOKEN_URL = "https://sts.{UNIVERSE_DOMAIN}/v1/token";
   static final String PROGRAMMATIC_METRICS_HEADER_VALUE = "programmatic";
 
   private final String transportFactoryClassName;
@@ -235,7 +235,13 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
     this.serviceAccountImpersonationUrl = builder.serviceAccountImpersonationUrl;
     this.clientId = builder.clientId;
     this.clientSecret = builder.clientSecret;
-    this.tokenUrl = builder.tokenUrl == null ? DEFAULT_TOKEN_URL : builder.tokenUrl;
+
+    if (builder.tokenUrl == null) {
+      this.tokenUrl = DEFAULT_TOKEN_URL.replace("{UNIVERSE_DOMAIN}", this.getUniverseDomain());
+    } else {
+      this.tokenUrl = builder.tokenUrl;
+    }
+
     this.scopes =
         (builder.scopes == null || builder.scopes.isEmpty())
             ? Arrays.asList(CLOUD_PLATFORM_SCOPE)
@@ -321,6 +327,17 @@ public void onFailure(Throwable exception) {
         });
   }
 
+  @Override
+  public String getUniverseDomain() {
+    try {
+      return super.getUniverseDomain();
+    } catch (IOException e) {
+      // Throwing an IOException would be a breaking change, so wrap it here.
+      // This should not happen for this credential type.
+      throw new IllegalStateException(e);
+    }
+  }
+
   @Override
   public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
     Map<String, List<String>> requestMetadata = super.getRequestMetadata(uri);
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java
index 986669c9c..9cefedb8c 100644
--- a/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java
+++ b/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java
@@ -565,6 +565,24 @@ public void constructor_builder_defaultTokenUrl() {
     assertEquals(STS_URL, credentials.getTokenUrl());
   }
 
+  @Test
+  public void constructor_builder_defaultTokenUrlwithUniverseDomain() {
+    HashMap<String, Object> credentialSource = new HashMap<>();
+    credentialSource.put("file", "file");
+
+    ExternalAccountCredentials credentials =
+        IdentityPoolCredentials.newBuilder()
+            .setHttpTransportFactory(transportFactory)
+            .setAudience(
+                "//iam.googleapis.com/locations/global/workforcePools/pool/providers/provider")
+            .setSubjectTokenType("subjectTokenType")
+            .setCredentialSource(new TestCredentialSource(credentialSource))
+            .setUniverseDomain("testdomain.org")
+            .build();
+
+    assertEquals("https://sts.testdomain.org/v1/token", credentials.getTokenUrl());
+  }
+
   @Test
   public void constructor_builder_subjectTokenTypeEnum() {
     HashMap<String, Object> credentialSource = new HashMap<>();