diff --git a/shadow-google-cloud-bigquery-helper/build.gradle b/shadow-google-cloud-bigquery-helper/build.gradle index 9b50911..80fc817 100644 --- a/shadow-google-cloud-bigquery-helper/build.gradle +++ b/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. diff --git a/shadow-google-cloud-bigquery-helper/gradle/dependency-locks/runtimeClasspath.lockfile b/shadow-google-cloud-bigquery-helper/gradle/dependency-locks/runtimeClasspath.lockfile index 93d94f9..f12541c 100644 --- a/shadow-google-cloud-bigquery-helper/gradle/dependency-locks/runtimeClasspath.lockfile +++ b/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 diff --git a/src/main/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplier.java b/src/main/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplier.java new file mode 100644 index 0000000..1aa25df --- /dev/null +++ b/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. + * + *

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(); + } + } +} diff --git a/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuth.java b/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuth.java index 6b73590..8d88d02 100644 --- a/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuth.java +++ b/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuth.java @@ -19,9 +19,7 @@ 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; @@ -29,17 +27,13 @@ public class WorkloadIdentityFederationAuth { private final Set scopes; public WorkloadIdentityFederationAuth( - String awsAccessKeyId, - String awsSecretAccessKey, - String awsSessionToken, + AwsSecurityCredentialsSupplier awsCredentialsSupplier, String awsRegion, String audience, String serviceAccountImpersonationUrl, String tokenUrl, Set 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; } } diff --git a/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationCredentials.java b/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationCredentials.java index 16a5a64..dd7af9d 100644 --- a/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationCredentials.java +++ b/src/main/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationCredentials.java @@ -24,13 +24,13 @@ public class WorkloadIdentityFederationCredentials extends GoogleCredentials { private final WorkloadIdentityFederationAuth auth; private static class CacheKey { - private final String awsAccessKeyId; + private final String awsRoleArn; private final String awsRegion; private final String audience; private final Set scopes; - CacheKey(String awsAccessKeyId, String awsRegion, String audience, Set scopes) { - this.awsAccessKeyId = awsAccessKeyId; + CacheKey(String awsRoleArn, String awsRegion, String audience, Set scopes) { + this.awsRoleArn = awsRoleArn; this.awsRegion = awsRegion; this.audience = audience; this.scopes = scopes; @@ -41,7 +41,7 @@ public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; CacheKey cacheKey = (CacheKey) o; - return Objects.equals(awsAccessKeyId, cacheKey.awsAccessKeyId) + return Objects.equals(awsRoleArn, cacheKey.awsRoleArn) && Objects.equals(awsRegion, cacheKey.awsRegion) && Objects.equals(audience, cacheKey.audience) && Objects.equals(scopes, cacheKey.scopes); @@ -49,39 +49,39 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(awsAccessKeyId, awsRegion, audience, scopes); + return Objects.hash(awsRoleArn, awsRegion, audience, scopes); } @Override public String toString() { return String.format( - "%s:%s:%s:%s", awsAccessKeyId, awsRegion, audience, String.join(",", scopes)); + "%s:%s:%s:%s", awsRoleArn, awsRegion, audience, String.join(",", scopes)); } } public static WorkloadIdentityFederationCredentials getOrCreateByFetchingToken( WorkloadIdentityFederationConfig wifConfig, Set scopes) throws IOException { JsonObject jsonConfig = parseConfig(wifConfig); + String audience = jsonConfig.get("audience").getAsString(); + CacheKey cacheKey = - new CacheKey( - wifConfig.getAwsAccessKeyId(), - wifConfig.getAwsRegion(), - jsonConfig.get("audience").getAsString(), - scopes); + new CacheKey(wifConfig.getAwsRoleArn(), wifConfig.getAwsRegion(), audience, scopes); + WorkloadIdentityFederationCredentials cached = cache.get(cacheKey); if (cached != null) { logger.debug("cache hit for cacheKey: {}", cacheKey); return cached; } logger.debug("cache miss for cacheKey: {}", cacheKey); + WorkloadIdentityFederationCredentials credentials = createByFetchingToken( - wifConfig.getAwsAccessKeyId(), - wifConfig.getAwsSecretAccessKey(), - wifConfig.getAwsSessionToken().orElse(null), + wifConfig.getAwsRoleArn(), + wifConfig.getAwsRoleSessionName(), wifConfig.getAwsRegion(), jsonConfig, scopes); + cache.put(cacheKey, credentials); return credentials; } @@ -93,25 +93,27 @@ private static JsonObject parseConfig(WorkloadIdentityFederationConfig wifConfig } private static WorkloadIdentityFederationCredentials createByFetchingToken( - String awsAccessKeyId, - String awsSecretAccessKey, - String awsSessionToken, + String awsRoleArn, + String awsRoleSessionName, String awsRegion, JsonObject jsonConfig, Set scopes) throws IOException { - logger.info("creating credentials by fetching token"); + logger.info("creating credentials using AssumeRole with role: {}", awsRoleArn); String tokenUrl = jsonConfig.has("token_url") ? jsonConfig.get("token_url").getAsString() : null; String serviceAccountImpersonationUrl = jsonConfig.has("service_account_impersonation_url") ? jsonConfig.get("service_account_impersonation_url").getAsString() : null; + + // Create the AWS role credentials supplier that handles AssumeRole and auto-refresh + AwsRoleCredentialsSupplier awsRoleCredentialsSupplier = + new AwsRoleCredentialsSupplier(awsRoleArn, awsRoleSessionName, awsRegion); + WorkloadIdentityFederationAuth auth = new WorkloadIdentityFederationAuth( - awsAccessKeyId, - awsSecretAccessKey, - awsSessionToken, + awsRoleCredentialsSupplier, awsRegion, jsonConfig.get("audience").getAsString(), serviceAccountImpersonationUrl, @@ -132,4 +134,9 @@ public AccessToken refreshAccessToken() throws IOException { logger.info("refreshing access token"); return auth.fetchAccessToken(); } + + /** Clear the cache. Useful for testing. */ + public static void clearCache() { + cache.clear(); + } } diff --git a/src/main/java/org/embulk/output/bigquery_java/config/WorkloadIdentityFederationConfig.java b/src/main/java/org/embulk/output/bigquery_java/config/WorkloadIdentityFederationConfig.java index 2c68521..f19e79f 100644 --- a/src/main/java/org/embulk/output/bigquery_java/config/WorkloadIdentityFederationConfig.java +++ b/src/main/java/org/embulk/output/bigquery_java/config/WorkloadIdentityFederationConfig.java @@ -1,6 +1,5 @@ package org.embulk.output.bigquery_java.config; -import java.util.Optional; import org.embulk.util.config.Config; import org.embulk.util.config.ConfigDefault; import org.embulk.util.config.Task; @@ -10,15 +9,12 @@ public interface WorkloadIdentityFederationConfig extends Task { @Config("config") LocalFile getConfig(); - @Config("aws_access_key_id") - String getAwsAccessKeyId(); + @Config("aws_role_arn") + String getAwsRoleArn(); - @Config("aws_secret_access_key") - String getAwsSecretAccessKey(); - - @Config("aws_session_token") - @ConfigDefault("null") - Optional getAwsSessionToken(); + @Config("aws_role_session_name") + @ConfigDefault("\"embulk-bigquery-output\"") + String getAwsRoleSessionName(); @Config("aws_region") @ConfigDefault("\"ap-northeast-1\"") diff --git a/src/test/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplierTest.java b/src/test/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplierTest.java new file mode 100644 index 0000000..45ece44 --- /dev/null +++ b/src/test/java/org/embulk/output/bigquery_java/AwsRoleCredentialsSupplierTest.java @@ -0,0 +1,70 @@ +package org.embulk.output.bigquery_java; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.mockito.Matchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.auth.oauth2.AwsSecurityCredentials; +import java.io.IOException; +import java.time.Instant; +import org.junit.Test; +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; + +public class AwsRoleCredentialsSupplierTest { + + private static final String TEST_ROLE_ARN = "arn:aws:iam::123456789012:role/test-role"; + private static final String TEST_SESSION_NAME = "test-session"; + private static final String TEST_REGION = "us-east-1"; + + @Test + public void testGetCredentials_refreshesWhenNeeded() throws IOException { + // Setup mock STS client + StsClient mockStsClient = mock(StsClient.class); + Credentials mockCredentials = + Credentials.builder() + .accessKeyId("test-access-key") + .secretAccessKey("test-secret-key") + .sessionToken("test-session-token") + .expiration(Instant.now().plusSeconds(3600)) + .build(); + AssumeRoleResponse mockResponse = + AssumeRoleResponse.builder().credentials(mockCredentials).build(); + when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))).thenReturn(mockResponse); + + // Create supplier with mock client + AwsRoleCredentialsSupplier supplier = + new AwsRoleCredentialsSupplier( + TEST_ROLE_ARN, TEST_SESSION_NAME, TEST_REGION, mockStsClient); + + // First call should trigger refresh + AwsSecurityCredentials credentials = supplier.getCredentials(null); + assertNotNull(credentials); + assertEquals("test-access-key", credentials.getAccessKeyId()); + assertEquals("test-secret-key", credentials.getSecretAccessKey()); + + // Second call should use cached credentials (no refresh) + AwsSecurityCredentials credentials2 = supplier.getCredentials(null); + assertNotNull(credentials2); + + // Verify assumeRole was called only once (cached for second call) + verify(mockStsClient, times(1)).assumeRole(any(AssumeRoleRequest.class)); + } + + @Test + public void testGetRegion_returnsConfiguredRegion() throws IOException { + StsClient mockStsClient = mock(StsClient.class); + AwsRoleCredentialsSupplier supplier = + new AwsRoleCredentialsSupplier( + TEST_ROLE_ARN, TEST_SESSION_NAME, TEST_REGION, mockStsClient); + + String region = supplier.getRegion(null); + assertEquals(TEST_REGION, region); + } +} diff --git a/src/test/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuthTest.java b/src/test/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuthTest.java new file mode 100644 index 0000000..5b8387a --- /dev/null +++ b/src/test/java/org/embulk/output/bigquery_java/WorkloadIdentityFederationAuthTest.java @@ -0,0 +1,153 @@ +package org.embulk.output.bigquery_java; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.google.auth.oauth2.AccessToken; +import com.google.auth.oauth2.AwsSecurityCredentials; +import com.google.auth.oauth2.AwsSecurityCredentialsSupplier; +import com.google.auth.oauth2.ExternalAccountSupplierContext; +import java.io.IOException; +import java.util.Date; +import java.util.HashSet; +import java.util.Set; +import org.junit.Test; + +public class WorkloadIdentityFederationAuthTest { + + private static final String TEST_AWS_REGION = "us-east-1"; + private static final String TEST_AUDIENCE = + "//iam.googleapis.com/projects/123/locations/global/workloadIdentityPools/pool/providers/provider"; + private static final String TEST_SERVICE_ACCOUNT_IMPERSONATION_URL = + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test@project.iam.gserviceaccount.com:generateAccessToken"; + + private Set createTestScopes() { + Set scopes = new HashSet<>(); + scopes.add("https://www.googleapis.com/auth/cloud-platform"); + return scopes; + } + + private AwsSecurityCredentialsSupplier createMockSupplier() { + return new AwsSecurityCredentialsSupplier() { + @Override + public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext context) + throws IOException { + return new AwsSecurityCredentials( + "assumed-access-key", "assumed-secret-key", "assumed-session-token"); + } + + @Override + public String getRegion(ExternalAccountSupplierContext context) throws IOException { + return TEST_AWS_REGION; + } + }; + } + + @Test + public void testDirectAccessMode_returnsDirectFederatedToken() throws IOException { + // When serviceAccountImpersonationUrl is null, direct access mode should be used + AccessToken expectedToken = new AccessToken("federated-token", new Date()); + + WorkloadIdentityFederationAuth auth = + new TestableWorkloadIdentityFederationAuth( + createMockSupplier(), + TEST_AWS_REGION, + TEST_AUDIENCE, + null, // Direct access mode - no impersonation URL + null, + createTestScopes(), + expectedToken, + null); + + AccessToken result = auth.fetchAccessToken(); + + assertNotNull(result); + assertEquals("federated-token", result.getTokenValue()); + } + + @Test + public void testDirectAccessMode_withEmptyUrl_returnsDirectFederatedToken() throws IOException { + // When serviceAccountImpersonationUrl is empty string, direct access mode should be used + AccessToken expectedToken = new AccessToken("federated-token-empty", new Date()); + + WorkloadIdentityFederationAuth auth = + new TestableWorkloadIdentityFederationAuth( + createMockSupplier(), + TEST_AWS_REGION, + TEST_AUDIENCE, + "", // Empty string - should also use direct access mode + null, + createTestScopes(), + expectedToken, + null); + + AccessToken result = auth.fetchAccessToken(); + + assertNotNull(result); + assertEquals("federated-token-empty", result.getTokenValue()); + } + + @Test + public void testImpersonationMode_returnsImpersonatedToken() throws IOException { + // When serviceAccountImpersonationUrl is provided, impersonation mode should be used + AccessToken federatedToken = new AccessToken("federated-token", new Date()); + AccessToken impersonatedToken = new AccessToken("impersonated-token", new Date()); + + WorkloadIdentityFederationAuth auth = + new TestableWorkloadIdentityFederationAuth( + createMockSupplier(), + TEST_AWS_REGION, + TEST_AUDIENCE, + TEST_SERVICE_ACCOUNT_IMPERSONATION_URL, + null, + createTestScopes(), + federatedToken, + impersonatedToken); + + AccessToken result = auth.fetchAccessToken(); + + assertNotNull(result); + assertEquals("impersonated-token", result.getTokenValue()); + } + + /** + * Test subclass that allows overriding the token fetching methods to return mock tokens without + * making actual network calls. + */ + private static class TestableWorkloadIdentityFederationAuth + extends WorkloadIdentityFederationAuth { + + private final AccessToken mockFederatedToken; + private final AccessToken mockImpersonatedToken; + + public TestableWorkloadIdentityFederationAuth( + AwsSecurityCredentialsSupplier awsCredentialsSupplier, + String awsRegion, + String audience, + String serviceAccountImpersonationUrl, + String tokenUrl, + Set scopes, + AccessToken mockFederatedToken, + AccessToken mockImpersonatedToken) { + super( + awsCredentialsSupplier, + awsRegion, + audience, + serviceAccountImpersonationUrl, + tokenUrl, + scopes); + this.mockFederatedToken = mockFederatedToken; + this.mockImpersonatedToken = mockImpersonatedToken; + } + + @Override + protected AccessToken fetchFederatedToken() throws IOException { + return mockFederatedToken; + } + + @Override + protected AccessToken impersonateServiceAccount(AccessToken federatedToken) throws IOException { + return mockImpersonatedToken; + } + } +}