This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs-object-store.git
The following commit(s) were added to refs/heads/main by this push:
new ebfd02f AWS S3: Support STS endpoint, WebIdentity, RoleArn,
RoleSession configuration (#480)
ebfd02f is described below
commit ebfd02f0209e1cc3f5625815e3e5431ed0976d93
Author: Matt Friede <[email protected]>
AuthorDate: Fri Sep 19 10:11:06 2025 -0400
AWS S3: Support STS endpoint, WebIdentity, RoleArn, RoleSession
configuration (#480)
* Allow setting STS endpoint via env var
* Properly use AmazonS3Builder::credentials_from_env for
AssumeRoleWithWebIdentity auth flow
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
src/aws/builder.rs | 134 ++++++++++++++++++++++++++++++++++++++++++++++++++---
1 file changed, 128 insertions(+), 6 deletions(-)
diff --git a/src/aws/builder.rs b/src/aws/builder.rs
index 4dd41ff..06503ca 100644
--- a/src/aws/builder.rs
+++ b/src/aws/builder.rs
@@ -156,6 +156,14 @@ pub struct AmazonS3Builder {
container_credentials_full_uri: Option<String>,
/// Container authorization token file, see
<https://docs.aws.amazon.com/sdkref/latest/guide/feature-container-credentials.html>
container_authorization_token_file: Option<String>,
+ /// Web identity token file path for AssumeRoleWithWebIdentity
+ web_identity_token_file: Option<String>,
+ /// Role ARN to assume when using web identity token
+ role_arn: Option<String>,
+ /// Session name for web identity role assumption
+ role_session_name: Option<String>,
+ /// Custom STS endpoint for web identity token exchange
+ sts_endpoint: Option<String>,
/// Client options
client_options: ClientOptions,
/// Credentials
@@ -319,6 +327,34 @@ pub enum AmazonS3ConfigKey {
///
<https://docs.aws.amazon.com/sdkref/latest/guide/feature-container-credentials.html>
ContainerAuthorizationTokenFile,
+ /// Web identity token file path for AssumeRoleWithWebIdentity
+ ///
+ /// Supported keys:
+ /// - `aws_web_identity_token_file`
+ /// - `web_identity_token_file`
+ WebIdentityTokenFile,
+
+ /// Role ARN to assume when using web identity token
+ ///
+ /// Supported keys:
+ /// - `aws_role_arn`
+ /// - `role_arn`
+ RoleArn,
+
+ /// Session name for web identity role assumption
+ ///
+ /// Supported keys:
+ /// - `aws_role_session_name`
+ /// - `role_session_name`
+ RoleSessionName,
+
+ /// Custom STS endpoint for web identity token exchange
+ ///
+ /// Supported keys:
+ /// - `aws_endpoint_url_sts`
+ /// - `endpoint_url_sts`
+ StsEndpoint,
+
/// Configure how to provide `copy_if_not_exists`
///
/// See [`S3CopyIfNotExists`]
@@ -381,6 +417,10 @@ impl AsRef<str> for AmazonS3ConfigKey {
Self::ContainerCredentialsRelativeUri =>
"aws_container_credentials_relative_uri",
Self::ContainerCredentialsFullUri =>
"aws_container_credentials_full_uri",
Self::ContainerAuthorizationTokenFile =>
"aws_container_authorization_token_file",
+ Self::WebIdentityTokenFile => "aws_web_identity_token_file",
+ Self::RoleArn => "aws_role_arn",
+ Self::RoleSessionName => "aws_role_session_name",
+ Self::StsEndpoint => "aws_endpoint_url_sts",
Self::SkipSignature => "aws_skip_signature",
Self::CopyIfNotExists => "aws_copy_if_not_exists",
Self::ConditionalPut => "aws_conditional_put",
@@ -415,6 +455,12 @@ impl FromStr for AmazonS3ConfigKey {
"aws_container_credentials_relative_uri" =>
Ok(Self::ContainerCredentialsRelativeUri),
"aws_container_credentials_full_uri" =>
Ok(Self::ContainerCredentialsFullUri),
"aws_container_authorization_token_file" =>
Ok(Self::ContainerAuthorizationTokenFile),
+ "aws_web_identity_token_file" | "web_identity_token_file" => {
+ Ok(Self::WebIdentityTokenFile)
+ }
+ "aws_role_arn" | "role_arn" => Ok(Self::RoleArn),
+ "aws_role_session_name" | "role_session_name" =>
Ok(Self::RoleSessionName),
+ "aws_endpoint_url_sts" | "endpoint_url_sts" =>
Ok(Self::StsEndpoint),
"aws_skip_signature" | "skip_signature" => Ok(Self::SkipSignature),
"aws_copy_if_not_exists" | "copy_if_not_exists" =>
Ok(Self::CopyIfNotExists),
"aws_conditional_put" | "conditional_put" =>
Ok(Self::ConditionalPut),
@@ -458,6 +504,10 @@ impl AmazonS3Builder {
/// * `AWS_DEFAULT_REGION` -> region
/// * `AWS_ENDPOINT` -> endpoint
/// * `AWS_SESSION_TOKEN` -> token
+ /// * `AWS_WEB_IDENTITY_TOKEN_FILE` -> path to file containing web
identity token for AssumeRoleWithWebIdentity
+ /// * `AWS_ROLE_ARN` -> ARN of the role to assume when using web identity
token
+ /// * `AWS_ROLE_SESSION_NAME` -> optional session name for web identity
role assumption (defaults to "WebIdentitySession")
+ /// * `AWS_ENDPOINT_URL_STS` -> optional custom STS endpoint for web
identity token exchange (defaults to "https://sts.{region}.amazonaws.com";)
/// * `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI` ->
<https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html>
/// * `AWS_CONTAINER_CREDENTIALS_FULL_URI` ->
<https://docs.aws.amazon.com/sdkref/latest/guide/feature-container-credentials.html>
/// * `AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE` ->
<https://docs.aws.amazon.com/sdkref/latest/guide/feature-container-credentials.html>
@@ -543,6 +593,18 @@ impl AmazonS3Builder {
AmazonS3ConfigKey::ContainerAuthorizationTokenFile => {
self.container_authorization_token_file = Some(value.into());
}
+ AmazonS3ConfigKey::WebIdentityTokenFile => {
+ self.web_identity_token_file = Some(value.into());
+ }
+ AmazonS3ConfigKey::RoleArn => {
+ self.role_arn = Some(value.into());
+ }
+ AmazonS3ConfigKey::RoleSessionName => {
+ self.role_session_name = Some(value.into());
+ }
+ AmazonS3ConfigKey::StsEndpoint => {
+ self.sts_endpoint = Some(value.into());
+ }
AmazonS3ConfigKey::Client(key) => {
self.client_options = self.client_options.with_config(key,
value)
}
@@ -612,6 +674,10 @@ impl AmazonS3Builder {
AmazonS3ConfigKey::ContainerAuthorizationTokenFile => {
self.container_authorization_token_file.clone()
}
+ AmazonS3ConfigKey::WebIdentityTokenFile =>
self.web_identity_token_file.clone(),
+ AmazonS3ConfigKey::RoleArn => self.role_arn.clone(),
+ AmazonS3ConfigKey::RoleSessionName =>
self.role_session_name.clone(),
+ AmazonS3ConfigKey::StsEndpoint => self.sts_endpoint.clone(),
AmazonS3ConfigKey::SkipSignature =>
Some(self.skip_signature.to_string()),
AmazonS3ConfigKey::CopyIfNotExists => {
self.copy_if_not_exists.as_ref().map(ToString::to_string)
@@ -959,21 +1025,25 @@ impl AmazonS3Builder {
std::env::var("AWS_WEB_IDENTITY_TOKEN_FILE"),
std::env::var("AWS_ROLE_ARN"),
) {
- // TODO: Replace with `AmazonS3Builder::credentials_from_env`
debug!("Using WebIdentity credential provider");
- let session_name = std::env::var("AWS_ROLE_SESSION_NAME")
- .unwrap_or_else(|_| "WebIdentitySession".to_string());
+ let session_name = self
+ .role_session_name
+ .clone()
+ .unwrap_or_else(|| "WebIdentitySession".to_string());
- let endpoint = format!("https://sts.{region}.amazonaws.com";);
+ let endpoint = self
+ .sts_endpoint
+ .clone()
+ .unwrap_or_else(||
format!("https://sts.{region}.amazonaws.com";));
// Disallow non-HTTPs requests
let options = self.client_options.clone().with_allow_http(false);
let token = WebIdentityProvider {
- token_path,
+ token_path: token_path.clone(),
session_name,
- role_arn,
+ role_arn: role_arn.clone(),
endpoint,
};
@@ -1611,4 +1681,56 @@ mod tests {
"expected EKS provider but got: {debug_str}"
);
}
+
+ #[test]
+ fn test_builder_web_identity_with_config() {
+ let builder = AmazonS3Builder::new()
+ .with_bucket_name("some-bucket")
+ .with_config(
+ AmazonS3ConfigKey::WebIdentityTokenFile,
+ "/tmp/fake-token-file",
+ )
+ .with_config(
+ AmazonS3ConfigKey::RoleArn,
+ "arn:aws:iam::123456789012:role/test-role",
+ )
+ .with_config(AmazonS3ConfigKey::RoleSessionName, "TestSession")
+ .with_config(
+ AmazonS3ConfigKey::StsEndpoint,
+ "https://sts.us-west-2.amazonaws.com";,
+ );
+
+ assert_eq!(
+ builder
+ .get_config_value(&AmazonS3ConfigKey::WebIdentityTokenFile)
+ .unwrap(),
+ "/tmp/fake-token-file"
+ );
+ assert_eq!(
+ builder
+ .get_config_value(&AmazonS3ConfigKey::RoleArn)
+ .unwrap(),
+ "arn:aws:iam::123456789012:role/test-role"
+ );
+ assert_eq!(
+ builder
+ .get_config_value(&AmazonS3ConfigKey::RoleSessionName)
+ .unwrap(),
+ "TestSession"
+ );
+ assert_eq!(
+ builder
+ .get_config_value(&AmazonS3ConfigKey::StsEndpoint)
+ .unwrap(),
+ "https://sts.us-west-2.amazonaws.com";
+ );
+
+ let s3 = builder.build().expect("should build successfully");
+ let creds = &s3.client.config.credentials;
+ let debug_str = format!("{creds:?}");
+ assert!(
+ debug_str.contains("TokenCredentialProvider"),
+ "expected TokenCredentialProvider but got: {debug_str}"
+ );
+ }
}
