googleapis · lsirac · Apr 21, 2022 · Apr 21, 2022 · Apr 21, 2022
diff --git a/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java b/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java
@@ -99,7 +99,11 @@ abstract static class CredentialSource {
 
   protected transient HttpTransportFactory transportFactory;
 
-  @Nullable protected ImpersonatedCredentials impersonatedCredentials;
+  @Nullable protected final ImpersonatedCredentials impersonatedCredentials;
+
+  // Internal override for impersonated credentials. This is done to keep
+  // impersonatedCredentials final.
+  @Nullable private ImpersonatedCredentials impersonatedCredentialsOverride;
 
   private EnvironmentProvider environmentProvider;
 
@@ -196,7 +200,7 @@ protected ExternalAccountCredentials(
       validateServiceAccountImpersonationInfoUrl(serviceAccountImpersonationUrl);
     }
 
-    this.impersonatedCredentials = initializeImpersonatedCredentials();
+    this.impersonatedCredentials = buildImpersonatedCredentials();
   }
 
   /**
@@ -238,10 +242,10 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
       validateServiceAccountImpersonationInfoUrl(serviceAccountImpersonationUrl);
     }
 
-    this.impersonatedCredentials = initializeImpersonatedCredentials();
+    this.impersonatedCredentials = buildImpersonatedCredentials();
   }
 
-  protected ImpersonatedCredentials initializeImpersonatedCredentials() {
+  ImpersonatedCredentials buildImpersonatedCredentials() {
     if (serviceAccountImpersonationUrl == null) {
       return null;
     }
@@ -275,6 +279,10 @@ protected ImpersonatedCredentials initializeImpersonatedCredentials() {
         .build();
   }
 
+  void overrideImpersonatedCredentials(ImpersonatedCredentials credentials) {
+    this.impersonatedCredentialsOverride = credentials;
+  }
+
   @Override
   public void getRequestMetadata(
       URI uri, Executor executor, final RequestMetadataCallback callback) {
@@ -429,7 +437,10 @@ private static boolean isAwsCredential(Map<String, Object> credentialSource) {
   protected AccessToken exchangeExternalCredentialForAccessToken(
       StsTokenExchangeRequest stsTokenExchangeRequest) throws IOException {
     // Handle service account impersonation if necessary.
-    if (impersonatedCredentials != null) {
+    // Internal override takes priority.
+    if (impersonatedCredentialsOverride != null) {
+      return impersonatedCredentialsOverride.refreshAccessToken();
+    } else if (impersonatedCredentials != null) {
       return impersonatedCredentials.refreshAccessToken();
     }
 

diff --git a/oauth2_http/java/com/google/auth/oauth2/PluggableAuthCredentials.java b/oauth2_http/java/com/google/auth/oauth2/PluggableAuthCredentials.java
@@ -213,7 +213,7 @@ String getOutputFilePath() {
 
     // Re-initialize impersonated credentials as the handler hasn't been set yet when
     // this is called in the base class.
-    this.impersonatedCredentials = initializeImpersonatedCredentials();
+    overrideImpersonatedCredentials(buildImpersonatedCredentials());
   }
 
   @Override

diff --git a/oauth2_http/java/com/google/auth/oauth2/PluggableAuthHandler.java b/oauth2_http/java/com/google/auth/oauth2/PluggableAuthHandler.java
@@ -59,44 +59,6 @@
  */
 final class PluggableAuthHandler implements ExecutableHandler {
 
-  /** An interface for creating and managing a process. */
-  abstract static class InternalProcessBuilder {
-
-    abstract Map<String, String> environment();
-
-    abstract InternalProcessBuilder redirectErrorStream(boolean redirectErrorStream);
-
-    abstract Process start() throws IOException;
-  }
-
-  /**
-   * The default implementation that wraps {@link ProcessBuilder} for creating and managing a
-   * process.
-   */
-  static final class DefaultProcessBuilder extends InternalProcessBuilder {
-    ProcessBuilder processBuilder;
-
-    DefaultProcessBuilder(ProcessBuilder processBuilder) {
-      this.processBuilder = processBuilder;
-    }
-
-    @Override
-    Map<String, String> environment() {
-      return this.processBuilder.environment();
-    }
-
-    @Override
-    InternalProcessBuilder redirectErrorStream(boolean redirectErrorStream) {
-      this.processBuilder.redirectErrorStream(redirectErrorStream);
-      return this;
-    }
-
-    @Override
-    Process start() throws IOException {
-      return this.processBuilder.start();
-    }
-  }
-
   // The maximum supported version for the executable response.
   // The executable response always includes a version number that is used
   // to detect compatibility with the response and library verions.
@@ -288,4 +250,46 @@ InternalProcessBuilder getProcessBuilder(List<String> commandComponents) {
     }
     return new DefaultProcessBuilder(new ProcessBuilder(commandComponents));
   }
+
+  /**
+   * An interface for creating and managing a process.
+   *
+   * <p>ProcessBuilder is final and does not implement any interface. This class allows concrete
+   * implementations to be specified to test these changes.
+   */
+  abstract static class InternalProcessBuilder {
+
+    abstract Map<String, String> environment();
+
+    abstract InternalProcessBuilder redirectErrorStream(boolean redirectErrorStream);
+
+    abstract Process start() throws IOException;
+  }
+
+  /**
+   * A default implementation for {@link InternalProcessBuilder} that wraps {@link ProcessBuilder}.
+   */
+  static final class DefaultProcessBuilder extends InternalProcessBuilder {
+    ProcessBuilder processBuilder;
+
+    DefaultProcessBuilder(ProcessBuilder processBuilder) {
+      this.processBuilder = processBuilder;
+    }
+
+    @Override
+    Map<String, String> environment() {
+      return this.processBuilder.environment();
+    }
+
+    @Override
+    InternalProcessBuilder redirectErrorStream(boolean redirectErrorStream) {
+      this.processBuilder.redirectErrorStream(redirectErrorStream);
+      return this;
+    }
+
+    @Override
+    Process start() throws IOException {
+      return this.processBuilder.start();
+    }
+  }
 }
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/ExternalAccountCredentialsTest.java
@@ -47,6 +47,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.net.URI;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
@@ -571,6 +572,38 @@ void exchangeExternalCredentialForAccessToken_withServiceAccountImpersonation()
         transportFactory.transport.getServiceAccountAccessToken(), returnedToken.getTokenValue());
   }
 
+  @Test
+  void exchangeExternalCredentialForAccessToken_withServiceAccountImpersonationOverride()
+      throws IOException {
+    transportFactory.transport.setExpireTime(TestUtils.getDefaultExpireTime());
+
+    String serviceAccountEmail = "[email protected]";
+    ExternalAccountCredentials credential =
+        ExternalAccountCredentials.fromStream(
+            IdentityPoolCredentialsTest.writeIdentityPoolCredentialsStream(
+                transportFactory.transport.getStsUrl(),
+                transportFactory.transport.getMetadataUrl(),
+                transportFactory.transport.getServiceAccountImpersonationUrl()),
+            transportFactory);
+
+    // Override impersonated credentials.
+    ExternalAccountCredentials sourceCredentials =
+        IdentityPoolCredentials.newBuilder((IdentityPoolCredentials) credential)
+            .setServiceAccountImpersonationUrl(null)
+            .build();
+    credential.overrideImpersonatedCredentials(
+        new ImpersonatedCredentials.Builder(sourceCredentials, serviceAccountEmail)
+            .setScopes(new ArrayList<>(sourceCredentials.getScopes()))
+            .setHttpTransportFactory(transportFactory)
+            .build());
+
+    credential.exchangeExternalCredentialForAccessToken(
+        StsTokenExchangeRequest.newBuilder("credential", "subjectTokenType").build());
+
+    assertTrue(
+        transportFactory.transport.getRequests().get(2).getUrl().contains(serviceAccountEmail));
+  }
+
   @Test
   void exchangeExternalCredentialForAccessToken_throws() throws IOException {
     ExternalAccountCredentials credential =

diff --git a/oauth2_http/javatests/com/google/auth/oauth2/MockExternalAccountCredentialsTransport.java b/oauth2_http/javatests/com/google/auth/oauth2/MockExternalAccountCredentialsTransport.java
@@ -84,6 +84,8 @@ public class MockExternalAccountCredentialsTransport extends MockHttpTransport {
   static final String SERVICE_ACCOUNT_IMPERSONATION_URL =
       "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/[email protected]:generateAccessToken";
 
+  static final String IAM_ENDPOINT = "https://iamcredentials.googleapis.com";
+
   private Queue<Boolean> responseSequence = new ArrayDeque<>();
   private Queue<IOException> responseErrorSequence = new ArrayDeque<>();
   private Queue<String> refreshTokenSequence = new ArrayDeque<>();
@@ -193,7 +195,8 @@ public LowLevelHttpResponse execute() throws IOException {
                   .setContentType(Json.MEDIA_TYPE)
                   .setContent(response.toPrettyString());
             }
-            if (SERVICE_ACCOUNT_IMPERSONATION_URL.equals(url)) {
+
+            if (url.contains(IAM_ENDPOINT)) {
               GenericJson query =
                   OAuth2Utils.JSON_FACTORY
                       .createJsonParser(getContentAsString())