This is an automated email from the ASF dual-hosted git repository. github-actions[bot] pushed a commit to branch cherry-pick-8515fb54-to-branch-1.3 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit f06ca00b8a048c4e5808be8b9c132caf4b622485 Author: Qi Yu <[email protected]> AuthorDate: Wed Jun 24 01:06:51 2026 +0800 [#11734] feat(s3): introduce WebIdentityTokenSource SPI with file-based source (#11735) ### What changes were proposed in this pull request? Introduce a pluggable `WebIdentityTokenSource` SPI for obtaining the OIDC token used by AWS STS `AssumeRoleWithWebIdentity`, plus a built-in `file` source that reads the token from a configured file path (falling back to `AWS_WEB_IDENTITY_TOKEN_FILE`). `AwsIrsaCredentialGenerator` is refactored to retrieve the token through the SPI instead of reading the env-driven file directly. Existing IRSA behavior is unchanged. ### Why are the changes needed? Different deployments source the WebIdentity token differently (a file on disk for K8s/IRSA, OAuth2 client_credentials for VMs, JWT bearer for service accounts, ...). A pluggable SPI keeps the STS / scoped-policy logic generic and lets each deployment plug in the right token source. Fix: #11734 ### Does this PR introduce _any_ user-facing change? Yes, two new optional property keys: - `s3-web-identity-token-source` — selects the source (default `file`) - `s3-web-identity-token-file` — file path used by the `file` source (falls back to `AWS_WEB_IDENTITY_TOKEN_FILE` if absent) No existing keys or default behavior change; IRSA setups relying on `AWS_WEB_IDENTITY_TOKEN_FILE` keep working. ### How was this patch tested? - New unit tests under `bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/` cover configured file path resolution, re-read on each call, error paths, and source selection through the factory. - Refactored `TestAwsIrsaCredentialGenerator`. - `./gradlew :bundles:aws:test -PskipITs` passes. --- bundles/aws/build.gradle.kts | 2 + .../s3/credential/AwsIrsaCredentialGenerator.java | 175 ++++++++++--------- .../webidentity/FileWebIdentityTokenSource.java | 88 ++++++++++ .../webidentity/WebIdentityTokenSource.java | 52 ++++++ .../webidentity/WebIdentityTokenSourceConfig.java | 41 +++++ .../webidentity/WebIdentityTokenSources.java | 108 ++++++++++++ ...3.credential.webidentity.WebIdentityTokenSource | 19 +++ .../credential/TestAwsIrsaCredentialGenerator.java | 129 ++++++++++++++ .../TestFileWebIdentityTokenSource.java | 128 ++++++++++++++ .../webidentity/TestWebIdentityTokenSources.java | 187 +++++++++++++++++++++ 10 files changed, 854 insertions(+), 75 deletions(-) diff --git a/bundles/aws/build.gradle.kts b/bundles/aws/build.gradle.kts index 27ba900d5d..7b6d4f1ec1 100644 --- a/bundles/aws/build.gradle.kts +++ b/bundles/aws/build.gradle.kts @@ -46,6 +46,8 @@ dependencies { compileOnly(libs.hadoop3.client.api) testImplementation(libs.aws.iam) + testImplementation(libs.aws.policy) + testImplementation(libs.aws.sts) testImplementation(libs.junit.jupiter.api) testImplementation(libs.junit.jupiter.params) testRuntimeOnly(libs.junit.jupiter.engine) diff --git a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/AwsIrsaCredentialGenerator.java b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/AwsIrsaCredentialGenerator.java index 3c1641506e..c960da6900 100644 --- a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/AwsIrsaCredentialGenerator.java +++ b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/AwsIrsaCredentialGenerator.java @@ -20,15 +20,14 @@ package org.apache.gravitino.s3.credential; import java.io.IOException; import java.net.URI; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Paths; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Stream; import org.apache.commons.lang3.StringUtils; import org.apache.gravitino.credential.AwsIrsaCredential; @@ -36,9 +35,10 @@ import org.apache.gravitino.credential.CredentialContext; import org.apache.gravitino.credential.CredentialGenerator; import org.apache.gravitino.credential.PathBasedCredentialContext; import org.apache.gravitino.credential.config.S3CredentialConfig; +import org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource; +import org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSources; import software.amazon.awssdk.auth.credentials.AwsCredentials; import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; -import software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider; import software.amazon.awssdk.policybuilder.iam.IamConditionOperator; import software.amazon.awssdk.policybuilder.iam.IamEffect; import software.amazon.awssdk.policybuilder.iam.IamPolicy; @@ -47,6 +47,7 @@ import software.amazon.awssdk.policybuilder.iam.IamStatement; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleWithWebIdentityCredentialsProvider; import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest; import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityResponse; import software.amazon.awssdk.services.sts.model.Credentials; @@ -54,17 +55,18 @@ import software.amazon.awssdk.services.sts.model.Credentials; /** Generate AWS IRSA credentials according to the read and write paths. */ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCredential> { - private WebIdentityTokenFileCredentialsProvider baseCredentialsProvider; + private WebIdentityTokenSource tokenSource; private String roleArn; private int tokenExpireSecs; private String region; private String stsEndpoint; private boolean listLocationPrefix; + private StsClient basicModeStsClient; + private Map<String, StsAssumeRoleWithWebIdentityCredentialsProvider> basicModeProviders; @Override public void initialize(Map<String, String> properties) { - // Use WebIdentityTokenFileCredentialsProvider for base IRSA configuration - this.baseCredentialsProvider = WebIdentityTokenFileCredentialsProvider.create(); + this.tokenSource = WebIdentityTokenSources.create(properties); S3CredentialConfig s3CredentialConfig = new S3CredentialConfig(properties); this.roleArn = s3CredentialConfig.s3RoleArn(); @@ -72,32 +74,14 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr this.region = s3CredentialConfig.region(); this.stsEndpoint = s3CredentialConfig.stsEndpoint(); this.listLocationPrefix = s3CredentialConfig.listLocationPrefix(); + this.basicModeStsClient = createStsClient(); + this.basicModeProviders = new ConcurrentHashMap<>(); } @Override public AwsIrsaCredential generate(CredentialContext context) { if (!(context instanceof PathBasedCredentialContext)) { - // Fallback to original behavior for non-path-based contexts - AwsCredentials creds = baseCredentialsProvider.resolveCredentials(); - if (creds instanceof AwsSessionCredentials) { - AwsSessionCredentials sessionCreds = (AwsSessionCredentials) creds; - if (!sessionCreds.expirationTime().isPresent()) { - throw new IllegalStateException( - "AWS IRSA session credentials must include an expiration time for vended credential" - + " refresh"); - } - long expiration = sessionCreds.expirationTime().get().toEpochMilli(); - return new AwsIrsaCredential( - sessionCreds.accessKeyId(), - sessionCreds.secretAccessKey(), - sessionCreds.sessionToken(), - expiration); - } else { - throw new IllegalStateException( - "AWS IRSA credentials must be of type AwsSessionCredentials. " - + "Check your EKS/IRSA configuration. Got: " - + creds.getClass().getName()); - } + return toAwsIrsaCredential(createBasicModeCredentials(context.getUserName())); } PathBasedCredentialContext pathBasedCredentialContext = (PathBasedCredentialContext) context; @@ -107,11 +91,7 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr pathBasedCredentialContext.getReadPaths(), pathBasedCredentialContext.getWritePaths(), pathBasedCredentialContext.getUserName()); - return new AwsIrsaCredential( - s3Token.accessKeyId(), - s3Token.secretAccessKey(), - s3Token.sessionToken(), - s3Token.expiration().toEpochMilli()); + return toAwsIrsaCredential(s3Token); } private Credentials createCredentialsWithSessionPolicy( @@ -119,22 +99,61 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr validateInputParameters(readLocations, writeLocations, userName); IamPolicy sessionPolicy = createSessionPolicy(readLocations, writeLocations, region); - String webIdentityTokenFile = getValidatedWebIdentityTokenFile(); + return createCredentials(userName, Optional.of(sessionPolicy)); + } + + private Credentials createCredentials(String userName, Optional<IamPolicy> sessionPolicy) { String effectiveRoleArn = getValidatedRoleArn(roleArn); + String tokenContent = tokenSource.getToken(); + return assumeRoleWithWebIdentity(effectiveRoleArn, userName, tokenContent, sessionPolicy); + } - try { - String tokenContent = - new String(Files.readAllBytes(Paths.get(webIdentityTokenFile)), StandardCharsets.UTF_8); - if (StringUtils.isBlank(tokenContent)) { - throw new IllegalStateException( - "Web identity token file is empty: " + webIdentityTokenFile); - } + private AwsCredentials createBasicModeCredentials(String userName) { + return basicModeProviders + .computeIfAbsent(userName, this::createBasicModeCredentialsProvider) + .resolveCredentials(); + } + + private StsAssumeRoleWithWebIdentityCredentialsProvider createBasicModeCredentialsProvider( + String userName) { + return StsAssumeRoleWithWebIdentityCredentialsProvider.builder() + .stsClient(basicModeStsClient) + .refreshRequest( + () -> + createAssumeRoleWithWebIdentityRequest( + getValidatedRoleArn(roleArn), + userName, + tokenSource.getToken(), + Optional.empty())) + .build(); + } - return assumeRoleWithSessionPolicy(effectiveRoleArn, userName, tokenContent, sessionPolicy); - } catch (Exception e) { - throw new RuntimeException( - "Failed to create credentials with session policy for user: " + userName, e); + private AwsIrsaCredential toAwsIrsaCredential(Credentials credentials) { + return new AwsIrsaCredential( + credentials.accessKeyId(), + credentials.secretAccessKey(), + credentials.sessionToken(), + credentials.expiration().toEpochMilli()); + } + + private AwsIrsaCredential toAwsIrsaCredential(AwsCredentials credentials) { + if (!(credentials instanceof AwsSessionCredentials)) { + throw new IllegalStateException( + "AWS IRSA credentials must be of type AwsSessionCredentials. " + + "Check your EKS/IRSA configuration. Got: " + + credentials.getClass().getName()); } + AwsSessionCredentials sessionCredentials = (AwsSessionCredentials) credentials; + if (!sessionCredentials.expirationTime().isPresent()) { + throw new IllegalStateException( + "AWS IRSA session credentials must include an expiration time for vended credential" + + " refresh"); + } + return new AwsIrsaCredential( + sessionCredentials.accessKeyId(), + sessionCredentials.secretAccessKey(), + sessionCredentials.sessionToken(), + sessionCredentials.expirationTime().get().toEpochMilli()); } private IamPolicy createSessionPolicy( @@ -309,20 +328,6 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr } } - private String getValidatedWebIdentityTokenFile() { - String webIdentityTokenFile = System.getenv("AWS_WEB_IDENTITY_TOKEN_FILE"); - if (StringUtils.isBlank(webIdentityTokenFile)) { - throw new IllegalStateException( - "AWS_WEB_IDENTITY_TOKEN_FILE environment variable is not set. " - + "Ensure IRSA is properly configured in your EKS cluster."); - } - if (!Files.exists(Paths.get(webIdentityTokenFile))) { - throw new IllegalStateException( - "Web identity token file does not exist: " + webIdentityTokenFile); - } - return webIdentityTokenFile; - } - private String getValidatedRoleArn(String configRoleArn) { String effectiveRoleArn = StringUtils.isNotBlank(configRoleArn) ? configRoleArn : System.getenv("AWS_ROLE_ARN"); @@ -336,8 +341,18 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr return effectiveRoleArn; } - private Credentials assumeRoleWithSessionPolicy( - String roleArn, String userName, String webIdentityToken, IamPolicy sessionPolicy) { + private Credentials assumeRoleWithWebIdentity( + String roleArn, String userName, String webIdentityToken, Optional<IamPolicy> sessionPolicy) { + try (StsClient stsClient = createStsClient()) { + AssumeRoleWithWebIdentityResponse response = + stsClient.assumeRoleWithWebIdentity( + createAssumeRoleWithWebIdentityRequest( + roleArn, userName, webIdentityToken, sessionPolicy)); + return response.credentials(); + } + } + + private StsClient createStsClient() { StsClientBuilder stsBuilder = StsClient.builder(); if (StringUtils.isNotBlank(region)) { stsBuilder.region(Region.of(region)); @@ -345,22 +360,32 @@ public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCr if (StringUtils.isNotBlank(stsEndpoint)) { stsBuilder.endpointOverride(URI.create(stsEndpoint)); } + return stsBuilder.build(); + } - try (StsClient stsClient = stsBuilder.build()) { - AssumeRoleWithWebIdentityRequest request = - AssumeRoleWithWebIdentityRequest.builder() - .roleArn(roleArn) - .roleSessionName("gravitino_irsa_session_" + userName) - .durationSeconds(tokenExpireSecs) - .webIdentityToken(webIdentityToken) - .policy(sessionPolicy.toJson()) - .build(); - - AssumeRoleWithWebIdentityResponse response = stsClient.assumeRoleWithWebIdentity(request); - return response.credentials(); - } + private AssumeRoleWithWebIdentityRequest createAssumeRoleWithWebIdentityRequest( + String roleArn, String userName, String webIdentityToken, Optional<IamPolicy> sessionPolicy) { + AssumeRoleWithWebIdentityRequest.Builder requestBuilder = + AssumeRoleWithWebIdentityRequest.builder() + .roleArn(roleArn) + .roleSessionName("gravitino_irsa_session_" + userName) + .durationSeconds(tokenExpireSecs) + .webIdentityToken(webIdentityToken); + sessionPolicy.ifPresent(policy -> requestBuilder.policy(policy.toJson())); + return requestBuilder.build(); } @Override - public void close() throws IOException {} + public void close() throws IOException { + if (basicModeProviders != null) { + basicModeProviders.values().forEach(StsAssumeRoleWithWebIdentityCredentialsProvider::close); + basicModeProviders.clear(); + } + if (basicModeStsClient != null) { + basicModeStsClient.close(); + } + if (tokenSource != null) { + tokenSource.close(); + } + } } diff --git a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/FileWebIdentityTokenSource.java b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/FileWebIdentityTokenSource.java new file mode 100644 index 0000000000..5a49ea1d33 --- /dev/null +++ b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/FileWebIdentityTokenSource.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Map; +import org.apache.commons.lang3.StringUtils; + +/** + * A {@link WebIdentityTokenSource} that reads the token from a file on disk. Matches the AWS IRSA / + * Kubernetes projected service account token pattern. + * + * <p>The path is taken from {@code s3-web-identity-token-file}; if absent, the {@code + * AWS_WEB_IDENTITY_TOKEN_FILE} environment variable is used as a fallback so existing IRSA setups + * continue to work without configuration changes. + * + * <p>The token file is read on every call to {@link #getToken()} so that token rotations performed + * by the kubelet (or any external rotator) are picked up automatically. + */ +public class FileWebIdentityTokenSource implements WebIdentityTokenSource { + + public static final String NAME = "file"; + static final String AWS_WEB_IDENTITY_TOKEN_FILE_ENV = "AWS_WEB_IDENTITY_TOKEN_FILE"; + + private Path tokenFile; + + @Override + public String name() { + return NAME; + } + + @Override + public void initialize(Map<String, String> properties) { + String configured = properties.get(WebIdentityTokenSourceConfig.FILE_PATH); + String path = + StringUtils.isNotBlank(configured) + ? configured + : System.getenv(AWS_WEB_IDENTITY_TOKEN_FILE_ENV); + if (StringUtils.isBlank(path)) { + throw new IllegalStateException( + "No WebIdentity token file is configured. Set " + + WebIdentityTokenSourceConfig.FILE_PATH + + " or the " + + AWS_WEB_IDENTITY_TOKEN_FILE_ENV + + " environment variable."); + } + this.tokenFile = Paths.get(path); + } + + @Override + public String getToken() { + if (!Files.exists(tokenFile)) { + throw new IllegalStateException("WebIdentity token file does not exist: " + tokenFile); + } + try { + String token = new String(Files.readAllBytes(tokenFile), StandardCharsets.UTF_8).trim(); + if (StringUtils.isBlank(token)) { + throw new IllegalStateException("WebIdentity token file is empty: " + tokenFile); + } + return token; + } catch (IOException e) { + throw new IllegalStateException("Failed to read WebIdentity token file: " + tokenFile, e); + } + } + + @Override + public void close() {} +} diff --git a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSource.java b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSource.java new file mode 100644 index 0000000000..95e730a132 --- /dev/null +++ b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSource.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +import java.io.Closeable; +import java.util.Map; + +/** + * SPI for obtaining a WebIdentity token (an OIDC id_token or access_token) that can be exchanged + * with AWS STS {@code AssumeRoleWithWebIdentity} for temporary S3 credentials. + * + * <p>Implementations decide how the token is sourced (local file, OAuth flow, etc.) and are + * responsible for any caching and refresh behavior. + * + * <p>This SPI is scoped to AWS S3 credential vending (IRSA / {@code AssumeRoleWithWebIdentity}); it + * is not a generic cross-storage abstraction. Other storage backends (e.g. GCS, Azure) use + * different workload-identity mechanisms and are out of scope here. + */ +public interface WebIdentityTokenSource extends Closeable { + + /** + * The configuration name used to select this source via {@code s3-web-identity-token-source}. + * + * <p>Must be unique across all implementations on the classpath. + */ + String name(); + + /** Initialize the source with the credential provider's properties. */ + void initialize(Map<String, String> properties); + + /** + * Returns a current WebIdentity token. Implementations may cache and refresh internally; callers + * should invoke this for every credential request. + */ + String getToken(); +} diff --git a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSourceConfig.java b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSourceConfig.java new file mode 100644 index 0000000000..1174093f49 --- /dev/null +++ b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSourceConfig.java @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +/** + * Property keys shared by {@link WebIdentityTokenSource} implementations. + * + * <p>These keys are AWS S3-specific (prefixed {@code s3-}); the WebIdentity token source is scoped + * to AWS S3 credential vending and is not a generic cross-storage abstraction. + */ +public final class WebIdentityTokenSourceConfig { + + /** Selects which {@link WebIdentityTokenSource} implementation to use. */ + public static final String SOURCE = "s3-web-identity-token-source"; + + /** + * Path to a file containing the WebIdentity token. Used by {@link FileWebIdentityTokenSource}. + */ + public static final String FILE_PATH = "s3-web-identity-token-file"; + + /** Default source name used when {@link #SOURCE} is not set. */ + public static final String DEFAULT_SOURCE = FileWebIdentityTokenSource.NAME; + + private WebIdentityTokenSourceConfig() {} +} diff --git a/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSources.java b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSources.java new file mode 100644 index 0000000000..80f77a0490 --- /dev/null +++ b/bundles/aws/src/main/java/org/apache/gravitino/s3/credential/webidentity/WebIdentityTokenSources.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.ServiceLoader; +import org.apache.commons.lang3.StringUtils; + +/** Factory for resolving and initializing a {@link WebIdentityTokenSource} from properties. */ +public final class WebIdentityTokenSources { + + /** + * Resolve the {@link WebIdentityTokenSource} configured by {@code s3-web-identity-token-source} + * (defaults to {@code file}) and initialize it with the given properties. + * + * <p>Sources are discovered through {@link ServiceLoader}, so additional implementations can be + * registered by listing them under {@code + * META-INF/services/org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource}. + */ + public static WebIdentityTokenSource create(Map<String, String> properties) { + return create(properties, ServiceLoader.load(WebIdentityTokenSource.class)); + } + + /** + * Package-private overload that takes the candidate sources directly, enabling tests to exercise + * the resolution logic (in particular duplicate-name detection) without touching the {@link + * ServiceLoader} registry. + */ + static WebIdentityTokenSource create( + Map<String, String> properties, Iterable<WebIdentityTokenSource> candidates) { + String type = + properties.getOrDefault( + WebIdentityTokenSourceConfig.SOURCE, WebIdentityTokenSourceConfig.DEFAULT_SOURCE); + + List<String> available = new ArrayList<>(); + List<WebIdentityTokenSource> matched = new ArrayList<>(); + Map<String, WebIdentityTokenSource> sourcesByName = new HashMap<>(); + for (WebIdentityTokenSource candidate : candidates) { + String candidateName = candidate.name(); + if (StringUtils.isBlank(candidateName)) { + throw new IllegalStateException( + "WebIdentity token source " + + candidate.getClass().getName() + + " must return a non-blank name()."); + } + available.add(candidateName); + String normalizedName = candidateName.toLowerCase(Locale.ROOT); + WebIdentityTokenSource existing = sourcesByName.putIfAbsent(normalizedName, candidate); + if (existing != null) { + throw new IllegalStateException( + "Multiple WebIdentity token sources registered with name '" + + existing.name() + + "': " + + existing.getClass().getName() + + ", " + + candidate.getClass().getName() + + ". Each implementation must return a unique name()."); + } + if (type.equalsIgnoreCase(candidateName)) { + matched.add(candidate); + } + } + if (matched.isEmpty()) { + throw new IllegalArgumentException( + "Unknown WebIdentity token source: " + + type + + ". Available sources: " + + String.join(", ", available)); + } + if (matched.size() > 1) { + List<String> classes = new ArrayList<>(); + for (WebIdentityTokenSource source : matched) { + classes.add(source.getClass().getName()); + } + throw new IllegalStateException( + "Multiple WebIdentity token sources registered with name '" + + type + + "': " + + String.join(", ", classes) + + ". Each implementation must return a unique name()."); + } + WebIdentityTokenSource source = matched.get(0); + source.initialize(properties); + return source; + } + + private WebIdentityTokenSources() {} +} diff --git a/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource b/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource new file mode 100644 index 0000000000..7d96b37fe9 --- /dev/null +++ b/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource @@ -0,0 +1,19 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +org.apache.gravitino.s3.credential.webidentity.FileWebIdentityTokenSource diff --git a/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/TestAwsIrsaCredentialGenerator.java b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/TestAwsIrsaCredentialGenerator.java new file mode 100644 index 0000000000..dbe75f2d42 --- /dev/null +++ b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/TestAwsIrsaCredentialGenerator.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.URLDecoder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Instant; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.gravitino.credential.AwsIrsaCredential; +import org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSourceConfig; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class TestAwsIrsaCredentialGenerator { + + private HttpServer server; + private AtomicReference<String> requestBody; + private AtomicInteger requestCount; + + @BeforeEach + void setUp() throws IOException { + requestBody = new AtomicReference<>(); + requestCount = new AtomicInteger(); + server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext("/", this::handleAssumeRoleWithWebIdentity); + server.start(); + } + + @AfterEach + void tearDown() { + server.stop(0); + } + + @Test + void basicModeUsesConfiguredTokenSourceWithoutSessionPolicy(@TempDir Path dir) + throws IOException { + try (AwsIrsaCredentialGenerator generator = new AwsIrsaCredentialGenerator()) { + generator.initialize(createProperties(dir)); + + AwsIrsaCredential credential = generator.generate(() -> "test-user"); + + assertEquals("access-key", credential.accessKeyId()); + assertEquals("secret-key", credential.secretAccessKey()); + assertEquals("session-token", credential.sessionToken()); + } + + String body = URLDecoder.decode(requestBody.get(), StandardCharsets.UTF_8.name()); + assertTrue(body.contains("WebIdentityToken=configured-token")); + assertTrue(body.contains("RoleSessionName=gravitino_irsa_session_test-user")); + assertFalse(body.contains("Policy=")); + } + + @Test + void basicModeReusesCachedStsCredentials(@TempDir Path dir) throws IOException { + try (AwsIrsaCredentialGenerator generator = new AwsIrsaCredentialGenerator()) { + generator.initialize(createProperties(dir)); + + generator.generate(() -> "test-user"); + generator.generate(() -> "test-user"); + } + + assertEquals(1, requestCount.get()); + } + + private Map<String, String> createProperties(Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "configured-token".getBytes(StandardCharsets.UTF_8)); + + Map<String, String> properties = new HashMap<>(); + properties.put(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString()); + properties.put("s3-role-arn", "arn:aws:iam::123456789012:role/test-role"); + properties.put("s3-region", "us-east-1"); + properties.put( + "s3-token-service-endpoint", + String.format("http://127.0.0.1:%s";, server.getAddress().getPort())); + return properties; + } + + private void handleAssumeRoleWithWebIdentity(HttpExchange exchange) throws IOException { + requestCount.incrementAndGet(); + requestBody.set(new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8)); + byte[] response = + ("<AssumeRoleWithWebIdentityResponse xmlns=\"https://sts.amazonaws.com/doc/2011-06-15/\";>" + + "<AssumeRoleWithWebIdentityResult><Credentials>" + + "<AccessKeyId>access-key</AccessKeyId>" + + "<SecretAccessKey>secret-key</SecretAccessKey>" + + "<SessionToken>session-token</SessionToken>" + + "<Expiration>" + + Instant.now().plusSeconds(3600) + + "</Expiration>" + + "</Credentials></AssumeRoleWithWebIdentityResult>" + + "</AssumeRoleWithWebIdentityResponse>") + .getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, response.length); + exchange.getResponseBody().write(response); + exchange.close(); + } +} diff --git a/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestFileWebIdentityTokenSource.java b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestFileWebIdentityTokenSource.java new file mode 100644 index 0000000000..7d8ac9bcf3 --- /dev/null +++ b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestFileWebIdentityTokenSource.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import org.apache.commons.lang3.StringUtils; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class TestFileWebIdentityTokenSource { + + @Test + void nameIsFile() { + assertEquals("file", new FileWebIdentityTokenSource().name()); + } + + @Test + void readsTokenFromConfiguredPath(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "abc123".getBytes(StandardCharsets.UTF_8)); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString())); + + assertEquals("abc123", source.getToken()); + } + + @Test + void trimsTrailingWhitespace(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "abc123\n".getBytes(StandardCharsets.UTF_8)); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString())); + + assertEquals("abc123", source.getToken()); + } + + @Test + void rereadsOnEachCallSoRotatedTokensArePickedUp(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "first".getBytes(StandardCharsets.UTF_8)); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString())); + assertEquals("first", source.getToken()); + + Files.write(tokenFile, "second".getBytes(StandardCharsets.UTF_8)); + assertEquals("second", source.getToken()); + } + + @Test + void initializeFailsWhenNeitherPropertyNorEnvIsSet() { + assumeTrue( + StringUtils.isBlank( + System.getenv(FileWebIdentityTokenSource.AWS_WEB_IDENTITY_TOKEN_FILE_ENV))); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + + IllegalStateException error = + assertThrows(IllegalStateException.class, () -> source.initialize(Collections.emptyMap())); + assertTrue(error.getMessage().contains(WebIdentityTokenSourceConfig.FILE_PATH)); + } + + @Test + void getTokenFailsWhenFileMissing(@TempDir Path dir) { + Path tokenFile = dir.resolve("missing-token"); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString())); + + IllegalStateException error = assertThrows(IllegalStateException.class, source::getToken); + assertTrue(error.getMessage().contains("does not exist")); + } + + @Test + void getTokenFailsWhenFileIsEmpty(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("empty-token"); + Files.write(tokenFile, new byte[0]); + + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString())); + + IllegalStateException error = assertThrows(IllegalStateException.class, source::getToken); + assertTrue(error.getMessage().contains("empty")); + } + + @Test + void getTokenWrapsReadFailureAsIllegalStateException(@TempDir Path dir) { + FileWebIdentityTokenSource source = new FileWebIdentityTokenSource(); + source.initialize( + Collections.singletonMap(WebIdentityTokenSourceConfig.FILE_PATH, dir.toString())); + + IllegalStateException error = assertThrows(IllegalStateException.class, source::getToken); + assertTrue(error.getMessage().contains("Failed to read WebIdentity token file")); + assertTrue(error.getCause() instanceof IOException); + } +} diff --git a/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestWebIdentityTokenSources.java b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestWebIdentityTokenSources.java new file mode 100644 index 0000000000..56b960894a --- /dev/null +++ b/bundles/aws/src/test/java/org/apache/gravitino/s3/credential/webidentity/TestWebIdentityTokenSources.java @@ -0,0 +1,187 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.s3.credential.webidentity; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class TestWebIdentityTokenSources { + + @Test + void defaultsToFileSource(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "tok".getBytes(StandardCharsets.UTF_8)); + + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString()); + + WebIdentityTokenSource source = WebIdentityTokenSources.create(props); + assertInstanceOf(FileWebIdentityTokenSource.class, source); + assertEquals("tok", source.getToken()); + } + + @Test + void selectsSourceByName(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "tok".getBytes(StandardCharsets.UTF_8)); + + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "file"); + props.put(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString()); + + WebIdentityTokenSource source = WebIdentityTokenSources.create(props); + assertInstanceOf(FileWebIdentityTokenSource.class, source); + } + + @Test + void selectsSourceByNameIgnoringCase(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "tok".getBytes(StandardCharsets.UTF_8)); + + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "FILE"); + props.put(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString()); + + WebIdentityTokenSource source = WebIdentityTokenSources.create(props); + assertInstanceOf(FileWebIdentityTokenSource.class, source); + } + + @Test + void throwsForUnknownSource() { + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "no-such-source"); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> WebIdentityTokenSources.create(props)); + assertTrue(error.getMessage().contains("no-such-source")); + } + + @Test + void blankSourcePropertyIsRejected(@TempDir Path dir) throws IOException { + Path tokenFile = dir.resolve("token"); + Files.write(tokenFile, "tok".getBytes(StandardCharsets.UTF_8)); + + // The default applies only when the key is absent; an explicitly blank value + // is treated as an unknown source rather than silently falling back. + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, ""); + props.put(WebIdentityTokenSourceConfig.FILE_PATH, tokenFile.toString()); + + assertThrows(IllegalArgumentException.class, () -> WebIdentityTokenSources.create(props)); + } + + @Test + void throwsWhenMultipleSourcesShareTheSameName() { + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "dup"); + + WebIdentityTokenSource first = new NamedNoopSource("dup"); + WebIdentityTokenSource second = new NamedNoopSource("dup"); + + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> WebIdentityTokenSources.create(props, Arrays.asList(first, second))); + assertTrue(error.getMessage().contains("Multiple WebIdentity token sources")); + assertTrue(error.getMessage().contains("dup")); + assertTrue(error.getMessage().contains(NamedNoopSource.class.getName())); + } + + @Test + void throwsWhenMultipleSourcesShareTheSameNameIgnoringCase() { + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "dup"); + + WebIdentityTokenSource first = new NamedNoopSource("dup"); + WebIdentityTokenSource second = new NamedNoopSource("DUP"); + + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> WebIdentityTokenSources.create(props, Arrays.asList(first, second))); + assertTrue(error.getMessage().contains("Multiple WebIdentity token sources")); + assertTrue(error.getMessage().contains("dup")); + } + + @Test + void throwsWhenUnselectedSourcesShareTheSameName() { + Map<String, String> props = new HashMap<>(); + props.put(WebIdentityTokenSourceConfig.SOURCE, "selected"); + + WebIdentityTokenSource selected = new NamedNoopSource("selected"); + WebIdentityTokenSource firstDuplicate = new NamedNoopSource("dup"); + WebIdentityTokenSource secondDuplicate = new NamedNoopSource("dup"); + + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + WebIdentityTokenSources.create( + props, Arrays.asList(selected, firstDuplicate, secondDuplicate))); + assertTrue(error.getMessage().contains("Multiple WebIdentity token sources")); + assertTrue(error.getMessage().contains("dup")); + } + + @Test + void throwsWhenSourceNameIsBlank() { + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + WebIdentityTokenSources.create( + new HashMap<>(), Arrays.asList(new NamedNoopSource(null)))); + assertTrue(error.getMessage().contains("must return a non-blank name")); + } + + private static final class NamedNoopSource implements WebIdentityTokenSource { + private final String name; + + NamedNoopSource(String name) { + this.name = name; + } + + @Override + public String name() { + return name; + } + + @Override + public void initialize(Map<String, String> properties) {} + + @Override + public String getToken() { + return "noop"; + } + + @Override + public void close() {} + } +}
