ITWorkloadIdentityFederationTest.java

/*
 * Copyright 2021 Google LLC
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *    * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *    * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *
 *    * Neither the name of Google LLC nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package com.google.auth.oauth2;

import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

import com.google.api.client.http.ByteArrayContent;
import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpContent;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpRequestFactory;
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.UrlEncodedContent;
import com.google.api.client.http.javanet.NetHttpTransport;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.JsonObjectParser;
import com.google.api.client.json.gson.GsonFactory;
import com.google.api.client.json.webtoken.JsonWebSignature;
import com.google.api.client.json.webtoken.JsonWebToken;
import com.google.api.client.util.GenericData;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.auth.http.HttpTransportFactory;
import com.google.auth.mtls.MtlsHttpTransportFactory;
import com.google.auth.mtls.X509Provider;
import com.google.auth.oauth2.ExternalAccountCredentials.SubjectTokenTypes;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.security.MessageDigest;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.util.Base64;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

/**
 * Integration tests for Workload Identity Federation.
 *
 * <p>The only requirements for this test suite to run is to set the environment variable
 * GOOGLE_APPLICATION_CREDENTIALS to point to the same service account keys used in the setup script
 * (workloadidentityfederation-setup). These tests call GCS to get bucket information. The bucket
 * name must be provided through the GCS_BUCKET environment variable. The mTLS tests use {@code
 * testresources/mtls/certificate_config.json}.
 */
final class ITWorkloadIdentityFederationTest {

  // Copy output from workloadidentityfederation-setup.
  private static final String AUDIENCE_PREFIX =
      "//iam.googleapis.com/projects/1016721519174/locations/global/workloadIdentityPools/pool-1/providers/";
  private static final String AWS_ROLE_NAME = "ci-java-test";
  private static final String AWS_ROLE_ARN = "arn:aws:iam::027472800722:role/ci-java-test";

  private static final String AWS_AUDIENCE = AUDIENCE_PREFIX + "aws-1";
  private static final String OIDC_AUDIENCE = AUDIENCE_PREFIX + "oidc-1";
  private static final String OIDC_JWT_AUDIENCE = AUDIENCE_PREFIX + "oidc-jwt-1";
  private static final String OIDC_JWT_ISSUER =
      "https://storage.googleapis.com/example-oidc-issuer";
  private static final String DEFAULT_MTLS_CERT_CONFIG_PATH =
      "testresources/mtls/certificate_config.json";

  private static volatile String syncedOidcJwtKeyId;

  private String clientEmail;
  private String clientId;
  private String privateKeyId;
  private PrivateKey privateKey;

  @BeforeEach
  void setup() throws IOException {
    GenericJson keys = getServiceAccountKeyFileAsJson();
    clientEmail = (String) keys.get("client_email");
    clientId = (String) keys.get("client_id");
    privateKeyId = (String) keys.get("private_key_id");
    privateKey = OAuth2Utils.privateKeyFromPkcs8((String) keys.get("private_key"));
  }

  /**
   * IdentityPoolCredentials (OIDC provider): Uses the service account to generate a Google ID token
   * using the iamcredentials generateIdToken API. This will use the service account client ID as
   * the sub field of the token. This OIDC token will be used as the external subject token to be
   * exchanged for a GCP access token via GCP STS endpoint and then to impersonate the original
   * service account key. Retrieves the OIDC token from a file.
   */
  @Test
  void identityPoolCredentials() throws IOException {
    IdentityPoolCredentials identityPoolCredentials =
        (IdentityPoolCredentials)
            ExternalAccountCredentials.fromJson(
                buildIdentityPoolCredentialConfig(), OAuth2Utils.HTTP_TRANSPORT_FACTORY);

    callGcs(identityPoolCredentials);
  }

  /**
   * AwsCredentials (AWS provider): Uses the service account keys to generate a Google ID token
   * using the iamcredentials generateIdToken API. Exchanges the OIDC ID token for AWS security keys
   * using AWS STS AssumeRoleWithWebIdentity API. These values will be set as AWS environment
   * variables to simulate an AWS VM. The Auth library can now read these variables and create a
   * signed request to AWS GetCallerIdentity. This will be used as the external subject token to be
   * exchanged for a GCP access token via GCP STS endpoint and then to impersonate the original
   * service account key.
   */
  @Test
  void awsCredentials() throws Exception {
    String idToken = generateGoogleIdToken(AWS_AUDIENCE);

    String url =
        String.format(
            "https://sts.amazonaws.com/?Action=AssumeRoleWithWebIdentity"
                + "&Version=2011-06-15&DurationSeconds=3600&RoleSessionName=%s"
                + "&RoleArn=%s&WebIdentityToken=%s",
            AWS_ROLE_NAME, AWS_ROLE_ARN, idToken);

    HttpRequestFactory requestFactory = new NetHttpTransport().createRequestFactory();
    HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(url));

    JsonObjectParser parser = new JsonObjectParser(GsonFactory.getDefaultInstance());
    request.setParser(parser);

    HttpResponse response = request.execute();
    String rawXml = response.parseAsString();

    String awsAccessKeyId = getXmlValueByTagName(rawXml, "AccessKeyId");
    String awsSecretAccessKey = getXmlValueByTagName(rawXml, "SecretAccessKey");
    String awsSessionToken = getXmlValueByTagName(rawXml, "SessionToken");

    AwsCredentials awsCredentialWithoutEnvProvider =
        (AwsCredentials)
            AwsCredentials.fromJson(buildAwsCredentialConfig(), OAuth2Utils.HTTP_TRANSPORT_FACTORY);
    TestEnvironmentProvider testEnvironmentProvider = new TestEnvironmentProvider();
    testEnvironmentProvider
        .setEnv("AWS_ACCESS_KEY_ID", awsAccessKeyId)
        .setEnv("AWS_SECRET_ACCESS_KEY", awsSecretAccessKey)
        .setEnv("AWS_SESSION_TOKEN", awsSessionToken)
        .setEnv("AWS_REGION", "us-east-2");

    AwsCredentials awsCredential =
        AwsCredentials.newBuilder(awsCredentialWithoutEnvProvider)
            .setEnvironmentProvider(testEnvironmentProvider)
            .build();

    callGcs(awsCredential);
  }

  /**
   * AwsCredentials (AWS Provider): Uses the service account keys to generate a Google ID token
   * using the iamcredentials generateIdToken API. Exchanges the OIDC ID token for AWS security keys
   * using AWS STS AssumeRoleWithWebIdentity API. These values will be returned as a
   * AwsSecurityCredentials object and returned by a Supplier. The Auth library can now call get()
   * from the supplier and create a signed request to AWS GetCallerIdentity. This will be used as
   * the external subject token to be exchanged for a GCP access token via GCP STS endpoint and then
   * to impersonate the original service account key.
   */
  @Test
  void awsCredentials_withProgrammaticAuth() throws Exception {
    String idToken = generateGoogleIdToken(AWS_AUDIENCE);

    String url =
        String.format(
            "https://sts.amazonaws.com/?Action=AssumeRoleWithWebIdentity"
                + "&Version=2011-06-15&DurationSeconds=3600&RoleSessionName=%s"
                + "&RoleArn=%s&WebIdentityToken=%s",
            AWS_ROLE_NAME, AWS_ROLE_ARN, idToken);

    HttpRequestFactory requestFactory = new NetHttpTransport().createRequestFactory();
    HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(url));

    JsonObjectParser parser = new JsonObjectParser(GsonFactory.getDefaultInstance());
    request.setParser(parser);

    HttpResponse response = request.execute();
    String rawXml = response.parseAsString();

    String awsAccessKeyId = getXmlValueByTagName(rawXml, "AccessKeyId");
    String awsSecretAccessKey = getXmlValueByTagName(rawXml, "SecretAccessKey");
    String awsSessionToken = getXmlValueByTagName(rawXml, "SessionToken");

    AwsSecurityCredentials credentials =
        new AwsSecurityCredentials(awsAccessKeyId, awsSecretAccessKey, awsSessionToken);

    AwsSecurityCredentialsSupplier provider =
        new ITAwsSecurityCredentialsProvider("us-east-2", credentials);
    AwsCredentials awsCredential =
        AwsCredentials.newBuilder()
            .setAwsSecurityCredentialsSupplier(provider)
            .setSubjectTokenType(SubjectTokenTypes.AWS4)
            .setAudience(AWS_AUDIENCE)
            .setServiceAccountImpersonationUrl(
                String.format(
                    "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
                    clientEmail))
            .setHttpTransportFactory(OAuth2Utils.HTTP_TRANSPORT_FACTORY)
            .build();

    callGcs(awsCredential);
  }

  /**
   * PluggableCredential (OIDC provider): Uses the service account to generate a Google ID token
   * using the iamcredentials generateIdToken API. This will use the service account client ID as
   * the sub field of the token. This OIDC token will be used as the external subject token to be
   * exchanged for a GCP access token via GCP STS endpoint and then to impersonate the original
   * service account key. Runs an executable to get the OIDC token.
   */
  @Test
  void pluggableAuthCredentials() throws IOException {
    PluggableAuthCredentials pluggableAuthCredentials =
        (PluggableAuthCredentials)
            ExternalAccountCredentials.fromJson(
                buildPluggableCredentialConfig(), OAuth2Utils.HTTP_TRANSPORT_FACTORY);

    callGcs(pluggableAuthCredentials);
  }

  /**
   * Sets the service account impersonation object in configuration JSON with a non-default value
   * for token_lifetime_seconds and validates that the lifetime is used for the access token.
   */
  @Test
  void identityPoolCredentials_withServiceAccountImpersonationOptions() throws IOException {
    GenericJson identityPoolCredentialConfig = buildIdentityPoolCredentialConfig();
    Map<String, Object> map = new HashMap<String, Object>();
    map.put("token_lifetime_seconds", 2800);
    identityPoolCredentialConfig.put("service_account_impersonation", map);

    IdentityPoolCredentials identityPoolCredentials =
        (IdentityPoolCredentials)
            ExternalAccountCredentials.fromJson(
                identityPoolCredentialConfig, OAuth2Utils.HTTP_TRANSPORT_FACTORY);
    long maxExpirationTime = Instant.now().plusSeconds(2800 + 5).toEpochMilli();
    long minExpirationTime = Instant.now().plusSeconds(2800 - 5).toEpochMilli();

    callGcs(identityPoolCredentials);
    long tokenExpiry = identityPoolCredentials.getAccessToken().getExpirationTimeMillis();
    assertTrue(minExpirationTime <= tokenExpiry && tokenExpiry <= maxExpirationTime);
  }

  /**
   * IdentityPoolCredentials (OIDC provider): Uses the service account to generate a Google ID token
   * using the iamcredentials generateIdToken API. This will use the service account client ID as
   * the sub field of the token. This OIDC token will be used as the external subject token to be
   * exchanged for a GCP access token via GCP STS endpoint and then to impersonate the original
   * service account key. Retrieves the OIDC token from a Supplier that returns the subject token
   * when get() is called.
   */
  @Test
  void identityPoolCredentials_withProgrammaticAuth() throws IOException {

    IdentityPoolSubjectTokenSupplier tokenSupplier =
        (ExternalAccountSupplierContext context) -> {
          try {
            return generateGoogleIdToken(OIDC_AUDIENCE);
          } catch (IOException e) {
            throw new RuntimeException(e);
          }
        };

    IdentityPoolCredentials identityPoolCredentials =
        IdentityPoolCredentials.newBuilder()
            .setSubjectTokenSupplier(tokenSupplier)
            .setAudience(OIDC_AUDIENCE)
            .setSubjectTokenType(SubjectTokenTypes.JWT)
            .setServiceAccountImpersonationUrl(
                String.format(
                    "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
                    clientEmail))
            .setHttpTransportFactory(OAuth2Utils.HTTP_TRANSPORT_FACTORY)
            .build();

    callGcs(identityPoolCredentials);
  }

  /**
   * IdentityPoolCredentials (OIDC provider with certificate-bound workload and actor token): Uses
   * the service account private key to sign OIDC JWTs for subject and actor tokens with RFC 8705
   * certificate thumbprint binding (cnf.x5t#S256) and delegation (may_act) claims. Writes both
   * tokens to a temporary JSON file with subject_token and actor_token field names. Configures
   * certificate_config_location pointing to the certificate config. Exchanges the tokens over the
   * mTLS STS endpoint (https://sts.mtls.googleapis.com/v1/token) and calls GCS.
   */
  @Test
  void identityPoolCredentials_withCertificateBoundWorkloadAndActorToken() throws Exception {
    KeyStore keyStore = new X509Provider(DEFAULT_MTLS_CERT_CONFIG_PATH).getKeyStore();
    String certThumbprint = computeCertThumbprint(keyStore);
    String subjectToken =
        generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ true);
    String actorToken =
        generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ false);

    File tokenFile =
        File.createTempFile(
            "ITWorkloadIdentityFederation_cert_actor", /* suffix= */ null, /* directory= */ null);
    tokenFile.deleteOnExit();
    try {
      GenericJson tokenJson = new GenericJson();
      tokenJson.setFactory(OAuth2Utils.JSON_FACTORY);
      tokenJson.put("subject_token", subjectToken);
      tokenJson.put("actor_token", actorToken);

      OAuth2Utils.writeInputStreamToFile(
          new ByteArrayInputStream(tokenJson.toPrettyString().getBytes(StandardCharsets.UTF_8)),
          tokenFile.getAbsolutePath());

      GenericJson config = new GenericJson();
      config.put("type", "external_account");
      config.put("audience", OIDC_JWT_AUDIENCE);
      config.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
      config.put("actor_token_type", "urn:ietf:params:oauth:token-type:jwt");
      config.put("token_url", "https://sts.mtls.googleapis.com/v1/token");
      config.put(
          "service_account_impersonation_url",
          String.format(
              "https://iamcredentials.mtls.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
              clientEmail));

      GenericJson credentialSource = new GenericJson();
      credentialSource.put("file", tokenFile.getAbsolutePath());

      GenericJson format = new GenericJson();
      format.put("type", "json");
      format.put("subject_token_field_name", "subject_token");
      format.put("actor_token_field_name", "actor_token");
      credentialSource.put("format", format);

      GenericJson certificate = new GenericJson();
      certificate.put("certificate_config_location", DEFAULT_MTLS_CERT_CONFIG_PATH);
      credentialSource.put("certificate", certificate);

      config.put("credential_source", credentialSource);

      IdentityPoolCredentials identityPoolCredentials =
          (IdentityPoolCredentials)
              ExternalAccountCredentials.fromJson(config, OAuth2Utils.HTTP_TRANSPORT_FACTORY);

      HttpTransport mtlsTransport = new MtlsHttpTransportFactory(keyStore).create();
      callGcs(
          identityPoolCredentials,
          mtlsTransport,
          "https://storage.mtls.googleapis.com/storage/v1/b/");
    } finally {
      tokenFile.delete();
    }
  }

  /**
   * IdentityPoolCredentials (OIDC provider with programmatic mTLS and actor token): Uses the
   * service account private key to sign OIDC JWTs for subject and actor tokens via suppliers.
   * Configures mTLS transport using MtlsHttpTransportFactory with KeyStore loaded from the mTLS
   * certificate config. Exchanges the tokens over mTLS STS endpoint and calls GCS.
   */
  @Test
  void identityPoolCredentials_withProgrammaticMtlsAndActorToken() throws Exception {
    X509Provider x509Provider = new X509Provider(DEFAULT_MTLS_CERT_CONFIG_PATH);
    KeyStore keyStore = x509Provider.getKeyStore();
    String certThumbprint = computeCertThumbprint(keyStore);
    IdentityPoolSubjectTokenSupplier tokenSupplier =
        context ->
            generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ true);
    IdentityPoolActorTokenSupplier actorSupplier =
        context ->
            generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ false);

    HttpTransportFactory transportFactory = new MtlsHttpTransportFactory(keyStore);

    IdentityPoolCredentials credentials =
        IdentityPoolCredentials.newBuilder()
            .setSubjectTokenSupplier(tokenSupplier)
            .setActorTokenSupplier(actorSupplier)
            .setActorTokenType(SubjectTokenTypes.JWT.value)
            .setAudience(OIDC_JWT_AUDIENCE)
            .setSubjectTokenType(SubjectTokenTypes.JWT)
            .setTokenUrl("https://sts.mtls.googleapis.com/v1/token")
            .setServiceAccountImpersonationUrl(
                String.format(
                    "https://iamcredentials.mtls.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
                    clientEmail))
            .setHttpTransportFactory(transportFactory)
            .build();

    callGcs(
        credentials,
        transportFactory.create(),
        "https://storage.mtls.googleapis.com/storage/v1/b/");
  }

  /**
   * IdentityPoolCredentials (OIDC provider with certificate-bound workload and actor token, direct
   * STS): Exchanges tokens directly over the mTLS STS endpoint
   * (https://sts.mtls.googleapis.com/v1/token) without service account impersonation.
   */
  @Test
  void identityPoolCredentials_directSts_withCertificateBoundWorkloadAndActorToken()
      throws Exception {
    KeyStore keyStore = new X509Provider(DEFAULT_MTLS_CERT_CONFIG_PATH).getKeyStore();
    String certThumbprint = computeCertThumbprint(keyStore);
    String subjectToken =
        generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ true);
    String actorToken =
        generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ false);

    File tokenFile =
        File.createTempFile(
            "ITWorkloadIdentityFederation_direct_cert_actor",
            /* suffix= */ null,
            /* directory= */ null);
    tokenFile.deleteOnExit();
    try {
      GenericJson tokenJson = new GenericJson();
      tokenJson.setFactory(OAuth2Utils.JSON_FACTORY);
      tokenJson.put("subject_token", subjectToken);
      tokenJson.put("actor_token", actorToken);

      OAuth2Utils.writeInputStreamToFile(
          new ByteArrayInputStream(tokenJson.toPrettyString().getBytes(StandardCharsets.UTF_8)),
          tokenFile.getAbsolutePath());

      GenericJson config = new GenericJson();
      config.put("type", "external_account");
      config.put("audience", OIDC_JWT_AUDIENCE);
      config.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
      config.put("actor_token_type", "urn:ietf:params:oauth:token-type:jwt");
      config.put("token_url", "https://sts.mtls.googleapis.com/v1/token");

      GenericJson credentialSource = new GenericJson();
      credentialSource.put("file", tokenFile.getAbsolutePath());

      GenericJson format = new GenericJson();
      format.put("type", "json");
      format.put("subject_token_field_name", "subject_token");
      format.put("actor_token_field_name", "actor_token");
      credentialSource.put("format", format);

      GenericJson certificate = new GenericJson();
      certificate.put("certificate_config_location", DEFAULT_MTLS_CERT_CONFIG_PATH);
      credentialSource.put("certificate", certificate);

      config.put("credential_source", credentialSource);

      IdentityPoolCredentials identityPoolCredentials =
          (IdentityPoolCredentials)
              ExternalAccountCredentials.fromJson(config, OAuth2Utils.HTTP_TRANSPORT_FACTORY);

      AccessToken accessToken = identityPoolCredentials.refreshAccessToken();
      assertNotNull(accessToken);
      assertNotNull(accessToken.getTokenValue());
      assertFalse(accessToken.getTokenValue().isEmpty());
    } finally {
      tokenFile.delete();
    }
  }

  /**
   * IdentityPoolCredentials (OIDC provider with programmatic mTLS and actor token, direct STS):
   * Uses suppliers for subject and actor tokens, configuring MtlsHttpTransportFactory. Exchanges
   * directly over mTLS STS endpoint without service account impersonation.
   */
  @Test
  void identityPoolCredentials_directSts_withProgrammaticMtlsAndActorToken() throws Exception {
    X509Provider x509Provider = new X509Provider(DEFAULT_MTLS_CERT_CONFIG_PATH);
    KeyStore keyStore = x509Provider.getKeyStore();
    String certThumbprint = computeCertThumbprint(keyStore);
    IdentityPoolSubjectTokenSupplier tokenSupplier =
        context ->
            generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ true);
    IdentityPoolActorTokenSupplier actorSupplier =
        context ->
            generateSignedOidcJwt(OIDC_JWT_AUDIENCE, certThumbprint, /* includeMayAct= */ false);

    HttpTransportFactory transportFactory = new MtlsHttpTransportFactory(keyStore);

    IdentityPoolCredentials credentials =
        IdentityPoolCredentials.newBuilder()
            .setSubjectTokenSupplier(tokenSupplier)
            .setActorTokenSupplier(actorSupplier)
            .setActorTokenType(SubjectTokenTypes.JWT.value)
            .setAudience(OIDC_JWT_AUDIENCE)
            .setSubjectTokenType(SubjectTokenTypes.JWT)
            .setTokenUrl("https://sts.mtls.googleapis.com/v1/token")
            .setHttpTransportFactory(transportFactory)
            .build();

    AccessToken accessToken = credentials.refreshAccessToken();
    assertNotNull(accessToken);
    assertNotNull(accessToken.getTokenValue());
    assertFalse(accessToken.getTokenValue().isEmpty());
  }

  private GenericJson buildIdentityPoolCredentialConfig() throws IOException {
    String idToken = generateGoogleIdToken(OIDC_AUDIENCE);

    File file =
        File.createTempFile(
            "ITWorkloadIdentityFederation", /* suffix= */ null, /* directory= */ null);
    file.deleteOnExit();
    OAuth2Utils.writeInputStreamToFile(
        new ByteArrayInputStream(idToken.getBytes(StandardCharsets.UTF_8)), file.getAbsolutePath());

    GenericJson config = new GenericJson();
    config.put("type", "external_account");
    config.put("audience", OIDC_AUDIENCE);
    config.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
    config.put("token_url", "https://sts.googleapis.com/v1/token");
    config.put(
        "service_account_impersonation_url",
        String.format(
            "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
            clientEmail));

    GenericJson credentialSource = new GenericJson();
    credentialSource.put("file", file.getAbsolutePath());
    config.put("credential_source", credentialSource);

    return config;
  }

  private GenericJson buildPluggableCredentialConfig() throws IOException {
    String idToken = generateGoogleIdToken(OIDC_AUDIENCE);

    Instant expirationTime = Instant.now().plusSeconds(60 * 60);

    GenericJson executableJson = new GenericJson();
    executableJson.setFactory(OAuth2Utils.JSON_FACTORY);
    executableJson.put("success", true);
    executableJson.put("version", 1);
    executableJson.put("expirationTime", expirationTime.toEpochMilli());
    executableJson.put("token_type", "urn:ietf:params:oauth:token-type:jwt");
    executableJson.put("id_token", idToken);

    String fileContents =
        "#!/bin/bash\n"
            + "echo \""
            + executableJson.toPrettyString().replace("\"", "\\\"")
            + "\"\n";

    File file =
        File.createTempFile(
            "ITWorkloadIdentityFederation", /* suffix= */ null, /* directory= */ null);
    file.deleteOnExit();
    if (!file.setExecutable(true, true)) {
      throw new IOException("Unable to make script executable");
    }
    OAuth2Utils.writeInputStreamToFile(
        new ByteArrayInputStream(fileContents.getBytes(StandardCharsets.UTF_8)),
        file.getAbsolutePath());

    GenericJson config = new GenericJson();
    config.put("type", "external_account");
    config.put("audience", OIDC_AUDIENCE);
    config.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
    config.put("token_url", "https://sts.googleapis.com/v1/token");
    config.put(
        "service_account_impersonation_url",
        String.format(
            "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
            clientEmail));

    GenericJson credentialSource = new GenericJson();
    config.put("credential_source", credentialSource);

    GenericJson executableConfig = new GenericJson();
    credentialSource.put("executable", executableConfig);
    executableConfig.put("command", file.getAbsolutePath());

    return config;
  }

  private GenericJson buildAwsCredentialConfig() {
    GenericJson config = new GenericJson();
    config.put("type", "external_account");
    config.put("audience", AWS_AUDIENCE);
    config.put("subject_token_type", "urn:ietf:params:aws:token-type:aws4_request");
    config.put("token_url", "https://sts.googleapis.com/v1/token");
    config.put(
        "service_account_impersonation_url",
        String.format(
            "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken",
            clientEmail));

    GenericJson credentialSource = new GenericJson();
    credentialSource.put("environment_id", "aws1");
    credentialSource.put(
        "regional_cred_verification_url",
        "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15");
    config.put("credential_source", credentialSource);

    return config;
  }

  private String computeCertThumbprint(KeyStore keyStore) throws IOException {
    try {
      Enumeration<String> aliases = keyStore.aliases();
      while (aliases.hasMoreElements()) {
        String alias = aliases.nextElement();
        if (keyStore.isKeyEntry(alias)) {
          Certificate[] chain = keyStore.getCertificateChain(alias);
          Certificate cert =
              (chain != null && chain.length > 0) ? chain[0] : keyStore.getCertificate(alias);
          if (cert instanceof X509Certificate) {
            byte[] digest = MessageDigest.getInstance("SHA-256").digest(cert.getEncoded());
            return Base64.getUrlEncoder().withoutPadding().encodeToString(digest);
          }
        }
      }
    } catch (GeneralSecurityException e) {
      throw new IOException("Failed to compute X.509 certificate SHA-256 thumbprint.", e);
    }
    throw new IOException("No X509Certificate found in KeyStore.");
  }

  private String generateSignedOidcJwt(String audience, String x5tS256, boolean includeMayAct)
      throws IOException {
    ensureOidcJwtProviderJwksSynced();

    JsonWebSignature.Header header = new JsonWebSignature.Header();
    header.setAlgorithm("RS256");
    header.setType("JWT");
    header.setKeyId(privateKeyId);

    long nowSeconds = Instant.now().getEpochSecond();
    JsonWebToken.Payload payload = new JsonWebToken.Payload();
    payload.setIssuer(OIDC_JWT_ISSUER);
    payload.setSubject(clientId);
    payload.setAudience(audience);
    payload.setIssuedAtTimeSeconds(nowSeconds);
    payload.setExpirationTimeSeconds(nowSeconds + 3600);

    Map<String, Object> cnf = new HashMap<>();
    cnf.put("x5t#S256", x5tS256);
    payload.put("cnf", cnf);

    if (includeMayAct) {
      Map<String, Object> mayAct = new HashMap<>();
      mayAct.put("sub", clientId);
      mayAct.put("iss", OIDC_JWT_ISSUER);
      payload.put("may_act", mayAct);
    }

    try {
      return JsonWebSignature.signUsingRsaSha256(
          privateKey, OAuth2Utils.JSON_FACTORY, header, payload);
    } catch (GeneralSecurityException e) {
      throw new IOException("Error signing OIDC JWT with service account private key.", e);
    }
  }

  /**
   * Ensures the {@code oidc-jwt-1} Workload Identity Pool Provider's static {@code jwksJson}
   * includes the current service account {@code private_key_id}. When the integration test service
   * account rotates keys bi-monthly, the first test run automatically fetches the live public JWK
   * set for the service account and patches {@code oidc-jwt-1} using the service account's own
   * {@code roles/iam.workloadIdentityPoolAdmin} permission.
   */
  private void ensureOidcJwtProviderJwksSynced() throws IOException {
    if (privateKeyId.equals(syncedOidcJwtKeyId)) {
      return;
    }
    synchronized (ITWorkloadIdentityFederationTest.class) {
      if (privateKeyId.equals(syncedOidcJwtKeyId)) {
        return;
      }
      GoogleCredentials adminCredentials =
          GoogleCredentials.getApplicationDefault().createScoped(OAuth2Utils.CLOUD_PLATFORM_SCOPE);
      HttpTransport transport = new NetHttpTransport();
      HttpRequestFactory requestFactory =
          transport.createRequestFactory(new HttpCredentialsAdapter(adminCredentials));
      JsonObjectParser parser = new JsonObjectParser(GsonFactory.getDefaultInstance());

      String providerResource = OIDC_JWT_AUDIENCE.substring("//iam.googleapis.com/".length());
      String providerUrl = "https://iam.googleapis.com/v1/" + providerResource;

      HttpRequest getRequest = requestFactory.buildGetRequest(new GenericUrl(providerUrl));
      getRequest.setParser(parser);
      GenericData providerData;
      HttpResponse getResponse = getRequest.execute();
      try {
        providerData = getResponse.parseAs(GenericData.class);
      } finally {
        getResponse.disconnect();
      }

      Map<?, ?> oidc = (Map<?, ?>) providerData.get("oidc");
      String currentJwksJson = oidc != null ? (String) oidc.get("jwksJson") : null;
      if (currentJwksJson != null && currentJwksJson.contains(privateKeyId)) {
        syncedOidcJwtKeyId = privateKeyId;
        return;
      }

      HttpRequest jwksRequest =
          transport
              .createRequestFactory()
              .buildGetRequest(
                  new GenericUrl(
                      "https://www.googleapis.com/service_accounts/v1/jwk/" + clientEmail));
      String liveJwksJson;
      HttpResponse jwksResponse = jwksRequest.execute();
      try {
        liveJwksJson = jwksResponse.parseAsString();
      } finally {
        jwksResponse.disconnect();
      }

      GenericJson oidcPatch = new GenericJson();
      oidcPatch.setFactory(OAuth2Utils.JSON_FACTORY);
      oidcPatch.put("jwksJson", liveJwksJson);

      GenericJson patchBody = new GenericJson();
      patchBody.setFactory(OAuth2Utils.JSON_FACTORY);
      patchBody.put("oidc", oidcPatch);

      HttpContent content =
          new ByteArrayContent(
              "application/json", patchBody.toPrettyString().getBytes(StandardCharsets.UTF_8));
      HttpRequest patchRequest =
          requestFactory.buildPostRequest(
              new GenericUrl(providerUrl + "?updateMask=oidc.jwksJson"), content);
      patchRequest.getHeaders().set("X-HTTP-Method-Override", "PATCH");
      HttpResponse patchResponse = patchRequest.execute();
      try {
        assertTrue(patchResponse.isSuccessStatusCode());
      } finally {
        patchResponse.disconnect();
      }

      syncedOidcJwtKeyId = privateKeyId;
    }
  }

  private void callGcs(GoogleCredentials credentials) throws IOException {
    callGcs(credentials, new NetHttpTransport(), "https://storage.googleapis.com/storage/v1/b/");
  }

  private void callGcs(
      GoogleCredentials credentials, HttpTransport transport, String storageBaseUrl)
      throws IOException {
    String bucketName = System.getenv("GCS_BUCKET");
    if (bucketName == null) {
      fail("GCS bucket name not set through GCS_BUCKET env variable.");
    }

    String url = storageBaseUrl + bucketName;

    HttpCredentialsAdapter credentialsAdapter = new HttpCredentialsAdapter(credentials);
    HttpRequestFactory requestFactory = transport.createRequestFactory(credentialsAdapter);
    HttpRequest request = requestFactory.buildGetRequest(new GenericUrl(url));

    JsonObjectParser parser = new JsonObjectParser(GsonFactory.getDefaultInstance());
    request.setParser(parser);

    HttpResponse response = request.execute();
    try {
      assertTrue(response.isSuccessStatusCode());
    } finally {
      response.disconnect();
    }
  }

  /**
   * Generates a Google ID token using the iamcredentials generateIdToken API. <a
   * href="https://cloud.google.com/iam/docs/creating-short-lived-service-account-credentials#sa-credentials-oidc">...</a>
   */
  private String generateGoogleIdToken(String audience) throws IOException {
    GoogleCredentials googleCredentials =
        GoogleCredentials.getApplicationDefault().createScoped(OAuth2Utils.CLOUD_PLATFORM_SCOPE);

    HttpCredentialsAdapter credentialsAdapter = new HttpCredentialsAdapter(googleCredentials);
    HttpRequestFactory requestFactory =
        new NetHttpTransport().createRequestFactory(credentialsAdapter);

    String url =
        String.format(
            "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateIdToken",
            clientEmail);

    GenericData data = new GenericData();
    data.set("audience", audience);
    data.set("includeEmail", true);
    UrlEncodedContent content = new UrlEncodedContent(data);

    HttpRequest request = requestFactory.buildPostRequest(new GenericUrl(url), content);
    JsonObjectParser parser = new JsonObjectParser(GsonFactory.getDefaultInstance());
    request.setParser(parser);

    HttpResponse response = request.execute();
    GenericData responseData = response.parseAs(GenericData.class);
    return (String) responseData.get("token");
  }

  private GenericJson getServiceAccountKeyFileAsJson() throws IOException {
    String credentialsPath = System.getenv(DefaultCredentialsProvider.CREDENTIAL_ENV_VAR);
    JsonFactory jsonFactory = OAuth2Utils.JSON_FACTORY;
    JsonObjectParser parser = new JsonObjectParser(jsonFactory);
    return parser.parseAndClose(
        new FileInputStream(credentialsPath), StandardCharsets.UTF_8, GenericJson.class);
  }

  private String getXmlValueByTagName(String rawXml, String tagName) {
    int startIndex = rawXml.indexOf("<" + tagName + ">");
    int endIndex = rawXml.indexOf("</" + tagName + ">", startIndex);

    if (startIndex >= 0 && endIndex > startIndex) {
      return rawXml.substring(startIndex + tagName.length() + 2, endIndex);
    }
    return null;
  }

  private static class ITAwsSecurityCredentialsProvider implements AwsSecurityCredentialsSupplier {

    private final String region;
    private final AwsSecurityCredentials credentials;

    ITAwsSecurityCredentialsProvider(String region, AwsSecurityCredentials credentials) {
      this.region = region;
      this.credentials = credentials;
    }

    @Override
    public String getRegion(ExternalAccountSupplierContext context) {
      return this.region;
    }

    @Override
    public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext context) {
      return this.credentials;
    }
  }
}