andygrove commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4095970783
########## docs/source/user-guide/latest/s3-credential-providers.md: ########## @@ -100,6 +100,38 @@ Without the config set, no credential-related log lines appear at startup; nativ **Credentials silently going stale during long-running jobs.** When a vendor returns `expirationEpochMillis=0`, the bridge substitutes a 5-minute expiry before handing the credential to `opendal`, so `opendal`'s cache cannot hold a stale credential indefinitely. Returning a real expiry is preferred; the 5-minute fallback is a safety net, not a knob. +## EKS / IRSA: STS throttling protection (native Iceberg scan) + +This is automatic; there is nothing to configure to get the protection, and it does not involve a bridge class. It applies to the **native Iceberg scan**. The raw-Parquet path is unaffected: it uses the AWS SDK default chain, which already retries and stops on a provider error rather than downgrading. Review Comment: The guide and the design doc describe this as a scan-only change. `iceberg_write.rs:484` builds its `FileIO` through the same `load_file_io` with `AccessMode::Write`, so native Iceberg writes switch over too. Could both docs say reads and writes? Two other sentences don't match what happens. Line 112 says a throttle that outlasts the retries surfaces as a retryable error. What the reader actually gets is opendal's permanent `signing http request` error, and it's Spark's task retry that picks it up. Line 133 says a key without the `s3.` prefix is dropped. `CometScanRule` forwards the unfiltered FileIO property bag, and this page says so at lines 173 and 298, so a bare `comet.credential.webIdentity.enabled` does reach `catalog_properties`. It has no effect only because the lookup adds the prefix. The key-constant comment and the bare-key assertion in `iceberg_wiring_reads_s3_prefixed_keys` repeat the same claim. My earlier question was only whether the bare key arrived, and it does, through `fileIOProperties`. ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { + return None; + } + Some(Self { + role_arn, + token_file, + region: non_empty_env("AWS_REGION").or_else(|| non_empty_env("AWS_DEFAULT_REGION")), + max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS), DEFAULT_MAX_ATTEMPTS), + min_ttl: Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )), + max_jitter: Duration::from_secs(parse_setting( + resolve(KEY_JITTER_SECS), + DEFAULT_JITTER_SECS, + )), + }) + } + + fn entry_key(&self) -> EntryKey { + EntryKey { + role_arn: self.role_arn.clone(), + token_file: self.token_file.clone(), + region: self.region.clone(), + max_attempts: self.max_attempts, + min_ttl: self.min_ttl, + max_jitter: self.max_jitter, + } + } +} + +/// Process-wide cache key. A credential is shared per distinct identity AND resolved settings, so a +/// catalog that configures its own retry/refresh knobs gets its own entry with its own +/// configuration honored -- independent of which scan initializes first. Two callers with the same +/// identity and the same settings still share one entry (and one STS call). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct EntryKey { + role_arn: String, + token_file: String, + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential; `refresh_jitter` is drawn once per entry so each executor +/// refreshes at a slightly different time. `last_failure` coalesces a burst of readers that hit a +/// persistent failure into a single STS call, and remembers the real error so every waiter sees it. +#[derive(Debug)] +struct SharedEntry { + provider: Arc<dyn ProvideCredentials>, + cached: RwLock<Option<Credentials>>, + /// Single-flights refreshes so a burst of readers triggers exactly one STS call. + refresh_lock: tokio::sync::Mutex<()>, + /// When the last refresh failed and the error it produced. Waiters within `FAILURE_COOLDOWN` of + /// this replay that error without re-calling STS, so a failed burst costs one call rather than + /// one per reader and every reader sees the real cause (throttle vs bad token vs trust policy). + last_failure: RwLock<Option<(Instant, String)>>, + min_ttl: Duration, + refresh_jitter: Duration, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within + /// `min_ttl + refresh_jitter`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within_margin(cred) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within the refresh margin (`min_ttl + refresh_jitter`). A credential + /// with no reported expiry never does. + fn expires_within_margin(&self, cred: &Credentials) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter, + None => false, + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, replaying the recorded + /// error so callers bail out with the real cause instead of piling another assume-role call onto + /// a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let guard = self.last_failure.read().unwrap(); + let (at, err) = guard.as_ref()?; + (at.elapsed() < FAILURE_COOLDOWN) + .then(|| format!("{err} (backing off before retrying STS)")) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error + /// the error propagates -- we never fall back to a lower-privilege identity -- and is briefly + /// remembered so concurrent waiters do not each re-issue the same throttled call. + async fn credentials(&self) -> Result<Credentials, String> { + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + let _guard = self.refresh_lock.lock().await; + // Re-check: another task may have refreshed (or just failed) while we waited on the lock. + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + match self.provider.provide_credentials().await { + Ok(cred) => { + self.warn_if_immediately_stale(&cred); + *self.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + let err = format!("web-identity assume-role failed: {e}"); Review Comment: I don't think the error replay from my earlier comment gets the cause through. `{e}` on a `CredentialsError::ProviderError` always prints `an error occurred while loading credentials`. A throttle, a missing token file and a missing region all gave that same string when I tried them. `concurrent_failed_refresh_is_coalesced` only checks our own prefix, so it passes anyway. `aws_sdk_sts::error::DisplayErrorContext(&e)` walks the source chain and shows the `Throttling` code and the `Rate exceeded` message. The message also never reaches the task error. `ProvideCredentialChain` logs a provider `Err` at warn level and moves on, so the `Signer` reports `failed to load signing credential`. That warn line prints `{provider:?}`, and the derived `Debug` on `WebIdentityCredentialProvider` is about 30 KB because it includes the whole STS client config and the cached access key id. Could we use `DisplayErrorContext`, log the failure once ourselves when it's recorded, and give this type a compact manual `Debug` like `CometS3CredentialBridge` has? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { + return None; + } + Some(Self { + role_arn, + token_file, + region: non_empty_env("AWS_REGION").or_else(|| non_empty_env("AWS_DEFAULT_REGION")), + max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS), DEFAULT_MAX_ATTEMPTS), + min_ttl: Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )), + max_jitter: Duration::from_secs(parse_setting( + resolve(KEY_JITTER_SECS), + DEFAULT_JITTER_SECS, + )), + }) + } + + fn entry_key(&self) -> EntryKey { + EntryKey { + role_arn: self.role_arn.clone(), + token_file: self.token_file.clone(), + region: self.region.clone(), + max_attempts: self.max_attempts, + min_ttl: self.min_ttl, + max_jitter: self.max_jitter, + } + } +} + +/// Process-wide cache key. A credential is shared per distinct identity AND resolved settings, so a +/// catalog that configures its own retry/refresh knobs gets its own entry with its own +/// configuration honored -- independent of which scan initializes first. Two callers with the same +/// identity and the same settings still share one entry (and one STS call). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct EntryKey { + role_arn: String, + token_file: String, + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential; `refresh_jitter` is drawn once per entry so each executor +/// refreshes at a slightly different time. `last_failure` coalesces a burst of readers that hit a +/// persistent failure into a single STS call, and remembers the real error so every waiter sees it. +#[derive(Debug)] +struct SharedEntry { + provider: Arc<dyn ProvideCredentials>, + cached: RwLock<Option<Credentials>>, + /// Single-flights refreshes so a burst of readers triggers exactly one STS call. + refresh_lock: tokio::sync::Mutex<()>, + /// When the last refresh failed and the error it produced. Waiters within `FAILURE_COOLDOWN` of + /// this replay that error without re-calling STS, so a failed burst costs one call rather than + /// one per reader and every reader sees the real cause (throttle vs bad token vs trust policy). + last_failure: RwLock<Option<(Instant, String)>>, + min_ttl: Duration, + refresh_jitter: Duration, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within + /// `min_ttl + refresh_jitter`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within_margin(cred) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within the refresh margin (`min_ttl + refresh_jitter`). A credential + /// with no reported expiry never does. + fn expires_within_margin(&self, cred: &Credentials) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter, + None => false, + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, replaying the recorded + /// error so callers bail out with the real cause instead of piling another assume-role call onto + /// a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let guard = self.last_failure.read().unwrap(); + let (at, err) = guard.as_ref()?; + (at.elapsed() < FAILURE_COOLDOWN) + .then(|| format!("{err} (backing off before retrying STS)")) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error + /// the error propagates -- we never fall back to a lower-privilege identity -- and is briefly + /// remembered so concurrent waiters do not each re-issue the same throttled call. + async fn credentials(&self) -> Result<Credentials, String> { + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + let _guard = self.refresh_lock.lock().await; + // Re-check: another task may have refreshed (or just failed) while we waited on the lock. + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + match self.provider.provide_credentials().await { + Ok(cred) => { + self.warn_if_immediately_stale(&cred); + *self.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + let err = format!("web-identity assume-role failed: {e}"); + *self.last_failure.write().unwrap() = Some((Instant::now(), err.clone())); + Err(err) + } + } + } + + /// Warns once if a freshly fetched credential already falls inside the refresh margin -- a sign + /// `minTtlSeconds` is misconfigured larger than the STS session lifetime, which would make every + /// request refresh (the very burst this provider avoids). + fn warn_if_immediately_stale(&self, cred: &Credentials) { + static WARNED: OnceLock<()> = OnceLock::new(); + if self.expires_within_margin(cred) && WARNED.set(()).is_ok() { + log::warn!( + "A freshly fetched web-identity credential already falls within the {}s refresh \ + margin; comet.credential.webIdentity.minTtlSeconds may be larger than the STS \ + session lifetime, which forces a refresh on every request", + self.min_ttl.as_secs() + ); + } + } +} + +/// Registry of shared credential entries, one per identity, for the lifetime of the process. +/// +/// Process lifetime is the right scope for the same reason as the region cache in `s3.rs`: each +/// executor is dedicated to one Spark application, and there is a bounded set of assumed roles per +/// job. Entries are never evicted; the map stays proportional to the number of distinct roles. +fn registry() -> &'static std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>> { + static REGISTRY: OnceLock<std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>>> = + OnceLock::new(); + REGISTRY.get_or_init(|| std::sync::Mutex::new(HashMap::new())) +} + +/// Returns the shared entry for `cfg`, building the AWS SDK provider once if needed. The provider +/// is built outside the registry lock (it is async); a concurrent builder just loses the insert +/// race, which is harmless. +async fn shared_entry(cfg: &WebIdentityConfig) -> Arc<SharedEntry> { + let key = cfg.entry_key(); + if let Some(entry) = registry().lock().unwrap().get(&key).cloned() { + return entry; + } + + let provider = build_provider(cfg, None).await; + // Draw the refresh jitter once per entry (each distinct identity+settings key), so two + // executors -- or two catalogs with different tuning -- refresh at slightly different times and + // the cluster does not re-burst on a synchronized refresh. + let jitter = if cfg.max_jitter.is_zero() { + Duration::ZERO + } else { + Duration::from_secs(rand::rng().random_range(0..=cfg.max_jitter.as_secs())) + }; + let entry = Arc::new(SharedEntry { + provider, + cached: RwLock::new(None), + refresh_lock: tokio::sync::Mutex::new(()), + last_failure: RwLock::new(None), + min_ttl: cfg.min_ttl, + refresh_jitter: jitter, + }); + + let mut map = registry().lock().unwrap(); + Arc::clone(map.entry(key).or_insert(entry)) +} + +/// Builds the web-identity credential provider from the AWS SDK's fully-resolved config. +/// +/// The key move: we load a real `SdkConfig` (`aws_config::defaults(...).load()`), which resolves +/// region, FIPS, dual-stack, the profile, and any custom/profile STS endpoint with the SDK's normal +/// environment-then-profile precedence, and build the STS client from it. Because the client is +/// built from the resolved config rather than a hand-assembled one, there is no per-setting copying +/// to keep in sync -- every endpoint/region knob the SDK understands is honored. We only ever call +/// `AssumeRoleWithWebIdentity`, so there is no IMDS/instance-role fallback to downgrade to, and the +/// raised `RetryConfig` gives the throttle its retries. +/// +/// `http_override` lets tests drive the STS client through an in-memory stub; production passes +/// `None`. +async fn build_provider( + cfg: &WebIdentityConfig, + http_override: Option<SharedHttpClient>, +) -> Arc<dyn ProvideCredentials> { + let mut loader = aws_config::defaults(BehaviorVersion::latest()) + .retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts)); + if let Some(http) = http_override { + loader = loader.http_client(http); + } + let sdk = loader.load().await; + Arc::new(web_identity_provider_from( + cfg, + aws_sdk_sts::Client::new(&sdk), + )) +} + +/// Assembles the provider from an STS client. Split out so tests can supply a client built with an +/// in-memory HTTP stub while sharing the identity wiring with production. +fn web_identity_provider_from( + cfg: &WebIdentityConfig, + sts: aws_sdk_sts::Client, +) -> WebIdentityStsProvider { + WebIdentityStsProvider { + sts, + role_arn: cfg.role_arn.clone(), + token_file: cfg.token_file.clone(), + session_name: session_name(), + } +} + +/// STS `AssumeRoleWithWebIdentity` session name. Honors `AWS_ROLE_SESSION_NAME` first, matching the +/// default chain, so a trust policy conditioned on `sts:RoleSessionName` keeps working after the +/// take-over engages; otherwise falls back to a stable prefix plus a timestamp. +fn session_name() -> String { + if let Some(name) = non_empty_env("AWS_ROLE_SESSION_NAME") { + return name; + } + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + format!("comet-web-identity-{nanos}") +} + +/// A web-identity-only credential provider: it reads the projected token and calls STS +/// `AssumeRoleWithWebIdentity` on `sts`, and does nothing else. No credential chain, so a throttle +/// that outlasts the STS client's retries returns an error rather than a lower-privilege identity. +#[derive(Debug)] +struct WebIdentityStsProvider { + sts: aws_sdk_sts::Client, + role_arn: String, + token_file: String, + session_name: String, +} + +impl WebIdentityStsProvider { + async fn resolve(&self) -> Result<Credentials, CredentialsError> { + let token = std::fs::read_to_string(&self.token_file).map_err(|e| { + CredentialsError::provider_error(format!( + "reading web identity token file {}: {e}", + self.token_file + )) + })?; + let response = self + .sts + .assume_role_with_web_identity() + .role_arn(&self.role_arn) + .role_session_name(&self.session_name) + .web_identity_token(token.trim()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let creds = response.credentials().ok_or_else(|| { + CredentialsError::provider_error( + "STS AssumeRoleWithWebIdentity response had no credentials", + ) + })?; + let expiration = creds.expiration(); + let expiry = SystemTime::UNIX_EPOCH + .checked_add(Duration::new( + expiration.secs().max(0) as u64, + expiration.subsec_nanos(), + )) + .ok_or_else(|| { + CredentialsError::provider_error("STS credential expiry is out of range") + })?; + Ok(Credentials::new( + creds.access_key_id(), + creds.secret_access_key(), + Some(creds.session_token().to_string()), + Some(expiry), + "CometWebIdentity", + )) + } +} + +impl ProvideCredentials for WebIdentityStsProvider { + fn provide_credentials<'a>(&'a self) -> creds_future::ProvideCredentials<'a> + where + Self: 'a, + { + creds_future::ProvideCredentials::new(self.resolve()) + } +} + +/// The credential provider handed to opendal via `CustomAwsCredentialLoader` (the Iceberg path). +/// Holds only the cheap config plus a lazily resolved handle to the process-wide shared entry, so +/// the per-request path skips the registry lock after the first fetch. +#[derive(Debug)] +pub struct WebIdentityCredentialProvider { + config: WebIdentityConfig, + entry: tokio::sync::OnceCell<Arc<SharedEntry>>, +} + +impl WebIdentityCredentialProvider { + pub fn new(config: WebIdentityConfig) -> Self { + Self { + config, + entry: tokio::sync::OnceCell::new(), + } + } + + /// Resolves (once per provider) the shared entry for this identity. The entry itself is shared + /// process-wide via the registry; this just memoizes the lookup so repeated fetches avoid the + /// registry lock and the per-call `EntryKey` allocation. + async fn entry(&self) -> &Arc<SharedEntry> { + self.entry.get_or_init(|| shared_entry(&self.config)).await + } +} + +impl IcebergProvideCredential for WebIdentityCredentialProvider { + type Credential = IcebergAwsCredential; + + async fn provide_credential( + &self, + _ctx: &Context, + ) -> reqsign_core::Result<Option<Self::Credential>> { + let entry = self.entry().await; + let cred = entry + .credentials() + .await + .map_err(|e| ReqsignError::new(ReqsignErrorKind::CredentialInvalid, e))?; + + // Report the jittered refresh deadline (true expiry minus min_ttl minus jitter) as the + // expiry opendal caches against, so opendal refreshes when our own cache would, and each + // executor's refresh is spread out rather than synchronized. + let expires_in = match cred.expiry() { + Some(expiry) => { + let deadline = expiry + .checked_sub(entry.min_ttl + entry.refresh_jitter) + .unwrap_or(expiry); + Some(system_time_to_timestamp(deadline)?) + } + None => Some(Timestamp::now() + DEFAULT_EXPIRY_WHEN_UNKNOWN), + }; + + Ok(Some(IcebergAwsCredential { + access_key_id: cred.access_key_id().to_string(), + secret_access_key: cred.secret_access_key().to_string(), + session_token: cred.session_token().map(|s| s.to_string()), + expires_in, + })) + } +} + +/// Decides whether the Comet web-identity provider should take over credential resolution on the +/// Iceberg path. It does so only when the catalog names no explicit Comet provider class, has no +/// explicit credentials, and IRSA is detected; otherwise opendal keeps its default chain. `resolve` +/// reads a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) from the catalog property bag. +/// +/// It also stands aside for any credential source the default chain ranks ahead of web-identity: +/// static credentials in the environment (`AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY`) or a +/// configured profile (`AWS_PROFILE`, or a shared credentials / config file). opendal/reqsign +/// resolves Environment -> Profile -> WebIdentity, so taking over in those cases would silently +/// switch identity from the user's chosen source to the service-account role. The decision is +/// logged: `debug!` when the provider engages, `info!` (with the reason) when IRSA is detected but +/// we stand aside, so an operator can tell which branch a run took. +pub fn take_over_if_irsa<F>( + explicit_credentials: bool, + resolve: F, +) -> Option<WebIdentityCredentialProvider> +where + F: Fn(&str) -> Option<String>, +{ + if !irsa_present() { + // Not an IRSA environment; the default chain handles everything as before. No log: this is + // the common non-EKS case and would be pure noise. + return None; + } + let stand_aside_reason = if explicit_credentials { + Some("an explicit credential provider is configured") + } else if explicit_env_credentials() { + Some("static AWS credentials are set in the environment") + } else if configured_profile() { + Some("an AWS profile or config file is present") + } else { + None + }; + if let Some(reason) = stand_aside_reason { + log::info!("IRSA detected but the Comet web-identity provider is standing aside: {reason}"); Review Comment: This runs once per scan task and once per write task, because `load_file_io` builds the loader in `IcebergScanExec::execute_with_tasks` and in the write path, and native logging defaults to INFO. On EKS with a REST catalog that hands out `s3.access-key-id` and `s3.secret-access-key`, every task logs this line. So does every task for anyone who set `enabled=false`. `credential_bridge.rs` limits its missing-expiry warning to once per process for the same reason (`WARNED_MISSING_EXPIRY`). Could this do the same, once per reason? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { + return None; + } + Some(Self { + role_arn, + token_file, + region: non_empty_env("AWS_REGION").or_else(|| non_empty_env("AWS_DEFAULT_REGION")), + max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS), DEFAULT_MAX_ATTEMPTS), + min_ttl: Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )), + max_jitter: Duration::from_secs(parse_setting( + resolve(KEY_JITTER_SECS), + DEFAULT_JITTER_SECS, + )), + }) + } + + fn entry_key(&self) -> EntryKey { + EntryKey { + role_arn: self.role_arn.clone(), + token_file: self.token_file.clone(), + region: self.region.clone(), + max_attempts: self.max_attempts, + min_ttl: self.min_ttl, + max_jitter: self.max_jitter, + } + } +} + +/// Process-wide cache key. A credential is shared per distinct identity AND resolved settings, so a +/// catalog that configures its own retry/refresh knobs gets its own entry with its own +/// configuration honored -- independent of which scan initializes first. Two callers with the same +/// identity and the same settings still share one entry (and one STS call). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct EntryKey { + role_arn: String, + token_file: String, + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential; `refresh_jitter` is drawn once per entry so each executor +/// refreshes at a slightly different time. `last_failure` coalesces a burst of readers that hit a +/// persistent failure into a single STS call, and remembers the real error so every waiter sees it. +#[derive(Debug)] +struct SharedEntry { + provider: Arc<dyn ProvideCredentials>, + cached: RwLock<Option<Credentials>>, + /// Single-flights refreshes so a burst of readers triggers exactly one STS call. + refresh_lock: tokio::sync::Mutex<()>, + /// When the last refresh failed and the error it produced. Waiters within `FAILURE_COOLDOWN` of + /// this replay that error without re-calling STS, so a failed burst costs one call rather than + /// one per reader and every reader sees the real cause (throttle vs bad token vs trust policy). + last_failure: RwLock<Option<(Instant, String)>>, + min_ttl: Duration, + refresh_jitter: Duration, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within + /// `min_ttl + refresh_jitter`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within_margin(cred) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within the refresh margin (`min_ttl + refresh_jitter`). A credential + /// with no reported expiry never does. + fn expires_within_margin(&self, cred: &Credentials) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter, + None => false, + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, replaying the recorded + /// error so callers bail out with the real cause instead of piling another assume-role call onto + /// a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let guard = self.last_failure.read().unwrap(); + let (at, err) = guard.as_ref()?; + (at.elapsed() < FAILURE_COOLDOWN) + .then(|| format!("{err} (backing off before retrying STS)")) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error + /// the error propagates -- we never fall back to a lower-privilege identity -- and is briefly + /// remembered so concurrent waiters do not each re-issue the same throttled call. + async fn credentials(&self) -> Result<Credentials, String> { + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + let _guard = self.refresh_lock.lock().await; + // Re-check: another task may have refreshed (or just failed) while we waited on the lock. + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + match self.provider.provide_credentials().await { + Ok(cred) => { + self.warn_if_immediately_stale(&cred); + *self.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + let err = format!("web-identity assume-role failed: {e}"); + *self.last_failure.write().unwrap() = Some((Instant::now(), err.clone())); + Err(err) + } + } + } + + /// Warns once if a freshly fetched credential already falls inside the refresh margin -- a sign + /// `minTtlSeconds` is misconfigured larger than the STS session lifetime, which would make every + /// request refresh (the very burst this provider avoids). + fn warn_if_immediately_stale(&self, cred: &Credentials) { + static WARNED: OnceLock<()> = OnceLock::new(); + if self.expires_within_margin(cred) && WARNED.set(()).is_ok() { + log::warn!( + "A freshly fetched web-identity credential already falls within the {}s refresh \ + margin; comet.credential.webIdentity.minTtlSeconds may be larger than the STS \ + session lifetime, which forces a refresh on every request", + self.min_ttl.as_secs() + ); + } + } +} + +/// Registry of shared credential entries, one per identity, for the lifetime of the process. +/// +/// Process lifetime is the right scope for the same reason as the region cache in `s3.rs`: each +/// executor is dedicated to one Spark application, and there is a bounded set of assumed roles per +/// job. Entries are never evicted; the map stays proportional to the number of distinct roles. +fn registry() -> &'static std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>> { + static REGISTRY: OnceLock<std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>>> = + OnceLock::new(); + REGISTRY.get_or_init(|| std::sync::Mutex::new(HashMap::new())) +} + +/// Returns the shared entry for `cfg`, building the AWS SDK provider once if needed. The provider +/// is built outside the registry lock (it is async); a concurrent builder just loses the insert +/// race, which is harmless. +async fn shared_entry(cfg: &WebIdentityConfig) -> Arc<SharedEntry> { + let key = cfg.entry_key(); + if let Some(entry) = registry().lock().unwrap().get(&key).cloned() { + return entry; + } + + let provider = build_provider(cfg, None).await; + // Draw the refresh jitter once per entry (each distinct identity+settings key), so two + // executors -- or two catalogs with different tuning -- refresh at slightly different times and + // the cluster does not re-burst on a synchronized refresh. + let jitter = if cfg.max_jitter.is_zero() { + Duration::ZERO + } else { + Duration::from_secs(rand::rng().random_range(0..=cfg.max_jitter.as_secs())) + }; + let entry = Arc::new(SharedEntry { + provider, + cached: RwLock::new(None), + refresh_lock: tokio::sync::Mutex::new(()), + last_failure: RwLock::new(None), + min_ttl: cfg.min_ttl, + refresh_jitter: jitter, + }); + + let mut map = registry().lock().unwrap(); + Arc::clone(map.entry(key).or_insert(entry)) +} + +/// Builds the web-identity credential provider from the AWS SDK's fully-resolved config. +/// +/// The key move: we load a real `SdkConfig` (`aws_config::defaults(...).load()`), which resolves +/// region, FIPS, dual-stack, the profile, and any custom/profile STS endpoint with the SDK's normal +/// environment-then-profile precedence, and build the STS client from it. Because the client is +/// built from the resolved config rather than a hand-assembled one, there is no per-setting copying +/// to keep in sync -- every endpoint/region knob the SDK understands is honored. We only ever call +/// `AssumeRoleWithWebIdentity`, so there is no IMDS/instance-role fallback to downgrade to, and the +/// raised `RetryConfig` gives the throttle its retries. +/// +/// `http_override` lets tests drive the STS client through an in-memory stub; production passes +/// `None`. +async fn build_provider( + cfg: &WebIdentityConfig, + http_override: Option<SharedHttpClient>, +) -> Arc<dyn ProvideCredentials> { + let mut loader = aws_config::defaults(BehaviorVersion::latest()) Review Comment: What happens on a web-identity pod with no `AWS_REGION` or `AWS_DEFAULT_REGION`? When I tried it with no profile and no instance metadata service, the STS client sent no request at all and failed with the generic message. reqsign's web-identity provider falls back to the global `sts.amazonaws.com` in that case, unless `AWS_STS_REGIONAL_ENDPOINTS` is `regional` (see `sts_endpoint` in `reqsign-aws-core`). So a setup that works today, for example OIDC federation from a non-EKS cluster with `client.region` set on the catalog, would fail every read. The EKS webhook injects `AWS_REGION`, so could a missing region in the environment be another reason to stand aside in `take_over_if_irsa`? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { Review Comment: `parse_setting` uses `str::parse::<bool>`, which only accepts lowercase `true` and `false`. So `enabled=False` or `FALSE` falls back to the default and the take-over stays on, with nothing in the logs. Since this is the only opt-out, could it compare case-insensitively like `s3.rs:245` and iceberg-rust's `is_truthy` do, and warn when a value doesn't parse? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { + return None; + } + Some(Self { + role_arn, + token_file, + region: non_empty_env("AWS_REGION").or_else(|| non_empty_env("AWS_DEFAULT_REGION")), + max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS), DEFAULT_MAX_ATTEMPTS), + min_ttl: Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )), + max_jitter: Duration::from_secs(parse_setting( + resolve(KEY_JITTER_SECS), + DEFAULT_JITTER_SECS, + )), + }) + } + + fn entry_key(&self) -> EntryKey { + EntryKey { + role_arn: self.role_arn.clone(), + token_file: self.token_file.clone(), + region: self.region.clone(), + max_attempts: self.max_attempts, + min_ttl: self.min_ttl, + max_jitter: self.max_jitter, + } + } +} + +/// Process-wide cache key. A credential is shared per distinct identity AND resolved settings, so a +/// catalog that configures its own retry/refresh knobs gets its own entry with its own +/// configuration honored -- independent of which scan initializes first. Two callers with the same +/// identity and the same settings still share one entry (and one STS call). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct EntryKey { + role_arn: String, + token_file: String, + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential; `refresh_jitter` is drawn once per entry so each executor +/// refreshes at a slightly different time. `last_failure` coalesces a burst of readers that hit a +/// persistent failure into a single STS call, and remembers the real error so every waiter sees it. +#[derive(Debug)] +struct SharedEntry { + provider: Arc<dyn ProvideCredentials>, + cached: RwLock<Option<Credentials>>, + /// Single-flights refreshes so a burst of readers triggers exactly one STS call. + refresh_lock: tokio::sync::Mutex<()>, + /// When the last refresh failed and the error it produced. Waiters within `FAILURE_COOLDOWN` of + /// this replay that error without re-calling STS, so a failed burst costs one call rather than + /// one per reader and every reader sees the real cause (throttle vs bad token vs trust policy). + last_failure: RwLock<Option<(Instant, String)>>, + min_ttl: Duration, + refresh_jitter: Duration, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within + /// `min_ttl + refresh_jitter`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within_margin(cred) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within the refresh margin (`min_ttl + refresh_jitter`). A credential + /// with no reported expiry never does. + fn expires_within_margin(&self, cred: &Credentials) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter, + None => false, + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, replaying the recorded + /// error so callers bail out with the real cause instead of piling another assume-role call onto + /// a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let guard = self.last_failure.read().unwrap(); + let (at, err) = guard.as_ref()?; + (at.elapsed() < FAILURE_COOLDOWN) + .then(|| format!("{err} (backing off before retrying STS)")) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error + /// the error propagates -- we never fall back to a lower-privilege identity -- and is briefly + /// remembered so concurrent waiters do not each re-issue the same throttled call. + async fn credentials(&self) -> Result<Credentials, String> { + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + let _guard = self.refresh_lock.lock().await; + // Re-check: another task may have refreshed (or just failed) while we waited on the lock. + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + match self.provider.provide_credentials().await { + Ok(cred) => { + self.warn_if_immediately_stale(&cred); + *self.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure.write().unwrap() = None; + Ok(cred) + } + Err(e) => { Review Comment: When a refresh fails, this returns the error even if the cached credential still has minutes left. The margin is 300 to 360 seconds, so an STS throttle at refresh time fails every read on the executor while we hold a credential that would have worked. With a cached credential 240 seconds from expiry and a failing provider, `credentials()` returns `Err`. sunchao noted this in the first review and left it because the old Parquet cache did the same, but that comparison is gone now. Could the error path return the cached credential while it's still valid past reqsign's 120 second cache margin plus the 10 second signing headroom, and still record the failure for the cooldown? The expiry reported in that path would need to be the real one, or the `Signer` rejects it for the reason in my other comment. ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1263 @@ +// 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. + +//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential provider for the native S3 +//! paths. +//! +//! Why this exists: on EKS with IRSA the native reader assumes the app role by calling STS +//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x many cores starting +//! together) STS throttles that call. opendal's default reqsign chain (used by the Iceberg path +//! when no Comet provider class is set) does NOT retry the throttle and silently downgrades to the +//! EC2/EKS node instance role, which lacks bucket access -> every read then fails with a hard S3 +//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md. +//! +//! This provider fixes all three parts of that failure: +//! 1. Retry on throttle. It builds an STS client from the AWS SDK's fully-resolved `SdkConfig` +//! (`aws_config::defaults(...).load()`) with a raised `RetryConfig`, and calls +//! `AssumeRoleWithWebIdentity` on it. Because the client comes from the resolved config, it +//! honors region, FIPS, dual-stack and any profile/custom STS endpoint the SDK would -- +//! there is no hand-assembled config to drift. `max_attempts` is configurable. +//! 2. No silent downgrade. It only ever calls `AssumeRoleWithWebIdentity` -- there is no +//! credential chain and no IMDS/instance-role fallback -- so a throttle that outlasts the +//! retries surfaces as an error instead of a wrong-identity credential. +//! 3. Shared, jittered cache. One assumed-role credential is cached per process, keyed by +//! identity (role_arn, token_file, region) and the resolved retry/refresh settings, and shared +//! across all reader threads and scans that resolve to the same key. Refresh fires ahead of +//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide refreshes do not +//! synchronize into another burst; a failed refresh is briefly remembered so a throttled burst +//! costs one STS call rather than one per reader. +//! +//! It is wired into the Iceberg scan path (`iceberg_common::build_s3_credential_loader`), which is +//! where the reported failure occurs: opendal's default reqsign chain is the one that downgrades to +//! the node role. The raw-Parquet path is left on the AWS SDK default chain, which already retries +//! and stops on a provider error rather than downgrading. The provider is exposed to opendal as +//! reqsign's `ProvideCredential` via `CustomAwsCredentialLoader`, mirroring +//! `credential_bridge::CometS3CredentialBridge`. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aws_config::retry::RetryConfig; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::future as creds_future; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use rand::RngExt; +use reqsign_core::time::Timestamp; +use reqsign_core::{ + Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind, + ProvideCredential as IcebergProvideCredential, +}; + +use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN; + +/// EKS-projected env vars that signal IRSA is in effect. Both must be present. +const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN"; + +/// Config keys in their bare form. On the Iceberg path they are resolved under the `s3.` prefix in +/// the catalog property bag (e.g. `s3.comet.credential.webIdentity.enabled`), matching the existing +/// `s3.comet.credential.provider.class` SPI key. The `s3.` prefix is required: that is how a catalog +/// property reaches iceberg-rust's FileIO property bag; a bare, unprefixed key would be dropped and +/// the opt-out would silently have no effect. +const KEY_ENABLED: &str = "comet.credential.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.credential.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.credential.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.credential.webIdentity.refreshJitterSeconds"; + +const DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; +const DEFAULT_JITTER_SECS: u64 = 60; + +/// After a refresh exhausts its STS retries and fails, waiters within this window get the failure +/// without each firing their own assume-role call. Bounds STS pressure during a sustained throttle +/// (one call per entry per window instead of one per reader) while still letting the credential +/// recover shortly after. Kept short: the SDK has already spent its retry budget by the time we +/// record a failure. +const FAILURE_COOLDOWN: Duration = Duration::from_secs(1); + +/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the expensive AWS SDK +/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`. +#[derive(Clone, Debug)] +pub struct WebIdentityConfig { + role_arn: String, + token_file: String, + /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key. The STS client's + /// actual region (and endpoint) comes from the resolved `SdkConfig`. + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +impl WebIdentityConfig { + /// Returns a config only when IRSA is in effect (both env vars present) and the feature is + /// enabled. `resolve` looks up a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) in whichever config + /// bag the caller owns -- the Iceberg catalog bag or the Parquet `fs.s3a.*` bag -- so the two + /// scan paths share one detection routine without sharing a config-key scheme. Returns `None` + /// when IRSA is not detected or the feature is disabled. + pub fn detect_with<F>(resolve: F) -> Option<Self> + where + F: Fn(&str) -> Option<String>, + { + let token_file = non_empty_env(ENV_TOKEN_FILE)?; + let role_arn = non_empty_env(ENV_ROLE_ARN)?; + if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) { + return None; + } + Some(Self { + role_arn, + token_file, + region: non_empty_env("AWS_REGION").or_else(|| non_empty_env("AWS_DEFAULT_REGION")), + max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS), DEFAULT_MAX_ATTEMPTS), + min_ttl: Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )), + max_jitter: Duration::from_secs(parse_setting( + resolve(KEY_JITTER_SECS), + DEFAULT_JITTER_SECS, + )), + }) + } + + fn entry_key(&self) -> EntryKey { + EntryKey { + role_arn: self.role_arn.clone(), + token_file: self.token_file.clone(), + region: self.region.clone(), + max_attempts: self.max_attempts, + min_ttl: self.min_ttl, + max_jitter: self.max_jitter, + } + } +} + +/// Process-wide cache key. A credential is shared per distinct identity AND resolved settings, so a +/// catalog that configures its own retry/refresh knobs gets its own entry with its own +/// configuration honored -- independent of which scan initializes first. Two callers with the same +/// identity and the same settings still share one entry (and one STS call). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct EntryKey { + role_arn: String, + token_file: String, + region: Option<String>, + max_attempts: u32, + min_ttl: Duration, + max_jitter: Duration, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential; `refresh_jitter` is drawn once per entry so each executor +/// refreshes at a slightly different time. `last_failure` coalesces a burst of readers that hit a +/// persistent failure into a single STS call, and remembers the real error so every waiter sees it. +#[derive(Debug)] +struct SharedEntry { + provider: Arc<dyn ProvideCredentials>, + cached: RwLock<Option<Credentials>>, + /// Single-flights refreshes so a burst of readers triggers exactly one STS call. + refresh_lock: tokio::sync::Mutex<()>, + /// When the last refresh failed and the error it produced. Waiters within `FAILURE_COOLDOWN` of + /// this replay that error without re-calling STS, so a failed burst costs one call rather than + /// one per reader and every reader sees the real cause (throttle vs bad token vs trust policy). + last_failure: RwLock<Option<(Instant, String)>>, + min_ttl: Duration, + refresh_jitter: Duration, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within + /// `min_ttl + refresh_jitter`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within_margin(cred) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within the refresh margin (`min_ttl + refresh_jitter`). A credential + /// with no reported expiry never does. + fn expires_within_margin(&self, cred: &Credentials) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter, + None => false, + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, replaying the recorded + /// error so callers bail out with the real cause instead of piling another assume-role call onto + /// a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let guard = self.last_failure.read().unwrap(); + let (at, err) = guard.as_ref()?; + (at.elapsed() < FAILURE_COOLDOWN) + .then(|| format!("{err} (backing off before retrying STS)")) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error + /// the error propagates -- we never fall back to a lower-privilege identity -- and is briefly + /// remembered so concurrent waiters do not each re-issue the same throttled call. + async fn credentials(&self) -> Result<Credentials, String> { + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + let _guard = self.refresh_lock.lock().await; + // Re-check: another task may have refreshed (or just failed) while we waited on the lock. + if let Some(cred) = self.fresh() { + return Ok(cred); + } + if let Some(err) = self.in_failure_cooldown() { + return Err(err); + } + match self.provider.provide_credentials().await { + Ok(cred) => { + self.warn_if_immediately_stale(&cred); + *self.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + let err = format!("web-identity assume-role failed: {e}"); + *self.last_failure.write().unwrap() = Some((Instant::now(), err.clone())); + Err(err) + } + } + } + + /// Warns once if a freshly fetched credential already falls inside the refresh margin -- a sign + /// `minTtlSeconds` is misconfigured larger than the STS session lifetime, which would make every + /// request refresh (the very burst this provider avoids). + fn warn_if_immediately_stale(&self, cred: &Credentials) { + static WARNED: OnceLock<()> = OnceLock::new(); + if self.expires_within_margin(cred) && WARNED.set(()).is_ok() { + log::warn!( + "A freshly fetched web-identity credential already falls within the {}s refresh \ + margin; comet.credential.webIdentity.minTtlSeconds may be larger than the STS \ + session lifetime, which forces a refresh on every request", + self.min_ttl.as_secs() + ); + } + } +} + +/// Registry of shared credential entries, one per identity, for the lifetime of the process. +/// +/// Process lifetime is the right scope for the same reason as the region cache in `s3.rs`: each +/// executor is dedicated to one Spark application, and there is a bounded set of assumed roles per +/// job. Entries are never evicted; the map stays proportional to the number of distinct roles. +fn registry() -> &'static std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>> { + static REGISTRY: OnceLock<std::sync::Mutex<HashMap<EntryKey, Arc<SharedEntry>>>> = + OnceLock::new(); + REGISTRY.get_or_init(|| std::sync::Mutex::new(HashMap::new())) +} + +/// Returns the shared entry for `cfg`, building the AWS SDK provider once if needed. The provider +/// is built outside the registry lock (it is async); a concurrent builder just loses the insert +/// race, which is harmless. +async fn shared_entry(cfg: &WebIdentityConfig) -> Arc<SharedEntry> { + let key = cfg.entry_key(); + if let Some(entry) = registry().lock().unwrap().get(&key).cloned() { + return entry; + } + + let provider = build_provider(cfg, None).await; + // Draw the refresh jitter once per entry (each distinct identity+settings key), so two + // executors -- or two catalogs with different tuning -- refresh at slightly different times and + // the cluster does not re-burst on a synchronized refresh. + let jitter = if cfg.max_jitter.is_zero() { + Duration::ZERO + } else { + Duration::from_secs(rand::rng().random_range(0..=cfg.max_jitter.as_secs())) + }; + let entry = Arc::new(SharedEntry { + provider, + cached: RwLock::new(None), + refresh_lock: tokio::sync::Mutex::new(()), + last_failure: RwLock::new(None), + min_ttl: cfg.min_ttl, + refresh_jitter: jitter, + }); + + let mut map = registry().lock().unwrap(); + Arc::clone(map.entry(key).or_insert(entry)) +} + +/// Builds the web-identity credential provider from the AWS SDK's fully-resolved config. +/// +/// The key move: we load a real `SdkConfig` (`aws_config::defaults(...).load()`), which resolves +/// region, FIPS, dual-stack, the profile, and any custom/profile STS endpoint with the SDK's normal +/// environment-then-profile precedence, and build the STS client from it. Because the client is +/// built from the resolved config rather than a hand-assembled one, there is no per-setting copying +/// to keep in sync -- every endpoint/region knob the SDK understands is honored. We only ever call +/// `AssumeRoleWithWebIdentity`, so there is no IMDS/instance-role fallback to downgrade to, and the +/// raised `RetryConfig` gives the throttle its retries. +/// +/// `http_override` lets tests drive the STS client through an in-memory stub; production passes +/// `None`. +async fn build_provider( + cfg: &WebIdentityConfig, + http_override: Option<SharedHttpClient>, +) -> Arc<dyn ProvideCredentials> { + let mut loader = aws_config::defaults(BehaviorVersion::latest()) + .retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts)); + if let Some(http) = http_override { + loader = loader.http_client(http); + } + let sdk = loader.load().await; + Arc::new(web_identity_provider_from( + cfg, + aws_sdk_sts::Client::new(&sdk), + )) +} + +/// Assembles the provider from an STS client. Split out so tests can supply a client built with an +/// in-memory HTTP stub while sharing the identity wiring with production. +fn web_identity_provider_from( + cfg: &WebIdentityConfig, + sts: aws_sdk_sts::Client, +) -> WebIdentityStsProvider { + WebIdentityStsProvider { + sts, + role_arn: cfg.role_arn.clone(), + token_file: cfg.token_file.clone(), + session_name: session_name(), + } +} + +/// STS `AssumeRoleWithWebIdentity` session name. Honors `AWS_ROLE_SESSION_NAME` first, matching the +/// default chain, so a trust policy conditioned on `sts:RoleSessionName` keeps working after the +/// take-over engages; otherwise falls back to a stable prefix plus a timestamp. +fn session_name() -> String { + if let Some(name) = non_empty_env("AWS_ROLE_SESSION_NAME") { + return name; + } + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + format!("comet-web-identity-{nanos}") +} + +/// A web-identity-only credential provider: it reads the projected token and calls STS +/// `AssumeRoleWithWebIdentity` on `sts`, and does nothing else. No credential chain, so a throttle +/// that outlasts the STS client's retries returns an error rather than a lower-privilege identity. +#[derive(Debug)] +struct WebIdentityStsProvider { + sts: aws_sdk_sts::Client, + role_arn: String, + token_file: String, + session_name: String, +} + +impl WebIdentityStsProvider { + async fn resolve(&self) -> Result<Credentials, CredentialsError> { + let token = std::fs::read_to_string(&self.token_file).map_err(|e| { + CredentialsError::provider_error(format!( + "reading web identity token file {}: {e}", + self.token_file + )) + })?; + let response = self + .sts + .assume_role_with_web_identity() + .role_arn(&self.role_arn) + .role_session_name(&self.session_name) + .web_identity_token(token.trim()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let creds = response.credentials().ok_or_else(|| { + CredentialsError::provider_error( + "STS AssumeRoleWithWebIdentity response had no credentials", + ) + })?; + let expiration = creds.expiration(); + let expiry = SystemTime::UNIX_EPOCH + .checked_add(Duration::new( + expiration.secs().max(0) as u64, + expiration.subsec_nanos(), + )) + .ok_or_else(|| { + CredentialsError::provider_error("STS credential expiry is out of range") + })?; + Ok(Credentials::new( + creds.access_key_id(), + creds.secret_access_key(), + Some(creds.session_token().to_string()), + Some(expiry), + "CometWebIdentity", + )) + } +} + +impl ProvideCredentials for WebIdentityStsProvider { + fn provide_credentials<'a>(&'a self) -> creds_future::ProvideCredentials<'a> + where + Self: 'a, + { + creds_future::ProvideCredentials::new(self.resolve()) + } +} + +/// The credential provider handed to opendal via `CustomAwsCredentialLoader` (the Iceberg path). +/// Holds only the cheap config plus a lazily resolved handle to the process-wide shared entry, so +/// the per-request path skips the registry lock after the first fetch. +#[derive(Debug)] +pub struct WebIdentityCredentialProvider { + config: WebIdentityConfig, + entry: tokio::sync::OnceCell<Arc<SharedEntry>>, +} + +impl WebIdentityCredentialProvider { + pub fn new(config: WebIdentityConfig) -> Self { + Self { + config, + entry: tokio::sync::OnceCell::new(), + } + } + + /// Resolves (once per provider) the shared entry for this identity. The entry itself is shared + /// process-wide via the registry; this just memoizes the lookup so repeated fetches avoid the + /// registry lock and the per-call `EntryKey` allocation. + async fn entry(&self) -> &Arc<SharedEntry> { + self.entry.get_or_init(|| shared_entry(&self.config)).await + } +} + +impl IcebergProvideCredential for WebIdentityCredentialProvider { + type Credential = IcebergAwsCredential; + + async fn provide_credential( + &self, + _ctx: &Context, + ) -> reqsign_core::Result<Option<Self::Credential>> { + let entry = self.entry().await; + let cred = entry + .credentials() + .await + .map_err(|e| ReqsignError::new(ReqsignErrorKind::CredentialInvalid, e))?; + + // Report the jittered refresh deadline (true expiry minus min_ttl minus jitter) as the + // expiry opendal caches against, so opendal refreshes when our own cache would, and each + // executor's refresh is spread out rather than synchronized. + let expires_in = match cred.expiry() { Review Comment: I think there's a window before every refresh where all signing fails. `provide_credential` reports `expiry - min_ttl - jitter` as `expires_in`, and `fresh()` keeps returning the cached credential right up to that same instant. reqsign's `Signer` only caches a credential while its `expires_in` is more than 120s out (`Credential::is_valid` in `reqsign-aws-core`). After that it calls the loader on every request, and it requires `expires_in > now + 10s` before signing (`validate_refreshed_credential` in `reqsign-core`, `CREDENTIAL_OPERATION_HEADROOM` in `reqsign-aws-v4`). So for the last 10 seconds before our deadline we hand back a credential that reqsign rejects with `refreshed signing credential expires before the requested operation deadline`. opendal wraps that as a permanent `signing http request` error, so `RetryLayer` doesn't retry it. I reproduced this by signing through a real `Signer` and `RequestSigner`. With a cached credential whose true expiry is `min_ttl + 5s` away, signing fails and no STS call is made. With 60s of headroom it succeeds. With a one-hour session that's a 10 second outage roughly every 55 minutes on every executor, for native Iceberg reads and writes. The same arithmetic puts the deadline in the past when `minTtlSeconds` is at or above the session length, so every request fails. `minTtlSeconds=0` has the same window at the true expiry, and the one-time warning doesn't fire for it. Could we make sure we never return a credential whose reported expiry is inside reqsign's margins, either by reporting the real STS expiry like the bridge does, or by adding reqsign's 120s to the jittered deadline we report? Could we also add a test that signs through `reqsign_core::Signer`? `reqsign-aws-v4` is already in the lock, so it would just be a dev-dependency. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
