feat: ImpersonatedCredentials to support universe domain for idtoken and signblob (#1566)

zhumin8 · web-flow · commit adc2ff3dcabb · 2025-01-21T13:11:48.000-05:00
follow up to #1528. idtoken and sign flow are tested E2E according to TPC test guide for sa-to-sa impersonation.
diff --git a/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java
@@ -95,9 +95,6 @@ public class ComputeEngineCredentials extends GoogleCredentials
 
   static final String DEFAULT_METADATA_SERVER_URL = "http://metadata.google.internal";
 
-  static final String SIGN_BLOB_URL_FORMAT =
-      "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:signBlob";
-
   // Note: the explicit `timeout` and `tries` below is a workaround. The underlying
   // issue is that resolving an unknown host on some networks will take
   // 20-30 seconds; making this timeout short fixes the issue, but
@@ -675,11 +672,18 @@ public byte[] sign(byte[] toSign) {
     try {
       String account = getAccount();
       return IamUtils.sign(
-          account, this, transportFactory.create(), toSign, Collections.<String, Object>emptyMap());
+          account,
+          this,
+          this.getUniverseDomain(),
+          transportFactory.create(),
+          toSign,
+          Collections.<String, Object>emptyMap());
     } catch (SigningException ex) {
       throw ex;
     } catch (RuntimeException ex) {
       throw new SigningException("Signing failed", ex);
+    } catch (IOException ex) {
+      throw new SigningException("Failed to sign: Error obtaining universe domain", ex);
     }
   }
 
diff --git a/oauth2_http/java/com/google/auth/oauth2/IamUtils.java b/oauth2_http/java/com/google/auth/oauth2/IamUtils.java
@@ -62,10 +62,14 @@
  * features like signing.
  */
 class IamUtils {
-  private static final String SIGN_BLOB_URL_FORMAT =
-      "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:signBlob";
-  private static final String ID_TOKEN_URL_FORMAT =
-      "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateIdToken";
+
+  // IAM credentials endpoints are to be formatted with universe domain and client email.
+  static final String IAM_ID_TOKEN_ENDPOINT_FORMAT =
+      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateIdToken";
+  static final String IAM_ACCESS_TOKEN_ENDPOINT_FORMAT =
+      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateAccessToken";
+  static final String IAM_SIGN_BLOB_ENDPOINT_FORMAT =
+      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:signBlob";
   private static final String PARSE_ERROR_MESSAGE = "Error parsing error message response. ";
   private static final String PARSE_ERROR_SIGNATURE = "Error parsing signature response. ";
 
@@ -88,6 +92,7 @@ class IamUtils {
   static byte[] sign(
       String serviceAccountEmail,
       Credentials credentials,
+      String universeDomain,
       HttpTransport transport,
       byte[] toSign,
       Map<String, ?> additionalFields) {
@@ -97,7 +102,12 @@ static byte[] sign(
     String signature;
     try {
       signature =
-          getSignature(serviceAccountEmail, base64.encode(toSign), additionalFields, factory);
+          getSignature(
+              serviceAccountEmail,
+              universeDomain,
+              base64.encode(toSign),
+              additionalFields,
+              factory);
     } catch (IOException ex) {
       throw new ServiceAccountSigner.SigningException("Failed to sign the provided bytes", ex);
     }
@@ -106,11 +116,13 @@ static byte[] sign(
 
   private static String getSignature(
       String serviceAccountEmail,
+      String universeDomain,
       String bytes,
       Map<String, ?> additionalFields,
       HttpRequestFactory factory)
       throws IOException {
-    String signBlobUrl = String.format(SIGN_BLOB_URL_FORMAT, serviceAccountEmail);
+    String signBlobUrl =
+        String.format(IAM_SIGN_BLOB_ENDPOINT_FORMAT, universeDomain, serviceAccountEmail);
     GenericUrl genericUrl = new GenericUrl(signBlobUrl);
 
     GenericData signRequest = new GenericData();
@@ -193,10 +205,12 @@ static IdToken getIdToken(
       String targetAudience,
       boolean includeEmail,
       Map<String, ?> additionalFields,
-      CredentialTypeForMetrics credentialTypeForMetrics)
+      CredentialTypeForMetrics credentialTypeForMetrics,
+      String universeDomain)
       throws IOException {
 
-    String idTokenUrl = String.format(ID_TOKEN_URL_FORMAT, serviceAccountEmail);
+    String idTokenUrl =
+        String.format(IAM_ID_TOKEN_ENDPOINT_FORMAT, universeDomain, serviceAccountEmail);
     GenericUrl genericUrl = new GenericUrl(idTokenUrl);
 
     GenericData idTokenRequest = new GenericData();
diff --git a/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java
@@ -345,12 +345,19 @@ public void setTransportFactory(HttpTransportFactory httpTransportFactory) {
    */
   @Override
   public byte[] sign(byte[] toSign) {
-    return IamUtils.sign(
-        getAccount(),
-        sourceCredentials,
-        transportFactory.create(),
-        toSign,
-        ImmutableMap.of("delegates", this.delegates));
+    try {
+      return IamUtils.sign(
+          getAccount(),
+          sourceCredentials,
+          getUniverseDomain(),
+          transportFactory.create(),
+          toSign,
+          ImmutableMap.of("delegates", this.delegates));
+    } catch (IOException ex) {
+      // Throwing an IOException would be a breaking change, so wrap it here.
+      // This should not happen for this credential type.
+      throw new SigningException("Failed to sign: Error obtaining universe domain", ex);
+    }
   }
 
   /**
@@ -525,7 +532,7 @@ public AccessToken refreshAccessToken() throws IOException {
         this.iamEndpointOverride != null
             ? this.iamEndpointOverride
             : String.format(
-                OAuth2Utils.IAM_ACCESS_TOKEN_ENDPOINT_FORMAT,
+                IamUtils.IAM_ACCESS_TOKEN_ENDPOINT_FORMAT,
                 getUniverseDomain(),
                 this.targetPrincipal);
 
@@ -593,7 +600,8 @@ public IdToken idTokenWithAudience(String targetAudience, List<IdTokenProvider.O
         targetAudience,
         includeEmail,
         ImmutableMap.of("delegates", this.delegates),
-        getMetricsCredentialType());
+        getMetricsCredentialType(),
+        getUniverseDomain());
   }
 
   @Override
diff --git a/oauth2_http/java/com/google/auth/oauth2/OAuth2Utils.java b/oauth2_http/java/com/google/auth/oauth2/OAuth2Utils.java
@@ -77,15 +77,6 @@ class OAuth2Utils {
   static final String TOKEN_TYPE_TOKEN_EXCHANGE = "urn:ietf:params:oauth:token-type:token-exchange";
   static final String GRANT_TYPE_JWT_BEARER = "urn:ietf:params:oauth:grant-type:jwt-bearer";
 
-  // generateIdToken endpoint is to be formatted with universe domain and client email
-  static final String IAM_ID_TOKEN_ENDPOINT_FORMAT =
-      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateIdToken";
-
-  static final String IAM_ACCESS_TOKEN_ENDPOINT_FORMAT =
-      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateAccessToken";
-  static final String SIGN_BLOB_ENDPOINT_FORMAT =
-      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:signBlob";
-
   static final URI TOKEN_SERVER_URI = URI.create("https://oauth2.googleapis.com/token");
 
   static final URI TOKEN_REVOKE_URI = URI.create("https://oauth2.googleapis.com/revoke");
diff --git a/oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java
@@ -636,8 +636,7 @@ private IdToken getIdTokenIamEndpoint(String targetAudience) throws IOException
     // `getUniverseDomain()` throws an IOException that would need to be caught
     URI iamIdTokenUri =
         URI.create(
-            String.format(
-                OAuth2Utils.IAM_ID_TOKEN_ENDPOINT_FORMAT, getUniverseDomain(), clientEmail));
+            String.format(IamUtils.IAM_ID_TOKEN_ENDPOINT_FORMAT, getUniverseDomain(), clientEmail));
     HttpRequest request = buildIdTokenRequest(iamIdTokenUri, transportFactory, content);
     // Use the Access Token from the SSJWT to request the ID Token from IAM Endpoint
     request.setHeaders(new HttpHeaders().set(AuthHttpConstants.AUTHORIZATION, accessToken));
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java
@@ -66,6 +66,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
+import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -645,7 +646,7 @@ public LowLevelHttpResponse execute() throws IOException {
   }
 
   @Test
-  public void sign_sameAs() throws IOException {
+  public void sign_sameAs() {
     MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
     String defaultAccountEmail = "mail@mail.com";
     byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
@@ -659,21 +660,36 @@ public void sign_sameAs() throws IOException {
   }
 
   @Test
-  public void sign_getAccountFails() throws IOException {
+  public void sign_getUniverseException() {
+    MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
+
+    String defaultAccountEmail = "mail@mail.com";
+    transportFactory.transport.setServiceAccountEmail(defaultAccountEmail);
+    ComputeEngineCredentials credentials =
+        ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
+
+    transportFactory.transport.setRequestStatusCode(501);
+    Assert.assertThrows(IOException.class, credentials::getUniverseDomain);
+
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
+    SigningException signingException =
+        Assert.assertThrows(SigningException.class, () -> credentials.sign(expectedSignature));
+    assertEquals("Failed to sign: Error obtaining universe domain", signingException.getMessage());
+  }
+
+  @Test
+  public void sign_getAccountFails() {
     MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
     byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
 
     transportFactory.transport.setSignature(expectedSignature);
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-    try {
-      credentials.sign(expectedSignature);
-      fail("Should not be able to use credential without exception.");
-    } catch (SigningException ex) {
-      assertNotNull(ex.getMessage());
-      assertNotNull(ex.getCause());
-    }
+    SigningException exception =
+        Assert.assertThrows(SigningException.class, () -> credentials.sign(expectedSignature));
+    assertNotNull(exception.getMessage());
+    assertNotNull(exception.getCause());
   }
 
   @Test
@@ -705,15 +721,13 @@ public LowLevelHttpResponse execute() throws IOException {
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-    try {
-      byte[] bytes = {0xD, 0xE, 0xA, 0xD};
-      credentials.sign(bytes);
-      fail("Signing should have failed");
-    } catch (SigningException e) {
-      assertEquals("Failed to sign the provided bytes", e.getMessage());
-      assertNotNull(e.getCause());
-      assertTrue(e.getCause().getMessage().contains("403"));
-    }
+    byte[] bytes = {0xD, 0xE, 0xA, 0xD};
+
+    SigningException exception =
+        Assert.assertThrows(SigningException.class, () -> credentials.sign(bytes));
+    assertEquals("Failed to sign the provided bytes", exception.getMessage());
+    assertNotNull(exception.getCause());
+    assertTrue(exception.getCause().getMessage().contains("403"));
   }
 
   @Test
@@ -745,15 +759,13 @@ public LowLevelHttpResponse execute() throws IOException {
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-    try {
-      byte[] bytes = {0xD, 0xE, 0xA, 0xD};
-      credentials.sign(bytes);
-      fail("Signing should have failed");
-    } catch (SigningException e) {
-      assertEquals("Failed to sign the provided bytes", e.getMessage());
-      assertNotNull(e.getCause());
-      assertTrue(e.getCause().getMessage().contains("500"));
-    }
+    byte[] bytes = {0xD, 0xE, 0xA, 0xD};
+
+    SigningException exception =
+        Assert.assertThrows(SigningException.class, () -> credentials.sign(bytes));
+    assertEquals("Failed to sign the provided bytes", exception.getMessage());
+    assertNotNull(exception.getCause());
+    assertTrue(exception.getCause().getMessage().contains("500"));
   }
 
   @Test
@@ -778,14 +790,11 @@ public LowLevelHttpResponse execute() throws IOException {
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-    try {
-      credentials.refreshAccessToken();
-      fail("Should have failed");
-    } catch (IOException e) {
-      assertTrue(e.getCause().getMessage().contains("503"));
-      assertTrue(e instanceof GoogleAuthException);
-      assertTrue(((GoogleAuthException) e).isRetryable());
-    }
+    IOException exception =
+        Assert.assertThrows(IOException.class, () -> credentials.refreshAccessToken());
+    assertTrue(exception.getCause().getMessage().contains("503"));
+    assertTrue(exception instanceof GoogleAuthException);
+    assertTrue(((GoogleAuthException) exception).isRetryable());
   }
 
   @Test
@@ -818,12 +827,9 @@ public LowLevelHttpResponse execute() throws IOException {
       ComputeEngineCredentials credentials =
           ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-      try {
-        credentials.refreshAccessToken();
-        fail("Should have failed");
-      } catch (IOException e) {
-        assertFalse(e instanceof GoogleAuthException);
-      }
+      IOException exception =
+          Assert.assertThrows(IOException.class, () -> credentials.refreshAccessToken());
+      assertFalse(exception instanceof GoogleAuthException);
     }
   }
 
@@ -993,15 +999,13 @@ public LowLevelHttpResponse execute() throws IOException {
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
-    try {
-      byte[] bytes = {0xD, 0xE, 0xA, 0xD};
-      credentials.sign(bytes);
-      fail("Signing should have failed");
-    } catch (SigningException e) {
-      assertEquals("Failed to sign the provided bytes", e.getMessage());
-      assertNotNull(e.getCause());
-      assertTrue(e.getCause().getMessage().contains("Empty content"));
-    }
+    byte[] bytes = {0xD, 0xE, 0xA, 0xD};
+
+    SigningException exception =
+        Assert.assertThrows(SigningException.class, () -> credentials.sign(bytes));
+    assertEquals("Failed to sign the provided bytes", exception.getMessage());
+    assertNotNull(exception.getCause());
+    assertTrue(exception.getCause().getMessage().contains("Empty content"));
   }
 
   @Test
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/IamUtilsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/IamUtilsTest.java
@@ -36,6 +36,7 @@
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.client.http.HttpStatusCodes;
+import com.google.auth.Credentials;
 import com.google.auth.ServiceAccountSigner;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
@@ -60,6 +61,7 @@ public void setup() throws IOException {
     // token
     credentials = Mockito.mock(ServiceAccountCredentials.class);
     Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+    Mockito.when(credentials.getUniverseDomain()).thenReturn("googleapis.com");
   }
 
   @Test
@@ -76,6 +78,7 @@ public void sign_success_noRetry() {
         IamUtils.sign(
             CLIENT_EMAIL,
             credentials,
+            Credentials.GOOGLE_DEFAULT_UNIVERSE,
             transportFactory.getTransport(),
             expectedSignature,
             ImmutableMap.of());
@@ -107,6 +110,7 @@ public void sign_retryTwoTimes_success() {
         IamUtils.sign(
             CLIENT_EMAIL,
             credentials,
+            Credentials.GOOGLE_DEFAULT_UNIVERSE,
             transportFactory.getTransport(),
             expectedSignature,
             ImmutableMap.of());
@@ -143,6 +147,7 @@ public void sign_retryThreeTimes_success() {
         IamUtils.sign(
             CLIENT_EMAIL,
             credentials,
+            Credentials.GOOGLE_DEFAULT_UNIVERSE,
             transportFactory.getTransport(),
             expectedSignature,
             ImmutableMap.of());
@@ -185,6 +190,7 @@ public void sign_retryThreeTimes_exception() {
                 IamUtils.sign(
                     CLIENT_EMAIL,
                     credentials,
+                    Credentials.GOOGLE_DEFAULT_UNIVERSE,
                     transportFactory.getTransport(),
                     expectedSignature,
                     ImmutableMap.of()));
@@ -220,6 +226,7 @@ public void sign_4xxError_noRetry_exception() {
                 IamUtils.sign(
                     CLIENT_EMAIL,
                     credentials,
+                    Credentials.GOOGLE_DEFAULT_UNIVERSE,
                     transportFactory.getTransport(),
                     expectedSignature,
                     ImmutableMap.of()));
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/MockIAMCredentialsServiceTransport.java b/oauth2_http/javatests/com/google/auth/oauth2/MockIAMCredentialsServiceTransport.java
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/MockMetadataServerTransport.java b/oauth2_http/javatests/com/google/auth/oauth2/MockMetadataServerTransport.java
