sunchao commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4066405556
########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1015 @@ +// 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 uses the AWS SDK `WebIdentityTokenCredentialsProvider`, whose STS +//! client retries throttling with exponential backoff + jitter. `max_attempts` is +//! configurable (default higher than the SDK's default of 3). +//! 2. No silent downgrade. The provider is web-identity ONLY -- there is no IMDS/instance-role +//! fallback -- so a transient throttle surfaces as a retryable 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::provider_config::ProviderConfig; +use aws_config::retry::RetryConfig; +use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider; +use aws_config::BehaviorVersion; +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 actual region + /// resolution is done by the AWS SDK's `with_default_region`. + 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 a web-identity-only AWS SDK credential provider with STS retry raised to +/// `cfg.max_attempts`. No IMDS/instance-role fallback is wired in, so a throttle that outlasts the +/// retries errors instead of downgrading. +async fn build_provider(cfg: &WebIdentityConfig) -> Arc<dyn ProvideCredentials> { + let provider_config = base_provider_config(cfg).await; + let provider = WebIdentityTokenCredentialsProvider::builder() + .configure(&provider_config) + .build(); + Arc::new(provider) +} + +/// The `ProviderConfig` shared by the production build and the tests: it resolves region, FIPS and +/// dual-stack the same way the AWS default chain would (environment then profile), then raises the +/// STS retry budget. Forwarding `use_fips` / `use_dual_stack` matters because a bare +/// `with_default_region()` leaves them unset -- the STS client would then ignore +/// `AWS_USE_FIPS_ENDPOINT` / `AWS_USE_DUALSTACK_ENDPOINT` and hit the standard endpoint, which can +/// break access from restricted networks. Tests attach an in-memory HTTP client to the result. +async fn base_provider_config(cfg: &WebIdentityConfig) -> ProviderConfig { + let sdk = aws_config::defaults(BehaviorVersion::latest()).load().await; + ProviderConfig::without_region() + .with_region(sdk.region().cloned()) + .with_use_fips(sdk.use_fips()) + .with_use_dual_stack(sdk.use_dual_stack()) + .with_retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts)) Review Comment: ### Correctness [P2] Preserve profile-configured STS endpoint URLs This creates a fresh `ProviderConfig`, so the profile parsed by `sdk.load()` is discarded. In the pinned SDK, `ProviderConfig::client_config()` uses only its already-parsed profile when constructing the STS service config. With `AWS_PROFILE` unset, no shared credentials file, and an STS `endpoint_url` configured in `AWS_CONFIG_FILE` or `~/.aws/config` (for example through `[default] services = private` and `[services private] sts = ...`), `configured_profile()` permits takeover but this client uses the standard STS endpoint. The previous Parquet default chain honors that profile endpoint. This breaks credential acquisition where only the configured STS endpoint is reachable. Could the provider preserve profile-derived endpoint configuration and add a config-only endpoint regression test? The FIPS/dual-stack copies do fix their separate flags, and endpoint environment variables still work. -- 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]
