feat: support ability to set universe domain in ServiceAccountJwtAccessCredentials

Author: mpeddada1

oauth2_http/java/com/google/auth/oauth2/ServiceAccountJwtAccessCredentials.java

@@ -88,6 +88,7 @@ public class ServiceAccountJwtAccessCredentials extends Credentials
   private final String privateKeyId;
   private final URI defaultAudience;
   private final String quotaProjectId;
+  private final String universeDomain;
 
   private transient LoadingCache<JwtClaims, JwtCredentials> credentialsCache;
 
@@ -103,8 +104,8 @@ public class ServiceAccountJwtAccessCredentials extends Credentials
    * @param privateKeyId Private key identifier for the service account. May be null.
    */
   private ServiceAccountJwtAccessCredentials(
-      String clientId, String clientEmail, PrivateKey privateKey, String privateKeyId) {
-    this(clientId, clientEmail, privateKey, privateKeyId, null, null);
+          String clientId, String clientEmail, PrivateKey privateKey, String privateKeyId) {
+    this(clientId, clientEmail, privateKey, privateKeyId, null, null,Credentials.GOOGLE_DEFAULT_UNIVERSE);
   }
 
   /**
@@ -115,21 +116,28 @@ private ServiceAccountJwtAccessCredentials(
    * @param privateKey RSA private key object for the service account.
    * @param privateKeyId Private key identifier for the service account. May be null.
    * @param defaultAudience Audience to use if not provided by transport. May be null.
+   * @param universeDomain Universe domain
    */
   private ServiceAccountJwtAccessCredentials(
       String clientId,
       String clientEmail,
       PrivateKey privateKey,
       String privateKeyId,
       URI defaultAudience,
-      String quotaProjectId) {
+      String quotaProjectId,
+      String universeDomain) {
     this.clientId = clientId;
     this.clientEmail = Preconditions.checkNotNull(clientEmail);
     this.privateKey = Preconditions.checkNotNull(privateKey);
     this.privateKeyId = privateKeyId;
     this.defaultAudience = defaultAudience;
     this.credentialsCache = createCache();
     this.quotaProjectId = quotaProjectId;
+    if (universeDomain == null || universeDomain.trim().isEmpty()) {
+      this.universeDomain = Credentials.GOOGLE_DEFAULT_UNIVERSE;
+    } else {
+      this.universeDomain = universeDomain;
+    }
   }
 
   /**
@@ -160,6 +168,9 @@ static ServiceAccountJwtAccessCredentials fromJson(Map<String, Object> json, URI
     String privateKeyPkcs8 = (String) json.get("private_key");
     String privateKeyId = (String) json.get("private_key_id");
     String quoataProjectId = (String) json.get("quota_project_id");
+    String rawUniverseDomain = (String) json.get("universe_domain");
+    String resolvedUniverseDomain = (rawUniverseDomain == null) ? Credentials.GOOGLE_DEFAULT_UNIVERSE : rawUniverseDomain;
+
     if (clientId == null
         || clientEmail == null
         || privateKeyPkcs8 == null
@@ -169,9 +180,10 @@ static ServiceAccountJwtAccessCredentials fromJson(Map<String, Object> json, URI
               + "expecting  'client_id', 'client_email', 'private_key' and 'private_key_id'.");
     }
     return ServiceAccountJwtAccessCredentials.fromPkcs8(
-        clientId, clientEmail, privateKeyPkcs8, privateKeyId, defaultAudience, quoataProjectId);
+        clientId, clientEmail, privateKeyPkcs8, privateKeyId, defaultAudience, quoataProjectId, resolvedUniverseDomain);
   }
 
+
   /**
    * Factory using PKCS#8 for the private key.
    *
@@ -207,7 +219,7 @@ public static ServiceAccountJwtAccessCredentials fromPkcs8(
       URI defaultAudience)
       throws IOException {
     return ServiceAccountJwtAccessCredentials.fromPkcs8(
-        clientId, clientEmail, privateKeyPkcs8, privateKeyId, defaultAudience, null);
+        clientId, clientEmail, privateKeyPkcs8, privateKeyId, defaultAudience, null, Credentials.GOOGLE_DEFAULT_UNIVERSE);
   }
 
   static ServiceAccountJwtAccessCredentials fromPkcs8(
@@ -216,11 +228,12 @@ static ServiceAccountJwtAccessCredentials fromPkcs8(
       String privateKeyPkcs8,
       String privateKeyId,
       URI defaultAudience,
-      String quotaProjectId)
+      String quotaProjectId,
+      String universeDomain)
       throws IOException {
     PrivateKey privateKey = OAuth2Utils.privateKeyFromPkcs8(privateKeyPkcs8);
     return new ServiceAccountJwtAccessCredentials(
-        clientId, clientEmail, privateKey, privateKeyId, defaultAudience, quotaProjectId);
+        clientId, clientEmail, privateKey, privateKeyId, defaultAudience, quotaProjectId, universeDomain);
   }
 
   /**
@@ -352,22 +365,21 @@ public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException
                 + "defaultAudience to be specified");
       }
     }
-
     try {
       JwtClaims defaultClaims =
-          JwtClaims.newBuilder()
-              .setAudience(uri.toString())
-              .setIssuer(clientEmail)
-              .setSubject(clientEmail)
-              .build();
+              JwtClaims.newBuilder()
+                      .setAudience(uri.toString())
+                      .setIssuer(clientEmail)
+                      .setSubject(clientEmail)
+                      .build();
       JwtCredentials credentials = credentialsCache.get(defaultClaims);
       Map<String, List<String>> requestMetadata = credentials.getRequestMetadata(uri);
       return addQuotaProjectIdToRequestMetadata(quotaProjectId, requestMetadata);
     } catch (ExecutionException e) {
       Throwables.propagateIfPossible(e.getCause(), IOException.class);
       // Should never happen
       throw new IllegalStateException(
-          "generateJwtAccess threw an unexpected checked exception", e.getCause());
+              "generateJwtAccess threw an unexpected checked exception", e.getCause());
 
     } catch (UncheckedExecutionException e) {
       Throwables.throwIfUnchecked(e);
@@ -399,6 +411,10 @@ public final String getPrivateKeyId() {
     return privateKeyId;
   }
 
+  public final String getUniverseDomain() {
+    return universeDomain;
+  }
+
   @Override
   public String getAccount() {
     return getClientEmail();
@@ -474,6 +490,7 @@ public static class Builder {
     private String privateKeyId;
     private URI defaultAudience;
     private String quotaProjectId;
+    private String universeDomain;
 
     protected Builder() {}
 
@@ -484,6 +501,7 @@ protected Builder(ServiceAccountJwtAccessCredentials credentials) {
       this.privateKeyId = credentials.privateKeyId;
       this.defaultAudience = credentials.defaultAudience;
       this.quotaProjectId = credentials.quotaProjectId;
+      this.universeDomain = credentials.universeDomain;
     }
 
     @CanIgnoreReturnValue
@@ -522,6 +540,13 @@ public Builder setQuotaProjectId(String quotaProjectId) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder setUniverseDomain(String universeDomain) {
+      this.universeDomain = universeDomain;
+      return this;
+    }
+
+
     public String getClientId() {
       return clientId;
     }
@@ -546,9 +571,13 @@ public String getQuotaProjectId() {
       return quotaProjectId;
     }
 
+    public String getUniverseDomain() {
+      return universeDomain;
+    }
+
     public ServiceAccountJwtAccessCredentials build() {
       return new ServiceAccountJwtAccessCredentials(
-          clientId, clientEmail, privateKey, privateKeyId, defaultAudience, quotaProjectId);
+          clientId, clientEmail, privateKey, privateKeyId, defaultAudience, quotaProjectId, universeDomain);
     }
   }
 }

oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountJwtAccessCredentialsTest.java

@@ -31,6 +31,7 @@
 
 package com.google.auth.oauth2;
 
+import static com.google.auth.Credentials.GOOGLE_DEFAULT_UNIVERSE;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -111,6 +112,7 @@ public void constructor_allParameters_constructs() throws IOException {
     assertEquals(privateKey, credentials.getPrivateKey());
     assertEquals(SA_PRIVATE_KEY_ID, credentials.getPrivateKeyId());
     assertEquals(QUOTA_PROJECT, credentials.getQuotaProjectId());
+    assertEquals(GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
   }
 
   @Test
@@ -829,6 +831,20 @@ public void onFailure(Throwable exception) {
     assertTrue("Should have run onSuccess() callback", success.get());
   }
 
+  @Test
+  public void getUniverseDomain_defaultUniverse() throws IOException {
+    PrivateKey privateKey = OAuth2Utils.privateKeyFromPkcs8(SA_PRIVATE_KEY_PKCS8);
+    ServiceAccountJwtAccessCredentials credentials =
+            ServiceAccountJwtAccessCredentials.newBuilder()
+                    .setClientId(SA_CLIENT_ID)
+                    .setClientEmail(SA_CLIENT_EMAIL)
+                    .setPrivateKey(privateKey)
+                    .setPrivateKeyId(SA_PRIVATE_KEY_ID)
+                    .setDefaultAudience(URI.create("default-audience"))
+                    .build();
+    assertEquals(GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
+  }
+
   private void verifyJwtAccess(
       Map<String, List<String>> metadata,
       String expectedEmail,