feat: Async Refresh for Regional Access Boundaries

oauth2_http/java/com/google/auth/http/HttpCredentialsAdapter.java

@@ -39,6 +39,8 @@
 import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
 import com.google.api.client.util.Preconditions;
 import com.google.auth.Credentials;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.RegionalAccessBoundary;
 import java.io.IOException;
 import java.net.URI;
 import java.util.ArrayList;
@@ -61,6 +63,8 @@ public class HttpCredentialsAdapter
   private static final Pattern INVALID_TOKEN_ERROR =
       Pattern.compile("\\s*error\\s*=\\s*\"?invalid_token\"?");
 
+  private static final String STALE_RAB_ERROR_MESSAGE = "stale regional access boundary";
+
   private final Credentials credentials;
 
   /**
@@ -119,6 +123,12 @@ public void initialize(HttpRequest request) throws IOException {
    */
   @Override
   public boolean handleResponse(HttpRequest request, HttpResponse response, boolean supportsRetry) {
+    if (response.getStatusCode() == HttpStatusCodes.STATUS_CODE_BAD_REQUEST) {
+      if (handleStaleRegionalAccessBoundaryError(request, response)) {
+        return true;
+      }
+    }
+
     boolean refreshToken = false;
     boolean bearer = false;
 
@@ -152,4 +162,43 @@ public boolean handleResponse(HttpRequest request, HttpResponse response, boolea
     }
     return false;
   }
+
+  private boolean handleStaleRegionalAccessBoundaryError(
+      HttpRequest request, HttpResponse response) {
+    if (!(credentials instanceof GoogleCredentials)) {
+      return false;
+    }
+    GoogleCredentials googleCredentials = (GoogleCredentials) credentials;
+
+    // Only check for stale RAB error if we actually sent the header.
+    if (request.getHeaders().get(RegionalAccessBoundary.HEADER_KEY) == null) {
+      return false;
+    }
+
+    // Skip check for STS and IAM Credentials endpoints as per design.
+    String url = request.getUrl().toString();
+    if (url.contains("sts.googleapis.com") || url.contains("iamcredentials.googleapis.com")) {
+      return false;
+    }
+
+    try {
+      // Check for the stale regional access boundary error message in the response body.
+      // Note: This consumes the response stream.
+      String content = response.parseAsString();
+      if (content != null && content.toLowerCase().contains(STALE_RAB_ERROR_MESSAGE)) {
+        URI uri = null;
+        if (request.getUrl() != null) {
+          uri = request.getUrl().toURI();
+        }
+        googleCredentials.reactiveRefreshRegionalAccessBoundary(
+            uri, googleCredentials.getAccessToken());
+        // Re-initialize headers (this will remove the stale header since cache is cleared)
+        initialize(request);
+        return true;
+      }
+    } catch (Exception e) {
+      LOGGER.log(Level.FINE, "Error while checking for stale regional access boundary", e);
+    }
+    return false;
+  }
 }

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

@@ -83,7 +83,7 @@
  * <p>These credentials use the IAM API to sign data. See {@link #sign(byte[])} for more details.
  */
 public class ComputeEngineCredentials extends GoogleCredentials
-    implements ServiceAccountSigner, IdTokenProvider, TrustBoundaryProvider {
+    implements ServiceAccountSigner, IdTokenProvider, RegionalAccessBoundaryProvider {
 
   static final String METADATA_RESPONSE_EMPTY_CONTENT_ERROR_MESSAGE =
       "Empty content from metadata token server request.";
@@ -386,11 +386,7 @@ public AccessToken refreshAccessToken() throws IOException {
     int expiresInSeconds =
         OAuth2Utils.validateInt32(responseData, "expires_in", PARSE_ERROR_PREFIX);
     long expiresAtMilliseconds = clock.currentTimeMillis() + expiresInSeconds * 1000;
-    AccessToken newAccessToken = new AccessToken(accessToken, new Date(expiresAtMilliseconds));
-
-    refreshTrustBoundary(newAccessToken, transportFactory);
-
-    return newAccessToken;
+    return new AccessToken(accessToken, new Date(expiresAtMilliseconds));
   }
 
   /**
@@ -694,6 +690,11 @@ public static Builder newBuilder() {
    *
    * @throws RuntimeException if the default service account cannot be read
    */
+  @Override
+  HttpTransportFactory getTransportFactory() {
+    return transportFactory;
+  }
+
   @Override
   // todo(#314) getAccount should not throw a RuntimeException
   public String getAccount() {
@@ -709,11 +710,9 @@ public String getAccount() {
 
   @InternalApi
   @Override
-  public String getTrustBoundaryUrl() throws IOException {
+  public String getRegionalAccessBoundaryUrl() throws IOException {
     return String.format(
-        OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
-        getUniverseDomain(),
-        getAccount());
+        OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT, getAccount());
   }
 
   /**

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

@@ -80,7 +80,7 @@
  * </pre>
  */
 public class ExternalAccountAuthorizedUserCredentials extends GoogleCredentials
-    implements TrustBoundaryProvider {
+    implements RegionalAccessBoundaryProvider {
 
   private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
 
@@ -214,28 +214,28 @@ public AccessToken refreshAccessToken() throws IOException {
       this.refreshToken = refreshToken;
     }
 
-    AccessToken newAccessToken =
-        AccessToken.newBuilder()
-            .setExpirationTime(expiresAtMilliseconds)
-            .setTokenValue(accessToken)
-            .build();
-
-    refreshTrustBoundary(newAccessToken, transportFactory);
-    return newAccessToken;
+    return AccessToken.newBuilder()
+        .setExpirationTime(expiresAtMilliseconds)
+        .setTokenValue(accessToken)
+        .build();
   }
 
   @InternalApi
   @Override
-  public String getTrustBoundaryUrl() throws IOException {
+  public String getRegionalAccessBoundaryUrl() 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. "
               + "Refer: https://docs.cloud.google.com/iam/docs/principal-identifiers");
     }
     String poolId = matcher.group("pool");
-    return String.format(
-        IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL, getUniverseDomain(), poolId);
+    return String.format(IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL, poolId);
+  }
+
+  @Override
+  HttpTransportFactory getTransportFactory() {
+    return transportFactory;
   }
 
   @Nullable

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

@@ -69,7 +69,7 @@
  * account impersonation.
  */
 public abstract class ExternalAccountCredentials extends GoogleCredentials
-    implements TrustBoundaryProvider {
+    implements RegionalAccessBoundaryProvider {
 
   private static final long serialVersionUID = 8049126194174465023L;
 
@@ -532,11 +532,7 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
       this.impersonatedCredentials = this.buildImpersonatedCredentials();
     }
     if (this.impersonatedCredentials != null) {
-      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;
+      return this.impersonatedCredentials.refreshAccessToken();
     }
 
     StsRequestHandler.Builder requestHandler =
@@ -565,9 +561,7 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
     }
 
     StsTokenExchangeResponse response = requestHandler.build().exchangeToken();
-    AccessToken accessToken = response.getAccessToken();
-    refreshTrustBoundary(accessToken, transportFactory);
-    return accessToken;
+    return response.getAccessToken();
   }
 
   /**
@@ -581,6 +575,11 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
    */
   public abstract String retrieveSubjectToken() throws IOException;
 
+  @Override
+  HttpTransportFactory getTransportFactory() {
+    return transportFactory;
+  }
+
   public String getAudience() {
     return audience;
   }
@@ -626,14 +625,18 @@ public String getServiceAccountEmail() {
 
   @InternalApi
   @Override
-  public String getTrustBoundaryUrl() {
+  public String getRegionalAccessBoundaryUrl() throws IOException {
+    if (getServiceAccountEmail() != null) {
+      return String.format(
+          OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
+          getServiceAccountEmail());
+    }
+
     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);
+          OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKFORCE_POOL, poolId);
     }
 
     Matcher workloadMatcher = WORKLOAD_AUDIENCE_PATTERN.matcher(getAudience());
@@ -642,7 +645,6 @@ public String getTrustBoundaryUrl() {
       String poolId = workloadMatcher.group("pool");
       return String.format(
           OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_WORKLOAD_POOL,
-          getUniverseDomain(),
           projectNumber,
           poolId);
     }