zakariya-s commented on code in PR #2932: URL: https://github.com/apache/iceberg-rust/pull/2932#discussion_r3690419391
########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,1193 @@ +// 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. + +//! Refresh of vended storage credentials against 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, and +//! the returned [`StorageCredential`] enum lets the storage adapter enforce the +//! expected backend-specific type. This preserves Java's per-cloud refresh +//! policy while supporting mixed-cloud tables through a resolving FileIO. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, first add its credential type to Iceberg's storage API and +//! teach the storage adapter to consume it. Then write its `parse_*` function, +//! add a `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + AWS_REFRESH_CREDENTIALS_ENABLED, AWS_REFRESH_CREDENTIALS_ENDPOINT, + GCS_REFRESH_CREDENTIALS_ENABLED, GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, + GCS_TOKEN_EXPIRES_AT, GcsCredential, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, + S3_SESSION_TOKEN_EXPIRES_AT_MS, S3Credential, StorageCredential, StorageCredentialKind, + StorageCredentialProvider, +}; +use iceberg::{Error, ErrorKind, Result}; +use rand::Rng; +use reqwest::{Method, StatusCode, Url}; +use tokio::sync::Mutex; + +use crate::REST_CATALOG_PROP_SCAN_PLAN_ID; +use crate::client::{HttpClient, deserialize_unexpected_catalog_error}; +use crate::types::LoadCredentialsResponse; + +/// 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 to opt out; refresh is enabled unless this is `"false"`. + enabled_key: &'static str, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: fn(&HashMap<String, String>) -> Result<StorageCredential>, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + }; + // TODO: Azure (ADLS) is not yet supported: opendal 0.57's Azdls builder exposes no + // custom credential-provider hook, and reqsign's SAS-token credential has no + // expiry, so reqsign-based refresh isn't possible. + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP]; + + /// The backend that serves `location`, by its URL scheme, or `None` if no + /// supported backend matches (in which case static credentials are used + /// as-is, as before). + fn for_location(location: &str) -> Option<&'static Self> { + let scheme = scheme_of(location)?; + Self::SUPPORTED + .iter() + .find(|cloud| cloud.schemes.contains(&scheme.as_str())) + } + + fn matches_credential_prefix(&self, prefix: &str) -> bool { + scheme_of(prefix).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 failure backoff while a cached credential remains usable. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A vended credential paired with the storage-location prefix it applies to. +#[derive(Clone)] +struct CachedEntry { + prefix: String, + credential: StorageCredential, + /// When this entry becomes eligible for prefetch. `None` means it does not + /// expire and therefore never needs proactive refresh. + refresh_at: Option<SystemTime>, +} + +impl CachedEntry { + fn new(prefix: String, credential: StorageCredential, jitter_prefetch: bool) -> Self { + let refresh_at = credential + .expires_at + .map(|expires_at| prefetch_time(expires_at, jitter_prefetch)); + Self { + prefix, + credential, + refresh_at, + } + } + + /// 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(prefix: String, credential: StorageCredential, jitter_prefetch: bool) -> Self { + let due = credential.expires_at.is_some_and(|expires_at| { + SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= expires_at) + }); + let mut entry = Self::new(prefix, credential, jitter_prefetch); + if due { + entry.refresh_at = Some(UNIX_EPOCH); + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + self.refresh_at.is_none_or(|refresh_at| now < refresh_at) + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + self.credential + .expires_at + .is_none_or(|expires_at| now < expires_at) + } +} + +/// Cached credentials plus failure-backoff state. +struct CacheState { + entries: Vec<CachedEntry>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex<CacheState>, + /// 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<()>, +} + +/// Fetches and refreshes vended credentials from a REST catalog's table +/// credentials endpoint. +/// +/// Each cloud cache is seeded with the credential from the initial table +/// properties (when complete) and re-fetched from its endpoint as it +/// nears expiry. +pub(crate) struct RestVendedCredentialProvider { + client: Arc<HttpClient>, + /// Optional scan-plan identifier. + plan_id: Option<String>, + /// Independently configured endpoint and cache for each backing cloud. + clouds: Vec<ConfiguredCloud>, +} + +impl RestVendedCredentialProvider { + fn new(client: Arc<HttpClient>, plan_id: Option<String>, clouds: Vec<ConfiguredCloud>) -> Self { + Self { + client, + plan_id, + clouds, + } + } + + /// Fetch fresh credentials from the catalog's credentials endpoint. + async fn fetch(&self, configured: &ConfiguredCloud) -> Result<Vec<CachedEntry>> { + let mut request = self.client.request(Method::GET, &configured.endpoint); + if let Some(plan_id) = &self.plan_id { + request = request.query(&[("planId", plan_id)]); + } + let request = request.build()?; + let response = self.client.query_catalog(request).await?; + + match response.status() { + StatusCode::OK => { + let parsed: LoadCredentialsResponse = response.json().await?; + parsed + .storage_credentials + .into_iter() + .filter(|sc| configured.cloud.matches_credential_prefix(&sc.prefix)) + .map(|sc| { + (configured.cloud.parse_credential)(&sc.config).map(|credential| { + CachedEntry::new( + sc.prefix, + credential, + configured.cloud.jitter_prefetch, + ) + }) + }) + .collect() + } + _ => Err(deserialize_unexpected_catalog_error( + response, + self.client.disable_header_redaction(), + ) + .await), + } + } + + async fn refresh_credential( + &self, + configured: &ConfiguredCloud, + path: &str, + fallback: Option<CachedEntry>, + ) -> Result<StorageCredential> { + let refreshed = self.fetch(configured).await.and_then(|entries| { + let credential = longest_prefix_match(&entries, path) + .filter(|entry| entry.is_unexpired(SystemTime::now())) + .map(|entry| entry.credential.clone()) + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("no unexpired vended credential matches storage location: {path}"), + ) + })?; + Ok((entries, credential)) + }); + + match refreshed { + Ok((entries, credential)) => { + let mut cache = configured.cache.lock().await; + cache.entries = entries; + cache.consecutive_failures = 0; + cache.retry_not_before = None; + Ok(credential) + } + Err(fetch_error) => { + let mut cache = configured.cache.lock().await; + cache.consecutive_failures = cache.consecutive_failures.saturating_add(1); + cache.retry_not_before = + Instant::now().checked_add(failure_backoff(cache.consecutive_failures)); + + // Graceful degradation: while the cached credential remains + // usable, serve it and retry after jittered backoff. Expired + // credentials are never served. + fallback + .filter(|entry| entry.is_unexpired(SystemTime::now())) + .map(|entry| entry.credential) + .ok_or(fetch_error) + } + } + } +} + +impl std::fmt::Debug for RestVendedCredentialProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RestVendedCredentialProvider") + .field("configured_clouds", &self.clouds.len()) + .finish_non_exhaustive() + } +} + +enum CacheDecision { + Use(StorageCredential), + Refresh(Option<CachedEntry>), +} + +async fn cache_decision(configured: &ConfiguredCloud, path: &str) -> CacheDecision { + let cache = configured.cache.lock().await; + let current = longest_prefix_match(&cache.entries, path).cloned(); + let now = SystemTime::now(); + + if let Some(entry) = current.as_ref().filter(|entry| entry.is_fresh(now)) { + return CacheDecision::Use(entry.credential.clone()); + } + + if cache + .retry_not_before + .is_some_and(|retry_at| Instant::now() < retry_at) + && let Some(entry) = current.as_ref().filter(|entry| entry.is_unexpired(now)) + { + return CacheDecision::Use(entry.credential.clone()); + } + + CacheDecision::Refresh(current) +} + +#[async_trait] +impl StorageCredentialProvider for RestVendedCredentialProvider { + fn supports_path(&self, path: &str) -> bool { + CloudRefresh::for_location(path).is_some_and(|cloud| { + self.clouds + .iter() + .any(|configured| configured.cloud.endpoint_key == cloud.endpoint_key) + }) + } + + async fn load_credential(&self, path: &str) -> Result<StorageCredential> { + let cloud = CloudRefresh::for_location(path).ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!("no credential refresh implementation for storage location: {path}"), + ) + })?; + let configured = self + .clouds + .iter() + .find(|configured| configured.cloud.endpoint_key == cloud.endpoint_key) + .ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!("credential refresh is not configured for storage location: {path}"), + ) + })?; + + let current = match cache_decision(configured, path).await { + CacheDecision::Use(credential) => return Ok(credential), + CacheDecision::Refresh(current) => current, + }; + + // One caller refreshes, while concurrent callers immediately keep using the + // unexpired cached credential. With no usable credential, callers wait + // for the in-flight refresh instead. + let usable = current + .as_ref() + .filter(|entry| entry.is_unexpired(SystemTime::now())); + let _refresh_guard = if let Some(entry) = usable { + match configured.refresh.try_lock() { + Ok(guard) => guard, + Err(_) => return Ok(entry.credential.clone()), + } + } else { + configured.refresh.lock().await + }; + + // Another caller may have completed a refresh between our cache check + // and acquiring the single-flight guard. + let current = match cache_decision(configured, path).await { + CacheDecision::Use(credential) => return Ok(credential), + CacheDecision::Refresh(current) => current, + }; + + self.refresh_credential(configured, path, current).await + } +} + +/// Select the credential whose prefix is the longest match for `path`. +fn longest_prefix_match<'a>(entries: &'a [CachedEntry], path: &str) -> Option<&'a CachedEntry> { + entries + .iter() + .filter(|entry| path.starts_with(&entry.prefix)) + .max_by_key(|entry| entry.prefix.len()) +} + +/// Compute a successful credential's prefetch time. +fn prefetch_time(expires_at: SystemTime, jitter: bool) -> SystemTime { + let base = expires_at.checked_sub(REFRESH_BUFFER).unwrap_or(UNIX_EPOCH); + if !jitter { + return base; + } + + let jitter_window = REFRESH_BUFFER.saturating_sub(MIN_REFRESH_BUFFER); + let jitter_millis = rand::rng().random_range(0..jitter_window.as_millis() as u64); + base.checked_add(Duration::from_millis(jitter_millis)) + .unwrap_or(base) +} + +/// Equal-jitter exponential backoff. The random lower half avoids both hot +/// retry loops and synchronized retries across clients. +fn failure_backoff(consecutive_failures: u32) -> Duration { + let exponent = consecutive_failures.saturating_sub(1).min(5); + let ceiling = INITIAL_FAILURE_BACKOFF + .checked_mul(1 << exponent) + .unwrap_or(MAX_FAILURE_BACKOFF) + .min(MAX_FAILURE_BACKOFF); + let ceiling_millis = ceiling.as_millis() as u64; + let floor_millis = ceiling_millis / 2; + Duration::from_millis(rand::rng().random_range(floor_millis..=ceiling_millis)) +} + +/// Build a credential provider from a table's properties, +/// or `None` when no supported cloud advertises an enabled refresh endpoint. +/// +/// `base_uri` is the catalog URI, used to resolve a relative endpoint. +pub(crate) fn build_vended_credential_provider( + client: Arc<HttpClient>, + base_uri: &str, + props: &HashMap<String, String>, +) -> Result<Option<Arc<dyn StorageCredentialProvider>>> { + let clouds = CloudRefresh::SUPPORTED + .iter() + .filter_map(|cloud| { + // Refresh is enabled by default and invalid booleans disable refresh. + let enabled = props + .get(cloud.enabled_key) + .is_none_or(|value| value.parse().unwrap_or(false)); + if !enabled { + return None; + } + + let endpoint = resolve_endpoint( + base_uri, + props + .get(cloud.endpoint_key) + .filter(|value| !value.is_empty())?, + ); + let entries = (cloud.parse_credential)(props) + .ok() + .map(|credential| { + vec![CachedEntry::seed( + // Flat table properties carry no prefix. This cache is + // already cloud-specific, so an empty prefix is safe. + String::new(), + credential, + cloud.jitter_prefetch, + )] + }) + .unwrap_or_default(); + + Some(ConfiguredCloud { + cloud, + endpoint, + cache: Mutex::new(CacheState { + entries, + consecutive_failures: 0, + retry_not_before: None, + }), + refresh: Mutex::new(()), + }) + }) + .collect::<Vec<_>>(); + + if clouds.is_empty() { + return Ok(None); + } + + let table_client = Arc::new(client.for_table(base_uri, props.clone())?); + let plan_id = props.get(REST_CATALOG_PROP_SCAN_PLAN_ID).cloned(); + Ok(Some(Arc::new(RestVendedCredentialProvider::new( + table_client, + plan_id, + clouds, + )))) +} + +/// Resolve a possibly-relative refresh endpoint against the catalog base URI. +fn resolve_endpoint(base_uri: &str, endpoint: &str) -> String { Review Comment: The behaviour here was done to match Java -- 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]
