feat: Adds support for user defined subject token suppliers in AWSCredentials and IdentityPoolCredentials

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

@@ -43,7 +43,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -62,16 +61,12 @@ public class AwsCredentials extends ExternalAccountCredentials {
 
   private static final long serialVersionUID = -3670131891574618105L;
 
-  @Nullable private final AwsSecurityCredentialsProvider awsSecurityCredentialsProvider;
+  @Nullable private final AwsSecurityCredentialsSupplier awsSecurityCredentialsSupplier;
   // Regional credential verification url override. This needs to be its own value so we can
   // correctly pass it to a builder.
   @Nullable private final String regionalCredentialVerificationUrlOverride;
   @Nullable private final String regionalCredentialVerificationUrl;
-  @Nullable private final String awsRegion;
-
-  // Boolean to tell if a Supplier was provided to build the credential. This is needed to tell
-  // whether the supplier should be passed when creating a new builder from this credential.
-  private final boolean builtWithSupplier;
+  private final String metricsHeaderValue;
 
   /** Internal constructor. See {@link AwsCredentials.Builder}. */
   AwsCredentials(Builder builder) {
@@ -87,7 +82,6 @@ public class AwsCredentials extends ExternalAccountCredentials {
     }
 
     AwsCredentialSource credentialSource = (AwsCredentialSource) builder.credentialSource;
-    this.awsRegion = builder.awsRegion;
     // Set regional credential verification url override if provided.
     this.regionalCredentialVerificationUrlOverride =
         builder.regionalCredentialVerificationUrlOverride;
@@ -104,18 +98,13 @@ public class AwsCredentials extends ExternalAccountCredentials {
     // If user has provided a security credential supplier, use that to retrieve the AWS security
     // credentials.
     if (builder.awsSecurityCredentialsSupplier != null) {
-      this.awsSecurityCredentialsProvider =
-          new ProgrammaticAwsSecurityCredentialsProvider(
-              builder.awsSecurityCredentialsSupplier, this.awsRegion);
-      this.builtWithSupplier = true;
+      this.awsSecurityCredentialsSupplier = builder.awsSecurityCredentialsSupplier;
+      this.metricsHeaderValue = PROGRAMMATIC_METRICS_HEADER_VALUE;
     } else {
-      this.awsSecurityCredentialsProvider =
-          new InternalAwsSecurityCredentialsProvider(
-              credentialSource,
-              this.getEnvironmentProvider(),
-              this.transportFactory,
-              builder.awsRegion);
-      this.builtWithSupplier = false;
+      this.awsSecurityCredentialsSupplier =
+          new InternalAwsSecurityCredentialsSupplier(
+              credentialSource, this.getEnvironmentProvider(), this.transportFactory);
+      this.metricsHeaderValue = AWS_METRICS_HEADER_VALUE;
     }
   }
 
@@ -139,9 +128,9 @@ public String retrieveSubjectToken() throws IOException {
 
     // The targeted region is required to generate the signed request. The regional
     // endpoint must also be used.
-    String region = awsSecurityCredentialsProvider.getRegion();
+    String region = awsSecurityCredentialsSupplier.getRegion();
 
-    AwsSecurityCredentials credentials = awsSecurityCredentialsProvider.getCredentials();
+    AwsSecurityCredentials credentials = awsSecurityCredentialsSupplier.getCredentials();
 
     // Generate the signed request to the AWS STS GetCallerIdentity API.
     Map<String, String> headers = new HashMap<>();
@@ -168,7 +157,7 @@ public GoogleCredentials createScoped(Collection<String> newScopes) {
 
   @Override
   String getCredentialSourceType() {
-    return this.awsSecurityCredentialsProvider.getMetricsHeaderValue();
+    return this.metricsHeaderValue;
   }
 
   private String buildSubjectToken(AwsRequestSignature signature)
@@ -205,8 +194,8 @@ String getEnv(String name) {
   }
 
   @VisibleForTesting
-  AwsSecurityCredentialsProvider getAwsSecurityCredentialsProvider() {
-    return this.awsSecurityCredentialsProvider;
+  AwsSecurityCredentialsSupplier getAwsSecurityCredentialsSupplier() {
+    return this.awsSecurityCredentialsSupplier;
   }
 
   @Nullable
@@ -238,52 +227,35 @@ public static AwsCredentials.Builder newBuilder(AwsCredentials awsCredentials) {
 
   public static class Builder extends ExternalAccountCredentials.Builder {
 
-    private Supplier<AwsSecurityCredentials> awsSecurityCredentialsSupplier;
-
-    private String awsRegion;
+    private AwsSecurityCredentialsSupplier awsSecurityCredentialsSupplier;
 
     private String regionalCredentialVerificationUrlOverride;
 
     Builder() {}
 
     Builder(AwsCredentials credentials) {
       super(credentials);
-      this.awsRegion = credentials.awsRegion;
-      if (credentials.builtWithSupplier) {
-        this.awsSecurityCredentialsSupplier =
-            credentials.awsSecurityCredentialsProvider.getSupplier();
+      if (this.credentialSource == null) {
+        this.awsSecurityCredentialsSupplier = credentials.awsSecurityCredentialsSupplier;
       }
       this.regionalCredentialVerificationUrlOverride =
           credentials.regionalCredentialVerificationUrlOverride;
     }
 
     /**
      * Sets the AWS security credentials supplier. The supplier should return a valid {@code
-     * AwsSecurityCredentials} object. An AWS region also is required when using a supplier.
+     * AwsSecurityCredentials} object and a valid AWS region.
      *
-     * @param awsSecurityCredentialsSupplier the supplier method to be called.
+     * @param awsSecurityCredentialsSupplier the supplier to use.
      * @return this {@code Builder} object
      */
     @CanIgnoreReturnValue
     public Builder setAwsSecurityCredentialsSupplier(
-        Supplier<AwsSecurityCredentials> awsSecurityCredentialsSupplier) {
+        AwsSecurityCredentialsSupplier awsSecurityCredentialsSupplier) {
       this.awsSecurityCredentialsSupplier = awsSecurityCredentialsSupplier;
       return this;
     }
 
-    /**
-     * Sets the AWS region. Required when using an AWS Security Credentials Supplier. If set, will
-     * override any region obtained via environment variables or the metadata endpoint.
-     *
-     * @param awsRegion the aws region to set.
-     * @return this {@code Builder} object
-     */
-    @CanIgnoreReturnValue
-    public Builder setAwsRegion(String awsRegion) {
-      this.awsRegion = awsRegion;
-      return this;
-    }
-
     /**
      * Sets the AWS regional credential verification URL. If set, will override any credential
      * verification URL provided in the credential source. If not set, the credential verification

oauth2_http/java/com/google/auth/oauth2/AwsSecurityCredentialsProvider.java → oauth2_http/java/com/google/auth/oauth2/AwsSecurityCredentialsSupplier.java

@@ -33,41 +33,26 @@
 
 import java.io.IOException;
 import java.io.Serializable;
-import java.util.function.Supplier;
 
 /**
- * Provider for retrieving AWS Security credentials for {@Link AwsCredentials} to exchange for GCP
+ * Supplier for retrieving AWS Security credentials for {@Link AwsCredentials} to exchange for GCP
  * access tokens.
  */
-abstract class AwsSecurityCredentialsProvider implements Serializable {
+public interface AwsSecurityCredentialsSupplier extends Serializable {
 
   /**
    * Gets the AWS region to use.
    *
    * @return the AWS region that should be used for the credential.
    * @throws IOException
    */
-  abstract String getRegion() throws IOException;
+  String getRegion() throws IOException;
 
   /**
    * Gets AWS security credentials.
    *
    * @return valid AWS security credentials that can be exchanged for a GCP access token.
    * @throws IOException
    */
-  abstract AwsSecurityCredentials getCredentials() throws IOException;
-
-  /**
-   * Gets the metrics header value that should be used for the sts request.
-   *
-   * @return the metrics header value.
-   */
-  abstract String getMetricsHeaderValue();
-
-  /**
-   * Gets the Aws security credential supplier.
-   *
-   * @return the Supplier used to retrieve the AWS security credentials.
-   */
-  abstract Supplier<AwsSecurityCredentials> getSupplier();
+  AwsSecurityCredentials getCredentials() throws IOException;
 }

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

@@ -476,8 +476,8 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
     if (this.shouldBuildImpersonatedCredential()) {
       this.impersonatedCredentials = this.buildImpersonatedCredentials();
     }
-    if (impersonatedCredentials != null) {
-      return impersonatedCredentials.refreshAccessToken();
+    if (this.impersonatedCredentials != null) {
+      return this.impersonatedCredentials.refreshAccessToken();
     }
 
     StsRequestHandler.Builder requestHandler =

oauth2_http/java/com/google/auth/oauth2/FileIdentityPoolSubjectTokenProvider.java → oauth2_http/java/com/google/auth/oauth2/FileIdentityPoolSubjectTokenSupplier.java

@@ -31,21 +31,26 @@
 
 package com.google.auth.oauth2;
 
-import static com.google.auth.oauth2.IdentityPoolCredentials.FILE_METRICS_HEADER_VALUE;
-
+import com.google.api.client.json.GenericJson;
+import com.google.api.client.json.JsonObjectParser;
+import com.google.auth.oauth2.IdentityPoolCredentialSource.CredentialFormatType;
+import com.google.common.io.CharStreams;
+import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.LinkOption;
 import java.nio.file.Paths;
-import java.util.function.Supplier;
 
 /**
  * Internal provider for retrieving subject tokens for {@Link IdentityPoolCredentials} to exchange
  * for GCP access tokens via a local file.
  */
-class FileIdentityPoolSubjectTokenProvider extends IdentityPoolSubjectTokenProvider {
+class FileIdentityPoolSubjectTokenSupplier implements IdentityPoolSubjectTokenSupplier {
 
   private final long serialVersionUID = 2475549052347431992L;
 
@@ -56,32 +61,12 @@ class FileIdentityPoolSubjectTokenProvider extends IdentityPoolSubjectTokenProvi
    *
    * @param credentialSource the credential source to use.
    */
-  FileIdentityPoolSubjectTokenProvider(IdentityPoolCredentialSource credentialSource) {
+  FileIdentityPoolSubjectTokenSupplier(IdentityPoolCredentialSource credentialSource) {
     this.credentialSource = credentialSource;
   }
 
   @Override
-  String getSubjectToken() throws IOException {
-    return this.retrieveSubjectTokenFromCredentialFile();
-  }
-
-  @Override
-  String getMetricsHeaderValue() {
-    return FILE_METRICS_HEADER_VALUE;
-  }
-
-  @Override
-  Supplier<String> getSupplier() {
-    return () -> {
-      try {
-        return this.getSubjectToken();
-      } catch (IOException e) {
-        throw new RuntimeException(e);
-      }
-    };
-  }
-
-  private String retrieveSubjectTokenFromCredentialFile() throws IOException {
+  public String getSubjectToken() throws IOException {
     String credentialFilePath = this.credentialSource.credentialLocation;
     if (!Files.exists(Paths.get(credentialFilePath), LinkOption.NOFOLLOW_LINKS)) {
       throw new IOException(
@@ -95,4 +80,22 @@ private String retrieveSubjectTokenFromCredentialFile() throws IOException {
           "Error when attempting to read the subject token from the credential file.", e);
     }
   }
+
+  static String parseToken(InputStream inputStream, IdentityPoolCredentialSource credentialSource)
+      throws IOException {
+    if (credentialSource.credentialFormatType == CredentialFormatType.TEXT) {
+      BufferedReader reader =
+          new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
+      return CharStreams.toString(reader);
+    }
+
+    JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
+    GenericJson fileContents =
+        parser.parseAndClose(inputStream, StandardCharsets.UTF_8, GenericJson.class);
+
+    if (!fileContents.containsKey(credentialSource.subjectTokenFieldName)) {
+      throw new IOException("Invalid subject token field name. No subject token was found.");
+    }
+    return (String) fileContents.get(credentialSource.subjectTokenFieldName);
+  }
 }

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

@@ -38,7 +38,6 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Map;
-import java.util.function.Supplier;
 
 /**
  * Url-sourced, file-sourced, or user provided supplier method-sourced external account credentials.
@@ -51,10 +50,8 @@ public class IdentityPoolCredentials extends ExternalAccountCredentials {
   static final String URL_METRICS_HEADER_VALUE = "url";
   private static final long serialVersionUID = 2471046175477275881L;
 
-  private final IdentityPoolSubjectTokenProvider subjectTokenProvider;
-  // Boolean to tell if a Supplier was provided to build the credential. This is needed to tell
-  // whether the supplier should be passed when creating a new builder from this credential.
-  private final boolean builtWithSupplier;
+  private final IdentityPoolSubjectTokenSupplier subjectTokenSupplier;
+  private final String metricsHeaderValue;
 
   /** Internal constructor. See {@link Builder}. */
   IdentityPoolCredentials(Builder builder) {
@@ -72,17 +69,16 @@ public class IdentityPoolCredentials extends ExternalAccountCredentials {
           "A subjectTokenSupplier or a credentialSource must be provided.");
     }
     if (builder.subjectTokenSupplier != null) {
-      this.subjectTokenProvider =
-          new ProgrammaticIdentityPoolSubjectTokenProvider(builder.subjectTokenSupplier);
-      this.builtWithSupplier = true;
+      this.subjectTokenSupplier = builder.subjectTokenSupplier;
+      this.metricsHeaderValue = PROGRAMMATIC_METRICS_HEADER_VALUE;
     } else if (credentialSource.credentialSourceType
         == IdentityPoolCredentialSource.IdentityPoolCredentialSourceType.FILE) {
-      this.subjectTokenProvider = new FileIdentityPoolSubjectTokenProvider(credentialSource);
-      this.builtWithSupplier = false;
+      this.subjectTokenSupplier = new FileIdentityPoolSubjectTokenSupplier(credentialSource);
+      this.metricsHeaderValue = FILE_METRICS_HEADER_VALUE;
     } else {
-      this.subjectTokenProvider =
-          new UrlIdentityPoolSubjectTokenProvider(credentialSource, this.transportFactory);
-      this.builtWithSupplier = false;
+      this.subjectTokenSupplier =
+          new UrlIdentityPoolSubjectTokenSupplier(credentialSource, this.transportFactory);
+      this.metricsHeaderValue = URL_METRICS_HEADER_VALUE;
     }
   }
 
@@ -103,17 +99,17 @@ public AccessToken refreshAccessToken() throws IOException {
 
   @Override
   public String retrieveSubjectToken() throws IOException {
-    return this.subjectTokenProvider.getSubjectToken();
+    return this.subjectTokenSupplier.getSubjectToken();
   }
 
   @Override
   String getCredentialSourceType() {
-    return this.subjectTokenProvider.getMetricsHeaderValue();
+    return this.metricsHeaderValue;
   }
 
   @VisibleForTesting
-  IdentityPoolSubjectTokenProvider getIdentityPoolSubjectTokenProvider() {
-    return this.subjectTokenProvider;
+  IdentityPoolSubjectTokenSupplier getIdentityPoolSubjectTokenSupplier() {
+    return this.subjectTokenSupplier;
   }
 
   /** Clones the IdentityPoolCredentials with the specified scopes. */
@@ -133,25 +129,25 @@ public static Builder newBuilder(IdentityPoolCredentials identityPoolCredentials
 
   public static class Builder extends ExternalAccountCredentials.Builder {
 
-    private Supplier<String> subjectTokenSupplier;
+    private IdentityPoolSubjectTokenSupplier subjectTokenSupplier;
 
     Builder() {}
 
     Builder(IdentityPoolCredentials credentials) {
       super(credentials);
-      if (credentials.builtWithSupplier) {
-        this.setSubjectTokenSupplier(credentials.subjectTokenProvider.getSupplier());
+      if (this.credentialSource == null) {
+        this.subjectTokenSupplier = credentials.subjectTokenSupplier;
       }
     }
 
     /**
      * Sets the subject token supplier. The supplier should return a valid subject token string.
      *
-     * @param subjectTokenSupplier the supplier method to be called.
+     * @param subjectTokenSupplier the supplier to use.
      * @return this {@code Builder} object
      */
     @CanIgnoreReturnValue
-    public Builder setSubjectTokenSupplier(Supplier<String> subjectTokenSupplier) {
+    public Builder setSubjectTokenSupplier(IdentityPoolSubjectTokenSupplier subjectTokenSupplier) {
       this.subjectTokenSupplier = subjectTokenSupplier;
       return this;
     }