sunchao commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4100561993
########## native/core/src/cloud/s3/web_identity.rs: ########## @@ -0,0 +1,1491 @@ +// 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 cache. One assumed-role credential is cached per process, keyed by identity +//! (role_arn, token_file, region) and the resolved settings, and shared across all reader +//! threads and scans that resolve to the same key -- so a startup burst makes one STS call per +//! executor rather than one per reader thread. A failed refresh keeps serving the still-valid +//! cached credential and 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_sdk_sts::error::DisplayErrorContext; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential; +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. A bare key without the `s3.` prefix still reaches +/// the catalog bag (Comet forwards the unfiltered FileIO properties), but the lookup below adds the +/// prefix, so only the `s3.`-spelled key takes 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 DEFAULT_ENABLED: bool = true; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; +const DEFAULT_MIN_TTL_SECS: u64 = 300; + +/// reqsign's signer treats a credential as needing refresh once it is within 120s of its reported +/// expiry (`Credential::is_valid` in reqsign-aws-v4) and refuses to sign within 10s of it +/// (`CREDENTIAL_OPERATION_HEADROOM`). We must (a) report the real STS expiry so the signer never +/// sees a credential that is nominally inside those margins, and (b) refresh our own cache at or +/// before the signer's 120s point so that when the signer asks us to reload it gets a fresh +/// credential. So `min_ttl` is floored to this value. +const REQSIGN_REFRESH_MARGIN: Duration = Duration::from_secs(120); + +/// 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`. The STS client's region comes from the resolved + /// `SdkConfig`; we also require it to be present before taking over (see `take_over_if_irsa`), + /// because a web-identity STS client with no region silently fails. + region: Option<String>, + max_attempts: u32, + /// Refresh margin for our own cache. Floored to `REQSIGN_REFRESH_MARGIN` so our cache refreshes + /// at or before the point reqsign asks the loader to reload, avoiding a signing dead zone. + min_ttl: 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 the catalog + /// property bag. 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_enabled(resolve(KEY_ENABLED)) { + return None; + } + let min_ttl = Duration::from_secs(parse_setting( + resolve(KEY_MIN_TTL_SECS), + DEFAULT_MIN_TTL_SECS, + )) + .max(REQSIGN_REFRESH_MARGIN); + 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, + }) + } + + 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, + } + } +} + +/// 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, +} + +/// The shared, cached credential for one identity. `provider` resolves credentials via STS; +/// `cached` holds the last credential. `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, +} + +impl SharedEntry { + /// Returns the cached credential if it is still fresh, i.e. it does not expire within `min_ttl`. + fn fresh(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + if self.expires_within(cred, self.min_ttl) { + None + } else { + Some(cred.clone()) + } + } + + /// True if `cred` expires within `margin` from now. A credential with no reported expiry never + /// does. + fn expires_within(&self, cred: &Credentials, margin: Duration) -> bool { + match cred.expiry() { + Some(expiry) => expiry <= SystemTime::now() + margin, + 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)")) + } + + /// The cached credential if it is still safely signable -- outside reqsign's refresh margin -- + /// even though it is inside our own (larger) refresh margin. Used to keep serving reads when a + /// refresh fails but the current credential still has real headroom. + fn still_signable(&self) -> Option<Credentials> { + let guard = self.cached.read().unwrap(); + let cred = guard.as_ref()?; + (!self.expires_within(cred, REQSIGN_REFRESH_MARGIN)).then(|| cred.clone()) + } + + /// Fetches a fresh credential, refreshing from STS at most once at a time. On a refresh error we + /// keep serving the cached credential while it is still safely signable; only once it is too + /// close to expiry does the error propagate. We never fall back to a lower-privilege identity. + /// The failure is recorded 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 self.still_signable().ok_or(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 self.still_signable().ok_or(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: {}", + DisplayErrorContext(&e) + ); + log::warn!("Comet web-identity credential refresh failed: {err}"); + *self.last_failure.write().unwrap() = Some((Instant::now(), err.clone())); + // A refresh failure while the current credential is still safely signable must not + // fail reads that would have worked; keep serving it and let the cooldown throttle + // retries. + self.still_signable().ok_or(err) + } + } + } + + /// Warns once if a freshly fetched credential already falls inside our 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(cred, self.min_ttl) && 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; + 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, + }); + + 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: [P2] Could the takeover preserve `AWS_STS_REGIONAL_ENDPOINTS=legacy`, or stand aside when it is explicitly requested? With both IRSA variables set, `AWS_REGION=us-west-2`, and no explicit credentials/profile, takeover engages. The previous reqsign provider honors `legacy` and calls `sts.amazonaws.com`, but this SDK configuration ignores that setting and calls `sts.us-west-2.amazonaws.com`. An executor whose STS egress permits only the configured global endpoint therefore loses credential acquisition for native Iceberg reads and writes after upgrading. Preserve the endpoint choice and add a regression test through `build_provider`. Evidence: An offline in-memory HTTP probe compared the pinned reqsign web-identity provider with this head’s production `build_provider` using the same synthetic IRSA configuration. It printed `legacy endpoint: base host=sts.amazonaws.com; head uri=https://sts.us-west-2.amazonaws.com/`. The probe also confirmed `take_over_if_irsa(false, |_| None)` returns a provider. Reproduce with `cargo test --offline --manifest-path /tmp/comet6025-dbb0-probe/Cargo.toml review_legacy_sts_endpoint_comparison -- --nocapture`. All harness dependencies match versions in the PR lockfile. No AWS request was sent. -- 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]
