parthchandra commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4096798096
########## 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: You're right, this was a real ~10s outage every hour. Fixed by reporting the credential's real expiry instead of an early one, so reqsign never sees it as "about to expire." Our own cache now refreshes at the 120s mark (min_ttl floored to 120), so reqsign always gets a fresh credential when it reloads. Added a test that signs through a real reqsign Signer to prove there's no dead zone. ########## 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: Fixed. The "standing aside" info line is now logged once per process per reason, like the existing once-per-process warning in credential_bridge.rs. ########## 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: Fixed. The docs now say reads and writes. The retry sentence now says the throttle becomes an error that Spark's task retry picks up. And the prefix note now says a bare key does reach the reader but isn't looked up, rather than being "dropped." ########## 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: Good catch. We now stand aside when there's no AWS_REGION/AWS_DEFAULT_REGION, so a non-EKS setup that relies on the default global-STS-endpoint fallback keeps working. EKS sets the region, so the take-over still runs there. Added a test. ########## 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: Fixed. enabled is now parsed case-insensitively, so False/FALSE works, and an unrecognized value warns once instead of silently leaving the feature on. Added a test. ########## 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: Fixed all three. We now format the STS error with DisplayErrorContext so the real cause (like Throttling) shows, log it once when it happens, and gave the provider a small hand-written Debug that prints only the role, not the whole STS client and key. ########## 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: Fixed. If a refresh fails but the cached credential is still safely usable, we now keep serving it instead of failing the read. We only return an error once the credential is too close to expiry. -- 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]
