laskoviymishka commented on code in PR #2932: URL: https://github.com/apache/iceberg-rust/pull/2932#discussion_r4326462087
########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// 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. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) Review Comment: This drops expired entries before it picks the longest prefix. So once a narrow `s3://bucket/table` credential expires while a broader entry is still valid, `cache_decision` serves the broader one and never refreshes the narrow prefix. The catalog scoped that prefix separately for a reason. After an idle gap longer than the narrow credential's lifetime we'd sign with a credential that may carry different permissions, and the caller gets a 403 with no refresh attempted. I'd pick the longest covering prefix whether or not it has expired, and return `Refresh` if it has. Java's `clientForStoragePath` uses the same longest-prefix rule without the fallback. `unexpired_match_selects_longest_unexpired_prefix` currently pins the fallback as intended, so that test flips. ########## crates/catalog/rest/src/types.rs: ########## @@ -244,6 +255,27 @@ pub struct StorageCredential { pub config: HashMap<String, String>, } +impl std::fmt::Debug for StorageCredential { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("StorageCredential") + .field("prefix", &self.prefix) + .field("config_keys", &self.config.keys().collect::<Vec<_>>()) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +/// Response from the table credentials endpoint +/// (`GET /v1/{prefix}/namespaces/{namespace}/tables/{table}/credentials`). +/// +/// Returns freshly vended storage credentials so clients can refresh temporary +/// credentials before they expire. +pub struct LoadCredentialsResponse { Review Comment: This becomes part of the crate's public API, with a `pub` field and five derives. Its only consumer is the private `credential` module, which only needs `Deserialize`. I'd make it `pub(crate)` and drop the extra derives. Once it ships, we're committed to keeping its shape stable. ########## crates/iceberg/src/io/file_io.rs: ########## @@ -123,14 +146,38 @@ impl FileIO { /// /// All storage configuration properties are included in the serialized representation. These /// properties may contain credentials or other sensitive values, so the returned bytes must be - /// protected in transit and at rest by the application embedding this crate. + /// protected in transit and at rest by the application embedding this crate. A serialized + /// credential provider may likewise carry catalog authentication and vended credentials. /// /// Storage factories are serialized through [`typetag`](https://docs.rs/typetag). Third-party /// factories must use `#[typetag::serde]` on their [`StorageFactory`] implementation. + /// + /// A credential provider is serialized as the + /// [`StorageCredentialProviderFactory`] returned by + /// [`StorageCredentialProvider::factory`], and rebuilt on deserialization. A provider that + /// cannot be rebuilt in another process is left out with a warning, as with + /// [`FileIO::without_credential_provider`]: the deserialized `FileIO` uses the credentials in + /// its configuration, which are not refreshed. Storage factories that hold their own + /// credential sources, such as a custom AWS credential loader, still fail to serialize, as + /// leaving them out would leave no credentials at all. pub fn serialize_all(&self) -> Result<Vec<u8>> { + let credential_provider = self.credential_provider.as_ref().and_then(|provider| { + provider + .factory() + .and_then(|factory| Ok(serde_json::to_value(factory)?)) + .inspect_err(|error| { + static WARNED: std::sync::Once = std::sync::Once::new(); Review Comment: Following up on @mbutrovich's thread about warning instead of erroring: a process-wide `Once` means that only the first `FileIO` to lose its provider gets a warning. On a long-lived worker, every later one downgrades silently, and the first sign of trouble is a 403 hours later. I'd dedupe per factory type or per error instead of with one static. Better still, add a way for callers to check that the provider survived, such as `FileIO::has_credential_provider()`. ########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// 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. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) + .max_by_key(|(_, entry)| entry.credential.prefix().len()) + .map(|(index, _)| index) + } + + /// Whether a failed prefix covering `path` is narrower than `credential`'s. + fn masks(&self, path: &str, credential: Option<&StorageCredential>) -> bool { + let prefix_len = credential.map(|credential| credential.prefix().len()); + self.failed_prefixes.iter().any(|failed| { + prefix_covers(failed, path) && prefix_len.is_none_or(|len| failed.len() > len) + }) + } + + /// Cache `entry`, replacing any entry with the same prefix, and return its + /// index. + fn insert(&mut self, entry: CachedEntry) -> usize { + let prefix = entry.credential.prefix().to_owned(); + self.entries + .retain(|cached| cached.credential.prefix() != prefix); + self.entries.push(entry); + self.entries.len() - 1 + } + + fn record_success(&mut self) { + self.consecutive_failures = 0; + self.retry_not_before = None; + self.last_failure = None; + } + + fn record_failure(&mut self, error: &Error) { + self.consecutive_failures = self.consecutive_failures.saturating_add(1); + self.retry_not_before = + Instant::now().checked_add(failure_backoff(self.consecutive_failures)); + self.last_failure = Some(error.to_string()); + } + + /// Replace cached credentials with the fetched ones, prefix by prefix, + /// and record the prefixes whose credential was invalid in this response. + /// + /// An unexpired cached credential survives when the response carries no + /// valid replacement for its prefix, so an absent or malformed entry never + /// evicts a usable credential. Returns the prefixes the response replaced. + fn merge(&mut self, fetched: Vec<CachedEntry>, errors: &[CredentialError]) -> HashSet<String> { + let fetched_prefixes = fetched + .iter() + .map(|entry| entry.credential.prefix().to_owned()) + .collect::<HashSet<_>>(); + // A response names every prefix of the table, so a prefix it omits is + // no longer scoped separately: a fetched credential that covers it + // replaces its cached one. A prefix whose credential is invalid keeps + // its cached one. + let failed_prefixes = errors + .iter() + .map(|error| error.prefix.clone()) + .collect::<HashSet<_>>(); + self.entries.retain(|entry| { + let prefix = entry.credential.prefix(); + failed_prefixes.contains(prefix) + || !fetched + .iter() + .any(|fetched| fetched.credential.covers(prefix)) + }); + self.entries.extend(fetched); + self.failed_prefixes = failed_prefixes; + fetched_prefixes + } +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex<CacheState>, + /// Initial property credentials whose scope cannot be represented by a + /// single leading URI prefix, keyed according to the cloud strategy. + keyed_seeds: HashMap<String, CachedEntry>, + /// Only one caller fetches at a time. The cache lock is deliberately + /// separate so other callers can keep using an unexpired credential while + /// the refresh is in flight. + refresh: Mutex<()>, +} + +impl ConfiguredCloud { + fn new( + cloud: &'static CloudRefresh, + endpoint: String, + entries: Vec<CachedEntry>, + keyed_seeds: HashMap<String, CachedEntry>, + ) -> Self { + Self { + cloud, + endpoint, + cache: Mutex::new(CacheState::new(entries)), + keyed_seeds, + refresh: Mutex::new(()), + } + } + + /// Configure `cloud` from the FileIO properties, or `None` when they + /// advertise no enabled refresh endpoint for it. + fn configure( + cloud: &'static CloudRefresh, + base_uri: &str, + props: &HashMap<String, String>, + ) -> Option<Self> { + let enabled = props + .get(cloud.enabled_key) + .is_none_or(|value| value.eq_ignore_ascii_case("true")) + && cloud + .required_key + .is_none_or(|key| props.get(key).is_some_and(|value| !value.is_empty())); + let endpoint = props + .get(cloud.endpoint_key) + .filter(|endpoint| enabled && !endpoint.is_empty()) + .map(|endpoint| resolve_endpoint(base_uri, endpoint))?; + + let mut entries = Vec::new(); + let mut keyed_seeds = HashMap::new(); + match &cloud.seed_strategy { + SeedStrategy::Flat { prefixes } => { + entries.extend(prefixes.iter().filter_map(|prefix| { + (cloud.parse_credential)(props, prefix.to_string()) + .ok() + .map(|vended| CachedEntry::seed(vended, cloud.jitter_prefetch)) + })); + } + SeedStrategy::Keyed(strategy) => { + keyed_seeds = (strategy.parse_credentials)(props) + .into_iter() + .map(|(key, credential)| { + (key, CachedEntry::seed(credential, cloud.jitter_prefetch)) + }) + .collect(); + } + } + Some(Self::new(cloud, endpoint, entries, keyed_seeds)) + } + + fn keyed_seed_for_path(&self, path: &str) -> Result<Option<CachedEntry>> { + let SeedStrategy::Keyed(strategy) = &self.cloud.seed_strategy else { + return Ok(None); + }; + let resolved = (strategy.resolve_path)(path)?; + let Some(seed) = resolved + .keys + .iter() + .find_map(|key| self.keyed_seeds.get(key)) + else { + return Ok(None); + }; + let credential = StorageCredential::new(resolved.scope, seed.credential.config().clone()); + // The scope is normalized, e.g. with a lowercase scheme, so it may not + // cover the path as written; the path then fetches its credential. + Ok(credential.covers(path).then_some(CachedEntry { + credential, + expires_at: seed.expires_at, + refresh_at: seed.refresh_at, + })) + } +} + +/// Serializable recipe for rebuilding a [`RestVendedCredentialProvider`] from +/// the FileIO properties, mirroring how Java rebuilds its provider on workers. +#[derive(Clone, Serialize, Deserialize)] +pub(crate) struct RestVendedCredentialProviderFactory { + /// Catalog URI, used to resolve relative refresh endpoints. + catalog_uri: String, + table: TableIdent, + /// The `rest.auth.type` the catalog resolved, so a rebuilt provider uses + /// the same auth manager even when the merged FileIO properties would + /// infer another one. + auth_type: String, + /// Whether `auth_type` was inferred rather than configured. Like Java, + /// a table `token` then authenticates a catalog without auth. + auth_type_inferred: bool, + /// The auth-related part of the unmerged config returned by the table + /// endpoint, from which the auth manager derives a table session. Keeping + /// it separate prevents local FileIO overrides from masking table auth. + table_auth: HashMap<String, String>, +} + +/// Table config keys that may override authentication in a table session. +const TABLE_AUTH_KEYS: &[&str] = &["token"]; + +impl RestVendedCredentialProviderFactory { + pub(crate) fn new( + catalog_uri: impl Into<String>, + table: TableIdent, + auth_type: impl Into<String>, + auth_type_inferred: bool, + table_config: &HashMap<String, String>, + ) -> Self { + Self { + catalog_uri: catalog_uri.into(), + table, + auth_type: auth_type.into(), + auth_type_inferred, + table_auth: table_config + .iter() + .filter(|(key, _)| TABLE_AUTH_KEYS.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + } + } + + /// The table token that authenticates credential requests of a catalog + /// without auth. Java's credential providers infer OAuth2 from a `token` + /// unless `rest.auth.type` is configured. + fn bearer_table_token(&self) -> Option<&str> { + (self.auth_type == AUTH_TYPE_NONE && self.auth_type_inferred) + .then(|| self.table_auth.get("token").map(String::as_str)) + .flatten() + } + + /// Connect to the catalog from the FileIO properties, as the catalog would. + async fn connect(&self, props: &HashMap<String, String>) -> Result<HttpClient> { + let mut config_props = props.clone(); + config_props.insert( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + self.auth_type.clone(), + ); + let config = RestCatalogConfig::builder() + .uri(self.catalog_uri.clone()) + .props(config_props) + .build(); + let auth_manager = load_auth_manager(&config)?; + let client = HttpClient::new(&config)?; + let session = auth_manager + .catalog_session(&client.without_auth_session(), &config.auth_props()) + .await?; + table_client( + &client.with_auth_session(session), + auth_manager.as_ref(), + &self.table, + props, + &self.table_auth, + self.bearer_table_token(), + ) + .await + } +} + +impl std::fmt::Debug for RestVendedCredentialProviderFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RestVendedCredentialProviderFactory") + .field("catalog_uri", &self.catalog_uri) + .field("table", &self.table) + .finish_non_exhaustive() + } +} + +#[typetag::serde(name = "RestVendedCredentialProviderFactory")] +impl StorageCredentialProviderFactory for RestVendedCredentialProviderFactory { + fn build(&self, config: &StorageConfig) -> Result<Arc<dyn StorageCredentialProvider>> { + let provider = + RestVendedCredentialProvider::configure(self, config.props()).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + "FileIO configuration no longer configures vended credentials", + ) + })?; + Ok(Arc::new(RestVendedCredentialProvider { + factory: Some(self.clone()), + ..provider + })) + } +} + +/// Derive the table-scoped client used for credential requests. +/// +/// A `bearer_table_token` authenticates instead of a session from the auth +/// manager. +async fn table_client( + catalog_client: &HttpClient, + auth_manager: &dyn AuthManager, + table: &TableIdent, + props: &HashMap<String, String>, + table_config: &HashMap<String, String>, + bearer_table_token: Option<&str>, +) -> Result<HttpClient> { + let session = match bearer_table_token { + Some(token) => static_token_session(token), + None => { + auth_manager + .table_session( + &catalog_client.without_auth_session(), + table, + table_config, + catalog_client.auth_session(), + ) + .await? + } + }; + catalog_client.for_table(props, session) +} + +/// Serves vended credentials for a table, refreshing them from the REST +/// catalog's table credentials endpoint. +/// +/// Each cloud cache is seeded with the credentials from the initial +/// table properties (when complete) and re-fetched from its endpoint as they +/// near expiry. +pub(crate) struct RestVendedCredentialProvider { + /// Table-scoped catalog client. Set when the catalog builds the provider, + /// and connected on first refresh after deserialization. + client: OnceCell<HttpClient>, + /// Rebuilds this provider in another process. `None` when the catalog + /// authentication cannot be rebuilt from properties. + factory: Option<RestVendedCredentialProviderFactory>, + /// Effective FileIO properties, which supply catalog connection settings. + props: HashMap<String, String>, + /// Query parameters sent with every credentials request. + query_params: Vec<(&'static str, String)>, + /// Independently configured endpoint and cache for each backing cloud. + clouds: Vec<ConfiguredCloud>, +} + +impl RestVendedCredentialProvider { + /// Configure a provider without a catalog connection, or `None` when no + /// supported cloud advertises an enabled refresh endpoint. + fn configure( + factory: &RestVendedCredentialProviderFactory, + props: &HashMap<String, String>, + ) -> Option<Self> { + let clouds = CloudRefresh::SUPPORTED + .iter() + .filter_map(|cloud| ConfiguredCloud::configure(cloud, &factory.catalog_uri, props)) + .collect::<Vec<_>>(); + (!clouds.is_empty()).then(|| Self { + client: OnceCell::new(), + factory: None, + props: props.clone(), + query_params: credentials_query_params(props), + clouds, + }) + } + + fn configured_cloud_for_location(&self, location: &str) -> Option<&ConfiguredCloud> { + self.clouds + .iter() + .find(|configured| configured.cloud.matches_location(location)) + } + + async fn client(&self) -> Result<&HttpClient> { + self.client + .get_or_try_init(|| async { + let factory = self.factory.as_ref().ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "vended credential provider has no catalog connection", + ) + })?; + factory.connect(&self.props).await + }) + .await + } + + /// Fetch fresh credentials from the catalog's credentials endpoint. + async fn fetch(&self, configured: &ConfiguredCloud) -> Result<ParsedCredentials> { + let cloud = configured.cloud; + let client = self.client().await?; + let url = credentials_url(&configured.endpoint, &self.query_params)?; + let request = HttpRequest::build(client.request(Method::GET, url))?; + let response = client.query_catalog(request).await?; + + if !response.status().is_success() { + return Err(unexpected_catalog_error_without_body( + response, + client.disable_header_redaction(), + )); + } + + // Credential responses contain secrets, and serde's errors quote the + // values they reject, so only the error's category and position are + // reported. + let parsed: LoadCredentialsResponse = + serde_json::from_slice(response.body()).map_err(|error| { + Error::new( + ErrorKind::Unexpected, + format!( + "failed to parse vended credential response: {:?} error at line {}, \ + column {}", + error.classify(), + error.line(), + error.column() + ), + ) + })?; + let now = SystemTime::now(); + let mut entries = Vec::new(); + let mut errors = Vec::new(); + for credential in parsed + .storage_credentials + .into_iter() + .filter(|credential| cloud.matches_location(&credential.prefix)) + { + let prefix = credential.prefix; + let parsed = (cloud.parse_credential)(&credential.config, prefix.clone()) + .map(|vended| CachedEntry::new(vended, cloud.jitter_prefetch)) + .and_then(|entry| { + if entry.is_unexpired(now) { + Ok(entry) + } else { + Err(Error::new( + ErrorKind::DataInvalid, + "invalid vended credential: credential is already expired", + )) + } + }); + match parsed { + Ok(entry) => entries.push(entry), + Err(error) => errors.push(CredentialError { prefix, error }), + } + } + Ok(ParsedCredentials { entries, errors }) + } + + async fn refresh_credential( + &self, + configured: &ConfiguredCloud, + path: &str, + fallback: Option<CachedEntry>, + ) -> Result<StorageCredential> { + let fetched = self.fetch(configured).await; + let mut cache = configured.cache.lock().await; + let now = SystemTime::now(); + cache.entries.retain(|entry| entry.is_unexpired(now)); + + let (fetched_prefixes, failure) = match fetched { + Ok(ParsedCredentials { entries, errors }) => { + let fetched_prefixes = cache.merge(entries, &errors); + // The most specific error for `path`. + let failure = errors + .into_iter() + .filter(|error| prefix_covers(&error.prefix, path)) + .max_by_key(|error| error.prefix.len()) + .map(|error| error.error); + (Some(fetched_prefixes), failure) + } + Err(error) => (None, Some(error)), + }; + + let selected = cache.unexpired_match(path, now); + // A keyed seed is only cached once served as the fallback. + let seed = selected + .is_none() + .then(|| { + fallback + .clone() + .filter(|fallback| fallback.is_unexpired(now)) + }) + .flatten(); + let candidate = selected + .map(|index| &cache.entries[index].credential) Review Comment: `cache.entries[index]` appears six times between here and line 795, plus once in `cache_decision`. It's correct today because nothing changes `entries` between `unexpired_match` and those uses. But `insert` does a `retain` and `merge` reorders the vector, so one reordered edit later turns this into a panic on the credential path. I'd use `get`/`get_mut` and return an error, or have `unexpired_match` return the entry itself instead of an index. ########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// 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. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) + .max_by_key(|(_, entry)| entry.credential.prefix().len()) + .map(|(index, _)| index) + } + + /// Whether a failed prefix covering `path` is narrower than `credential`'s. + fn masks(&self, path: &str, credential: Option<&StorageCredential>) -> bool { + let prefix_len = credential.map(|credential| credential.prefix().len()); + self.failed_prefixes.iter().any(|failed| { + prefix_covers(failed, path) && prefix_len.is_none_or(|len| failed.len() > len) + }) + } + + /// Cache `entry`, replacing any entry with the same prefix, and return its + /// index. + fn insert(&mut self, entry: CachedEntry) -> usize { + let prefix = entry.credential.prefix().to_owned(); + self.entries + .retain(|cached| cached.credential.prefix() != prefix); + self.entries.push(entry); + self.entries.len() - 1 + } + + fn record_success(&mut self) { + self.consecutive_failures = 0; + self.retry_not_before = None; + self.last_failure = None; + } + + fn record_failure(&mut self, error: &Error) { + self.consecutive_failures = self.consecutive_failures.saturating_add(1); + self.retry_not_before = + Instant::now().checked_add(failure_backoff(self.consecutive_failures)); + self.last_failure = Some(error.to_string()); + } + + /// Replace cached credentials with the fetched ones, prefix by prefix, + /// and record the prefixes whose credential was invalid in this response. + /// + /// An unexpired cached credential survives when the response carries no + /// valid replacement for its prefix, so an absent or malformed entry never + /// evicts a usable credential. Returns the prefixes the response replaced. + fn merge(&mut self, fetched: Vec<CachedEntry>, errors: &[CredentialError]) -> HashSet<String> { + let fetched_prefixes = fetched + .iter() + .map(|entry| entry.credential.prefix().to_owned()) + .collect::<HashSet<_>>(); + // A response names every prefix of the table, so a prefix it omits is + // no longer scoped separately: a fetched credential that covers it + // replaces its cached one. A prefix whose credential is invalid keeps + // its cached one. + let failed_prefixes = errors + .iter() + .map(|error| error.prefix.clone()) + .collect::<HashSet<_>>(); + self.entries.retain(|entry| { + let prefix = entry.credential.prefix(); + failed_prefixes.contains(prefix) + || !fetched + .iter() + .any(|fetched| fetched.credential.covers(prefix)) + }); + self.entries.extend(fetched); + self.failed_prefixes = failed_prefixes; + fetched_prefixes + } +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex<CacheState>, + /// Initial property credentials whose scope cannot be represented by a + /// single leading URI prefix, keyed according to the cloud strategy. + keyed_seeds: HashMap<String, CachedEntry>, + /// Only one caller fetches at a time. The cache lock is deliberately + /// separate so other callers can keep using an unexpired credential while + /// the refresh is in flight. + refresh: Mutex<()>, +} + +impl ConfiguredCloud { + fn new( + cloud: &'static CloudRefresh, + endpoint: String, + entries: Vec<CachedEntry>, + keyed_seeds: HashMap<String, CachedEntry>, + ) -> Self { + Self { + cloud, + endpoint, + cache: Mutex::new(CacheState::new(entries)), + keyed_seeds, + refresh: Mutex::new(()), + } + } + + /// Configure `cloud` from the FileIO properties, or `None` when they + /// advertise no enabled refresh endpoint for it. + fn configure( + cloud: &'static CloudRefresh, + base_uri: &str, + props: &HashMap<String, String>, + ) -> Option<Self> { + let enabled = props + .get(cloud.enabled_key) + .is_none_or(|value| value.eq_ignore_ascii_case("true")) + && cloud + .required_key + .is_none_or(|key| props.get(key).is_some_and(|value| !value.is_empty())); + let endpoint = props + .get(cloud.endpoint_key) + .filter(|endpoint| enabled && !endpoint.is_empty()) + .map(|endpoint| resolve_endpoint(base_uri, endpoint))?; + + let mut entries = Vec::new(); + let mut keyed_seeds = HashMap::new(); + match &cloud.seed_strategy { + SeedStrategy::Flat { prefixes } => { + entries.extend(prefixes.iter().filter_map(|prefix| { + (cloud.parse_credential)(props, prefix.to_string()) + .ok() Review Comment: A seed that fails `parse_s3_credential` is dropped here with no trace. The vended chain replaces the whole S3 credential chain in `s3_config_build`, so the static `s3.access-key-id`/`s3.secret-access-key` from the same config stop being used too. Take a catalog that advertises `client.refresh-credentials-endpoint` but sends keys without `s3.session-token-expires-at-ms`. That works on main today. Here the first file op goes to the endpoint, and if that call fails, every S3 op fails. I agree with not falling back to ambient AWS credentials. But the keys in `config` aren't ambient: the catalog vended them. I'd keep them as a seed with no expiry, which the trait doc already allows. At the very least I'd `tracing::warn!` the parse error so the unexpected network fetch has an explanation. ########## crates/storage/opendal/src/azdls.rs: ########## @@ -226,16 +372,24 @@ pub(crate) struct AzureStoragePath { } impl AzureStoragePath { + /// The host of the path, e.g. `account.dfs.core.windows.net`. + fn host(&self) -> String { + let service = match self.scheme { + AzureStorageScheme::Abfs | AzureStorageScheme::Abfss => "dfs", + AzureStorageScheme::Wasb | AzureStorageScheme::Wasbs => "blob", + }; + format!("{}.{service}.{}", self.account_name, self.endpoint_suffix) + } + /// Converts the AzureStoragePath into a full endpoint URL. /// /// This is possible because the path is fully qualified. + /// + /// Like Iceberg Java, the endpoint always uses TLS, also for the + /// non-secure schemes: SAS tokens are query parameters and must not be + /// sent over a plaintext connection unless the user configures one. fn as_endpoint(&self) -> String { - format!( - "{}://{}.dfs.{}", - self.scheme.as_http_scheme(), - self.account_name, - self.endpoint_suffix - ) + format!("https://{}.dfs.{}";, self.account_name, self.endpoint_suffix) Review Comment: This switches `abfs://`/`wasb://` from `http` to `https` for everyone, and `requires_tls` changes which configured endpoints are accepted. Keeping SAS tokens off plaintext connections and matching Java are good reasons. But it's a user-visible change unrelated to refresh, so I'd move it into the ADLS piece and call it out in the changelog. ########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// 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. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) + .max_by_key(|(_, entry)| entry.credential.prefix().len()) + .map(|(index, _)| index) + } + + /// Whether a failed prefix covering `path` is narrower than `credential`'s. + fn masks(&self, path: &str, credential: Option<&StorageCredential>) -> bool { + let prefix_len = credential.map(|credential| credential.prefix().len()); + self.failed_prefixes.iter().any(|failed| { + prefix_covers(failed, path) && prefix_len.is_none_or(|len| failed.len() > len) + }) + } + + /// Cache `entry`, replacing any entry with the same prefix, and return its + /// index. + fn insert(&mut self, entry: CachedEntry) -> usize { + let prefix = entry.credential.prefix().to_owned(); + self.entries + .retain(|cached| cached.credential.prefix() != prefix); + self.entries.push(entry); + self.entries.len() - 1 + } + + fn record_success(&mut self) { + self.consecutive_failures = 0; + self.retry_not_before = None; + self.last_failure = None; + } + + fn record_failure(&mut self, error: &Error) { + self.consecutive_failures = self.consecutive_failures.saturating_add(1); + self.retry_not_before = + Instant::now().checked_add(failure_backoff(self.consecutive_failures)); + self.last_failure = Some(error.to_string()); + } + + /// Replace cached credentials with the fetched ones, prefix by prefix, + /// and record the prefixes whose credential was invalid in this response. + /// + /// An unexpired cached credential survives when the response carries no + /// valid replacement for its prefix, so an absent or malformed entry never + /// evicts a usable credential. Returns the prefixes the response replaced. + fn merge(&mut self, fetched: Vec<CachedEntry>, errors: &[CredentialError]) -> HashSet<String> { + let fetched_prefixes = fetched + .iter() + .map(|entry| entry.credential.prefix().to_owned()) + .collect::<HashSet<_>>(); + // A response names every prefix of the table, so a prefix it omits is + // no longer scoped separately: a fetched credential that covers it + // replaces its cached one. A prefix whose credential is invalid keeps + // its cached one. + let failed_prefixes = errors + .iter() + .map(|error| error.prefix.clone()) + .collect::<HashSet<_>>(); + self.entries.retain(|entry| { + let prefix = entry.credential.prefix(); + failed_prefixes.contains(prefix) + || !fetched + .iter() + .any(|fetched| fetched.credential.covers(prefix)) + }); + self.entries.extend(fetched); + self.failed_prefixes = failed_prefixes; + fetched_prefixes + } +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex<CacheState>, + /// Initial property credentials whose scope cannot be represented by a + /// single leading URI prefix, keyed according to the cloud strategy. + keyed_seeds: HashMap<String, CachedEntry>, + /// Only one caller fetches at a time. The cache lock is deliberately + /// separate so other callers can keep using an unexpired credential while + /// the refresh is in flight. + refresh: Mutex<()>, +} + +impl ConfiguredCloud { + fn new( + cloud: &'static CloudRefresh, + endpoint: String, + entries: Vec<CachedEntry>, + keyed_seeds: HashMap<String, CachedEntry>, + ) -> Self { + Self { + cloud, + endpoint, + cache: Mutex::new(CacheState::new(entries)), + keyed_seeds, + refresh: Mutex::new(()), + } + } + + /// Configure `cloud` from the FileIO properties, or `None` when they + /// advertise no enabled refresh endpoint for it. + fn configure( + cloud: &'static CloudRefresh, + base_uri: &str, + props: &HashMap<String, String>, + ) -> Option<Self> { + let enabled = props + .get(cloud.enabled_key) + .is_none_or(|value| value.eq_ignore_ascii_case("true")) + && cloud + .required_key + .is_none_or(|key| props.get(key).is_some_and(|value| !value.is_empty())); + let endpoint = props + .get(cloud.endpoint_key) + .filter(|endpoint| enabled && !endpoint.is_empty()) + .map(|endpoint| resolve_endpoint(base_uri, endpoint))?; + + let mut entries = Vec::new(); + let mut keyed_seeds = HashMap::new(); + match &cloud.seed_strategy { + SeedStrategy::Flat { prefixes } => { + entries.extend(prefixes.iter().filter_map(|prefix| { + (cloud.parse_credential)(props, prefix.to_string()) + .ok() + .map(|vended| CachedEntry::seed(vended, cloud.jitter_prefetch)) + })); + } + SeedStrategy::Keyed(strategy) => { + keyed_seeds = (strategy.parse_credentials)(props) + .into_iter() + .map(|(key, credential)| { + (key, CachedEntry::seed(credential, cloud.jitter_prefetch)) + }) + .collect(); + } + } + Some(Self::new(cloud, endpoint, entries, keyed_seeds)) + } + + fn keyed_seed_for_path(&self, path: &str) -> Result<Option<CachedEntry>> { + let SeedStrategy::Keyed(strategy) = &self.cloud.seed_strategy else { + return Ok(None); + }; + let resolved = (strategy.resolve_path)(path)?; + let Some(seed) = resolved + .keys + .iter() + .find_map(|key| self.keyed_seeds.get(key)) + else { + return Ok(None); + }; + let credential = StorageCredential::new(resolved.scope, seed.credential.config().clone()); + // The scope is normalized, e.g. with a lowercase scheme, so it may not + // cover the path as written; the path then fetches its credential. + Ok(credential.covers(path).then_some(CachedEntry { + credential, + expires_at: seed.expires_at, + refresh_at: seed.refresh_at, + })) + } +} + +/// Serializable recipe for rebuilding a [`RestVendedCredentialProvider`] from +/// the FileIO properties, mirroring how Java rebuilds its provider on workers. +#[derive(Clone, Serialize, Deserialize)] +pub(crate) struct RestVendedCredentialProviderFactory { + /// Catalog URI, used to resolve relative refresh endpoints. + catalog_uri: String, + table: TableIdent, + /// The `rest.auth.type` the catalog resolved, so a rebuilt provider uses + /// the same auth manager even when the merged FileIO properties would + /// infer another one. + auth_type: String, + /// Whether `auth_type` was inferred rather than configured. Like Java, + /// a table `token` then authenticates a catalog without auth. + auth_type_inferred: bool, + /// The auth-related part of the unmerged config returned by the table + /// endpoint, from which the auth manager derives a table session. Keeping + /// it separate prevents local FileIO overrides from masking table auth. + table_auth: HashMap<String, String>, +} + +/// Table config keys that may override authentication in a table session. +const TABLE_AUTH_KEYS: &[&str] = &["token"]; + +impl RestVendedCredentialProviderFactory { + pub(crate) fn new( + catalog_uri: impl Into<String>, + table: TableIdent, + auth_type: impl Into<String>, + auth_type_inferred: bool, + table_config: &HashMap<String, String>, + ) -> Self { + Self { + catalog_uri: catalog_uri.into(), + table, + auth_type: auth_type.into(), + auth_type_inferred, + table_auth: table_config + .iter() + .filter(|(key, _)| TABLE_AUTH_KEYS.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + } + } + + /// The table token that authenticates credential requests of a catalog + /// without auth. Java's credential providers infer OAuth2 from a `token` + /// unless `rest.auth.type` is configured. + fn bearer_table_token(&self) -> Option<&str> { + (self.auth_type == AUTH_TYPE_NONE && self.auth_type_inferred) + .then(|| self.table_auth.get("token").map(String::as_str)) + .flatten() + } + + /// Connect to the catalog from the FileIO properties, as the catalog would. + async fn connect(&self, props: &HashMap<String, String>) -> Result<HttpClient> { + let mut config_props = props.clone(); + config_props.insert( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + self.auth_type.clone(), + ); + let config = RestCatalogConfig::builder() + .uri(self.catalog_uri.clone()) + .props(config_props) + .build(); + let auth_manager = load_auth_manager(&config)?; + let client = HttpClient::new(&config)?; + let session = auth_manager + .catalog_session(&client.without_auth_session(), &config.auth_props()) + .await?; + table_client( + &client.with_auth_session(session), + auth_manager.as_ref(), + &self.table, + props, + &self.table_auth, + self.bearer_table_token(), + ) + .await + } +} + +impl std::fmt::Debug for RestVendedCredentialProviderFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RestVendedCredentialProviderFactory") + .field("catalog_uri", &self.catalog_uri) + .field("table", &self.table) + .finish_non_exhaustive() + } +} + +#[typetag::serde(name = "RestVendedCredentialProviderFactory")] +impl StorageCredentialProviderFactory for RestVendedCredentialProviderFactory { + fn build(&self, config: &StorageConfig) -> Result<Arc<dyn StorageCredentialProvider>> { + let provider = + RestVendedCredentialProvider::configure(self, config.props()).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + "FileIO configuration no longer configures vended credentials", + ) + })?; + Ok(Arc::new(RestVendedCredentialProvider { + factory: Some(self.clone()), + ..provider + })) + } +} + +/// Derive the table-scoped client used for credential requests. +/// +/// A `bearer_table_token` authenticates instead of a session from the auth +/// manager. +async fn table_client( + catalog_client: &HttpClient, + auth_manager: &dyn AuthManager, + table: &TableIdent, + props: &HashMap<String, String>, + table_config: &HashMap<String, String>, + bearer_table_token: Option<&str>, +) -> Result<HttpClient> { + let session = match bearer_table_token { + Some(token) => static_token_session(token), + None => { + auth_manager + .table_session( + &catalog_client.without_auth_session(), + table, + table_config, + catalog_client.auth_session(), + ) + .await? + } + }; + catalog_client.for_table(props, session) +} + +/// Serves vended credentials for a table, refreshing them from the REST +/// catalog's table credentials endpoint. +/// +/// Each cloud cache is seeded with the credentials from the initial +/// table properties (when complete) and re-fetched from its endpoint as they +/// near expiry. +pub(crate) struct RestVendedCredentialProvider { + /// Table-scoped catalog client. Set when the catalog builds the provider, + /// and connected on first refresh after deserialization. + client: OnceCell<HttpClient>, + /// Rebuilds this provider in another process. `None` when the catalog + /// authentication cannot be rebuilt from properties. + factory: Option<RestVendedCredentialProviderFactory>, + /// Effective FileIO properties, which supply catalog connection settings. + props: HashMap<String, String>, + /// Query parameters sent with every credentials request. + query_params: Vec<(&'static str, String)>, + /// Independently configured endpoint and cache for each backing cloud. + clouds: Vec<ConfiguredCloud>, +} + +impl RestVendedCredentialProvider { + /// Configure a provider without a catalog connection, or `None` when no + /// supported cloud advertises an enabled refresh endpoint. + fn configure( + factory: &RestVendedCredentialProviderFactory, + props: &HashMap<String, String>, + ) -> Option<Self> { + let clouds = CloudRefresh::SUPPORTED + .iter() + .filter_map(|cloud| ConfiguredCloud::configure(cloud, &factory.catalog_uri, props)) + .collect::<Vec<_>>(); + (!clouds.is_empty()).then(|| Self { + client: OnceCell::new(), + factory: None, + props: props.clone(), + query_params: credentials_query_params(props), + clouds, + }) + } + + fn configured_cloud_for_location(&self, location: &str) -> Option<&ConfiguredCloud> { + self.clouds + .iter() + .find(|configured| configured.cloud.matches_location(location)) + } + + async fn client(&self) -> Result<&HttpClient> { + self.client + .get_or_try_init(|| async { + let factory = self.factory.as_ref().ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "vended credential provider has no catalog connection", + ) + })?; + factory.connect(&self.props).await + }) + .await + } + + /// Fetch fresh credentials from the catalog's credentials endpoint. + async fn fetch(&self, configured: &ConfiguredCloud) -> Result<ParsedCredentials> { + let cloud = configured.cloud; + let client = self.client().await?; + let url = credentials_url(&configured.endpoint, &self.query_params)?; + let request = HttpRequest::build(client.request(Method::GET, url))?; + let response = client.query_catalog(request).await?; Review Comment: `fetch` has no timeout of its own, and as far as I can see the shared client doesn't set one either. Say the catalog accepts the connection and then stalls. The caller holding the single-flight guard (line 893) never returns. Every caller without a usable credential queues on `refresh.lock().await` behind it. No failure is ever recorded, so backoff never starts. Callers with a cached credential keep working until it expires, then they hang too, which is the long-running-writer case this PR is meant to fix. I'd wrap the fetch in `tokio::time::timeout`, record a timeout as a refresh failure, and add a test against an endpoint that never answers. ########## crates/storage/opendal/src/utils.rs: ########## @@ -27,3 +28,331 @@ pub(crate) fn from_opendal_error(e: opendal::Error) -> iceberg::Error { ) .with_source(e) } + +/// The non-empty value of `key` in a vended credential's config. +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +pub(crate) fn required_credential_property<'a>( + config: &'a std::collections::HashMap<String, String>, + key: &str, +) -> reqsign_core::Result<&'a str> { + config + .get(key) + .map(String::as_str) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + reqsign_core::Error::unexpected(format!("vended credential is missing {key}")) + }) +} + +/// The latest epoch millisecond `reqsign` timestamps can represent, +/// `9999-12-30T22:00:00Z`. +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +const MAX_TIMESTAMP_MILLIS: i64 = 253_402_207_200_000; + +/// The epoch-millisecond expiry under `key` in a vended credential's config, +/// if present, as the `reqsign` timestamp backend credentials use. Later +/// expiries, such as `Long.MAX_VALUE` for a credential that never expires, +/// are clamped to the latest representable timestamp. +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +pub(crate) fn credential_expiry( + config: &std::collections::HashMap<String, String>, + key: &str, +) -> reqsign_core::Result<Option<reqsign_core::time::Timestamp>> { + config + .get(key) + .filter(|value| !value.is_empty()) + .map(|value| { + value + .parse::<i64>() + .ok() + .and_then(|millis| { + reqsign_core::time::Timestamp::from_millisecond( + millis.min(MAX_TIMESTAMP_MILLIS), + ) + .ok() + }) + .ok_or_else(|| { + reqsign_core::Error::unexpected(format!( + "vended credential has an invalid {key}" + )) + }) + }) + .transpose() +} + +/// Validate that a provider's credential covers the location for which the +/// backend requested it. +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +pub(crate) fn validate_credential_prefix( + location: &str, + credential: &iceberg::io::StorageCredential, +) -> reqsign_core::Result<()> { + if credential.covers(location) { + Ok(()) + } else { + Err(reqsign_core::Error::unexpected(uncovered_location_message( + location, credential, + ))) + } +} + +/// The error message for a vended credential that does not cover `location`. +pub(crate) fn uncovered_location_message( + location: &str, + credential: &iceberg::io::StorageCredential, +) -> String { + format!( + "vended credential prefix {:?} does not cover storage location {location:?}", + credential.prefix() + ) +} + +/// The root of the storage location containing `path`, e.g. `s3://bucket/`. +pub(crate) fn storage_root(path: &str) -> iceberg::Result<String> { + let url = url::Url::parse(path)?; + Ok(format!("{}://{}/", url.scheme(), url.authority())) +} + +/// Loads vended credentials for one storage location on behalf of a backend's +/// `reqsign` credential provider. +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +pub(crate) struct VendedCredentialSource { + provider: std::sync::Arc<dyn iceberg::io::StorageCredentialProvider>, + /// Location handed to the provider: the path an operator serves, or the + /// scope shared by a bulk-delete batch. + location: String, +} + +#[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" +))] +impl VendedCredentialSource { + pub(crate) fn new( + provider: std::sync::Arc<dyn iceberg::io::StorageCredentialProvider>, + location: String, + ) -> Self { + Self { provider, location } + } + + /// Load a credential covering the location and convert its config with + /// `extract` into the backend's credential. + /// + /// Errors are logged here: reqsign's credential chain replaces them with an + /// error that only says no credential could be loaded, and reports the + /// cause only through the `log` crate. + pub(crate) async fn load<T>( + &self, + backend: &str, + extract: impl FnOnce(&std::collections::HashMap<String, String>) -> reqsign_core::Result<T>, + ) -> reqsign_core::Result<T> { + self.try_load(backend, extract).await.inspect_err(|error| { + tracing::warn!( Review Comment: While refresh is backed off, `load_credential` fails immediately and every object-store request ends up here, so a busy scan logs thousands of identical warnings per second. I'd move the `warn!` into the provider's `record_failure`, which fires once per failed refresh, and drop this one to `debug`. ########## crates/iceberg/src/io/storage/mod.rs: ########## @@ -139,4 +140,179 @@ pub trait StorageFactory: Debug + Send + Sync { /// A `Result` containing an `Arc<dyn Storage>` on success, or an error /// if the storage could not be created. fn build(&self, config: &StorageConfig) -> Result<Arc<dyn Storage>>; + + /// Build a new Storage instance, optionally supplying a credential provider + /// that the backend can call to obtain and refresh short-lived credentials. + /// + /// Backends that cannot use the provider ignore it and use the credentials + /// in `config`, as they would without one. The default does exactly that. + #[allow(unused_variables)] + fn build_with_credential_provider( Review Comment: Ignoring the provider by default is the right call, since it removes the wiring-order hazard @mbutrovich described. The cost is that a third-party `StorageFactory` keeps the seeded credentials and silently loses refresh, which shows up as a 403 when they expire. I'd `tracing::warn!` here when `credential_provider` is `Some`; this runs once per storage build, not per request. I'd also rename it to `_credential_provider` instead of using `#[allow(unused_variables)]`. ########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// 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. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) + .max_by_key(|(_, entry)| entry.credential.prefix().len()) + .map(|(index, _)| index) + } + + /// Whether a failed prefix covering `path` is narrower than `credential`'s. + fn masks(&self, path: &str, credential: Option<&StorageCredential>) -> bool { + let prefix_len = credential.map(|credential| credential.prefix().len()); + self.failed_prefixes.iter().any(|failed| { + prefix_covers(failed, path) && prefix_len.is_none_or(|len| failed.len() > len) + }) + } + + /// Cache `entry`, replacing any entry with the same prefix, and return its + /// index. + fn insert(&mut self, entry: CachedEntry) -> usize { + let prefix = entry.credential.prefix().to_owned(); + self.entries + .retain(|cached| cached.credential.prefix() != prefix); + self.entries.push(entry); + self.entries.len() - 1 + } + + fn record_success(&mut self) { + self.consecutive_failures = 0; + self.retry_not_before = None; + self.last_failure = None; + } + + fn record_failure(&mut self, error: &Error) { + self.consecutive_failures = self.consecutive_failures.saturating_add(1); + self.retry_not_before = + Instant::now().checked_add(failure_backoff(self.consecutive_failures)); + self.last_failure = Some(error.to_string()); + } + + /// Replace cached credentials with the fetched ones, prefix by prefix, + /// and record the prefixes whose credential was invalid in this response. + /// + /// An unexpired cached credential survives when the response carries no + /// valid replacement for its prefix, so an absent or malformed entry never + /// evicts a usable credential. Returns the prefixes the response replaced. + fn merge(&mut self, fetched: Vec<CachedEntry>, errors: &[CredentialError]) -> HashSet<String> { + let fetched_prefixes = fetched + .iter() + .map(|entry| entry.credential.prefix().to_owned()) + .collect::<HashSet<_>>(); + // A response names every prefix of the table, so a prefix it omits is + // no longer scoped separately: a fetched credential that covers it + // replaces its cached one. A prefix whose credential is invalid keeps + // its cached one. + let failed_prefixes = errors + .iter() + .map(|error| error.prefix.clone()) + .collect::<HashSet<_>>(); + self.entries.retain(|entry| { + let prefix = entry.credential.prefix(); + failed_prefixes.contains(prefix) + || !fetched + .iter() + .any(|fetched| fetched.credential.covers(prefix)) + }); + self.entries.extend(fetched); + self.failed_prefixes = failed_prefixes; + fetched_prefixes + } +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex<CacheState>, + /// Initial property credentials whose scope cannot be represented by a + /// single leading URI prefix, keyed according to the cloud strategy. + keyed_seeds: HashMap<String, CachedEntry>, + /// Only one caller fetches at a time. The cache lock is deliberately + /// separate so other callers can keep using an unexpired credential while + /// the refresh is in flight. + refresh: Mutex<()>, +} + +impl ConfiguredCloud { + fn new( + cloud: &'static CloudRefresh, + endpoint: String, + entries: Vec<CachedEntry>, + keyed_seeds: HashMap<String, CachedEntry>, + ) -> Self { + Self { + cloud, + endpoint, + cache: Mutex::new(CacheState::new(entries)), + keyed_seeds, + refresh: Mutex::new(()), + } + } + + /// Configure `cloud` from the FileIO properties, or `None` when they + /// advertise no enabled refresh endpoint for it. + fn configure( + cloud: &'static CloudRefresh, + base_uri: &str, + props: &HashMap<String, String>, + ) -> Option<Self> { + let enabled = props + .get(cloud.enabled_key) + .is_none_or(|value| value.eq_ignore_ascii_case("true")) + && cloud + .required_key + .is_none_or(|key| props.get(key).is_some_and(|value| !value.is_empty())); + let endpoint = props + .get(cloud.endpoint_key) + .filter(|endpoint| enabled && !endpoint.is_empty()) + .map(|endpoint| resolve_endpoint(base_uri, endpoint))?; + + let mut entries = Vec::new(); + let mut keyed_seeds = HashMap::new(); + match &cloud.seed_strategy { + SeedStrategy::Flat { prefixes } => { + entries.extend(prefixes.iter().filter_map(|prefix| { + (cloud.parse_credential)(props, prefix.to_string()) + .ok() + .map(|vended| CachedEntry::seed(vended, cloud.jitter_prefetch)) + })); + } + SeedStrategy::Keyed(strategy) => { + keyed_seeds = (strategy.parse_credentials)(props) + .into_iter() + .map(|(key, credential)| { + (key, CachedEntry::seed(credential, cloud.jitter_prefetch)) + }) + .collect(); + } + } + Some(Self::new(cloud, endpoint, entries, keyed_seeds)) + } + + fn keyed_seed_for_path(&self, path: &str) -> Result<Option<CachedEntry>> { + let SeedStrategy::Keyed(strategy) = &self.cloud.seed_strategy else { + return Ok(None); + }; + let resolved = (strategy.resolve_path)(path)?; + let Some(seed) = resolved + .keys + .iter() + .find_map(|key| self.keyed_seeds.get(key)) + else { + return Ok(None); + }; + let credential = StorageCredential::new(resolved.scope, seed.credential.config().clone()); + // The scope is normalized, e.g. with a lowercase scheme, so it may not + // cover the path as written; the path then fetches its credential. + Ok(credential.covers(path).then_some(CachedEntry { + credential, + expires_at: seed.expires_at, + refresh_at: seed.refresh_at, + })) + } +} + +/// Serializable recipe for rebuilding a [`RestVendedCredentialProvider`] from +/// the FileIO properties, mirroring how Java rebuilds its provider on workers. +#[derive(Clone, Serialize, Deserialize)] +pub(crate) struct RestVendedCredentialProviderFactory { + /// Catalog URI, used to resolve relative refresh endpoints. + catalog_uri: String, + table: TableIdent, + /// The `rest.auth.type` the catalog resolved, so a rebuilt provider uses + /// the same auth manager even when the merged FileIO properties would + /// infer another one. + auth_type: String, + /// Whether `auth_type` was inferred rather than configured. Like Java, + /// a table `token` then authenticates a catalog without auth. + auth_type_inferred: bool, + /// The auth-related part of the unmerged config returned by the table + /// endpoint, from which the auth manager derives a table session. Keeping + /// it separate prevents local FileIO overrides from masking table auth. + table_auth: HashMap<String, String>, +} + +/// Table config keys that may override authentication in a table session. +const TABLE_AUTH_KEYS: &[&str] = &["token"]; + +impl RestVendedCredentialProviderFactory { + pub(crate) fn new( + catalog_uri: impl Into<String>, + table: TableIdent, + auth_type: impl Into<String>, + auth_type_inferred: bool, + table_config: &HashMap<String, String>, + ) -> Self { + Self { + catalog_uri: catalog_uri.into(), + table, + auth_type: auth_type.into(), + auth_type_inferred, + table_auth: table_config + .iter() + .filter(|(key, _)| TABLE_AUTH_KEYS.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + } + } + + /// The table token that authenticates credential requests of a catalog + /// without auth. Java's credential providers infer OAuth2 from a `token` + /// unless `rest.auth.type` is configured. + fn bearer_table_token(&self) -> Option<&str> { + (self.auth_type == AUTH_TYPE_NONE && self.auth_type_inferred) + .then(|| self.table_auth.get("token").map(String::as_str)) + .flatten() + } + + /// Connect to the catalog from the FileIO properties, as the catalog would. + async fn connect(&self, props: &HashMap<String, String>) -> Result<HttpClient> { + let mut config_props = props.clone(); + config_props.insert( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + self.auth_type.clone(), + ); + let config = RestCatalogConfig::builder() + .uri(self.catalog_uri.clone()) + .props(config_props) + .build(); + let auth_manager = load_auth_manager(&config)?; + let client = HttpClient::new(&config)?; + let session = auth_manager + .catalog_session(&client.without_auth_session(), &config.auth_props()) + .await?; + table_client( + &client.with_auth_session(session), + auth_manager.as_ref(), + &self.table, + props, + &self.table_auth, + self.bearer_table_token(), + ) + .await + } +} + +impl std::fmt::Debug for RestVendedCredentialProviderFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RestVendedCredentialProviderFactory") + .field("catalog_uri", &self.catalog_uri) + .field("table", &self.table) + .finish_non_exhaustive() + } +} + +#[typetag::serde(name = "RestVendedCredentialProviderFactory")] +impl StorageCredentialProviderFactory for RestVendedCredentialProviderFactory { + fn build(&self, config: &StorageConfig) -> Result<Arc<dyn StorageCredentialProvider>> { + let provider = + RestVendedCredentialProvider::configure(self, config.props()).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + "FileIO configuration no longer configures vended credentials", + ) + })?; + Ok(Arc::new(RestVendedCredentialProvider { + factory: Some(self.clone()), + ..provider + })) + } +} + +/// Derive the table-scoped client used for credential requests. +/// +/// A `bearer_table_token` authenticates instead of a session from the auth +/// manager. +async fn table_client( + catalog_client: &HttpClient, + auth_manager: &dyn AuthManager, + table: &TableIdent, + props: &HashMap<String, String>, + table_config: &HashMap<String, String>, + bearer_table_token: Option<&str>, +) -> Result<HttpClient> { + let session = match bearer_table_token { + Some(token) => static_token_session(token), + None => { + auth_manager + .table_session( + &catalog_client.without_auth_session(), + table, + table_config, + catalog_client.auth_session(), + ) + .await? + } + }; + catalog_client.for_table(props, session) +} + +/// Serves vended credentials for a table, refreshing them from the REST +/// catalog's table credentials endpoint. +/// +/// Each cloud cache is seeded with the credentials from the initial +/// table properties (when complete) and re-fetched from its endpoint as they +/// near expiry. +pub(crate) struct RestVendedCredentialProvider { + /// Table-scoped catalog client. Set when the catalog builds the provider, + /// and connected on first refresh after deserialization. + client: OnceCell<HttpClient>, + /// Rebuilds this provider in another process. `None` when the catalog + /// authentication cannot be rebuilt from properties. + factory: Option<RestVendedCredentialProviderFactory>, + /// Effective FileIO properties, which supply catalog connection settings. + props: HashMap<String, String>, + /// Query parameters sent with every credentials request. + query_params: Vec<(&'static str, String)>, + /// Independently configured endpoint and cache for each backing cloud. + clouds: Vec<ConfiguredCloud>, +} + +impl RestVendedCredentialProvider { + /// Configure a provider without a catalog connection, or `None` when no + /// supported cloud advertises an enabled refresh endpoint. + fn configure( + factory: &RestVendedCredentialProviderFactory, + props: &HashMap<String, String>, + ) -> Option<Self> { + let clouds = CloudRefresh::SUPPORTED + .iter() + .filter_map(|cloud| ConfiguredCloud::configure(cloud, &factory.catalog_uri, props)) + .collect::<Vec<_>>(); + (!clouds.is_empty()).then(|| Self { + client: OnceCell::new(), + factory: None, + props: props.clone(), + query_params: credentials_query_params(props), + clouds, + }) + } + + fn configured_cloud_for_location(&self, location: &str) -> Option<&ConfiguredCloud> { + self.clouds + .iter() + .find(|configured| configured.cloud.matches_location(location)) + } + + async fn client(&self) -> Result<&HttpClient> { + self.client + .get_or_try_init(|| async { + let factory = self.factory.as_ref().ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "vended credential provider has no catalog connection", + ) + })?; + factory.connect(&self.props).await + }) + .await + } + + /// Fetch fresh credentials from the catalog's credentials endpoint. + async fn fetch(&self, configured: &ConfiguredCloud) -> Result<ParsedCredentials> { + let cloud = configured.cloud; + let client = self.client().await?; + let url = credentials_url(&configured.endpoint, &self.query_params)?; + let request = HttpRequest::build(client.request(Method::GET, url))?; + let response = client.query_catalog(request).await?; + + if !response.status().is_success() { + return Err(unexpected_catalog_error_without_body( + response, + client.disable_header_redaction(), + )); + } + + // Credential responses contain secrets, and serde's errors quote the + // values they reject, so only the error's category and position are + // reported. + let parsed: LoadCredentialsResponse = + serde_json::from_slice(response.body()).map_err(|error| { + Error::new( + ErrorKind::Unexpected, + format!( + "failed to parse vended credential response: {:?} error at line {}, \ + column {}", + error.classify(), + error.line(), + error.column() + ), + ) + })?; + let now = SystemTime::now(); + let mut entries = Vec::new(); + let mut errors = Vec::new(); + for credential in parsed + .storage_credentials + .into_iter() + .filter(|credential| cloud.matches_location(&credential.prefix)) + { + let prefix = credential.prefix; + let parsed = (cloud.parse_credential)(&credential.config, prefix.clone()) + .map(|vended| CachedEntry::new(vended, cloud.jitter_prefetch)) + .and_then(|entry| { + if entry.is_unexpired(now) { + Ok(entry) + } else { + Err(Error::new( + ErrorKind::DataInvalid, + "invalid vended credential: credential is already expired", + )) + } + }); + match parsed { + Ok(entry) => entries.push(entry), + Err(error) => errors.push(CredentialError { prefix, error }), + } + } + Ok(ParsedCredentials { entries, errors }) + } + + async fn refresh_credential( + &self, + configured: &ConfiguredCloud, + path: &str, + fallback: Option<CachedEntry>, + ) -> Result<StorageCredential> { + let fetched = self.fetch(configured).await; + let mut cache = configured.cache.lock().await; + let now = SystemTime::now(); + cache.entries.retain(|entry| entry.is_unexpired(now)); + + let (fetched_prefixes, failure) = match fetched { + Ok(ParsedCredentials { entries, errors }) => { + let fetched_prefixes = cache.merge(entries, &errors); + // The most specific error for `path`. + let failure = errors + .into_iter() + .filter(|error| prefix_covers(&error.prefix, path)) + .max_by_key(|error| error.prefix.len()) + .map(|error| error.error); + (Some(fetched_prefixes), failure) + } + Err(error) => (None, Some(error)), + }; + + let selected = cache.unexpired_match(path, now); + // A keyed seed is only cached once served as the fallback. + let seed = selected + .is_none() + .then(|| { + fallback + .clone() + .filter(|fallback| fallback.is_unexpired(now)) + }) + .flatten(); + let candidate = selected + .map(|index| &cache.entries[index].credential) + .or(seed.as_ref().map(|seed| &seed.credential)); + if cache.masks(path, candidate) { + let error = failure.unwrap_or_else(|| masked_credential_error(path)); + cache.record_failure(&error); + return Err(error); + } + + let refreshed = selected.filter(|&index| { + fetched_prefixes + .as_ref() + .is_some_and(|fetched| fetched.contains(cache.entries[index].credential.prefix())) + }); + if let Some(index) = refreshed { + let entry = &mut cache.entries[index]; + // A catalog may return the same credential until shortly before + // it expires. Treat one that is no newer like no answer, so the + // checks stay few as expiry approaches. + if fallback + .as_ref() + .is_some_and(|previous| entry.expires_at <= previous.expires_at) + { + entry.refresh_at = recheck_time(now, entry.expires_at); + } + cache.record_success(); + return Ok(cache.entries[index].credential.clone()); + } + + let selected = selected.or_else(|| seed.map(|seed| cache.insert(seed))); + let Some(index) = selected else { + let error = failure.unwrap_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("no unexpired vended credential matches storage location: {path}"), + ) + }); + cache.record_failure(&error); Review Comment: Recording a cloud-wide failure here treats "this path isn't vended" the same as "the catalog is failing". Once an endpoint is configured, `supports_path` (line 862) returns true for every `s3*` location. So a read outside the table's prefixes, like an external file in another bucket, triggers a fetch, ends up here, and backs off refresh for the whole cloud for up to 30s. During that window, a covered path whose credential has just expired gets the backoff error instead of a refresh. When the fetch succeeded and nothing covers `path`, I'd return the error without touching `consecutive_failures`/`retry_not_before`. Backoff should only apply to failed fetches and failed parses. ########## crates/iceberg/src/io/storage/mod.rs: ########## @@ -139,4 +140,179 @@ pub trait StorageFactory: Debug + Send + Sync { /// A `Result` containing an `Arc<dyn Storage>` on success, or an error /// if the storage could not be created. fn build(&self, config: &StorageConfig) -> Result<Arc<dyn Storage>>; + + /// Build a new Storage instance, optionally supplying a credential provider + /// that the backend can call to obtain and refresh short-lived credentials. + /// + /// Backends that cannot use the provider ignore it and use the credentials + /// in `config`, as they would without one. The default does exactly that. + #[allow(unused_variables)] + fn build_with_credential_provider( + &self, + config: &StorageConfig, + credential_provider: Option<Arc<dyn StorageCredentialProvider>>, + ) -> Result<Arc<dyn Storage>> { + self.build(config) + } +} + +/// Supplies fresh, backend-specific storage credentials on demand. +/// +/// A catalog that vends temporary credentials implements this trait so that +/// storage backends can re-fetch credentials as they approach expiry instead +/// of failing once the initial token's TTL runs out. +/// +/// # Debug output +/// +/// [`FileIO`](crate::io::FileIO) and storage implementations print the +/// provider in their `Debug` output, so the provider's `Debug` implementation +/// must not expose credentials or other secrets. +/// +/// # Caching +/// +/// [`load_credential`](Self::load_credential) may be called very frequently. +/// Implementations must cache internally and only re-fetch when the current +/// credential is at or near expiry; otherwise every object-store request could +/// trigger a call back to the catalog. +#[async_trait] +pub trait StorageCredentialProvider: Debug + Send + Sync { + /// Return whether this provider supplies credentials for `path`. + /// + /// Backends replace their own credential chain with the provider only for + /// paths it supports, so a provider must return `false` for the storage it + /// does not serve, such as other schemes behind a resolving storage. + fn supports_path(&self, path: &str) -> bool; + + /// Load a fresh credential for the storage location identified by `path`. + /// + /// `path` is an absolute location: usually the file being accessed, such + /// as `s3://bucket/warehouse/db/table/data/file.parquet`, but for a bulk + /// delete the location shared by a batch, which may be a storage root + /// (`s3://bucket/`) or the prefix of a previously returned credential + /// (`s3://bucket/warehouse/db/table`). Providers that vend distinct + /// credentials per location prefix use it to select the most specific + /// match, which must [cover](StorageCredential::covers) `path`. + /// + /// Backends read the credential's expiry from its config, such as + /// `s3.session-token-expires-at-ms`, and load a new one before it expires. + /// A credential without an expiry is used for as long as the backend lives. + async fn load_credential(&self, path: &str) -> Result<StorageCredential>; Review Comment: The trait doc says this "may be called very frequently", yet each call returns an owned `StorageCredential`. So every file op, per-path delete key and signer refresh clones the config map and its secrets. The trait is new and public, so I'd return `Arc<StorageCredential>` now. Changing it later would break every implementor. -- 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]
