andygrove commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4073736586
########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1149 @@ +// 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. +//! +//! The same struct is exposed as both `object_store::CredentialProvider` (raw Parquet path) and +//! reqsign's `ProvideCredential` (Iceberg via opendal / `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 async_trait::async_trait; +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 iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +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 read from the Iceberg catalog property bag. A non-`s3.`/`client.` prefix keeps them +/// from being forwarded into opendal's S3 config (see `iceberg_common::STORAGE_PROPERTY_PREFIXES`). +const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.s3.credentials.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.s3.credentials.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` is the AWS SDK web-identity provider +/// built once; `cached` holds the last credential; `refresh_jitter` is drawn once per process so +/// each executor refreshes at a slightly different time. `last_failure_at` coalesces a burst of +/// readers that hit a persistent throttle into a single STS call. +#[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. Waiters within `FAILURE_COOLDOWN` of this get the failure + /// without re-calling STS, so a failed burst costs one call rather than one per reader. + last_failure_at: RwLock<Option<Instant>>, + 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()?; + match cred.expiry() { + Some(expiry) => { + if expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter { + None + } else { + Some(cred.clone()) + } + } + // No expiry reported: keep it. Web-identity credentials normally carry one. + None => Some(cred.clone()), + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, so callers can bail out + /// instead of piling another assume-role call onto a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { Review Comment: This returns a fixed string, so the real STS error only reaches whichever reader lost the race. In the burst this PR is about, that means most readers see "backing off before retrying STS" with no way to tell a throttle from a bad token or a trust-policy rejection. Could `last_failure_at` hold `Option<(Instant, String)>` and replay the recorded error alongside the backoff note? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1149 @@ +// 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. +//! +//! The same struct is exposed as both `object_store::CredentialProvider` (raw Parquet path) and +//! reqsign's `ProvideCredential` (Iceberg via opendal / `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 async_trait::async_trait; +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 iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +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 read from the Iceberg catalog property bag. A non-`s3.`/`client.` prefix keeps them +/// from being forwarded into opendal's S3 config (see `iceberg_common::STORAGE_PROPERTY_PREFIXES`). +const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.s3.credentials.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.s3.credentials.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` is the AWS SDK web-identity provider +/// built once; `cached` holds the last credential; `refresh_jitter` is drawn once per process so +/// each executor refreshes at a slightly different time. `last_failure_at` coalesces a burst of +/// readers that hit a persistent throttle into a single STS call. +#[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. Waiters within `FAILURE_COOLDOWN` of this get the failure + /// without re-calling STS, so a failed burst costs one call rather than one per reader. + last_failure_at: RwLock<Option<Instant>>, + 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()?; + match cred.expiry() { + Some(expiry) => { + if expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter { + None + } else { + Some(cred.clone()) + } + } + // No expiry reported: keep it. Web-identity credentials normally carry one. + None => Some(cred.clone()), + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, so callers can bail out + /// instead of piling another assume-role call onto a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let at = (*self.last_failure_at.read().unwrap())?; + (at.elapsed() < FAILURE_COOLDOWN).then(|| { + "web-identity credential refresh failed recently; backing off before retrying STS" + .to_string() + }) + } + + /// 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.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure_at.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + *self.last_failure_at.write().unwrap() = Some(Instant::now()); + Err(format!("web-identity assume-role failed: {e}")) + } + } + } +} + +/// 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).await; + // Draw the refresh jitter once. subsec_nanos at build time differs across processes, so this + // seeds a per-executor offset even before rand is consulted. + 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_at: 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. +async fn build_provider(cfg: &WebIdentityConfig) -> Arc<dyn ProvideCredentials> { + let sdk = aws_config::defaults(BehaviorVersion::latest()) + .retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts)) + .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. STS requires one; it is informational only, so a +/// stable prefix plus a timestamp keeps sessions distinguishable in CloudTrail. +fn session_name() -> String { Review Comment: The default chain reads `AWS_ROLE_SESSION_NAME` and only generates a name when it is unset (`aws-config-1.12.0/src/web_identity_token.rs:78,138`). This always generates one. Could `session_name()` check that env var first and fall back to the generated name? The reason I care is that trust policies conditioned on `sts:RoleSessionName` are fairly common, and since the take-over is on by default, an executor that was satisfying that condition before would start getting `AccessDenied` after an upgrade. That is the same class of surprise the other stand-asides were added to prevent, and it is the one case they do not cover. ########## native/core/src/execution/operators/iceberg_common.rs: ########## @@ -170,7 +171,18 @@ fn build_s3_credential_loader( .map(|s| s.trim()) .filter(|s| !s.is_empty()) else { - return Ok(None); + // No explicit Comet provider class. On EKS/IRSA, take over credential resolution with the + // Comet web-identity provider (retry on STS throttle, no node-role downgrade, shared + // jittered cache) instead of leaving it to opendal's default reqsign chain, which + // downgrades to the node instance role under throttling. Non-IRSA setups (static keys, + // env, profile) keep the default chain. We also defer to any credentials the user + // configured explicitly in the catalog (static keys or an assume-role arn) -- explicit + // config always wins, same as a named provider class does. + let explicit = has_explicit_s3_credentials(catalog_properties); + return Ok( + take_over_if_irsa(explicit, |key| catalog_properties.get(key).cloned()) Review Comment: How did you verify that the Iceberg spelling arrives? On the Parquet side I can follow it, since `NativeConfig.extractObjectStoreOptions` forwards every `fs.s3a.*` key verbatim. Here a bare `comet.s3.credentials.webIdentity.*` catalog property can only reach `catalog_properties` through `fileIOProperties`, because `hadoopToIcebergS3Properties` drops unmapped suffixes. The one existing Comet catalog property is spelled `s3.comet.credential.provider.class` (`iceberg_common.rs:34`) and is covered end to end in `CometS3CredentialBridgeSuite:53`, so the only evidence we have points the other way. The comment on the key constants says the non-`s3.` prefix keeps them out of opendal's S3 config, but `s3.comet.credential.provider.class` is already forwarded there and causes no trouble, so I do not think that reason holds up. Would you either look up `s3.{key}` here and match the existing convention, or add a case to `CometS3CredentialBridgeSuite` showing a bare key arriving? This matters more to me than a naming preference, because `enabled=false` is the only way to turn this off and the Iceberg path is where the reported failure lives. If the key does not arrive, nobody can opt out. While the spelling is still open, could we also settle `comet.s3.credentials.` against the existing `comet.credential.`? Singular, and no `s3` segment under `fs.s3a.`. These are permanent surface and easier to change now than later. ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1149 @@ +// 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. +//! +//! The same struct is exposed as both `object_store::CredentialProvider` (raw Parquet path) and +//! reqsign's `ProvideCredential` (Iceberg via opendal / `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 async_trait::async_trait; +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 iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +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 read from the Iceberg catalog property bag. A non-`s3.`/`client.` prefix keeps them +/// from being forwarded into opendal's S3 config (see `iceberg_common::STORAGE_PROPERTY_PREFIXES`). +const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.s3.credentials.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.s3.credentials.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` is the AWS SDK web-identity provider +/// built once; `cached` holds the last credential; `refresh_jitter` is drawn once per process so +/// each executor refreshes at a slightly different time. `last_failure_at` coalesces a burst of +/// readers that hit a persistent throttle into a single STS call. +#[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. Waiters within `FAILURE_COOLDOWN` of this get the failure + /// without re-calling STS, so a failed burst costs one call rather than one per reader. + last_failure_at: RwLock<Option<Instant>>, + 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()?; + match cred.expiry() { + Some(expiry) => { + if expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter { + None + } else { + Some(cred.clone()) + } + } + // No expiry reported: keep it. Web-identity credentials normally carry one. + None => Some(cred.clone()), + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, so callers can bail out + /// instead of piling another assume-role call onto a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let at = (*self.last_failure_at.read().unwrap())?; + (at.elapsed() < FAILURE_COOLDOWN).then(|| { + "web-identity credential refresh failed recently; backing off before retrying STS" + .to_string() + }) + } + + /// 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.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure_at.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + *self.last_failure_at.write().unwrap() = Some(Instant::now()); + Err(format!("web-identity assume-role failed: {e}")) + } + } + } +} + +/// 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).await; + // Draw the refresh jitter once. subsec_nanos at build time differs across processes, so this + // seeds a per-executor offset even before rand is consulted. + 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_at: 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. +async fn build_provider(cfg: &WebIdentityConfig) -> Arc<dyn ProvideCredentials> { + let sdk = aws_config::defaults(BehaviorVersion::latest()) + .retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts)) + .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. STS requires one; it is informational only, so a +/// stable prefix plus a timestamp keeps sessions distinguishable in CloudTrail. +fn session_name() -> String { + 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 `object_store` (Parquet) and, via +/// `CustomAwsCredentialLoader`, to opendal (Iceberg). 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 + } +} + +#[async_trait] +impl CredentialProvider for WebIdentityCredentialProvider { + type Credential = AwsCredential; + + async fn get_credential(&self) -> object_store::Result<Arc<AwsCredential>> { + let cred = + self.entry() + .await + .credentials() + .await + .map_err(|e| object_store::Error::Generic { + store: "S3", + source: e.into(), + })?; + Ok(Arc::new(AwsCredential { + key_id: cred.access_key_id().to_string(), + secret_key: cred.secret_access_key().to_string(), + token: cred.session_token().map(|s| s.to_string()), + })) + } +} + +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. It does +/// so only when the caller has no explicit credentials configured and IRSA is detected; otherwise +/// the caller keeps its default chain. `resolve` reads a bare setting key (e.g. `KEY_MAX_ATTEMPTS`) +/// from whichever config bag the caller owns. Both scan paths share this one decision. +/// +/// 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). Both the AWS SDK +/// default chain (Parquet) and opendal/reqsign (Iceberg) resolve Environment -> Profile -> +/// WebIdentity, so taking over in those cases would silently switch identity from the user's chosen +/// source to the service-account role -- and would also drop profile-configured settings such as a +/// custom STS endpoint that a hand-built `ProviderConfig` cannot reconstruct here. +pub fn take_over_if_irsa<F>( + explicit_credentials: bool, + resolve: F, +) -> Option<WebIdentityCredentialProvider> +where + F: Fn(&str) -> Option<String>, +{ + if explicit_credentials || explicit_env_credentials() || configured_profile() { + return None; + } + WebIdentityConfig::detect_with(resolve).map(WebIdentityCredentialProvider::new) +} + +/// True if explicit static credentials are present in the environment. These outrank web-identity +/// in every default chain, so the take-over must not shadow them. +fn explicit_env_credentials() -> bool { + non_empty_env("AWS_ACCESS_KEY_ID").is_some() && non_empty_env("AWS_SECRET_ACCESS_KEY").is_some() +} + +/// True if a profile is configured that the default chain would consult ahead of web-identity: +/// `AWS_PROFILE` is set, or a shared credentials file (`AWS_SHARED_CREDENTIALS_FILE`, else +/// `~/.aws/credentials`) or a config file (`AWS_CONFIG_FILE`, else `~/.aws/config`) exists. We defer +/// to the default chain in all of these because it resolves profile credentials AND +/// profile-configured settings (region, endpoint URLs, FIPS/dual-stack) that a hand-built +/// `ProviderConfig` cannot reconstruct here. Conservative by design: standing aside just falls back +/// to the pre-existing default-chain behavior, so it is never worse than before. On an EKS/IRSA pod +/// none of these are normally present, so the take-over still applies there. +fn configured_profile() -> bool { Review Comment: Could we log the take-over decision somewhere? There are no log statements anywhere in this module, while the analogous credential fallback in `iceberg_common.rs:198` warns and `s3.rs` uses `debug!` and `error!` throughout. The case that made me notice is this function. An `~/.aws/config` that only sets a region makes it return true, but `ProfileFileCredentialsProvider` would have returned `CredentialsNotLoaded` for that file and the chain would have continued on to web identity anyway. So a base image that ships a region-only config file gets none of this protection, and the symptom is the same 403 storm with nothing in the logs to say why the fix did not engage. I am not asking you to narrow the check. Standing aside is safe and I think the conservative choice is right. I would just like an operator to be able to find out which branch they took. A `debug!` when the provider installs, and an `info!` naming the reason when IRSA is detected but we stand aside, would do it. ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1149 @@ +// 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. +//! +//! The same struct is exposed as both `object_store::CredentialProvider` (raw Parquet path) and +//! reqsign's `ProvideCredential` (Iceberg via opendal / `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 async_trait::async_trait; +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 iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +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 read from the Iceberg catalog property bag. A non-`s3.`/`client.` prefix keeps them +/// from being forwarded into opendal's S3 config (see `iceberg_common::STORAGE_PROPERTY_PREFIXES`). +const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.s3.credentials.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.s3.credentials.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( Review Comment: `max_attempts` goes through `parse_u32`, which rejects zero, but `min_ttl` uses plain `parse_setting` with no bounds. Two things fall out of that. `minTtlSeconds=0` is accepted and defeats the point of the refresh margin, since a credential can then be handed out with milliseconds left. And any value at or above the STS session lifetime makes `fresh()` return `None` on every call, so every request takes the refresh lock and issues its own `AssumeRoleWithWebIdentity`. The `expires_in` reported to opendal goes the same way, with the deadline landing in the past. The defaults have a wide margin so this only bites on misconfiguration, but it is a documented knob with no stated upper bound, and the failure mode is a worse version of the problem this PR fixes. Could it clamp, or warn once when a freshly fetched credential is already outside the freshness window? ########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1149 @@ +// 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. +//! +//! The same struct is exposed as both `object_store::CredentialProvider` (raw Parquet path) and +//! reqsign's `ProvideCredential` (Iceberg via opendal / `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 async_trait::async_trait; +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 iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +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 read from the Iceberg catalog property bag. A non-`s3.`/`client.` prefix keeps them +/// from being forwarded into opendal's S3 config (see `iceberg_common::STORAGE_PROPERTY_PREFIXES`). +const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled"; +const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts"; +const KEY_MIN_TTL_SECS: &str = "comet.s3.credentials.webIdentity.minTtlSeconds"; +const KEY_JITTER_SECS: &str = "comet.s3.credentials.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` is the AWS SDK web-identity provider +/// built once; `cached` holds the last credential; `refresh_jitter` is drawn once per process so +/// each executor refreshes at a slightly different time. `last_failure_at` coalesces a burst of +/// readers that hit a persistent throttle into a single STS call. +#[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. Waiters within `FAILURE_COOLDOWN` of this get the failure + /// without re-calling STS, so a failed burst costs one call rather than one per reader. + last_failure_at: RwLock<Option<Instant>>, + 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()?; + match cred.expiry() { + Some(expiry) => { + if expiry <= SystemTime::now() + self.min_ttl + self.refresh_jitter { + None + } else { + Some(cred.clone()) + } + } + // No expiry reported: keep it. Web-identity credentials normally carry one. + None => Some(cred.clone()), + } + } + + /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`, so callers can bail out + /// instead of piling another assume-role call onto a throttled STS. + fn in_failure_cooldown(&self) -> Option<String> { + let at = (*self.last_failure_at.read().unwrap())?; + (at.elapsed() < FAILURE_COOLDOWN).then(|| { + "web-identity credential refresh failed recently; backing off before retrying STS" + .to_string() + }) + } + + /// 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.cached.write().unwrap() = Some(cred.clone()); + *self.last_failure_at.write().unwrap() = None; + Ok(cred) + } + Err(e) => { + *self.last_failure_at.write().unwrap() = Some(Instant::now()); + Err(format!("web-identity assume-role failed: {e}")) + } + } + } +} + +/// 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).await; + // Draw the refresh jitter once. subsec_nanos at build time differs across processes, so this Review Comment: This comment describes an earlier version. There is no `subsec_nanos` in the jitter draw any more, it is just `rand::rng().random_range(..)`. The `SharedEntry` doc on line 167 also says `refresh_jitter` is drawn once per process when it is drawn once per entry. That amounts to the same thing for a single identity, but the key now includes the tuning knobs, so two catalogs with different settings get different jitters. Could you fix both? -- 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]
