feat: Add AWS Role Chaining support for Workload Identity Federation

shadow-google-cloud-bigquery-helper/build.gradle

@@ -52,6 +52,14 @@ dependencies {
     // Newer google-auth-library for AwsCredentials, ImpersonatedCredentials support
     compile "com.google.auth:google-auth-library-oauth2-http:1.41.0"
     compile "com.google.http-client:google-http-client-gson:1.45.3"
+
+    // AWS SDK for STS (AssumeRole for Role Chaining)
+    compile platform("software.amazon.awssdk:bom:2.29.51")
+    compile ("software.amazon.awssdk:sts") {
+        exclude group: "software.amazon.awssdk", module: "netty-nio-client"
+        exclude group: "software.amazon.awssdk", module: "apache-client"
+    }
+    compile "software.amazon.awssdk:url-connection-client"
 }
 
 // Relocate Guava and Jackson packages since they are incompatible from Embulk's.

shadow-google-cloud-bigquery-helper/gradle/dependency-locks/runtimeClasspath.lockfile

@@ -74,7 +74,35 @@ org.codehaus.mojo:animal-sniffer-annotations:1.21
 org.conscrypt:conscrypt-openjdk-uber:2.5.1
 org.json:json:20200518
 org.jspecify:jspecify:1.0.0
+org.reactivestreams:reactive-streams:1.0.4
 org.slf4j:jcl-over-slf4j:1.7.12
-org.slf4j:slf4j-api:1.7.25
+org.slf4j:slf4j-api:1.7.36
 org.threeten:threeten-extra:1.7.0
 org.threeten:threetenbp:1.6.0
+software.amazon.awssdk:annotations:2.29.51
+software.amazon.awssdk:auth:2.29.51
+software.amazon.awssdk:aws-core:2.29.51
+software.amazon.awssdk:aws-query-protocol:2.29.51
+software.amazon.awssdk:bom:2.29.51
+software.amazon.awssdk:checksums-spi:2.29.51
+software.amazon.awssdk:checksums:2.29.51
+software.amazon.awssdk:endpoints-spi:2.29.51
+software.amazon.awssdk:http-auth-aws-eventstream:2.29.51
+software.amazon.awssdk:http-auth-aws:2.29.51
+software.amazon.awssdk:http-auth-spi:2.29.51
+software.amazon.awssdk:http-auth:2.29.51
+software.amazon.awssdk:http-client-spi:2.29.51
+software.amazon.awssdk:identity-spi:2.29.51
+software.amazon.awssdk:json-utils:2.29.51
+software.amazon.awssdk:metrics-spi:2.29.51
+software.amazon.awssdk:profiles:2.29.51
+software.amazon.awssdk:protocol-core:2.29.51
+software.amazon.awssdk:regions:2.29.51
+software.amazon.awssdk:retries-spi:2.29.51
+software.amazon.awssdk:retries:2.29.51
+software.amazon.awssdk:sdk-core:2.29.51
+software.amazon.awssdk:sts:2.29.51
+software.amazon.awssdk:third-party-jackson-core:2.29.51
+software.amazon.awssdk:url-connection-client:2.29.51
+software.amazon.awssdk:utils:2.29.51
+software.amazon.eventstream:eventstream:1.0.1

src/main/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplier.java

@@ -0,0 +1,135 @@
+package org.embulk.output.bigquery_java;
+
+import com.google.auth.oauth2.AwsSecurityCredentials;
+import com.google.auth.oauth2.AwsSecurityCredentialsSupplier;
+import com.google.auth.oauth2.ExternalAccountSupplierContext;
+import java.io.IOException;
+import java.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
+import software.amazon.awssdk.services.sts.model.Credentials;
+
+/**
+ * Supplies AWS security credentials by assuming an IAM role. This class handles the AssumeRole
+ * operation and automatic credential refresh when credentials are about to expire.
+ *
+ * <p>This is designed for AWS Role Chaining scenarios where the base credentials (from IRSA, ECS
+ * Task Role, or environment variables) need to assume a middle role for Workload Identity
+ * Federation.
+ */
+public class AwsRoleCredentialsSupplier implements AwsSecurityCredentialsSupplier {
+  private static final Logger logger = LoggerFactory.getLogger(AwsRoleCredentialsSupplier.class);
+
+  /** Default session duration for AssumeRole (1 hour - maximum for role chaining) */
+  private static final int SESSION_DURATION_SECONDS = 3600;
+
+  /** Refresh credentials 5 minutes before expiration */
+  private static final int REFRESH_THRESHOLD_SECONDS = 300;
+
+  private final String roleArn;
+  private final String sessionName;
+  private final String region;
+  private final StsClient stsClient;
+
+  private Credentials currentCredentials;
+  private Instant expirationTime;
+
+  /**
+   * Creates a new AwsRoleCredentialsSupplier.
+   *
+   * @param roleArn The ARN of the IAM role to assume
+   * @param sessionName The session name for the assumed role session
+   * @param region The AWS region for STS calls
+   */
+  public AwsRoleCredentialsSupplier(String roleArn, String sessionName, String region) {
+    this.roleArn = roleArn;
+    this.sessionName = sessionName;
+    this.region = region;
+    this.stsClient =
+        StsClient.builder()
+            .region(Region.of(region))
+            .credentialsProvider(DefaultCredentialsProvider.create())
+            .build();
+    logger.debug(
+        "AwsRoleCredentialsSupplier created for role: {}, region: {}, session: {}",
+        roleArn,
+        region,
+        sessionName);
+  }
+
+  /**
+   * Constructor for testing - allows injecting a mock StsClient.
+   *
+   * @param roleArn The ARN of the IAM role to assume
+   * @param sessionName The session name for the assumed role session
+   * @param region The AWS region for STS calls
+   * @param stsClient The STS client to use (for testing)
+   */
+  AwsRoleCredentialsSupplier(
+      String roleArn, String sessionName, String region, StsClient stsClient) {
+    this.roleArn = roleArn;
+    this.sessionName = sessionName;
+    this.region = region;
+    this.stsClient = stsClient;
+  }
+
+  @Override
+  public synchronized AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext context)
+      throws IOException {
+    refreshIfNeeded();
+    return new AwsSecurityCredentials(
+        currentCredentials.accessKeyId(),
+        currentCredentials.secretAccessKey(),
+        currentCredentials.sessionToken());
+  }
+
+  @Override
+  public String getRegion(ExternalAccountSupplierContext context) throws IOException {
+    return region;
+  }
+
+  private boolean shouldRefresh() {
+    if (currentCredentials == null || expirationTime == null) {
+      return true;
+    }
+    // Refresh if we're within the threshold of expiration
+    Instant refreshThreshold = Instant.now().plusSeconds(REFRESH_THRESHOLD_SECONDS);
+    return refreshThreshold.isAfter(expirationTime);
+  }
+
+  private void refreshIfNeeded() throws IOException {
+    if (!shouldRefresh()) {
+      logger.debug("Using cached credentials, expires at: {}", expirationTime);
+      return;
+    }
+
+    logger.info("Refreshing AWS credentials by assuming role: {}", roleArn);
+    try {
+      AssumeRoleResponse response =
+          stsClient.assumeRole(
+              AssumeRoleRequest.builder()
+                  .roleArn(roleArn)
+                  .roleSessionName(sessionName)
+                  .durationSeconds(SESSION_DURATION_SECONDS)
+                  .build());
+
+      currentCredentials = response.credentials();
+      expirationTime = currentCredentials.expiration();
+      logger.info("AWS credentials refreshed, new expiration: {}", expirationTime);
+    } catch (Exception e) {
+      throw new IOException("Failed to assume role: " + roleArn, e);
+    }
+  }
+
+  /** Closes the underlying STS client. */
+  public void close() {
+    if (stsClient != null) {
+      stsClient.close();
+    }
+  }
+}

src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuth.java

@@ -19,27 +19,21 @@ public class WorkloadIdentityFederationAuth {
       LoggerFactory.getLogger(WorkloadIdentityFederationAuth.class);
   private static final int TOKEN_LIFETIME_SECONDS = 3600;
 
-  private final String awsAccessKeyId;
-  private final String awsSecretAccessKey;
-  private final String awsSessionToken;
+  private final AwsSecurityCredentialsSupplier awsCredentialsSupplier;
   private final String awsRegion;
   private final String audience;
   private final String serviceAccountImpersonationUrl;
   private final String tokenUrl;
   private final Set<String> scopes;
 
   public WorkloadIdentityFederationAuth(
-      String awsAccessKeyId,
-      String awsSecretAccessKey,
-      String awsSessionToken,
+      AwsSecurityCredentialsSupplier awsCredentialsSupplier,
       String awsRegion,
       String audience,
       String serviceAccountImpersonationUrl,
       String tokenUrl,
       Set<String> scopes) {
-    this.awsAccessKeyId = awsAccessKeyId;
-    this.awsSecretAccessKey = awsSecretAccessKey;
-    this.awsSessionToken = awsSessionToken;
+    this.awsCredentialsSupplier = awsCredentialsSupplier;
     this.awsRegion = awsRegion;
     this.audience = audience;
     this.serviceAccountImpersonationUrl = serviceAccountImpersonationUrl;
@@ -49,21 +43,13 @@ public WorkloadIdentityFederationAuth(
 
   public AccessToken fetchAccessToken() throws IOException {
     AccessToken federatedToken = fetchFederatedToken();
-
-    // Direct resource access (no service account impersonation)
-    if (serviceAccountImpersonationUrl == null || serviceAccountImpersonationUrl.isEmpty()) {
+    if (serviceAccountImpersonationUrl != null && !serviceAccountImpersonationUrl.isEmpty()) {
+      // Service account impersonation mode
+      return impersonateServiceAccount(federatedToken);
+    } else {
+      // Direct access mode (no impersonation)
       return federatedToken;
     }
-
-    return impersonateServiceAccount(federatedToken);
-  }
-
-  private static String formatAccessToken(AccessToken token) {
-    String expireTime =
-        token.getExpirationTime() != null
-            ? token.getExpirationTime().toInstant().toString()
-            : "null";
-    return "expireTime: " + expireTime;
   }
 
   private String getServiceAccountEmail() {
@@ -79,16 +65,17 @@ private String getServiceAccountEmail() {
   }
 
   // https://docs.cloud.google.com/iam/docs/reference/sts/rest/v1/TopLevel/token
-  private AccessToken fetchFederatedToken() throws IOException {
+  // Protected for testing - allows subclasses to provide mock tokens
+  protected AccessToken fetchFederatedToken() throws IOException {
     logger.debug("fetching federated token using AwsCredentials");
 
-    // Create AWS security credentials supplier
-    AwsSecurityCredentialsSupplier supplier =
+    // Create AWS security credentials supplier that uses our supplier and provides region
+    final AwsSecurityCredentialsSupplier supplierWithRegion =
         new AwsSecurityCredentialsSupplier() {
           @Override
           public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext context)
               throws IOException {
-            return new AwsSecurityCredentials(awsAccessKeyId, awsSecretAccessKey, awsSessionToken);
+            return awsCredentialsSupplier.getCredentials(context);
           }
 
           @Override
@@ -100,7 +87,7 @@ public String getRegion(ExternalAccountSupplierContext context) throws IOExcepti
     // Build AwsCredentials using the supplier
     AwsCredentials awsCredentials =
         AwsCredentials.newBuilder()
-            .setAwsSecurityCredentialsSupplier(supplier)
+            .setAwsSecurityCredentialsSupplier(supplierWithRegion)
             .setAudience(audience)
             .setTokenUrl(tokenUrl)
             .setSubjectTokenType("urn:ietf:params:aws:token-type:aws4_request")
@@ -110,13 +97,14 @@ public String getRegion(ExternalAccountSupplierContext context) throws IOExcepti
     awsCredentials.refresh();
     AccessToken accessToken = awsCredentials.getAccessToken();
 
-    logger.debug("federated token obtained, {}", formatAccessToken(accessToken));
+    logger.debug("federated token obtained");
 
     return accessToken;
   }
 
   // https://cloud.google.com/iam/docs/reference/credentials/rest/v1/projects.serviceAccounts/generateAccessToken
-  private AccessToken impersonateServiceAccount(AccessToken federatedToken) throws IOException {
+  // Protected for testing - allows subclasses to provide mock impersonation
+  protected AccessToken impersonateServiceAccount(AccessToken federatedToken) throws IOException {
     String serviceAccountEmail = getServiceAccountEmail();
     logger.debug("impersonating service account: {}", serviceAccountEmail);
 
@@ -131,7 +119,7 @@ private AccessToken impersonateServiceAccount(AccessToken federatedToken) throws
 
     impersonatedCredentials.refresh();
     AccessToken accessToken = impersonatedCredentials.getAccessToken();
-    logger.debug("service account impersonation succeeded, {}", formatAccessToken(accessToken));
+    logger.debug("service account impersonation succeeded");
     return accessToken;
   }
 }