feat: Add TrustBoundaries support for ExternalAccounts.

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

@@ -31,7 +31,9 @@
 
 package com.google.auth.oauth2;
 
+import static com.google.auth.oauth2.OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL;
 import static com.google.auth.oauth2.OAuth2Utils.JSON_FACTORY;
+import static com.google.auth.oauth2.OAuth2Utils.WORKFORCE_AUDIENCE_PATTERN;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.api.client.http.GenericUrl;
@@ -44,6 +46,7 @@
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.util.GenericData;
 import com.google.api.client.util.Preconditions;
+import com.google.api.core.InternalApi;
 import com.google.auth.http.HttpTransportFactory;
 import com.google.common.base.MoreObjects;
 import com.google.common.io.BaseEncoding;
@@ -55,6 +58,7 @@
 import java.util.Date;
 import java.util.Map;
 import java.util.Objects;
+import java.util.regex.Matcher;
 import javax.annotation.Nullable;
 
 /**
@@ -75,12 +79,12 @@
  * }
  * </pre>
  */
-public class ExternalAccountAuthorizedUserCredentials extends GoogleCredentials {
+public class ExternalAccountAuthorizedUserCredentials extends GoogleCredentials
+    implements TrustBoundaryProvider {
 
   private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
 
   private static final long serialVersionUID = -2181779590486283287L;
-
   private final String transportFactoryClassName;
   private final String audience;
   private final String tokenUrl;
@@ -210,10 +214,27 @@ public AccessToken refreshAccessToken() throws IOException {
       this.refreshToken = refreshToken;
     }
 
-    return AccessToken.newBuilder()
-        .setExpirationTime(expiresAtMilliseconds)
-        .setTokenValue(accessToken)
-        .build();
+    AccessToken newAccessToken =
+        AccessToken.newBuilder()
+            .setExpirationTime(expiresAtMilliseconds)
+            .setTokenValue(accessToken)
+            .build();
+
+    refreshTrustBoundary(newAccessToken, transportFactory);
+    return newAccessToken;
+  }
+
+  @InternalApi
+  @Override
+  public String getTrustBoundaryUrl() throws IOException {
+    Matcher matcher = WORKFORCE_AUDIENCE_PATTERN.matcher(getAudience());
+    if (!matcher.matches()) {
+      throw new IllegalStateException(
+          "The provided audience is not in the correct format for a workforce pool.");
+    }
+    String poolId = matcher.group("pool");
+    return String.format(
+        IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL, getUniverseDomain(), poolId);
   }
 
   @Nullable

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

@@ -31,12 +31,15 @@
 
 package com.google.auth.oauth2;
 
+import static com.google.auth.oauth2.OAuth2Utils.WORKFORCE_AUDIENCE_PATTERN;
+import static com.google.auth.oauth2.OAuth2Utils.WORKLOAD_AUDIENCE_PATTERN;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.api.client.util.Data;
+import com.google.api.core.InternalApi;
 import com.google.auth.RequestMetadataCallback;
 import com.google.auth.http.HttpTransportFactory;
 import com.google.common.base.MoreObjects;
@@ -55,6 +58,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.Executor;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
@@ -64,7 +68,8 @@
  * <p>Handles initializing external credentials, calls to the Security Token Service, and service
  * account impersonation.
  */
-public abstract class ExternalAccountCredentials extends GoogleCredentials {
+public abstract class ExternalAccountCredentials extends GoogleCredentials
+    implements TrustBoundaryProvider {
 
   private static final long serialVersionUID = 8049126194174465023L;
 
@@ -527,7 +532,11 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
       this.impersonatedCredentials = this.buildImpersonatedCredentials();
     }
     if (this.impersonatedCredentials != null) {
-      return this.impersonatedCredentials.refreshAccessToken();
+      AccessToken accessToken = this.impersonatedCredentials.refreshAccessToken();
+      // After the impersonated credential refreshes, its trust boundary is
+      // also refreshed. That is the trust boundary we will use.
+      this.trustBoundary = this.impersonatedCredentials.getTrustBoundary();
+      return accessToken;
     }
 
     StsRequestHandler.Builder requestHandler =
@@ -556,7 +565,9 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
     }
 
     StsTokenExchangeResponse response = requestHandler.build().exchangeToken();
-    return response.getAccessToken();
+    AccessToken accessToken = response.getAccessToken();
+    refreshTrustBoundary(accessToken, transportFactory);
+    return accessToken;
   }
 
   /**
@@ -613,6 +624,33 @@ public String getServiceAccountEmail() {
     return ImpersonatedCredentials.extractTargetPrincipal(serviceAccountImpersonationUrl);
   }
 
+  @InternalApi
+  @Override
+  public String getTrustBoundaryUrl() {
+    Matcher workforceMatcher = WORKFORCE_AUDIENCE_PATTERN.matcher(getAudience());
+    if (workforceMatcher.matches()) {
+      String poolId = workforceMatcher.group("pool");
+      return String.format(
+          OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL,
+          getUniverseDomain(),
+          poolId);
+    }
+
+    Matcher workloadMatcher = WORKLOAD_AUDIENCE_PATTERN.matcher(getAudience());
+    if (workloadMatcher.matches()) {
+      String projectNumber = workloadMatcher.group("project");
+      String poolId = workloadMatcher.group("pool");
+      return String.format(
+          OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKLOAD_POOL,
+          getUniverseDomain(),
+          projectNumber,
+          poolId);
+    }
+
+    throw new IllegalStateException(
+        "The provided audience is not in a valid format for either a workload identity pool or a workforce pool.");
+  }
+
   @Nullable
   public String getClientId() {
     return clientId;

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

@@ -108,7 +108,7 @@ String getFileType() {
   private final String universeDomain;
   private final boolean isExplicitUniverseDomain;
 
-  private TrustBoundary trustBoundary;
+  TrustBoundary trustBoundary;
 
   protected final String quotaProjectId;
 
@@ -339,6 +339,10 @@ TrustBoundary getTrustBoundary() {
     return trustBoundary;
   }
 
+  protected void setTrustBoundary(TrustBoundary trustBoundary) {
+    this.trustBoundary = trustBoundary;
+  }
+
   /**
    * Refreshes the trust boundary by making a call to the trust boundary URL.
    *

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

@@ -68,6 +68,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.regex.Pattern;
 
 /**
  * Internal utilities for the com.google.auth.oauth2 namespace.
@@ -93,9 +94,6 @@ public class OAuth2Utils {
   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");
-
-  static final String IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT =
-      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s/allowedLocations";
   static final URI USER_AUTH_URI = URI.create("https://accounts.google.com/o/oauth2/auth");
 
   public static final String CLOUD_PLATFORM_SCOPE =
@@ -120,6 +118,22 @@ public class OAuth2Utils {
   static final double RETRY_MULTIPLIER = 2;
   static final int DEFAULT_NUMBER_OF_RETRIES = 3;
 
+  static final Pattern WORKFORCE_AUDIENCE_PATTERN =
+      Pattern.compile(
+          "^//iam.googleapis.com/locations/(?<location>[^/]+)/workforcePools/(?<pool>[^/]+)/providers/(?<provider>[^/]+)$");
+  static final Pattern WORKLOAD_AUDIENCE_PATTERN =
+      Pattern.compile(
+          "^//iam.googleapis.com/projects/(?<project>[^/]+)/locations/(?<location>[^/]+)/workloadIdentityPools/(?<pool>[^/]+)/providers/(?<provider>[^/]+)$");
+
+  static final String IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT =
+      "https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s/allowedLocations";
+
+  static final String IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL =
+      "https://iamcredentials.%s/v1/locations/global/workforcePools/%s/allowedLocations";
+
+  static final String IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKLOAD_POOL =
+      "https://iamcredentials.%s/v1/projects/%s/locations/global/workloadIdentityPools/%s/allowedLocations";
+
   // Includes expected server errors from Google token endpoint
   // Other 5xx codes are either not used or retries are unlikely to succeed
   public static final Set<Integer> TOKEN_ENDPOINT_RETRYABLE_STATUS_CODES =

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

@@ -184,9 +184,7 @@ static TrustBoundary refresh(
 
     // Add the cached trust boundary header, if available.
     if (cachedTrustBoundary != null) {
-      String headerValue =
-          cachedTrustBoundary.isNoOp() ? "" : cachedTrustBoundary.getEncodedLocations();
-      request.getHeaders().set(TRUST_BOUNDARY_KEY, headerValue);
+      request.getHeaders().set(TRUST_BOUNDARY_KEY, cachedTrustBoundary.getEncodedLocations());
     }
 
     // Add retry logic

oauth2_http/javatests/com/google/auth/TestUtils.java

@@ -42,6 +42,7 @@
 import com.google.api.client.json.gson.GsonFactory;
 import com.google.auth.http.AuthHttpConstants;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -55,6 +56,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.TimeZone;
 import javax.annotation.Nullable;
 
 /** Utilities for test code under com.google.auth. */
@@ -64,6 +66,9 @@ public class TestUtils {
       URI.create("https://auth.cloud.google/authorize");
   public static final URI WORKFORCE_IDENTITY_FEDERATION_TOKEN_SERVER_URI =
       URI.create("https://sts.googleapis.com/v1/oauthtoken");
+  public static final String TRUST_BOUNDARY_ENCODED_LOCATION = "0x800000";
+  public static final List<String> TRUST_BOUNDARY_LOCATIONS =
+      ImmutableList.of("us-central1", "us-central2");
 
   private static final JsonFactory JSON_FACTORY = GsonFactory.getDefaultInstance();
 
@@ -147,7 +152,9 @@ public static String getDefaultExpireTime() {
     Calendar calendar = Calendar.getInstance();
     calendar.setTime(new Date());
     calendar.add(Calendar.SECOND, 300);
-    return new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'").format(calendar.getTime());
+    SimpleDateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
+    dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
+    return dateFormat.format(calendar.getTime());
   }
 
   private TestUtils() {}

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

@@ -1399,4 +1399,34 @@ public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext cont
       return credentials;
     }
   }
+
+  @Test
+  public void testRefresh_trustBoundarySuccess() throws IOException {
+    TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
+    TrustBoundary.setEnvironmentProviderForTest(environmentProvider);
+    environmentProvider.setEnv("GOOGLE_AUTH_TRUST_BOUNDARY_ENABLE_EXPERIMENT", "1");
+
+    MockExternalAccountCredentialsTransportFactory transportFactory =
+        new MockExternalAccountCredentialsTransportFactory();
+
+    AwsSecurityCredentialsSupplier supplier =
+        new TestAwsSecurityCredentialsSupplier("test", programmaticAwsCreds, null, null);
+
+    AwsCredentials awsCredential =
+        AwsCredentials.newBuilder()
+            .setAwsSecurityCredentialsSupplier(supplier)
+            .setHttpTransportFactory(transportFactory)
+            .setAudience(
+                "//iam.googleapis.com/projects/12345/locations/global/workloadIdentityPools/pool/providers/provider")
+            .setTokenUrl(STS_URL)
+            .setSubjectTokenType("subjectTokenType")
+            .build();
+
+    awsCredential.refreshAccessToken();
+
+    TrustBoundary trustBoundary = awsCredential.getTrustBoundary();
+    assertNotNull(trustBoundary);
+    assertEquals(TestUtils.TRUST_BOUNDARY_ENCODED_LOCATION, trustBoundary.getEncodedLocations());
+    TrustBoundary.setEnvironmentProviderForTest(null);
+  }
 }

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

@@ -33,6 +33,7 @@
 
 import static com.google.auth.oauth2.ComputeEngineCredentials.METADATA_RESPONSE_EMPTY_CONTENT_ERROR_MESSAGE;
 import static com.google.auth.oauth2.ImpersonatedCredentialsTest.SA_CLIENT_EMAIL;
+import static com.google.auth.oauth2.TrustBoundary.TRUST_BOUNDARY_KEY;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -1159,15 +1160,17 @@ public void refresh_trustBoundarySuccess() throws IOException {
     String defaultAccountEmail = "[email protected]";
     MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
     TrustBoundary trustBoundary =
-        new TrustBoundary("0x80000", Collections.singletonList("us-central1"));
+        new TrustBoundary(
+            TestUtils.TRUST_BOUNDARY_ENCODED_LOCATION, TestUtils.TRUST_BOUNDARY_LOCATIONS);
     transportFactory.transport.setTrustBoundary(trustBoundary);
     transportFactory.transport.setServiceAccountEmail(defaultAccountEmail);
 
     ComputeEngineCredentials credentials =
         ComputeEngineCredentials.newBuilder().setHttpTransportFactory(transportFactory).build();
 
     Map<String, List<String>> headers = credentials.getRequestMetadata();
-    assertEquals(headers.get("x-allowed-locations"), Arrays.asList("0x80000"));
+    assertEquals(
+        headers.get(TRUST_BOUNDARY_KEY), Arrays.asList(TestUtils.TRUST_BOUNDARY_ENCODED_LOCATION));
   }
 
   @Test