Addressed PR comments. All refresh logic within Google Credentials. Use supportsTrustBoundary to identify whether a child class supports tb prior to refresh.

vverman · vverman · commit e6491ecb4a71 · 2025-10-10T13:37:25.000-07:00
diff --git a/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ComputeEngineCredentials.java
@@ -82,7 +82,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 {
 
   static final String METADATA_RESPONSE_EMPTY_CONTENT_ERROR_MESSAGE =
       "Empty content from metadata token server request.";
@@ -387,7 +387,7 @@ public AccessToken refreshAccessToken() throws IOException {
     long expiresAtMilliseconds = clock.currentTimeMillis() + expiresInSeconds * 1000;
     AccessToken newAccessToken = new AccessToken(accessToken, new Date(expiresAtMilliseconds));
 
-    refreshTrustBoundaries(newAccessToken);
+    refreshTrustBoundary(newAccessToken, getTrustBoundaryUrl(), transportFactory);
 
     return newAccessToken;
   }
@@ -707,12 +707,11 @@ public String getAccount() {
   }
 
   @Override
-  public HttpTransportFactory getTransportFactory() {
-    return transportFactory;
+  Boolean supportsTrustBoundary() {
+      return true;
   }
 
-  @Override
-  public String getTrustBoundaryUrl() throws IOException {
+  String getTrustBoundaryUrl() throws IOException {
     return String.format(
         OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
         getUniverseDomain(),
diff --git a/oauth2_http/java/com/google/auth/oauth2/GoogleCredentials.java b/oauth2_http/java/com/google/auth/oauth2/GoogleCredentials.java
@@ -341,6 +341,48 @@ public TrustBoundary getTrustBoundary() {
     return trustBoundary;
   }
 
+  Boolean supportsTrustBoundary() {
+    return false;
+  }
+
+  void refreshTrustBoundary(AccessToken newAccessToken, String trustBoundaryUrl, HttpTransportFactory transportFactory)
+      throws IOException {
+
+      if (!supportsTrustBoundary() || !TrustBoundary.isTrustBoundaryEnabled() || !isDefaultUniverseDomain()) {
+        return;
+      }
+
+      TrustBoundary cachedTrustBoundary;
+
+      synchronized (lock) {
+        // Do not refresh if the cached value is already NO_OP.
+        if (trustBoundary != null && trustBoundary.isNoOp()) {
+          return;
+        }
+        cachedTrustBoundary = trustBoundary;
+      }
+
+      TrustBoundary newTrustBoundary;
+      try {
+        newTrustBoundary =
+            TrustBoundary.refresh(
+                transportFactory, trustBoundaryUrl, newAccessToken, cachedTrustBoundary);
+      } catch (IOException e) {
+        // If refresh fails, check for a cached value.
+        if (cachedTrustBoundary == null) {
+          // No cached value, so fail hard.
+          throw new IOException(
+              "Failed to refresh trust boundary and no cached value is available.", e);
+        }
+        return;
+      }
+
+      // A lock is required to safely update the shared field.
+      synchronized (lock) {
+        trustBoundary = newTrustBoundary;
+      }
+  }
+
   /**
    * Gets the universe domain for the credential.
    *
@@ -372,7 +414,7 @@ protected boolean isExplicitUniverseDomain() {
    * @return true if universe domain equals to {@link Credentials#GOOGLE_DEFAULT_UNIVERSE}, false
    *     otherwise
    */
-  boolean isDefaultUniverseDomain() throws IOException {
+  public boolean isDefaultUniverseDomain() throws IOException {
     return getUniverseDomain().equals(Credentials.GOOGLE_DEFAULT_UNIVERSE);
   }
 
@@ -392,57 +434,6 @@ static Map<String, List<String>> addQuotaProjectIdToRequestMetadata(
     return Collections.unmodifiableMap(newRequestMetadata);
   }
 
-  /**
-   * Refreshes the trust boundary associated with these credentials.
-   *
-   * @param newAccessToken The newly refreshed access token to use for the trust boundary request.
-   * @throws IOException If the refresh fails and no cached trust boundary is available.
-   * @throws IllegalArgumentException If the provided access token is null or expired.
-   */
-  protected void refreshTrustBoundaries(AccessToken newAccessToken) throws IOException {
-    if (!(this instanceof TrustBoundaryProvider)) {
-      return;
-    }
-    if (!TrustBoundary.isTrustBoundaryEnabled() || !isDefaultUniverseDomain()) {
-      return;
-    }
-
-    TrustBoundaryProvider provider = (TrustBoundaryProvider) this;
-    TrustBoundary cachedTrustBoundary;
-
-    synchronized (lock) {
-      // Do not refresh if the cached value is already NO_OP.
-      if (this.trustBoundary != null && this.trustBoundary.isNoOp()) {
-        return;
-      }
-      cachedTrustBoundary = this.trustBoundary;
-    }
-
-    TrustBoundary newTrustBoundary;
-    try {
-      newTrustBoundary =
-          TrustBoundary.refresh(
-              provider.getTransportFactory(),
-              provider.getTrustBoundaryUrl(),
-              newAccessToken,
-              cachedTrustBoundary);
-    } catch (IOException e) {
-      // If refresh fails, check for a cached value.
-      if (cachedTrustBoundary == null) {
-        // No cached value, so fail hard.
-        throw new IOException(
-            "Failed to refresh trust boundary and no cached value is available.", e);
-      }
-
-      return;
-    }
-
-    // A lock is required to safely update the shared field.
-    synchronized (lock) {
-      this.trustBoundary = newTrustBoundary;
-    }
-  }
-
   @Override
   protected Map<String, List<String>> getAdditionalHeaders() {
     Map<String, List<String>> headers = new HashMap<>(super.getAdditionalHeaders());
diff --git a/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java
@@ -93,7 +93,7 @@
  * </pre>
  */
 public class ImpersonatedCredentials extends GoogleCredentials
-    implements ServiceAccountSigner, IdTokenProvider, TrustBoundaryProvider {
+    implements ServiceAccountSigner, IdTokenProvider {
 
   private static final long serialVersionUID = -2133257318957488431L;
   private static final String RFC3339 = "yyyy-MM-dd'T'HH:mm:ssX";
@@ -325,13 +325,12 @@ public GoogleCredentials getSourceCredentials() {
     return sourceCredentials;
   }
 
-  @Override
-  public HttpTransportFactory getTransportFactory() {
-    return transportFactory;
-  }
+    @Override
+    Boolean supportsTrustBoundary() {
+        return true;
+    }
 
-  @Override
-  public String getTrustBoundaryUrl() throws IOException {
+  String getTrustBoundaryUrl() throws IOException {
     return String.format(
         OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
         getUniverseDomain(),
@@ -342,6 +341,7 @@ int getLifetime() {
     return this.lifetime;
   }
 
+
   public void setTransportFactory(HttpTransportFactory httpTransportFactory) {
     this.transportFactory = httpTransportFactory;
   }
@@ -602,7 +602,7 @@ public AccessToken refreshAccessToken() throws IOException {
       Date date = format.parse(expireTime);
       AccessToken newAccessToken = new AccessToken(accessToken, date);
 
-      refreshTrustBoundaries(newAccessToken);
+      refreshTrustBoundary(newAccessToken, getTrustBoundaryUrl(), transportFactory);
 
       return newAccessToken;
     } catch (ParseException pe) {
diff --git a/oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java
@@ -89,7 +89,7 @@
  * <p>By default uses a JSON Web Token (JWT) to fetch access tokens.
  */
 public class ServiceAccountCredentials extends GoogleCredentials
-    implements ServiceAccountSigner, IdTokenProvider, JwtProvider, TrustBoundaryProvider {
+    implements ServiceAccountSigner, IdTokenProvider, JwtProvider {
 
   private static final long serialVersionUID = 7807543542681217978L;
   private static final String GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer";
@@ -582,7 +582,7 @@ public AccessToken refreshAccessToken() throws IOException {
     long expiresAtMilliseconds = clock.currentTimeMillis() + expiresInSeconds * 1000L;
     AccessToken newAccessToken = new AccessToken(accessToken, new Date(expiresAtMilliseconds));
 
-    refreshTrustBoundaries(newAccessToken);
+    refreshTrustBoundary(newAccessToken, getTrustBoundaryUrl(), transportFactory);
 
     return newAccessToken;
   }
@@ -823,13 +823,12 @@ public boolean getUseJwtAccessWithScope() {
     return useJwtAccessWithScope;
   }
 
-  @Override
-  public HttpTransportFactory getTransportFactory() {
-    return transportFactory;
-  }
+    @Override
+    Boolean supportsTrustBoundary() {
+        return true;
+    }
 
-  @Override
-  public String getTrustBoundaryUrl() throws IOException {
+  String getTrustBoundaryUrl() throws IOException {
     return String.format(
         OAuth2Utils.IAM_CREDENTIALS_ALLOWED_LOCATIONS_URL_FORMAT_SERVICE_ACCOUNT,
         getUniverseDomain(),
diff --git a/oauth2_http/java/com/google/auth/oauth2/TrustBoundary.java b/oauth2_http/java/com/google/auth/oauth2/TrustBoundary.java
@@ -46,6 +46,7 @@
 import com.google.auth.http.HttpTransportFactory;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Date;
@@ -168,9 +169,7 @@ static TrustBoundary refresh(
       AccessToken accessToken,
       @Nullable TrustBoundary cachedTrustBoundary)
       throws IOException {
-    if (accessToken == null) {
-      throw new IllegalArgumentException("The provided access token is null.");
-    }
+    Preconditions.checkNotNull(accessToken, "The provided access token is null.");
     if (accessToken.getExpirationTime() != null
         && accessToken.getExpirationTime().before(new Date())) {
       throw new IllegalArgumentException("The provided access token is expired.");
@@ -213,7 +212,8 @@ static TrustBoundary refresh(
     }
     String encodedLocations = json.getEncodedLocations();
     if (encodedLocations == null) {
-      throw new IOException("TrustBoundary: Malformed response from lookup endpoint.");
+      throw new IOException(
+          "TrustBoundary: Malformed response from lookup endpoint - `encodedLocations` was null.");
     }
     return new TrustBoundary(encodedLocations, json.getLocations());
   }
diff --git a/oauth2_http/java/com/google/auth/oauth2/TrustBoundaryProvider.java b/oauth2_http/java/com/google/auth/oauth2/TrustBoundaryProvider.java
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ComputeEngineCredentialsTest.java
@@ -129,6 +129,11 @@ public class ComputeEngineCredentialsTest extends BaseSerializationTest {
               })
           .collect(Collectors.toMap(data -> data[0], data -> data[1]));
 
+  @After
+  public void tearDown() {
+    TrustBoundary.setEnvironmentProviderForTest(null);
+  }
+
   @Test
   public void buildTokenUrlWithScopes_null_scopes() {
     ComputeEngineCredentials credentials =
@@ -397,7 +402,6 @@ public void getRequestMetadata_hasAccessToken() throws IOException {
     TestUtils.assertContainsBearerToken(metadata, ACCESS_TOKEN);
     // verify metrics header added and other header intact
     Map<String, List<String>> requestHeaders = transportFactory.transport.getRequest().getHeaders();
-    com.google.auth.oauth2.TestUtils.validateMetricsHeader(requestHeaders, "at", "mds");
     assertTrue(requestHeaders.containsKey("metadata-flavor"));
     assertTrue(requestHeaders.get("metadata-flavor").contains("Google"));
   }
@@ -423,11 +427,6 @@ public void getRequestMetadata_shouldInvalidateAccessTokenWhenScoped_newAccessTo
     TestUtils.assertNotContainsBearerToken(metadataForCopiedCredentials, ACCESS_TOKEN);
   }
 
-  @After
-  public void tearDown() {
-    TrustBoundary.setEnvironmentProviderForTest(null);
-  }
-
   @Test
   public void getRequestMetadata_missingServiceAccount_throws() {
     MockMetadataServerTransportFactory transportFactory = new MockMetadataServerTransportFactory();
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/GoogleCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/GoogleCredentialsTest.java
@@ -102,6 +102,11 @@ public class GoogleCredentialsTest extends BaseSerializationTest {
   private static final String GOOGLE_DEFAULT_UNIVERSE = "googleapis.com";
   private static final String TPC_UNIVERSE = "foo.bar";
 
+  @After
+  public void tearDown() {
+    TrustBoundary.setEnvironmentProviderForTest(null);
+  }
+
   @Test
   public void getApplicationDefault_nullTransport_throws() throws IOException {
     try {
@@ -936,11 +941,6 @@ public void getCredentialInfo_impersonatedServiceAccount() throws IOException {
         ImpersonatedCredentialsTest.IMPERSONATED_CLIENT_EMAIL, credentialInfo.get("Principal"));
   }
 
-  @After
-  public void tearDown() {
-    TrustBoundary.setEnvironmentProviderForTest(null);
-  }
-
   @Test
   public void trustBoundary_shouldNotCallLookupEndpointWhenDisabled() throws IOException {
     TestEnvironmentProvider environmentProvider = new TestEnvironmentProvider();
@@ -1079,9 +1079,8 @@ public void trustBoundary_refreshShouldThrowWhenNoValidAccessTokenIsPassed() thr
     try {
       credentials.getRequestMetadata();
       fail("Should have thrown an IOException.");
-    } catch (IOException e) {
-      assertEquals(
-          "Failed to refresh trust boundary and no cached value is available.", e.getMessage());
+    } catch (IllegalArgumentException e) {
+      assertEquals("The provided access token is expired.", e.getMessage());
     }
   }
 
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ImpersonatedCredentialsTest.java
