From 88c6632f4ab87aad70d27efb781a1ef594f88199 Mon Sep 17 00:00:00 2001
From: aeitzman <eitzman@google.com>
Date: Tue, 9 Apr 2024 12:54:02 -0700
Subject: [PATCH] fix: makes default token url universe aware

---
 .../oauth2/ExternalAccountCredentials.java     | 16 ++++++++++++++--
 .../oauth2/ExternalAccountCredentialsTest.java | 18 ++++++++++++++++++
 2 files changed, 32 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..f2e346d89 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,19 @@ 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){
+      try {
+        this.tokenUrl = DEFAULT_TOKEN_URL.replace("{UNIVERSE_DOMAIN}", this.getUniverseDomain());
+      } catch (IOException e) {
+        // Throwing an IOException would be a breaking change, so wrap it here.
+        throw new IllegalStateException(
+            "Error occurred when attempting to retrieve universe domain.", e);
+      }
+    } else {
+      this.tokenUrl = builder.tokenUrl;
+    }
+
     this.scopes =
         (builder.scopes == null || builder.scopes.isEmpty())
             ? Arrays.asList(CLOUD_PLATFORM_SCOPE)
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<>();