Adding testing functionality for requests made. Refactoring header pa…

…ssing logic

README.md

@@ -211,10 +211,11 @@ Where the following variables need to be substituted:
 
 This generates the configuration file in the specified output file.
 
-If you want to use IDMSv2, then below field needs to be added to credential_source
-section of credential configuration.
+If you want to use the AWS IMDSv2 flow, you can add the field below to the credential_source in your AWS ADC configuration file:
 "aws_session_token_url": "http://169.254.169.254/latest/api/token"
 
+The gcloud create-cred-config command will be updated to support this soon.
+
 You can now [use the Auth library](#using-external-identities) to call Google Cloud
 resources from AWS.
 

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

@@ -32,7 +32,9 @@
 package com.google.auth.oauth2;
 
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpContent;
 import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpMethods;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpRequestFactory;
 import com.google.api.client.http.HttpResponse;
@@ -49,6 +51,7 @@
 import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 
 /**
  * AWS credentials representing a third-party identity for calling Google APIs.
@@ -143,16 +146,34 @@ public AccessToken refreshAccessToken() throws IOException {
 
   @Override
   public String retrieveSubjectToken() throws IOException {
-    String awsSessionToken = null;
+    // AWS IDMSv2 introduced a requirement for a session token to be present
+    // with the requests made to metadata endpoints. This requirement is to help
+    // prevent SSRF attacks.
+    // Presence of "aws_session_token_url" in Credential Source of config file
+    // will trigger a flow with session token, else there will not be a session
+    // token with the metadata requests.
+    // Both flows work for IDMS v1 and v2. But if IDMSv2 is enabled, then if
+    // session token is not present, Unauthorized exception will be thrown.
+    Map<String, Object> metadataHeaders = new HashMap<>();
     if (awsCredentialSource.awsSessionTokenUrl != null) {
-      awsSessionToken = getAwsSessionToken(awsCredentialSource.awsSessionTokenUrl);
+      Map<String, Object> tokenRequestHeaders =
+          Map.of("x-aws-ec2-metadata-token-ttl-seconds", "21600");
+
+      String awsSessionToken =
+          retrieveResource(
+              awsCredentialSource.awsSessionTokenUrl,
+              "Session Token",
+              HttpMethods.PUT,
+              tokenRequestHeaders,
+              /*content =*/ null);
+      metadataHeaders.put("x-aws-ec2-metadata-token", awsSessionToken);
     }
 
     // The targeted region is required to generate the signed request. The regional
     // endpoint must also be used.
-    String region = getAwsRegion(awsSessionToken);
+    String region = getAwsRegion(metadataHeaders);
 
-    AwsSecurityCredentials credentials = getAwsSecurityCredentials(awsSessionToken);
+    AwsSecurityCredentials credentials = getAwsSecurityCredentials(metadataHeaders);
 
     // Generate the signed request to the AWS STS GetCallerIdentity API.
     Map<String, String> headers = new HashMap<>();
@@ -177,36 +198,32 @@ public GoogleCredentials createScoped(Collection<String> newScopes) {
     return new AwsCredentials((AwsCredentials.Builder) newBuilder(this).setScopes(newScopes));
   }
 
-  private String retrieveResource(String url, String resourceName, String sessionToken)
+  private String retrieveResource(String url, String resourceName, Map<String, Object> headers)
       throws IOException {
-    try {
-      HttpRequestFactory requestFactory = transportFactory.create().createRequestFactory();
-      HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(url));
-
-      if (sessionToken != null) {
-        HttpHeaders headers = request.getHeaders();
-        headers.set("X-aws-ec2-metadata-token", sessionToken);
-      }
-
-      HttpResponse response = request.execute();
-      return response.parseAsString();
-    } catch (IOException e) {
-      throw new IOException(String.format("Failed to retrieve AWS %s.", resourceName), e);
-    }
+    return retrieveResource(url, resourceName, HttpMethods.GET, headers, /*content =*/ null);
   }
 
-  private String getAwsSessionToken(String awsSessionTokenUrl) throws IOException {
+  private String retrieveResource(
+      String url,
+      String resourceName,
+      String requestMethod,
+      Map<String, Object> headers,
+      @Nullable HttpContent content)
+      throws IOException {
     try {
       HttpRequestFactory requestFactory = transportFactory.create().createRequestFactory();
       HttpRequest request =
-          requestFactory.buildPutRequest(new GenericUrl(awsSessionTokenUrl), null);
-      HttpHeaders headers = request.getHeaders();
-      headers.set("X-aws-ec2-metadata-token-ttl-seconds", "21600");
+          requestFactory.buildRequest(requestMethod, new GenericUrl(url), content);
+
+      HttpHeaders httpHeaders = request.getHeaders();
+      for (Map.Entry<String, Object> header : headers.entrySet()) {
+        httpHeaders.set(header.getKey(), header.getValue());
+      }
 
       HttpResponse response = request.execute();
       return response.parseAsString();
     } catch (IOException e) {
-      throw new IOException(String.format("Failed to fetch AWS Session Token"), e);
+      throw new IOException(String.format("Failed to retrieve AWS %s.", resourceName), e);
     }
   }
 
@@ -236,7 +253,7 @@ private String buildSubjectToken(AwsRequestSignature signature)
   }
 
   @VisibleForTesting
-  String getAwsRegion(String awsSessionToken) throws IOException {
+  String getAwsRegion(Map<String, Object> metadataHeaders) throws IOException {
     // For AWS Lambda, the region is retrieved through the AWS_REGION environment variable.
     String region = getEnvironmentProvider().getEnv("AWS_REGION");
     if (region != null) {
@@ -253,15 +270,16 @@ String getAwsRegion(String awsSessionToken) throws IOException {
           "Unable to determine the AWS region. The credential source does not contain the region URL.");
     }
 
-    region = retrieveResource(awsCredentialSource.regionUrl, "region", awsSessionToken);
+    region = retrieveResource(awsCredentialSource.regionUrl, "region", metadataHeaders);
 
     // There is an extra appended character that must be removed. If `us-east-1b` is returned,
     // we want `us-east-1`.
     return region.substring(0, region.length() - 1);
   }
 
   @VisibleForTesting
-  AwsSecurityCredentials getAwsSecurityCredentials(String awsSessionToken) throws IOException {
+  AwsSecurityCredentials getAwsSecurityCredentials(Map<String, Object> metadataHeaders)
+      throws IOException {
     // Check environment variables for credentials first.
     String accessKeyId = getEnvironmentProvider().getEnv("AWS_ACCESS_KEY_ID");
     String secretAccessKey = getEnvironmentProvider().getEnv("AWS_SECRET_ACCESS_KEY");
@@ -278,12 +296,12 @@ AwsSecurityCredentials getAwsSecurityCredentials(String awsSessionToken) throws
           "Unable to determine the AWS IAM role name. The credential source does not contain the"
               + " url field.");
     }
-    String roleName = retrieveResource(awsCredentialSource.url, "IAM role", awsSessionToken);
+    String roleName = retrieveResource(awsCredentialSource.url, "IAM role", metadataHeaders);
 
     // Retrieve the AWS security credentials by calling the endpoint specified by the credential
     // source.
     String awsCredentials =
-        retrieveResource(awsCredentialSource.url + "/" + roleName, "credentials", awsSessionToken);
+        retrieveResource(awsCredentialSource.url + "/" + roleName, "credentials", metadataHeaders);
 
     JsonParser parser = OAuth2Utils.JSON_FACTORY.createJsonParser(awsCredentials);
     GenericJson genericJson = parser.parseAndClose(GenericJson.class);