From 992bc103c8897689bd9e169a4162cd1a14f20ce9 Mon Sep 17 00:00:00 2001 From: Zakariya Stasa Date: Thu, 30 Jul 2026 22:25:12 +0100 Subject: [PATCH 1/4] feat(rest): support refreshing vended storage credentials --- Cargo.lock | 2 + crates/catalog/rest/Cargo.toml | 1 + crates/catalog/rest/public-api.txt | 15 + crates/catalog/rest/src/catalog.rs | 53 +- crates/catalog/rest/src/client.rs | 59 +- crates/catalog/rest/src/credential.rs | 1193 +++++++++++++++++++ crates/catalog/rest/src/lib.rs | 1 + crates/catalog/rest/src/types.rs | 38 +- crates/iceberg/public-api.txt | 43 + crates/iceberg/src/io/file_io.rs | 27 +- crates/iceberg/src/io/storage/config/gcs.rs | 6 + crates/iceberg/src/io/storage/config/mod.rs | 12 +- crates/iceberg/src/io/storage/config/s3.rs | 6 + crates/iceberg/src/io/storage/mod.rs | 102 +- crates/storage/opendal/Cargo.toml | 3 +- crates/storage/opendal/public-api.txt | 4 + crates/storage/opendal/src/gcs.rs | 119 +- crates/storage/opendal/src/lib.rs | 163 ++- crates/storage/opendal/src/resolving.rs | 38 +- crates/storage/opendal/src/s3.rs | 90 +- crates/storage/opendal/src/utils.rs | 22 + 21 files changed, 1934 insertions(+), 63 deletions(-) create mode 100644 crates/catalog/rest/src/credential.rs diff --git a/Cargo.lock b/Cargo.lock index b2269dc51f..8a359a3404 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3889,6 +3889,7 @@ dependencies = [ "iceberg_test_utils", "itertools 0.13.0", "mockito", + "rand 0.9.5", "reqwest 0.12.28", "serde", "serde_derive", @@ -4035,6 +4036,7 @@ dependencies = [ "opendal", "reqsign-aws-v4", "reqsign-core", + "reqsign-google", "reqwest 0.12.28", "serde", "tokio", diff --git a/crates/catalog/rest/Cargo.toml b/crates/catalog/rest/Cargo.toml index 247709efd4..d9526aa6bf 100644 --- a/crates/catalog/rest/Cargo.toml +++ b/crates/catalog/rest/Cargo.toml @@ -35,6 +35,7 @@ chrono = { workspace = true } http = { workspace = true } iceberg = { workspace = true } itertools = { workspace = true } +rand = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_derive = { workspace = true } diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 776b11c40a..c3e226fc1c 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -139,6 +139,20 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::ListTablesResponse pub fn iceberg_catalog_rest::ListTablesResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::ListTablesResponse pub fn iceberg_catalog_rest::ListTablesResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::LoadCredentialsResponse +pub iceberg_catalog_rest::LoadCredentialsResponse::storage_credentials: alloc::vec::Vec +impl core::clone::Clone for iceberg_catalog_rest::LoadCredentialsResponse +pub fn iceberg_catalog_rest::LoadCredentialsResponse::clone(&self) -> iceberg_catalog_rest::LoadCredentialsResponse +impl core::cmp::Eq for iceberg_catalog_rest::LoadCredentialsResponse +impl core::cmp::PartialEq for iceberg_catalog_rest::LoadCredentialsResponse +pub fn iceberg_catalog_rest::LoadCredentialsResponse::eq(&self, other: &iceberg_catalog_rest::LoadCredentialsResponse) -> bool +impl core::fmt::Debug for iceberg_catalog_rest::LoadCredentialsResponse +pub fn iceberg_catalog_rest::LoadCredentialsResponse::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::LoadCredentialsResponse +impl serde_core::ser::Serialize for iceberg_catalog_rest::LoadCredentialsResponse +pub fn iceberg_catalog_rest::LoadCredentialsResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::LoadCredentialsResponse +pub fn iceberg_catalog_rest::LoadCredentialsResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::LoadTableResult pub iceberg_catalog_rest::LoadTableResult::config: std::collections::hash::map::HashMap pub iceberg_catalog_rest::LoadTableResult::metadata: iceberg::spec::table_metadata::TableMetadata @@ -288,5 +302,6 @@ pub fn iceberg_catalog_rest::UpdateNamespacePropertiesResponse::serialize<__S>(& impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::UpdateNamespacePropertiesResponse pub fn iceberg_catalog_rest::UpdateNamespacePropertiesResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub const iceberg_catalog_rest::REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str +pub const iceberg_catalog_rest::REST_CATALOG_PROP_SCAN_PLAN_ID: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_URI: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_WAREHOUSE: &str diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 8642b32d22..1a400fd7da 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -54,6 +54,8 @@ pub const REST_CATALOG_PROP_URI: &str = "uri"; pub const REST_CATALOG_PROP_WAREHOUSE: &str = "warehouse"; /// Disable header redaction in error logs (defaults to false for security) pub const REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str = "disable-header-redaction"; +/// Identifier for a server-side scan plan associated with credential requests. +pub const REST_CATALOG_PROP_SCAN_PLAN_ID: &str = "rest.scan.plan-id"; const ICEBERG_REST_SPEC_VERSION: &str = "0.14.1"; const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -164,7 +166,7 @@ impl RestCatalogBuilder { } /// Rest catalog configuration. -#[derive(Clone, Debug, TypedBuilder)] +#[derive(Clone, TypedBuilder)] pub(crate) struct RestCatalogConfig { #[builder(default, setter(strip_option))] name: Option, @@ -181,6 +183,19 @@ pub(crate) struct RestCatalogConfig { client: Option, } +impl std::fmt::Debug for RestCatalogConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Catalog and table properties may contain OAuth credentials, delegated + // storage credentials, or arbitrary secret-bearing `header.*` values. + f.debug_struct("RestCatalogConfig") + .field("name", &self.name) + .field("uri", &self.uri) + .field("warehouse", &self.warehouse) + .field("property_keys", &self.props.keys().collect::>()) + .finish_non_exhaustive() + } +} + impl RestCatalogConfig { fn url_prefixed(&self, parts: &[&str]) -> String { [&self.uri, PATH_V1] @@ -358,7 +373,7 @@ impl RestCatalogConfig { #[derive(Debug)] struct RestContext { - client: HttpClient, + client: Arc, /// Runtime config is fetched from rest server and stored here. /// /// It's could be different from the user config. @@ -447,7 +462,7 @@ impl RestCatalog { Ok(RestContext { config, - client, + client: Arc::new(client), endpoints, }) }) @@ -510,17 +525,17 @@ impl RestCatalog { metadata_location: Option<&str>, extra_config: Option>, ) -> Result { - let mut props = self.context().await?.config.props.clone(); + let context = self.context().await?; + let mut props = context.config.props.clone(); if let Some(config) = extra_config { props.extend(config); } // If the warehouse is a logical identifier instead of a URL we don't want // to raise an exception - let warehouse_path = match self.context().await?.config.warehouse.as_deref() { + let warehouse_path = match context.config.warehouse.as_deref() { Some(url) if Url::parse(url).is_ok() => Some(url), - Some(_) => None, - None => None, + _ => None, }; if metadata_location.or(warehouse_path).is_none() { @@ -541,9 +556,20 @@ impl RestCatalog { ) })?; - let file_io = FileIOBuilder::new(factory).with_props(props).build(); + // If the catalog vends refreshable credentials for this table's storage, + // attach a provider so the backend re-fetches them before they expire. + let credential_provider = crate::credential::build_vended_credential_provider( + context.client.clone(), + &context.config.uri, + &props, + )?; + + let mut builder = FileIOBuilder::new(factory).with_props(props); + if let Some(provider) = credential_provider { + builder = builder.with_credential_provider(provider); + } - Ok(file_io) + Ok(builder.build()) } /// Invalidate the current token without generating a new one. On the next request, the client @@ -1048,7 +1074,14 @@ impl Catalog for RestCatalog { "Metadata location missing in `register_table` response!", ))?; - let file_io = self.load_file_io(Some(metadata_location), None).await?; + let config = response + .config + .into_iter() + .chain(self.user_config.props.clone()) + .collect(); + let file_io = self + .load_file_io(Some(metadata_location), Some(config)) + .await?; let mut table_builder = Table::builder() .identifier(table_ident.clone()) diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs index 07dc0620da..cc4b1c2f97 100644 --- a/crates/catalog/rest/src/client.rs +++ b/crates/catalog/rest/src/client.rs @@ -49,9 +49,13 @@ pub(crate) struct HttpClient { impl Debug for HttpClient { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // Omit the reqwest client: injected clients may carry secret default + // headers. Explicit headers use the same redaction policy as errors. f.debug_struct("HttpClient") - .field("client", &self.client) - .field("extra_headers", &self.extra_headers) + .field( + "extra_headers", + &format_headers_redacted(&self.extra_headers, self.disable_header_redaction), + ) .finish_non_exhaustive() } } @@ -71,6 +75,24 @@ impl HttpClient { }) } + /// Create a client for table-scoped resources while reusing this client's + /// underlying connection pool. + /// + /// A load-table response may supply a table token or `header.*` values that + /// must be used for subsequent table requests such as credential refresh. + pub(crate) fn for_table( + &self, + catalog_uri: &str, + props: HashMap, + ) -> Result { + let cfg = RestCatalogConfig::builder() + .uri(catalog_uri.to_string()) + .props(props) + .client(Some(self.client.clone())) + .build(); + Self::new(&cfg) + } + /// Update the http client with new configuration. /// /// If cfg carries new value, we will use cfg instead. @@ -156,7 +178,6 @@ impl HttpClient { ) .with_context("operation", "auth") .with_context("url", auth_url.to_string()) - .with_context("json", String::from_utf8_lossy(&text)) .with_source(e) })?) } else { @@ -170,7 +191,6 @@ impl HttpClient { .with_context("code", code.to_string()) .with_context("operation", "auth") .with_context("url", auth_url.to_string()) - .with_context("json", String::from_utf8_lossy(&text)) .with_source(e) })?; Err(Error::from(e)) @@ -278,29 +298,30 @@ pub(crate) async fn deserialize_catalog_response( let bytes = response.bytes().await?; serde_json::from_slice::(&bytes).map_err(|e| { + // Successful REST responses can contain OAuth tokens and delegated + // storage credentials. Never copy an unparseable response into an error. Error::new( ErrorKind::Unexpected, "Failed to parse response from rest catalog server", ) - .with_context("json", String::from_utf8_lossy(&bytes)) .with_source(e) }) } -/// Headers that contain sensitive information and should be excluded from logs. -const SENSITIVE_HEADERS: &[&str] = &[ - "authorization", - "proxy-authorization", - "set-cookie", - "cookie", - "x-api-key", - "x-auth-token", -]; - -/// Returns true if the header name is considered sensitive. +/// Returns true if the header may carry a secret. fn is_sensitive_header(name: &str) -> bool { let name_lower = name.to_lowercase(); - SENSITIVE_HEADERS.iter().any(|h| name_lower == *h) + [ + "auth", + "token", + "secret", + "key", + "password", + "cookie", + "credential", + ] + .iter() + .any(|pattern| name_lower.contains(pattern)) } /// Redacts sensitive headers and returns a debug-formatted string. @@ -386,6 +407,8 @@ mod tests { fn test_format_headers_redacted_filters_sensitive() { let mut headers = HeaderMap::new(); headers.insert("authorization", "Bearer secret-token".parse().unwrap()); + headers.insert("x-client-secret", "private-value".parse().unwrap()); + headers.insert("x-client-credential", "credential-value".parse().unwrap()); headers.insert("content-type", "application/json".parse().unwrap()); let result = format_headers_redacted(&headers, false); @@ -395,6 +418,8 @@ mod tests { assert!(result.contains("[REDACTED]")); // Sensitive value should NOT be present assert!(!result.contains("secret-token")); + assert!(!result.contains("private-value")); + assert!(!result.contains("credential-value")); // Non-sensitive header should be present with actual value assert!(result.contains("content-type")); assert!(result.contains("application/json")); diff --git a/crates/catalog/rest/src/credential.rs b/crates/catalog/rest/src/credential.rs new file mode 100644 index 0000000000..14363b3a71 --- /dev/null +++ b/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) -> Result, +} + +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, +} + +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, + consecutive_failures: u32, + retry_not_before: Option, +} + +struct ConfiguredCloud { + cloud: &'static CloudRefresh, + endpoint: String, + cache: Mutex, + /// 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, + /// Optional scan-plan identifier. + plan_id: Option, + /// Independently configured endpoint and cache for each backing cloud. + clouds: Vec, +} + +impl RestVendedCredentialProvider { + fn new(client: Arc, plan_id: Option, clouds: Vec) -> Self { + Self { + client, + plan_id, + clouds, + } + } + + /// Fetch fresh credentials from the catalog's credentials endpoint. + async fn fetch(&self, configured: &ConfiguredCloud) -> Result> { + 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, + ) -> Result { + 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), +} + +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 { + 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, + base_uri: &str, + props: &HashMap, +) -> Result>> { + 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::>(); + + 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 { + if endpoint.starts_with("http://") || endpoint.starts_with("https://") { + return endpoint.to_string(); + } + + let base = base_uri.trim_end_matches('/'); + let separator = if endpoint.starts_with('/') { "" } else { "/" }; + format!("{base}{separator}{endpoint}") +} + +/// The URL scheme of `location`, lowercased (e.g. `"s3"` for `s3://bucket/k`). +fn scheme_of(location: &str) -> Option { + Url::parse(location) + .ok() + .map(|url| url.scheme().to_string()) +} + +/// Parse a complete S3 credential returned by the credentials endpoint. +fn parse_s3_credential(config: &HashMap) -> Result { + let access_key_id = required_nonempty(config, S3_ACCESS_KEY_ID)?; + let secret_access_key = required_nonempty(config, S3_SECRET_ACCESS_KEY)?; + let session_token = required_nonempty(config, S3_SESSION_TOKEN)?; + let expires_at = required_epoch_millis(config, S3_SESSION_TOKEN_EXPIRES_AT_MS)?; + Ok(StorageCredential { + kind: StorageCredentialKind::S3(S3Credential { + access_key_id, + secret_access_key, + session_token: Some(session_token), + }), + expires_at: Some(expires_at), + }) +} + +/// Parse a complete GCS credential returned by the credentials endpoint. +fn parse_gcs_credential(config: &HashMap) -> Result { + let token = required_nonempty(config, GCS_TOKEN)?; + let expires_at = required_epoch_millis(config, GCS_TOKEN_EXPIRES_AT)?; + Ok(StorageCredential { + kind: StorageCredentialKind::Gcs(GcsCredential { token }), + expires_at: Some(expires_at), + }) +} + +fn required_nonempty(config: &HashMap, key: &str) -> Result { + config + .get(key) + .filter(|value| !value.is_empty()) + .cloned() + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("invalid vended credential: {key} is missing or empty"), + ) + }) +} + +fn required_epoch_millis(config: &HashMap, key: &str) -> Result { + let value = required_nonempty(config, key)?; + parse_epoch_millis(&value).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("invalid vended credential: {key} is not a valid epoch-millisecond timestamp"), + ) + }) +} + +/// Parse an epoch-millisecond timestamp into a [`SystemTime`]. +fn parse_epoch_millis(millis: &str) -> Option { + millis + .parse() + .ok() + .and_then(|millis| UNIX_EPOCH.checked_add(Duration::from_millis(millis))) +} + +#[cfg(test)] +mod tests { + use std::sync::{Barrier, mpsc}; + + use mockito::{Matcher, Server}; + + use super::*; + use crate::RestCatalogConfig; + + fn epoch_millis(time: SystemTime) -> String { + time.duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() + .to_string() + } + + fn s3_cred(access_key_id: &str, expires_at: Option) -> StorageCredential { + StorageCredential { + kind: StorageCredentialKind::S3(S3Credential { + access_key_id: access_key_id.to_string(), + secret_access_key: "secret".to_string(), + session_token: None, + }), + expires_at, + } + } + + fn s3_access_key_id(credential: &StorageCredential) -> &str { + match &credential.kind { + StorageCredentialKind::S3(s3) => &s3.access_key_id, + other => panic!("expected S3 credential, got {other:?}"), + } + } + + fn cached_s3(prefix: &str, access_key_id: &str, expires_at: Option) -> CachedEntry { + CachedEntry::new( + prefix.to_string(), + s3_cred(access_key_id, expires_at), + false, + ) + } + + #[test] + fn resolve_endpoint_matches_java_semantics() { + assert_eq!( + resolve_endpoint("https://catalog/", "https://other/creds"), + "https://other/creds" + ); + assert_eq!( + resolve_endpoint("https://catalog", "http://other/creds"), + "http://other/creds" + ); + assert_eq!( + resolve_endpoint("https://catalog", "v1/creds"), + "https://catalog/v1/creds" + ); + assert_eq!( + resolve_endpoint("https://catalog/", "v1/creds"), + "https://catalog/v1/creds" + ); + // All trailing slashes stripped from the base (Java stripTrailingSlash). + assert_eq!( + resolve_endpoint("https://catalog///", "/v1/creds"), + "https://catalog/v1/creds" + ); + // Existing leading slashes on the endpoint are preserved, not collapsed + // (matches Java's resolveEndpoint, which only prepends when absent). + assert_eq!( + resolve_endpoint("https://catalog/", "//v1/creds"), + "https://catalog//v1/creds" + ); + } + + #[test] + fn cloud_selection_uses_url_scheme() { + assert!(CloudRefresh::for_location("s3://b/k").is_some()); + assert!(CloudRefresh::for_location("S3://b/k").is_some()); + assert!(CloudRefresh::for_location("s3a://b/k").is_some()); + assert!(CloudRefresh::for_location("s3n://b/k").is_some()); + assert!(CloudRefresh::for_location("gs://b/k").is_some()); + assert!(CloudRefresh::for_location("gcs://b/k").is_some()); + assert!(CloudRefresh::for_location("not a url").is_none()); + // Azure not supported yet -> no provider, static creds used as-is. + assert!(CloudRefresh::for_location("abfss://fs@acct.dfs.core.windows.net/k").is_none()); + assert!(CloudRefresh::AWS.matches_credential_prefix("s3://bucket/path")); + assert!(!CloudRefresh::AWS.matches_credential_prefix("s3evil://bucket/path")); + } + + #[test] + fn parses_s3_credential() { + let mut config = HashMap::new(); + assert!(parse_s3_credential(&config).is_err()); + config.insert(S3_ACCESS_KEY_ID.to_string(), "AK".to_string()); + assert!(parse_s3_credential(&config).is_err()); + config.insert(S3_SECRET_ACCESS_KEY.to_string(), "SK".to_string()); + assert!(parse_s3_credential(&config).is_err()); + + config.insert(S3_SESSION_TOKEN.to_string(), "TOK".to_string()); + config.insert( + S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), + "1500".to_string(), + ); + let credential = parse_s3_credential(&config).unwrap(); + assert_eq!( + credential.expires_at, + Some(UNIX_EPOCH + Duration::from_millis(1500)) + ); + match credential.kind { + StorageCredentialKind::S3(s3) => assert_eq!(s3.session_token.as_deref(), Some("TOK")), + other => panic!("expected S3, got {other:?}"), + } + } + + #[test] + fn parse_gcs_requires_token_and_expiry() { + let mut config = HashMap::new(); + assert!(parse_gcs_credential(&config).is_err()); + config.insert(GCS_TOKEN.to_string(), "ya29.token".to_string()); + assert!(parse_gcs_credential(&config).is_err()); + + config.insert(GCS_TOKEN_EXPIRES_AT.to_string(), "2000".to_string()); + let credential = parse_gcs_credential(&config).unwrap(); + match &credential.kind { + StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token, "ya29.token"), + other => panic!("expected GCS, got {other:?}"), + } + assert_eq!( + credential.expires_at, + Some(UNIX_EPOCH + Duration::from_millis(2000)) + ); + } + + #[test] + fn cached_entry_freshness() { + let no_expiry = cached_s3("", "a", None); + let far = cached_s3( + "", + "a", + Some(SystemTime::now() + REFRESH_BUFFER + Duration::from_secs(60)), + ); + // Within the buffer but not yet expired: stale for a fast-path read, but + // still usable for graceful degradation. + let soon = cached_s3("", "a", Some(SystemTime::now() + Duration::from_secs(60))); + let past = cached_s3("", "a", Some(SystemTime::now() - Duration::from_secs(60))); + + let now = SystemTime::now(); + assert!(no_expiry.is_fresh(now)); + assert!(no_expiry.is_unexpired(now)); + assert!(far.is_fresh(now)); + assert!(!soon.is_fresh(now)); + assert!(soon.is_unexpired(now)); + assert!(!past.is_unexpired(now)); + } + + #[test] + fn prefetch_time_matches_cloud_policy() { + let expires_at = SystemTime::now() + Duration::from_secs(3600); + assert_eq!( + prefetch_time(expires_at, false), + expires_at - REFRESH_BUFFER + ); + + for _ in 0..100 { + let refresh_at = prefetch_time(expires_at, true); + assert!(refresh_at >= expires_at - REFRESH_BUFFER); + assert!(refresh_at < expires_at - MIN_REFRESH_BUFFER); + } + } + + #[test] + fn seed_inside_nominal_window_is_immediately_due() { + let credential = s3_cred( + "a", + Some(SystemTime::now() + REFRESH_BUFFER - Duration::from_secs(1)), + ); + let entry = CachedEntry::seed(String::new(), credential, true); + let now = SystemTime::now(); + assert!(!entry.is_fresh(now)); + assert!(entry.is_unexpired(now)); + } + + #[test] + fn failure_backoff_is_jittered_and_capped() { + for failures in 1_u32..=40 { + let ceiling = INITIAL_FAILURE_BACKOFF + .checked_mul(1_u32 << failures.saturating_sub(1).min(5)) + .unwrap_or(MAX_FAILURE_BACKOFF) + .min(MAX_FAILURE_BACKOFF); + for _ in 0..100 { + let backoff = failure_backoff(failures); + assert!(backoff >= ceiling / 2); + assert!(backoff <= ceiling); + } + } + } + + #[test] + fn longest_prefix_match_ignores_freshness() { + let far = Some(SystemTime::now() + REFRESH_BUFFER + Duration::from_secs(3600)); + let entries = vec![ + cached_s3("s3://bucket", "wide", far), + cached_s3("s3://bucket/warehouse/db", "narrow", far), + ]; + let got = longest_prefix_match(&entries, "s3://bucket/warehouse/db/t/f").unwrap(); + assert_eq!(s3_access_key_id(&got.credential), "narrow"); + assert!(longest_prefix_match(&entries, "s3://other/x").is_none()); + + let fresh = Some(SystemTime::now() + REFRESH_BUFFER + Duration::from_secs(60)); + let stale = Some(SystemTime::now() - Duration::from_secs(60)); + let entries = vec![ + cached_s3("s3://bucket", "wide", fresh), + cached_s3("s3://bucket/table", "narrow-stale", stale), + ]; + let selected = longest_prefix_match(&entries, "s3://bucket/table/f").unwrap(); + assert_eq!(s3_access_key_id(&selected.credential), "narrow-stale"); + assert!(!selected.is_fresh(SystemTime::now())); + } + + #[test] + fn provider_support_tracks_configured_cloud_endpoints() { + let config = RestCatalogConfig::builder() + .uri("http://cat".to_string()) + .build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let props = HashMap::from([( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/creds".to_string(), + )]); + + // The provider is configured independently of the table metadata scheme, + // but advertises support only for clouds whose endpoint is present. + let provider = build_vended_credential_provider(client.clone(), "http://cat", &props) + .unwrap() + .unwrap(); + assert!(provider.supports_path("s3://b/k")); + assert!(!provider.supports_path("abfss://fs@acct.dfs.core.windows.net/k")); + // No endpoint advertised. + assert!( + build_vended_credential_provider(client.clone(), "http://cat", &HashMap::new()) + .unwrap() + .is_none() + ); + // Explicitly disabled. + let disabled = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/creds".to_string(), + ), + ( + CloudRefresh::AWS.enabled_key.to_string(), + "false".to_string(), + ), + ]); + assert!( + build_vended_credential_provider(client, "http://cat", &disabled) + .unwrap() + .is_none() + ); + // Java's `Strings.isNullOrEmpty` check treats an empty endpoint as absent. + let empty = HashMap::from([(CloudRefresh::AWS.endpoint_key.to_string(), String::new())]); + let config = RestCatalogConfig::builder() + .uri("http://cat".to_string()) + .build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + assert!( + build_vended_credential_provider(client, "http://cat", &empty) + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn fetches_and_selects_vended_credential() { + let mut server = Server::new_async().await; + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let body = format!( + r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .match_query(Matcher::UrlEncoded( + "planId".to_string(), + "scan-plan-1".to_string(), + )) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + // No static creds -> no seed -> first load fetches from the endpoint. + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + ( + REST_CATALOG_PROP_SCAN_PLAN_ID.to_string(), + "scan-plan-1".to_string(), + ), + ]); + + let provider = build_vended_credential_provider(client, &server.url(), &props) + .expect("provider construction should succeed") + .expect("provider should be built"); + let credential = provider + .load_credential("s3://bucket/warehouse/f") + .await + .unwrap(); + + match credential.kind { + StorageCredentialKind::S3(s3) => { + assert_eq!(s3.access_key_id, "AK"); + assert_eq!(s3.session_token.as_deref(), Some("TOK")); + } + other => panic!("expected S3, got {other:?}"), + } + mock.assert_async().await; + } + + #[tokio::test] + async fn fresh_seed_is_served_without_fetching() { + let mut server = Server::new_async().await; + // Any call to the server is a failure: the fresh seed must be reused. + let mock = server + .mock("GET", Matcher::Any) + .expect(0) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), + (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), + (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), + (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expires), + ]); + + let provider = build_vended_credential_provider(client, &server.url(), &props) + .expect("provider construction should succeed") + .expect("provider should be built"); + let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); + + assert_eq!(s3_access_key_id(&credential), "SEED_AK"); + mock.assert_async().await; + } + + #[tokio::test] + async fn failed_refresh_is_backed_off_while_credential_is_unexpired() { + let mut server = Server::new_async().await; + // A refresh is due (the seed is within the buffer) but the catalog errors. + let mock = server + .mock("GET", Matcher::Any) + .expect(1) + .with_status(500) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + // Seeded credential is within the refresh buffer but not yet expired. + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), + (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), + (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), + (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expires), + ]); + + let provider = build_vended_credential_provider(client, &server.url(), &props) + .expect("provider construction should succeed") + .expect("provider should be built"); + // The first refresh fails, but the still-valid seed is served. Immediate + // follow-up operations stay inside the first jittered backoff window. + for _ in 0..3 { + let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); + assert_eq!(s3_access_key_id(&credential), "SEED_AK"); + } + mock.assert_async().await; + } + + #[tokio::test] + async fn concurrent_prefetch_serves_unexpired_credential_without_waiting() { + let mut server = Server::new_async().await; + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let body = format!( + r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"NEW_AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + ); + let (started_tx, started_rx) = mpsc::channel(); + let release = Arc::new(Barrier::new(2)); + let callback_release = Arc::clone(&release); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_chunked_body(move |writer| { + started_tx.send(()).unwrap(); + callback_release.wait(); + writer.write_all(body.as_bytes()) + }) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let seed_expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), + (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), + (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), + (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), seed_expires), + ]); + let provider = build_vended_credential_provider(client, &server.url(), &props) + .unwrap() + .unwrap(); + + let first_provider = Arc::clone(&provider); + let first = + tokio::spawn(async move { first_provider.load_credential("s3://bucket/x/f").await }); + tokio::task::spawn_blocking(move || { + started_rx.recv_timeout(Duration::from_secs(5)).unwrap() + }) + .await + .unwrap(); + + let second_provider = Arc::clone(&provider); + let second = + tokio::spawn(async move { second_provider.load_credential("s3://bucket/x/f").await }); + for _ in 0..100 { + if second.is_finished() { + break; + } + tokio::task::yield_now().await; + } + let completed_without_waiting = second.is_finished(); + release.wait(); + + assert!( + completed_without_waiting, + "a concurrent prefetch waited instead of using the unexpired credential" + ); + assert_eq!(s3_access_key_id(&second.await.unwrap().unwrap()), "SEED_AK"); + assert_eq!(s3_access_key_id(&first.await.unwrap().unwrap()), "NEW_AK"); + mock.assert_async().await; + } + + #[tokio::test] + async fn refresh_uses_table_scoped_token() { + let mut server = Server::new_async().await; + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let body = format!( + r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .match_header("authorization", "Bearer table-token") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + ("token".to_string(), "table-token".to_string()), + ]); + let provider = build_vended_credential_provider(client, &server.url(), &props) + .unwrap() + .unwrap(); + + provider.load_credential("s3://bucket/x/f").await.unwrap(); + mock.assert_async().await; + } + + #[tokio::test] + async fn one_provider_refreshes_multiple_clouds() { + let mut server = Server::new_async().await; + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let aws_body = format!( + r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + ); + let gcp_body = format!( + r#"{{"storage-credentials":[{{"prefix":"gs://bucket","config":{{"gcs.oauth2.token":"GCS","gcs.oauth2.token-expires-at":"{expires}"}}}}]}}"# + ); + let aws_mock = server + .mock("GET", "/v1/aws-credentials") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(aws_body) + .create_async() + .await; + let gcp_mock = server + .mock("GET", "/v1/gcp-credentials") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(gcp_body) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/aws-credentials".to_string(), + ), + ( + CloudRefresh::GCP.endpoint_key.to_string(), + "/v1/gcp-credentials".to_string(), + ), + ]); + let provider = build_vended_credential_provider(client, &server.url(), &props) + .unwrap() + .unwrap(); + + assert!(provider.supports_path("s3://bucket/x")); + assert!(provider.supports_path("gs://bucket/x")); + assert_eq!( + s3_access_key_id(&provider.load_credential("s3://bucket/x").await.unwrap()), + "AK" + ); + match provider + .load_credential("gs://bucket/x") + .await + .unwrap() + .kind + { + StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token, "GCS"), + other => panic!("expected GCS credential, got {other:?}"), + } + aws_mock.assert_async().await; + gcp_mock.assert_async().await; + } + + #[tokio::test] + async fn refresh_failure_never_serves_expired_credential() { + let mut server = Server::new_async().await; + let mock = server + .mock("GET", Matcher::Any) + .with_status(500) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + let expired = epoch_millis(SystemTime::now() - Duration::from_secs(60)); + let props = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + ), + (S3_ACCESS_KEY_ID.to_string(), "EXPIRED_AK".to_string()), + (S3_SECRET_ACCESS_KEY.to_string(), "EXPIRED_SK".to_string()), + (S3_SESSION_TOKEN.to_string(), "EXPIRED_TOK".to_string()), + (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expired), + ]); + + let provider = build_vended_credential_provider(client, &server.url(), &props) + .expect("provider construction should succeed") + .expect("provider should be built"); + + assert!(provider.load_credential("s3://bucket/x/f").await.is_err()); + mock.assert_async().await; + } + + #[tokio::test] + async fn sub_buffer_ttl_is_refetched_on_each_sequential_operation() { + let mut server = Server::new_async().await; + // The vended TTL (60s) is shorter than REFRESH_BUFFER, so every operation + // is eligible for refresh. This matches AWS CachedSupplier: successful + // results do not receive failure backoff. + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); + let body = format!( + r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .expect(3) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let config = RestCatalogConfig::builder().uri(server.url()).build(); + let client = Arc::new(HttpClient::new(&config).unwrap()); + // No static creds -> no seed -> each sequential load fetches because the + // returned credential is already inside its prefetch window. + let props = HashMap::from([( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/credentials".to_string(), + )]); + + let provider = build_vended_credential_provider(client, &server.url(), &props) + .expect("provider construction should succeed") + .expect("provider should be built"); + + for _ in 0..3 { + let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); + assert_eq!(s3_access_key_id(&credential), "AK"); + } + mock.assert_async().await; + } +} diff --git a/crates/catalog/rest/src/lib.rs b/crates/catalog/rest/src/lib.rs index 383728401f..a3f13742b8 100644 --- a/crates/catalog/rest/src/lib.rs +++ b/crates/catalog/rest/src/lib.rs @@ -53,6 +53,7 @@ mod catalog; mod client; +mod credential; mod endpoint; mod types; diff --git a/crates/catalog/rest/src/types.rs b/crates/catalog/rest/src/types.rs index 390521229e..8265ceb40b 100644 --- a/crates/catalog/rest/src/types.rs +++ b/crates/catalog/rest/src/types.rs @@ -101,7 +101,7 @@ impl From for Error { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Serialize, Deserialize)] pub(super) struct TokenResponse { pub(super) access_token: String, pub(super) token_type: String, @@ -205,7 +205,7 @@ pub struct RenameTableRequest { pub destination: TableIdent, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "kebab-case")] /// Result returned when a table is successfully loaded or created. /// @@ -231,7 +231,18 @@ pub struct LoadTableResult { pub storage_credentials: Option>, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +impl std::fmt::Debug for LoadTableResult { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadTableResult") + .field("metadata_location", &self.metadata_location) + .field("metadata", &self.metadata) + .field("config_keys", &self.config.keys().collect::>()) + .field("storage_credentials", &self.storage_credentials) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)] /// Storage credential for a specific location prefix. /// /// Indicates a storage location prefix where the credential is relevant. Clients should @@ -244,6 +255,27 @@ pub struct StorageCredential { pub config: HashMap, } +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::>()) + .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 { + /// Storage credentials, one entry per location prefix. + pub storage_credentials: Vec, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "kebab-case")] /// Request to create a new table in a namespace. diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 610675c83d..2771ffd998 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -689,6 +689,13 @@ pub fn iceberg::inspect::SnapshotsTable<'a>::new(table: &'a iceberg::table::Tabl pub async fn iceberg::inspect::SnapshotsTable<'a>::scan(&self) -> iceberg::Result pub fn iceberg::inspect::SnapshotsTable<'a>::schema(&self) -> iceberg::spec::Schema pub mod iceberg::io +pub enum iceberg::io::StorageCredentialKind +pub iceberg::io::StorageCredentialKind::Gcs(iceberg::io::GcsCredential) +pub iceberg::io::StorageCredentialKind::S3(iceberg::io::S3Credential) +impl core::clone::Clone for iceberg::io::StorageCredentialKind +pub fn iceberg::io::StorageCredentialKind::clone(&self) -> iceberg::io::StorageCredentialKind +impl core::fmt::Debug for iceberg::io::StorageCredentialKind +pub fn iceberg::io::StorageCredentialKind::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::io::AzdlsConfig pub iceberg::io::AzdlsConfig::account_key: core::option::Option pub iceberg::io::AzdlsConfig::account_name: core::option::Option @@ -739,6 +746,7 @@ impl iceberg::io::FileIOBuilder pub fn iceberg::io::FileIOBuilder::build(self) -> iceberg::io::FileIO pub fn iceberg::io::FileIOBuilder::config(&self) -> &iceberg::io::StorageConfig pub fn iceberg::io::FileIOBuilder::new(factory: alloc::sync::Arc) -> Self +pub fn iceberg::io::FileIOBuilder::with_credential_provider(self, provider: alloc::sync::Arc) -> Self pub fn iceberg::io::FileIOBuilder::with_prop(self, key: impl alloc::string::ToString, value: impl alloc::string::ToString) -> Self pub fn iceberg::io::FileIOBuilder::with_props(self, args: impl core::iter::traits::collect::IntoIterator) -> Self impl core::clone::Clone for iceberg::io::FileIOBuilder @@ -775,6 +783,12 @@ impl serde_core::ser::Serialize for iceberg::io::GcsConfig pub fn iceberg::io::GcsConfig::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::GcsConfig pub fn iceberg::io::GcsConfig::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg::io::GcsCredential +pub iceberg::io::GcsCredential::token: alloc::string::String +impl core::clone::Clone for iceberg::io::GcsCredential +pub fn iceberg::io::GcsCredential::clone(&self) -> iceberg::io::GcsCredential +impl core::fmt::Debug for iceberg::io::GcsCredential +pub fn iceberg::io::GcsCredential::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::io::HfConfig pub iceberg::io::HfConfig::endpoint: core::option::Option pub iceberg::io::HfConfig::revision: core::option::Option @@ -842,6 +856,7 @@ impl core::fmt::Debug for iceberg::io::LocalFsStorageFactory pub fn iceberg::io::LocalFsStorageFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::io::StorageFactory for iceberg::io::LocalFsStorageFactory pub fn iceberg::io::LocalFsStorageFactory::build(&self, _config: &iceberg::io::StorageConfig) -> iceberg::Result> +pub fn iceberg::io::LocalFsStorageFactory::build_with_credentials(&self, config: &iceberg::io::StorageConfig, credential_provider: core::option::Option>) -> iceberg::Result> impl serde_core::ser::Serialize for iceberg::io::LocalFsStorageFactory pub fn iceberg::io::LocalFsStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::LocalFsStorageFactory @@ -880,6 +895,7 @@ impl core::fmt::Debug for iceberg::io::MemoryStorageFactory pub fn iceberg::io::MemoryStorageFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::io::StorageFactory for iceberg::io::MemoryStorageFactory pub fn iceberg::io::MemoryStorageFactory::build(&self, _config: &iceberg::io::StorageConfig) -> iceberg::Result> +pub fn iceberg::io::MemoryStorageFactory::build_with_credentials(&self, config: &iceberg::io::StorageConfig, credential_provider: core::option::Option>) -> iceberg::Result> impl serde_core::ser::Serialize for iceberg::io::MemoryStorageFactory pub fn iceberg::io::MemoryStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::MemoryStorageFactory @@ -955,6 +971,14 @@ impl serde_core::ser::Serialize for iceberg::io::S3Config pub fn iceberg::io::S3Config::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::S3Config pub fn iceberg::io::S3Config::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg::io::S3Credential +pub iceberg::io::S3Credential::access_key_id: alloc::string::String +pub iceberg::io::S3Credential::secret_access_key: alloc::string::String +pub iceberg::io::S3Credential::session_token: core::option::Option +impl core::clone::Clone for iceberg::io::S3Credential +pub fn iceberg::io::S3Credential::clone(&self) -> iceberg::io::S3Credential +impl core::fmt::Debug for iceberg::io::S3Credential +pub fn iceberg::io::S3Credential::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::io::StorageConfig impl iceberg::io::StorageConfig pub fn iceberg::io::StorageConfig::from_props(props: std::collections::hash::map::HashMap) -> Self @@ -992,6 +1016,13 @@ impl serde_core::ser::Serialize for iceberg::io::StorageConfig pub fn iceberg::io::StorageConfig::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::StorageConfig pub fn iceberg::io::StorageConfig::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg::io::StorageCredential +pub iceberg::io::StorageCredential::expires_at: core::option::Option +pub iceberg::io::StorageCredential::kind: iceberg::io::StorageCredentialKind +impl core::clone::Clone for iceberg::io::StorageCredential +pub fn iceberg::io::StorageCredential::clone(&self) -> iceberg::io::StorageCredential +impl core::fmt::Debug for iceberg::io::StorageCredential +pub fn iceberg::io::StorageCredential::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub const iceberg::io::ADLS_ACCOUNT_KEY: &str pub const iceberg::io::ADLS_ACCOUNT_NAME: &str pub const iceberg::io::ADLS_AUTHORITY_HOST: &str @@ -1000,6 +1031,8 @@ pub const iceberg::io::ADLS_CLIENT_SECRET: &str pub const iceberg::io::ADLS_CONNECTION_STRING: &str pub const iceberg::io::ADLS_SAS_TOKEN: &str pub const iceberg::io::ADLS_TENANT_ID: &str +pub const iceberg::io::AWS_REFRESH_CREDENTIALS_ENABLED: &str +pub const iceberg::io::AWS_REFRESH_CREDENTIALS_ENDPOINT: &str pub const iceberg::io::CLIENT_REGION: &str pub const iceberg::io::GCS_ALLOW_ANONYMOUS: &str pub const iceberg::io::GCS_CREDENTIALS_JSON: &str @@ -1007,8 +1040,11 @@ pub const iceberg::io::GCS_DISABLE_CONFIG_LOAD: &str pub const iceberg::io::GCS_DISABLE_VM_METADATA: &str pub const iceberg::io::GCS_NO_AUTH: &str pub const iceberg::io::GCS_PROJECT_ID: &str +pub const iceberg::io::GCS_REFRESH_CREDENTIALS_ENABLED: &str +pub const iceberg::io::GCS_REFRESH_CREDENTIALS_ENDPOINT: &str pub const iceberg::io::GCS_SERVICE_PATH: &str pub const iceberg::io::GCS_TOKEN: &str +pub const iceberg::io::GCS_TOKEN_EXPIRES_AT: &str pub const iceberg::io::GCS_USER_PROJECT: &str pub const iceberg::io::HF_ENDPOINT: &str pub const iceberg::io::HF_REVISION: &str @@ -1028,6 +1064,7 @@ pub const iceberg::io::S3_PATH_STYLE_ACCESS: &str pub const iceberg::io::S3_REGION: &str pub const iceberg::io::S3_SECRET_ACCESS_KEY: &str pub const iceberg::io::S3_SESSION_TOKEN: &str +pub const iceberg::io::S3_SESSION_TOKEN_EXPIRES_AT_MS: &str pub const iceberg::io::S3_SSE_KEY: &str pub const iceberg::io::S3_SSE_MD5: &str pub const iceberg::io::S3_SSE_TYPE: &str @@ -1079,12 +1116,18 @@ pub fn iceberg::io::MemoryStorage::read<'life0, 'life1, 'async_trait>(&'life0 se pub fn iceberg::io::MemoryStorage::reader<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg::io::MemoryStorage::write<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str, bs: bytes::bytes::Bytes) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg::io::MemoryStorage::writer<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub trait iceberg::io::StorageCredentialProvider: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg::io::StorageCredentialProvider::load_credential<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg::io::StorageCredentialProvider::supports_path(&self, _path: &str) -> bool pub trait iceberg::io::StorageFactory: core::fmt::Debug + core::marker::Send + core::marker::Sync + typetag::Serialize + typetag::Deserialize pub fn iceberg::io::StorageFactory::build(&self, config: &iceberg::io::StorageConfig) -> iceberg::Result> +pub fn iceberg::io::StorageFactory::build_with_credentials(&self, config: &iceberg::io::StorageConfig, credential_provider: core::option::Option>) -> iceberg::Result> impl iceberg::io::StorageFactory for iceberg::io::LocalFsStorageFactory pub fn iceberg::io::LocalFsStorageFactory::build(&self, _config: &iceberg::io::StorageConfig) -> iceberg::Result> +pub fn iceberg::io::LocalFsStorageFactory::build_with_credentials(&self, config: &iceberg::io::StorageConfig, credential_provider: core::option::Option>) -> iceberg::Result> impl iceberg::io::StorageFactory for iceberg::io::MemoryStorageFactory pub fn iceberg::io::MemoryStorageFactory::build(&self, _config: &iceberg::io::StorageConfig) -> iceberg::Result> +pub fn iceberg::io::MemoryStorageFactory::build_with_credentials(&self, config: &iceberg::io::StorageConfig, credential_provider: core::option::Option>) -> iceberg::Result> pub mod iceberg::memory pub struct iceberg::memory::MemoryCatalog impl core::fmt::Debug for iceberg::memory::MemoryCatalog diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index cd0a4434c4..90d9e48785 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -22,7 +22,8 @@ use bytes::Bytes; use futures::{Stream, StreamExt}; use super::storage::{ - LocalFsStorageFactory, MemoryStorageFactory, Storage, StorageConfig, StorageFactory, + LocalFsStorageFactory, MemoryStorageFactory, Storage, StorageConfig, StorageCredentialProvider, + StorageFactory, }; use crate::Result; @@ -65,6 +66,8 @@ pub struct FileIO { config: StorageConfig, /// Factory for creating storage instances factory: Arc, + /// Optional provider of refreshable, backend-specific credentials + credential_provider: Option>, /// Cached storage instance (lazily initialized) storage: Arc>>, } @@ -77,6 +80,7 @@ impl FileIO { Self { config: StorageConfig::new(), factory: Arc::new(MemoryStorageFactory), + credential_provider: None, storage: Arc::new(OnceLock::new()), } } @@ -88,6 +92,7 @@ impl FileIO { Self { config: StorageConfig::new(), factory: Arc::new(LocalFsStorageFactory), + credential_provider: None, storage: Arc::new(OnceLock::new()), } } @@ -107,8 +112,11 @@ impl FileIO { return Ok(storage.clone()); } - // Build the storage - let storage = self.factory.build(&self.config)?; + // Build the storage, passing any credential provider so backends that + // support refreshable credentials can wire it into their operators. + let storage = self + .factory + .build_with_credentials(&self.config, self.credential_provider.clone())?; // Try to set it (another thread might have set it first) let _ = self.storage.set(storage.clone()); @@ -191,6 +199,8 @@ pub struct FileIOBuilder { factory: Arc, /// Storage configuration config: StorageConfig, + /// Optional provider of refreshable, backend-specific credentials + credential_provider: Option>, } impl FileIOBuilder { @@ -199,6 +209,7 @@ impl FileIOBuilder { Self { factory, config: StorageConfig::new(), + credential_provider: None, } } @@ -224,11 +235,21 @@ impl FileIOBuilder { &self.config } + /// Attach a provider of refreshable, backend-specific credentials. + pub fn with_credential_provider( + mut self, + provider: Arc, + ) -> Self { + self.credential_provider = Some(provider); + self + } + /// Builds [`FileIO`]. pub fn build(self) -> FileIO { FileIO { config: self.config, factory: self.factory, + credential_provider: self.credential_provider, storage: Arc::new(OnceLock::new()), } } diff --git a/crates/iceberg/src/io/storage/config/gcs.rs b/crates/iceberg/src/io/storage/config/gcs.rs index 5b11567f4d..1392e35fd1 100644 --- a/crates/iceberg/src/io/storage/config/gcs.rs +++ b/crates/iceberg/src/io/storage/config/gcs.rs @@ -41,6 +41,12 @@ pub const GCS_NO_AUTH: &str = "gcs.no-auth"; pub const GCS_CREDENTIALS_JSON: &str = "gcs.credentials-json"; /// Google Cloud Storage token. pub const GCS_TOKEN: &str = "gcs.oauth2.token"; +/// Epoch-millisecond timestamp at which the vended GCS OAuth2 token expires. +pub const GCS_TOKEN_EXPIRES_AT: &str = "gcs.oauth2.token-expires-at"; +/// Endpoint used to fetch and refresh vended GCS OAuth2 credentials. +pub const GCS_REFRESH_CREDENTIALS_ENDPOINT: &str = "gcs.oauth2.refresh-credentials-endpoint"; +/// Whether vended GCS OAuth2 credentials should be refreshed. Defaults to `true`. +pub const GCS_REFRESH_CREDENTIALS_ENABLED: &str = "gcs.oauth2.refresh-credentials-enabled"; /// Option to skip signing requests (e.g. for public buckets/folders). pub const GCS_ALLOW_ANONYMOUS: &str = "gcs.allow-anonymous"; /// Option to skip loading the credential from GCE metadata server. diff --git a/crates/iceberg/src/io/storage/config/mod.rs b/crates/iceberg/src/io/storage/config/mod.rs index 2350aab6dd..a79f871e8d 100644 --- a/crates/iceberg/src/io/storage/config/mod.rs +++ b/crates/iceberg/src/io/storage/config/mod.rs @@ -51,12 +51,22 @@ use serde::{Deserialize, Serialize}; /// which storage backend to use. The storage type is determined by the /// explicit factory selection. /// ``` -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)] +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct StorageConfig { /// Configuration properties for the storage backend props: HashMap, } +impl std::fmt::Debug for StorageConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Property values may hold vended credentials (e.g. `s3.secret-access-key`, + // `s3.session-token`) + f.debug_struct("StorageConfig") + .field("keys", &self.props.keys().collect::>()) + .finish_non_exhaustive() + } +} + impl StorageConfig { /// Create a new empty StorageConfig. pub fn new() -> Self { diff --git a/crates/iceberg/src/io/storage/config/s3.rs b/crates/iceberg/src/io/storage/config/s3.rs index 664e8637b8..eb8c9f54c5 100644 --- a/crates/iceberg/src/io/storage/config/s3.rs +++ b/crates/iceberg/src/io/storage/config/s3.rs @@ -35,10 +35,16 @@ pub const S3_ACCESS_KEY_ID: &str = "s3.access-key-id"; pub const S3_SECRET_ACCESS_KEY: &str = "s3.secret-access-key"; /// S3 session token (required when using temporary credentials). pub const S3_SESSION_TOKEN: &str = "s3.session-token"; +/// Epoch-millisecond timestamp at which the vended S3 session token expires. +pub const S3_SESSION_TOKEN_EXPIRES_AT_MS: &str = "s3.session-token-expires-at-ms"; /// S3 region. pub const S3_REGION: &str = "s3.region"; /// Region to use for the S3 client (takes precedence over [`S3_REGION`]). pub const CLIENT_REGION: &str = "client.region"; +/// Endpoint used to fetch and refresh vended AWS credentials. +pub const AWS_REFRESH_CREDENTIALS_ENDPOINT: &str = "client.refresh-credentials-endpoint"; +/// Whether vended AWS credentials should be refreshed. Defaults to `true`. +pub const AWS_REFRESH_CREDENTIALS_ENABLED: &str = "client.refresh-credentials-enabled"; /// S3 Path Style Access. pub const S3_PATH_STYLE_ACCESS: &str = "s3.path-style-access"; /// S3 Server Side Encryption Type. diff --git a/crates/iceberg/src/io/storage/mod.rs b/crates/iceberg/src/io/storage/mod.rs index 5276c7771f..b10bd380c0 100644 --- a/crates/iceberg/src/io/storage/mod.rs +++ b/crates/iceberg/src/io/storage/mod.rs @@ -23,6 +23,7 @@ mod memory; use std::fmt::Debug; use std::sync::Arc; +use std::time::SystemTime; use async_trait::async_trait; use bytes::Bytes; @@ -32,7 +33,7 @@ pub use local_fs::{LocalFsStorage, LocalFsStorageFactory}; pub use memory::{MemoryStorage, MemoryStorageFactory}; use super::{FileMetadata, FileRead, FileWrite, InputFile, OutputFile}; -use crate::Result; +use crate::{Error, ErrorKind, Result}; /// Trait for storage operations in Iceberg. /// @@ -139,4 +140,103 @@ pub trait StorageFactory: Debug + Send + Sync { /// A `Result` containing an `Arc` on success, or an error /// if the storage could not be created. fn build(&self, config: &StorageConfig) -> Result>; + + /// Build a new Storage instance, optionally supplying a credential provider + /// that the backend can call to obtain and refresh short-lived credentials. + fn build_with_credentials( + &self, + config: &StorageConfig, + credential_provider: Option>, + ) -> Result> { + if credential_provider.is_some() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Storage factory does not support refreshable credential providers", + )); + } + + 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. +/// +/// # Caching +/// +/// [`load_credential`](Self::load_credential) may be called very frequently — +/// the S3 backend, for example, rebuilds its operator (and therefore its +/// signer) on every file operation. Implementations must cache internally and +/// only re-fetch when the current credential is at or near expiry; otherwise +/// every object-store request would trigger a call back to the catalog. +#[async_trait] +pub trait StorageCredentialProvider: Debug + Send + Sync { + /// Return whether this provider has refresh configuration for `path`. + /// + /// Backends use this before replacing their normal credential chain. The + /// default is `true` for single-backend providers; multi-backend providers + /// should return `false` for schemes they do not configure. + fn supports_path(&self, _path: &str) -> bool { + true + } + + /// Load a fresh credential for the storage location identified by `path`. + /// + /// `path` is the absolute location being accessed (e.g. + /// `s3://bucket/warehouse/db/table/...`). Providers that vend distinct + /// credentials per location prefix use it to select the most specific + /// match. + async fn load_credential(&self, path: &str) -> Result; +} + +/// A vended storage credential together with when it expires. +#[derive(Clone, Debug)] +pub struct StorageCredential { + /// The backend-specific credential material. + pub kind: StorageCredentialKind, + /// When the credential expires, if known. `None` means non-expiring and + /// backends treat such a credential as always valid and never refresh it. + pub expires_at: Option, +} + +/// Backend-specific credential material. +#[derive(Clone, Debug)] +pub enum StorageCredentialKind { + /// Amazon S3 credentials. + S3(S3Credential), + /// Google Cloud Storage credentials. + Gcs(GcsCredential), +} + +/// Temporary Amazon S3 credentials. +#[derive(Clone)] +pub struct S3Credential { + /// AWS access key ID. + pub access_key_id: String, + /// AWS secret access key. + pub secret_access_key: String, + /// AWS session token, set for temporary (STS/vended) credentials. + pub session_token: Option, +} + +impl Debug for S3Credential { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("S3Credential").finish_non_exhaustive() + } +} + +/// Temporary Google Cloud Storage credentials (an OAuth2 access token). +#[derive(Clone)] +pub struct GcsCredential { + /// OAuth2 bearer token used to access GCS. + pub token: String, +} + +impl Debug for GcsCredential { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GcsCredential").finish_non_exhaustive() + } } diff --git a/crates/storage/opendal/Cargo.toml b/crates/storage/opendal/Cargo.toml index e43e7845b3..4cbf2892d9 100644 --- a/crates/storage/opendal/Cargo.toml +++ b/crates/storage/opendal/Cargo.toml @@ -41,7 +41,7 @@ opendal-all = [ opendal-azdls = ["opendal/services-azdls"] opendal-fs = ["opendal/services-fs"] -opendal-gcs = ["opendal/services-gcs"] +opendal-gcs = ["opendal/services-gcs", "reqsign-google", "reqsign-core"] opendal-hf = ["opendal/services-hf"] opendal-memory = ["opendal/services-memory"] opendal-oss = ["opendal/services-oss"] @@ -57,6 +57,7 @@ iceberg = { workspace = true } opendal = { workspace = true } reqsign-aws-v4 = { version = "3.0.0", optional = true } reqsign-core = { version = "3.0.0", optional = true } +reqsign-google = { version = "3.0.0", optional = true } serde = { workspace = true } typetag = { workspace = true } url = { workspace = true } diff --git a/crates/storage/opendal/public-api.txt b/crates/storage/opendal/public-api.txt index fc1ed7cf78..f420a891f1 100644 --- a/crates/storage/opendal/public-api.txt +++ b/crates/storage/opendal/public-api.txt @@ -6,6 +6,7 @@ pub iceberg_storage_opendal::OpenDalStorage::Azdls pub iceberg_storage_opendal::OpenDalStorage::Azdls::config: alloc::sync::Arc pub iceberg_storage_opendal::OpenDalStorage::Gcs pub iceberg_storage_opendal::OpenDalStorage::Gcs::config: alloc::sync::Arc +pub iceberg_storage_opendal::OpenDalStorage::Gcs::credential_provider: core::option::Option> pub iceberg_storage_opendal::OpenDalStorage::Hf pub iceberg_storage_opendal::OpenDalStorage::Hf::config: alloc::sync::Arc pub iceberg_storage_opendal::OpenDalStorage::LocalFs @@ -14,6 +15,7 @@ pub iceberg_storage_opendal::OpenDalStorage::Oss pub iceberg_storage_opendal::OpenDalStorage::Oss::config: alloc::sync::Arc pub iceberg_storage_opendal::OpenDalStorage::S3 pub iceberg_storage_opendal::OpenDalStorage::S3::config: alloc::sync::Arc +pub iceberg_storage_opendal::OpenDalStorage::S3::credential_provider: core::option::Option> pub iceberg_storage_opendal::OpenDalStorage::S3::customized_credential_load: core::option::Option impl core::clone::Clone for iceberg_storage_opendal::OpenDalStorage pub fn iceberg_storage_opendal::OpenDalStorage::clone(&self) -> iceberg_storage_opendal::OpenDalStorage @@ -50,6 +52,7 @@ impl core::fmt::Debug for iceberg_storage_opendal::OpenDalStorageFactory pub fn iceberg_storage_opendal::OpenDalStorageFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::io::storage::StorageFactory for iceberg_storage_opendal::OpenDalStorageFactory pub fn iceberg_storage_opendal::OpenDalStorageFactory::build(&self, config: &iceberg::io::storage::config::StorageConfig) -> iceberg::error::Result> +pub fn iceberg_storage_opendal::OpenDalStorageFactory::build_with_credentials(&self, config: &iceberg::io::storage::config::StorageConfig, credential_provider: core::option::Option>) -> iceberg::error::Result> impl serde_core::ser::Serialize for iceberg_storage_opendal::OpenDalStorageFactory pub fn iceberg_storage_opendal::OpenDalStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_opendal::OpenDalStorageFactory @@ -92,6 +95,7 @@ impl core::fmt::Debug for iceberg_storage_opendal::OpenDalResolvingStorageFactor pub fn iceberg_storage_opendal::OpenDalResolvingStorageFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::io::storage::StorageFactory for iceberg_storage_opendal::OpenDalResolvingStorageFactory pub fn iceberg_storage_opendal::OpenDalResolvingStorageFactory::build(&self, config: &iceberg::io::storage::config::StorageConfig) -> iceberg::error::Result> +pub fn iceberg_storage_opendal::OpenDalResolvingStorageFactory::build_with_credentials(&self, config: &iceberg::io::storage::config::StorageConfig, credential_provider: core::option::Option>) -> iceberg::error::Result> impl serde_core::ser::Serialize for iceberg_storage_opendal::OpenDalResolvingStorageFactory pub fn iceberg_storage_opendal::OpenDalResolvingStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_opendal::OpenDalResolvingStorageFactory diff --git a/crates/storage/opendal/src/gcs.rs b/crates/storage/opendal/src/gcs.rs index 47ca52ccc6..3bbaba2c5b 100644 --- a/crates/storage/opendal/src/gcs.rs +++ b/crates/storage/opendal/src/gcs.rs @@ -17,17 +17,20 @@ //! Google Cloud Storage properties use std::collections::HashMap; +use std::sync::Arc; use iceberg::io::{ GCS_ALLOW_ANONYMOUS, GCS_CREDENTIALS_JSON, GCS_DISABLE_CONFIG_LOAD, GCS_DISABLE_VM_METADATA, - GCS_NO_AUTH, GCS_SERVICE_PATH, GCS_TOKEN, + GCS_NO_AUTH, GCS_SERVICE_PATH, GCS_TOKEN, StorageCredentialKind, StorageCredentialProvider, }; use iceberg::{Error, ErrorKind, Result}; -use opendal::Operator; use opendal::services::GcsConfig; +use opendal::{Configurator, Operator}; +use reqsign_core::{Context, Error as ReqsignError, ProvideCredential, Result as ReqsignResult}; +use reqsign_google::{Credential as GoogleCredential, Token as GoogleToken}; use url::Url; -use crate::utils::{from_opendal_error, is_truthy}; +use crate::utils::{from_opendal_error, is_truthy, system_time_to_timestamp}; /// Parse iceberg properties to [`GcsConfig`]. pub(crate) fn gcs_config_parse(mut m: HashMap) -> Result { @@ -45,7 +48,9 @@ pub(crate) fn gcs_config_parse(mut m: HashMap) -> Result) -> Result Result { +pub(crate) fn gcs_config_build( + cfg: &GcsConfig, + credential_provider: &Option>, + path: &str, +) -> Result { let url = Url::parse(path)?; + if !matches!(url.scheme(), "gs" | "gcs") { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Invalid gcs url: {path}, expected gs:// or gcs://"), + )); + } let bucket = url.host_str().ok_or_else(|| { Error::new( ErrorKind::DataInvalid, @@ -82,7 +97,95 @@ pub(crate) fn gcs_config_build(cfg: &GcsConfig, path: &str) -> Result let mut cfg = cfg.clone(); cfg.bucket = bucket.to_string(); - Ok(Operator::from_config(cfg) - .map_err(from_opendal_error)? - .finish()) + + // When a catalog-supplied provider is present, make it the sole credential + // source. `reqsign_google` only prepends a custom provider (unlike S3, which + // replaces the chain) and its chain continues to the next provider on error, + // so without this a failed refresh would silently fall back to the stale seed + // token or ambient GCP credentials. + let credential_provider = credential_provider + .as_ref() + .filter(|provider| provider.supports_path(path)); + if credential_provider.is_some() { + if cfg.skip_signature { + return Err(Error::new( + ErrorKind::DataInvalid, + "Invalid GCS auth settings: anonymous access cannot be combined with refreshable credentials", + )); + } + cfg.token = None; + cfg.credential = None; + cfg.disable_vm_metadata = true; + cfg.disable_config_load = true; + } + + let mut builder = cfg.into_builder(); + + // A catalog-supplied provider re-fetches the vended OAuth2 token as it nears expiry + if let Some(provider) = credential_provider { + builder = builder.credential_provider(VendedGcsCredentialProvider::new( + Arc::clone(provider), + path.to_string(), + )); + } + + Ok(Operator::new(builder).map_err(from_opendal_error)?.finish()) +} + +/// Adapts a generic [`StorageCredentialProvider`] into a `reqsign` +/// [`ProvideCredential`], so the GCS signer can obtain and refresh vended OAuth2 +/// tokens. +struct VendedGcsCredentialProvider { + provider: Arc, + /// Absolute path this operator serves: handed back to the provider so it can + /// select the vended credential whose prefix best matches the location. + path: String, +} + +impl std::fmt::Debug for VendedGcsCredentialProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VendedGcsCredentialProvider") + .field("path", &self.path) + .finish_non_exhaustive() + } +} + +impl VendedGcsCredentialProvider { + fn new(provider: Arc, path: String) -> Self { + Self { provider, path } + } +} + +impl ProvideCredential for VendedGcsCredentialProvider { + type Credential = GoogleCredential; + + async fn provide_credential(&self, _ctx: &Context) -> ReqsignResult> { + let credential = self + .provider + .load_credential(&self.path) + .await + .map_err(|e| { + ReqsignError::unexpected(format!( + "failed to load vended GCS credential for {}", + self.path + )) + .with_source(e) + })?; + + let expires_at = credential + .expires_at + .map(system_time_to_timestamp) + .transpose()?; + match credential.kind { + StorageCredentialKind::Gcs(gcs) => { + Ok(Some(GoogleCredential::with_token(GoogleToken { + access_token: gcs.token, + expires_at, + }))) + } + _ => Err(ReqsignError::unexpected( + "GCS storage received a non-GCS credential from the provider", + )), + } + } } diff --git a/crates/storage/opendal/src/lib.rs b/crates/storage/opendal/src/lib.rs index 768d3ff77f..0f301dfc91 100644 --- a/crates/storage/opendal/src/lib.rs +++ b/crates/storage/opendal/src/lib.rs @@ -35,7 +35,7 @@ use futures::StreamExt; use futures::stream::BoxStream; use iceberg::io::{ FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageConfig, - StorageFactory, + StorageCredentialProvider, StorageFactory, }; use iceberg::{Error, ErrorKind, Result}; use opendal::Operator; @@ -135,8 +135,16 @@ pub enum OpenDalStorageFactory { #[typetag::serde(name = "OpenDalStorageFactory")] impl StorageFactory for OpenDalStorageFactory { - #[allow(unused_variables)] fn build(&self, config: &StorageConfig) -> Result> { + self.build_with_credentials(config, None) + } + + #[allow(unused_variables)] + fn build_with_credentials( + &self, + config: &StorageConfig, + credential_provider: Option>, + ) -> Result> { match self { #[cfg(feature = "opendal-memory")] OpenDalStorageFactory::Memory => { @@ -150,10 +158,12 @@ impl StorageFactory for OpenDalStorageFactory { } => Ok(Arc::new(OpenDalStorage::S3 { config: s3_config_parse(config.props().clone())?.into(), customized_credential_load: customized_credential_load.clone(), + credential_provider, })), #[cfg(feature = "opendal-gcs")] OpenDalStorageFactory::Gcs => Ok(Arc::new(OpenDalStorage::Gcs { config: gcs_config_parse(config.props().clone())?.into(), + credential_provider, })), #[cfg(feature = "opendal-oss")] OpenDalStorageFactory::Oss => Ok(Arc::new(OpenDalStorage::Oss { @@ -210,12 +220,18 @@ pub enum OpenDalStorage { /// Custom AWS credential loader. #[serde(skip)] customized_credential_load: Option, + /// Provider of refreshable vended credentials, supplied by the catalog. + #[serde(skip)] + credential_provider: Option>, }, /// GCS storage variant. #[cfg(feature = "opendal-gcs")] Gcs { /// GCS configuration. config: Arc, + /// Provider of refreshable vended credentials, supplied by the catalog. + #[serde(skip)] + credential_provider: Option>, }, /// OSS storage variant. #[cfg(feature = "opendal-oss")] @@ -287,8 +303,14 @@ impl OpenDalStorage { OpenDalStorage::S3 { config, customized_credential_load, + credential_provider, } => { - let op = s3_config_build(config, customized_credential_load, path)?; + let op = s3_config_build( + config, + customized_credential_load, + credential_provider, + path, + )?; let op_info = op.info(); // Use the URL scheme in the path for prefix matching. This enables @@ -310,9 +332,18 @@ impl OpenDalStorage { } } #[cfg(feature = "opendal-gcs")] - OpenDalStorage::Gcs { config } => { - let operator = gcs_config_build(config, path)?; - let prefix = format!("gs://{}/", operator.info().name()); + OpenDalStorage::Gcs { + config, + credential_provider, + } => { + let operator = gcs_config_build(config, credential_provider, path)?; + let url = url::Url::parse(path).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid gcs url: {path}: {e}"), + ) + })?; + let prefix = format!("{}://{}/", url.scheme(), operator.info().name()); if path.starts_with(&prefix) { (operator, &path[prefix.len()..]) } else { @@ -384,6 +415,27 @@ impl OpenDalStorage { } } + /// Returns whether `path` uses a path-scoped dynamic credential provider. + /// + /// Such paths cannot share an OpenDAL deleter unless the provider can expose + /// the credential scope that applies to each path. Process them independently + /// so bulk deletion remains bounded without crossing credential boundaries. + fn uses_dynamic_credentials(&self, path: &str) -> bool { + match self { + #[cfg(feature = "opendal-s3")] + OpenDalStorage::S3 { + credential_provider: Some(provider), + .. + } => provider.supports_path(path), + #[cfg(feature = "opendal-gcs")] + OpenDalStorage::Gcs { + credential_provider: Some(provider), + .. + } => provider.supports_path(path), + _ => false, + } + } + /// Extracts the relative path from an absolute path without building an operator. /// /// This is a lightweight alternative to [`create_operator`](Self::create_operator) for cases @@ -418,13 +470,19 @@ impl OpenDalStorage { #[cfg(feature = "opendal-gcs")] OpenDalStorage::Gcs { .. } => { let url = url::Url::parse(path)?; + if !matches!(url.scheme(), "gs" | "gcs") { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Invalid gcs url: {path}, expected gs:// or gcs://"), + )); + } let bucket = url.host_str().ok_or_else(|| { Error::new( ErrorKind::DataInvalid, format!("Invalid gcs url: {path}, missing bucket"), ) })?; - let prefix = format!("gs://{}/", bucket); + let prefix = format!("{}://{}/", url.scheme(), bucket); if path.starts_with(&prefix) { Ok(&path[prefix.len()..]) } else { @@ -553,6 +611,12 @@ impl Storage for OpenDalStorage { let mut deleters: HashMap = HashMap::new(); while let Some(path) = paths.next().await { + if self.uses_dynamic_credentials(&path) { + let (op, relative_path) = self.create_operator(&path)?; + op.delete(relative_path).await.map_err(from_opendal_error)?; + continue; + } + let bucket = self.batch_key_for_path(&path); let (relative_path, deleter) = match deleters.entry(bucket) { @@ -631,6 +695,18 @@ impl FileWrite for OpenDalWriter { mod tests { use super::*; + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + #[derive(Debug)] + struct AlwaysSupportedCredentialProvider; + + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + #[async_trait] + impl StorageCredentialProvider for AlwaysSupportedCredentialProvider { + async fn load_credential(&self, _path: &str) -> Result { + unreachable!("batch key selection must not load a credential") + } + } + #[cfg(feature = "opendal-memory")] #[test] fn test_default_memory_operator() { @@ -677,6 +753,7 @@ mod tests { let storage = OpenDalStorage::S3 { config: Arc::new(S3Config::default()), customized_credential_load: None, + credential_provider: None, }; // All S3-family schemes are accepted by the same storage instance. @@ -692,19 +769,73 @@ mod tests { } } + #[cfg(feature = "opendal-s3")] + #[test] + fn test_dynamic_credentials_use_isolated_deletes() { + let storage = OpenDalStorage::S3 { + config: Arc::new(S3Config::default()), + customized_credential_load: None, + credential_provider: Some(Arc::new(AlwaysSupportedCredentialProvider)), + }; + let first = "s3://bucket/table-a/file.parquet"; + let second = "s3://bucket/table-b/file.parquet"; + + assert!(storage.uses_dynamic_credentials(first)); + assert!(storage.uses_dynamic_credentials(second)); + assert_eq!(storage.batch_key_for_path(first), "bucket"); + assert_eq!(storage.batch_key_for_path(second), "bucket"); + } + + #[cfg(feature = "opendal-s3")] + #[test] + fn test_s3_rejects_anonymous_dynamic_credentials() { + let mut config = S3Config::default(); + config.skip_signature = true; + let storage = OpenDalStorage::S3 { + config: Arc::new(config), + customized_credential_load: None, + credential_provider: Some(Arc::new(AlwaysSupportedCredentialProvider)), + }; + + let error = storage + .create_operator(&"s3://bucket/file.parquet") + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::DataInvalid); + } + + #[cfg(feature = "opendal-gcs")] + #[test] + fn test_gcs_rejects_anonymous_dynamic_credentials() { + let mut config = GcsConfig::default(); + config.skip_signature = true; + let storage = OpenDalStorage::Gcs { + config: Arc::new(config), + credential_provider: Some(Arc::new(AlwaysSupportedCredentialProvider)), + }; + + let error = storage + .create_operator(&"gs://bucket/file.parquet") + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::DataInvalid); + } + #[cfg(feature = "opendal-gcs")] #[test] fn test_relativize_path_gcs() { let storage = OpenDalStorage::Gcs { config: Arc::new(GcsConfig::default()), + credential_provider: None, }; - assert_eq!( - storage - .relativize_path("gs://my-bucket/path/to/file.parquet") - .unwrap(), - "path/to/file.parquet" - ); + for scheme in ["gs", "gcs"] { + let path = format!("{scheme}://my-bucket/path/to/file.parquet"); + assert_eq!( + storage.relativize_path(&path).unwrap(), + "path/to/file.parquet" + ); + let (_, relative) = storage.create_operator(&path).unwrap(); + assert_eq!(relative, "path/to/file.parquet"); + } } #[cfg(feature = "opendal-gcs")] @@ -712,6 +843,7 @@ mod tests { fn test_relativize_path_gcs_invalid_scheme() { let storage = OpenDalStorage::Gcs { config: Arc::new(GcsConfig::default()), + credential_provider: None, }; assert!( @@ -719,6 +851,11 @@ mod tests { .relativize_path("s3://my-bucket/path/to/file.parquet") .is_err() ); + assert!( + storage + .create_operator(&"s3://my-bucket/path/to/file.parquet") + .is_err() + ); } #[cfg(feature = "opendal-oss")] diff --git a/crates/storage/opendal/src/resolving.rs b/crates/storage/opendal/src/resolving.rs index 86993220a8..0fad716e52 100644 --- a/crates/storage/opendal/src/resolving.rs +++ b/crates/storage/opendal/src/resolving.rs @@ -27,7 +27,7 @@ use futures::StreamExt; use futures::stream::BoxStream; use iceberg::io::{ FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageConfig, - StorageFactory, + StorageCredentialProvider, StorageFactory, }; use iceberg::{Error, ErrorKind, Result}; use serde::{Deserialize, Serialize}; @@ -85,6 +85,9 @@ fn build_storage_for_scheme( scheme: &'static str, props: &HashMap, #[cfg(feature = "opendal-s3")] customized_credential_load: &Option, + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] credential_provider: &Option< + Arc, + >, ) -> Result { match scheme { #[cfg(feature = "opendal-s3")] @@ -93,6 +96,7 @@ fn build_storage_for_scheme( Ok(OpenDalStorage::S3 { config: Arc::new(config), customized_credential_load: customized_credential_load.clone(), + credential_provider: credential_provider.clone(), }) } #[cfg(feature = "opendal-gcs")] @@ -100,6 +104,7 @@ fn build_storage_for_scheme( let config = crate::gcs::gcs_config_parse(props.clone())?; Ok(OpenDalStorage::Gcs { config: Arc::new(config), + credential_provider: credential_provider.clone(), }) } #[cfg(feature = "opendal-oss")] @@ -186,11 +191,22 @@ impl OpenDalResolvingStorageFactory { #[typetag::serde] impl StorageFactory for OpenDalResolvingStorageFactory { fn build(&self, config: &StorageConfig) -> Result> { + self.build_with_credentials(config, None) + } + + #[allow(unused_variables)] + fn build_with_credentials( + &self, + config: &StorageConfig, + credential_provider: Option>, + ) -> Result> { Ok(Arc::new(OpenDalResolvingStorage { props: config.props().clone(), storages: RwLock::new(HashMap::new()), #[cfg(feature = "opendal-s3")] customized_credential_load: self.customized_credential_load.clone(), + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + credential_provider, })) } } @@ -201,7 +217,7 @@ impl StorageFactory for OpenDalResolvingStorageFactory { /// Sub-storages are lazily created on first use for each scheme and cached /// for subsequent operations. Scheme aliases like `s3`/`s3a`/`s3n` map to /// the same canonical scheme, so they share a storage instance. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Serialize, Deserialize)] pub struct OpenDalResolvingStorage { /// Configuration properties shared across all backends. props: HashMap, @@ -212,6 +228,20 @@ pub struct OpenDalResolvingStorage { #[cfg(feature = "opendal-s3")] #[serde(skip)] customized_credential_load: Option, + /// Provider of refreshable vended credentials. + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + #[serde(skip)] + credential_provider: Option>, +} + +impl std::fmt::Debug for OpenDalResolvingStorage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // `props` can contain storage secrets, and a custom credential provider + // may carry secret state in its own Debug implementation + f.debug_struct("OpenDalResolvingStorage") + .field("property_keys", &self.props.keys().collect::>()) + .finish_non_exhaustive() + } } impl OpenDalResolvingStorage { @@ -247,6 +277,8 @@ impl OpenDalResolvingStorage { &self.props, #[cfg(feature = "opendal-s3")] &self.customized_credential_load, + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + &self.credential_provider, )?; let storage = Arc::new(storage); cache.insert(scheme, storage.clone()); @@ -334,6 +366,8 @@ mod tests { storages: RwLock::new(HashMap::new()), #[cfg(feature = "opendal-s3")] customized_credential_load: None, + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + credential_provider: None, } } diff --git a/crates/storage/opendal/src/s3.rs b/crates/storage/opendal/src/s3.rs index 4b3893b39c..8e57cdfdb6 100644 --- a/crates/storage/opendal/src/s3.rs +++ b/crates/storage/opendal/src/s3.rs @@ -22,7 +22,8 @@ use iceberg::io::{ CLIENT_REGION, S3_ACCESS_KEY_ID, S3_ALLOW_ANONYMOUS, S3_ASSUME_ROLE_ARN, S3_ASSUME_ROLE_EXTERNAL_ID, S3_ASSUME_ROLE_SESSION_NAME, S3_DISABLE_CONFIG_LOAD, S3_DISABLE_EC2_METADATA, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, S3_SECRET_ACCESS_KEY, - S3_SESSION_TOKEN, S3_SSE_KEY, S3_SSE_MD5, S3_SSE_TYPE, + S3_SESSION_TOKEN, S3_SSE_KEY, S3_SSE_MD5, S3_SSE_TYPE, StorageCredentialKind, + StorageCredentialProvider, }; use iceberg::{Error, ErrorKind, Result}; use opendal::services::S3Config; @@ -31,10 +32,13 @@ use opendal::{Configurator, Operator}; pub use reqsign_aws_v4::Credential as AwsCredential; /// Trait for types that can asynchronously supply [`AwsCredential`] to a [`CustomAwsCredentialLoader`]. pub use reqsign_core::ProvideCredential; -use reqsign_core::{ProvideCredentialChain, ProvideCredentialDyn}; +use reqsign_core::{ + Context, Error as ReqsignError, ProvideCredentialChain, ProvideCredentialDyn, + Result as ReqsignResult, +}; use url::Url; -use crate::utils::{from_opendal_error, is_truthy}; +use crate::utils::{from_opendal_error, is_truthy, system_time_to_timestamp}; /// Parse iceberg props to s3 config. pub(crate) fn s3_config_parse(mut m: HashMap) -> Result { @@ -129,6 +133,7 @@ pub(crate) fn s3_config_parse(mut m: HashMap) -> Result, + credential_provider: &Option>, path: &str, ) -> Result { let url = Url::parse(path)?; @@ -139,6 +144,19 @@ pub(crate) fn s3_config_build( ) })?; + // Preserve the existing custom-loader precedence: an explicitly configured loader + // is the sole source, otherwise install the catalog provider as a replacement chain so + // refresh failures cannot fall through to broader ambient AWS credentials. + let credential_provider = credential_provider + .as_ref() + .filter(|provider| provider.supports_path(path)); + if customized_credential_load.is_none() && credential_provider.is_some() && cfg.skip_signature { + return Err(Error::new( + ErrorKind::DataInvalid, + "Invalid S3 auth settings: anonymous access cannot be combined with refreshable credentials", + )); + } + let mut builder = cfg .clone() .into_builder() @@ -148,11 +166,75 @@ pub(crate) fn s3_config_build( if let Some(loader) = customized_credential_load { let chain = ProvideCredentialChain::new().push(Arc::clone(&loader.0)); builder = builder.credential_provider_chain(chain); + } else if let Some(provider) = credential_provider { + let chain = ProvideCredentialChain::new().push(VendedS3CredentialProvider::new( + Arc::clone(provider), + path.to_string(), + )); + builder = builder.credential_provider_chain(chain); } Ok(Operator::new(builder).map_err(from_opendal_error)?.finish()) } +/// Adapts a generic [`StorageCredentialProvider`] into a reqsign +/// [`ProvideCredential`], so the S3 signer can obtain and refresh vended +/// credentials. +struct VendedS3CredentialProvider { + provider: Arc, + /// Absolute path this operator serves; handed back to the provider so it can + /// select the vended credential whose prefix best matches the location. + path: String, +} + +impl std::fmt::Debug for VendedS3CredentialProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VendedS3CredentialProvider") + .field("path", &self.path) + .finish_non_exhaustive() + } +} + +impl VendedS3CredentialProvider { + fn new(provider: Arc, path: String) -> Self { + Self { provider, path } + } +} + +impl ProvideCredential for VendedS3CredentialProvider { + type Credential = AwsCredential; + + async fn provide_credential(&self, _ctx: &Context) -> ReqsignResult> { + let credential = self + .provider + .load_credential(&self.path) + .await + .map_err(|e| { + ReqsignError::unexpected(format!( + "failed to load vended S3 credential for {}", + self.path + )) + .with_source(e) + })?; + + let expires_in = credential + .expires_at + .map(system_time_to_timestamp) + .transpose()?; + match credential.kind { + StorageCredentialKind::S3(s3) => Ok(Some(AwsCredential { + access_key_id: s3.access_key_id, + secret_access_key: s3.secret_access_key, + session_token: s3.session_token, + expires_in, + })), + _ => Err(ReqsignError::unexpected( + "S3 storage received a non-S3 credential from the provider", + )), + } + } +} + /// Custom AWS credential loader. /// /// Wraps any [`ProvideCredential`] implementation for use with the S3 storage backend. @@ -176,7 +258,7 @@ impl std::fmt::Debug for CustomAwsCredentialLoader { impl CustomAwsCredentialLoader { /// Create a new custom AWS credential loader from any [`ProvideCredential`] implementation. pub fn new(provider: impl ProvideCredential + 'static) -> Self { - Self(Arc::new(provider) as Arc>) + Self(Arc::new(provider)) } } diff --git a/crates/storage/opendal/src/utils.rs b/crates/storage/opendal/src/utils.rs index 56f8c18059..3afd319f22 100644 --- a/crates/storage/opendal/src/utils.rs +++ b/crates/storage/opendal/src/utils.rs @@ -27,3 +27,25 @@ pub(crate) fn from_opendal_error(e: opendal::Error) -> iceberg::Error { ) .with_source(e) } + +/// Convert a [`SystemTime`](std::time::SystemTime) credential expiry into the +/// `reqsign` [`Timestamp`](reqsign_core::time::Timestamp) used on backend +/// credential types (e.g. `AwsCredential::expires_in`, `google::Token::expires_at`). +#[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] +pub(crate) fn system_time_to_timestamp( + time: std::time::SystemTime, +) -> reqsign_core::Result { + let millis = time + .duration_since(std::time::UNIX_EPOCH) + .map_err(|e| { + reqsign_core::Error::unexpected(format!( + "credential expiry precedes the UNIX epoch: {e}" + )) + })? + .as_millis(); + let millis = i64::try_from(millis).map_err(|_| { + reqsign_core::Error::unexpected("credential expiry overflows i64 milliseconds") + })?; + reqsign_core::time::Timestamp::from_millisecond(millis) + .map_err(|e| reqsign_core::Error::unexpected(format!("invalid credential expiry: {e}"))) +} From 26d7fcc6d299a8e814b6cceca480409b2c2cc3c9 Mon Sep 17 00:00:00 2001 From: Zakariya Stasa Date: Thu, 30 Jul 2026 23:08:08 +0100 Subject: [PATCH 2/4] chore: fix typo --- crates/catalog/rest/src/client.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs index cc4b1c2f97..8ac2bbfb45 100644 --- a/crates/catalog/rest/src/client.rs +++ b/crates/catalog/rest/src/client.rs @@ -299,7 +299,7 @@ pub(crate) async fn deserialize_catalog_response( serde_json::from_slice::(&bytes).map_err(|e| { // Successful REST responses can contain OAuth tokens and delegated - // storage credentials. Never copy an unparseable response into an error. + // storage credentials. Never copy an unparsable response into an error. Error::new( ErrorKind::Unexpected, "Failed to parse response from rest catalog server", From 8af37a53d8a5ec05bd523a8391ebd9c774bbafae Mon Sep 17 00:00:00 2001 From: Zakariya Stasa Date: Sun, 2 Aug 2026 17:36:00 +0100 Subject: [PATCH 3/4] refactor(rest): preserve vended credential scopes --- crates/catalog/rest/src/credential.rs | 142 ++++++++++++++------------ crates/iceberg/public-api.txt | 1 + crates/iceberg/src/io/storage/mod.rs | 8 +- crates/storage/opendal/src/gcs.rs | 6 +- crates/storage/opendal/src/s3.rs | 6 +- crates/storage/opendal/src/utils.rs | 18 ++++ 6 files changed, 111 insertions(+), 70 deletions(-) diff --git a/crates/catalog/rest/src/credential.rs b/crates/catalog/rest/src/credential.rs index 14363b3a71..457e90299f 100644 --- a/crates/catalog/rest/src/credential.rs +++ b/crates/catalog/rest/src/credential.rs @@ -74,7 +74,8 @@ struct CloudRefresh { /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. jitter_prefetch: bool, /// Parse a complete credential from catalog-supplied properties. - parse_credential: fn(&HashMap) -> Result, + parse_credential: + fn(config: &HashMap, prefix: Option) -> Result, } impl CloudRefresh { @@ -105,14 +106,13 @@ impl CloudRefresh { /// 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())) + .find(|cloud| cloud.matches_location(location)) } - fn matches_credential_prefix(&self, prefix: &str) -> bool { - scheme_of(prefix).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + fn matches_location(&self, location: &str) -> bool { + scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) } } @@ -130,10 +130,9 @@ 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. +/// A cached vended credential and its refresh schedule. #[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. @@ -141,12 +140,11 @@ struct CachedEntry { } impl CachedEntry { - fn new(prefix: String, credential: StorageCredential, jitter_prefetch: bool) -> Self { + fn new(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, } @@ -155,13 +153,13 @@ impl CachedEntry { /// 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 { + fn seed(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); + let mut entry = Self::new(credential, jitter_prefetch); if due { entry.refresh_at = Some(UNIX_EPOCH); } @@ -219,6 +217,12 @@ impl RestVendedCredentialProvider { } } + fn configured_cloud_for_location(&self, location: &str) -> Option<&ConfiguredCloud> { + self.clouds + .iter() + .find(|configured| configured.cloud.matches_location(location)) + } + /// Fetch fresh credentials from the catalog's credentials endpoint. async fn fetch(&self, configured: &ConfiguredCloud) -> Result> { let mut request = self.client.request(Method::GET, &configured.endpoint); @@ -234,15 +238,13 @@ impl RestVendedCredentialProvider { parsed .storage_credentials .into_iter() - .filter(|sc| configured.cloud.matches_credential_prefix(&sc.prefix)) + .filter(|sc| configured.cloud.matches_location(&sc.prefix)) .map(|sc| { - (configured.cloud.parse_credential)(&sc.config).map(|credential| { - CachedEntry::new( - sc.prefix, - credential, - configured.cloud.jitter_prefetch, - ) - }) + (configured.cloud.parse_credential)(&sc.config, Some(sc.prefix)).map( + |credential| { + CachedEntry::new(credential, configured.cloud.jitter_prefetch) + }, + ) }) .collect() } @@ -335,30 +337,22 @@ async fn cache_decision(configured: &ConfiguredCloud, path: &str) -> CacheDecisi #[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) - }) + self.configured_cloud_for_location(path).is_some() } async fn load_credential(&self, path: &str) -> Result { - let cloud = CloudRefresh::for_location(path).ok_or_else(|| { - Error::new( + if CloudRefresh::for_location(path).is_none() { + return Err(Error::new( ErrorKind::FeatureUnsupported, format!("no credential refresh implementation for storage location: {path}"), + )); + } + let configured = self.configured_cloud_for_location(path).ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!("credential refresh is not configured 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), @@ -395,8 +389,14 @@ impl StorageCredentialProvider for RestVendedCredentialProvider { 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()) + .filter(|entry| { + entry + .credential + .prefix + .as_deref() + .is_none_or(|prefix| path.starts_with(prefix)) + }) + .max_by_key(|entry| entry.credential.prefix.as_deref().map_or(0, str::len)) } /// Compute a successful credential's prefetch time. @@ -451,17 +451,9 @@ pub(crate) fn build_vended_credential_provider( .get(cloud.endpoint_key) .filter(|value| !value.is_empty())?, ); - let entries = (cloud.parse_credential)(props) + let entries = (cloud.parse_credential)(props, None) .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, - )] - }) + .map(|credential| vec![CachedEntry::seed(credential, cloud.jitter_prefetch)]) .unwrap_or_default(); Some(ConfiguredCloud { @@ -509,12 +501,16 @@ fn scheme_of(location: &str) -> Option { } /// Parse a complete S3 credential returned by the credentials endpoint. -fn parse_s3_credential(config: &HashMap) -> Result { +fn parse_s3_credential( + config: &HashMap, + prefix: Option, +) -> Result { let access_key_id = required_nonempty(config, S3_ACCESS_KEY_ID)?; let secret_access_key = required_nonempty(config, S3_SECRET_ACCESS_KEY)?; let session_token = required_nonempty(config, S3_SESSION_TOKEN)?; let expires_at = required_epoch_millis(config, S3_SESSION_TOKEN_EXPIRES_AT_MS)?; Ok(StorageCredential { + prefix, kind: StorageCredentialKind::S3(S3Credential { access_key_id, secret_access_key, @@ -525,10 +521,14 @@ fn parse_s3_credential(config: &HashMap) -> Result) -> Result { +fn parse_gcs_credential( + config: &HashMap, + prefix: Option, +) -> Result { let token = required_nonempty(config, GCS_TOKEN)?; let expires_at = required_epoch_millis(config, GCS_TOKEN_EXPIRES_AT)?; Ok(StorageCredential { + prefix, kind: StorageCredentialKind::Gcs(GcsCredential { token }), expires_at: Some(expires_at), }) @@ -581,8 +581,13 @@ mod tests { .to_string() } - fn s3_cred(access_key_id: &str, expires_at: Option) -> StorageCredential { + fn s3_cred( + prefix: Option<&str>, + access_key_id: &str, + expires_at: Option, + ) -> StorageCredential { StorageCredential { + prefix: prefix.map(str::to_string), kind: StorageCredentialKind::S3(S3Credential { access_key_id: access_key_id.to_string(), secret_access_key: "secret".to_string(), @@ -600,11 +605,7 @@ mod tests { } fn cached_s3(prefix: &str, access_key_id: &str, expires_at: Option) -> CachedEntry { - CachedEntry::new( - prefix.to_string(), - s3_cred(access_key_id, expires_at), - false, - ) + CachedEntry::new(s3_cred(Some(prefix), access_key_id, expires_at), false) } #[test] @@ -649,25 +650,26 @@ mod tests { assert!(CloudRefresh::for_location("not a url").is_none()); // Azure not supported yet -> no provider, static creds used as-is. assert!(CloudRefresh::for_location("abfss://fs@acct.dfs.core.windows.net/k").is_none()); - assert!(CloudRefresh::AWS.matches_credential_prefix("s3://bucket/path")); - assert!(!CloudRefresh::AWS.matches_credential_prefix("s3evil://bucket/path")); + assert!(CloudRefresh::AWS.matches_location("s3://bucket/path")); + assert!(!CloudRefresh::AWS.matches_location("s3evil://bucket/path")); } #[test] fn parses_s3_credential() { let mut config = HashMap::new(); - assert!(parse_s3_credential(&config).is_err()); + assert!(parse_s3_credential(&config, None).is_err()); config.insert(S3_ACCESS_KEY_ID.to_string(), "AK".to_string()); - assert!(parse_s3_credential(&config).is_err()); + assert!(parse_s3_credential(&config, None).is_err()); config.insert(S3_SECRET_ACCESS_KEY.to_string(), "SK".to_string()); - assert!(parse_s3_credential(&config).is_err()); + assert!(parse_s3_credential(&config, None).is_err()); config.insert(S3_SESSION_TOKEN.to_string(), "TOK".to_string()); config.insert( S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), "1500".to_string(), ); - let credential = parse_s3_credential(&config).unwrap(); + let credential = parse_s3_credential(&config, Some("s3://bucket".to_string())).unwrap(); + assert_eq!(credential.prefix.as_deref(), Some("s3://bucket")); assert_eq!( credential.expires_at, Some(UNIX_EPOCH + Duration::from_millis(1500)) @@ -681,12 +683,13 @@ mod tests { #[test] fn parse_gcs_requires_token_and_expiry() { let mut config = HashMap::new(); - assert!(parse_gcs_credential(&config).is_err()); + assert!(parse_gcs_credential(&config, None).is_err()); config.insert(GCS_TOKEN.to_string(), "ya29.token".to_string()); - assert!(parse_gcs_credential(&config).is_err()); + assert!(parse_gcs_credential(&config, None).is_err()); config.insert(GCS_TOKEN_EXPIRES_AT.to_string(), "2000".to_string()); - let credential = parse_gcs_credential(&config).unwrap(); + let credential = parse_gcs_credential(&config, Some("gs://bucket".to_string())).unwrap(); + assert_eq!(credential.prefix.as_deref(), Some("gs://bucket")); match &credential.kind { StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token, "ya29.token"), other => panic!("expected GCS, got {other:?}"), @@ -737,10 +740,11 @@ mod tests { #[test] fn seed_inside_nominal_window_is_immediately_due() { let credential = s3_cred( + None, "a", Some(SystemTime::now() + REFRESH_BUFFER - Duration::from_secs(1)), ); - let entry = CachedEntry::seed(String::new(), credential, true); + let entry = CachedEntry::seed(credential, true); let now = SystemTime::now(); assert!(!entry.is_fresh(now)); assert!(entry.is_unexpired(now)); @@ -770,6 +774,10 @@ mod tests { ]; let got = longest_prefix_match(&entries, "s3://bucket/warehouse/db/t/f").unwrap(); assert_eq!(s3_access_key_id(&got.credential), "narrow"); + assert_eq!( + got.credential.prefix.as_deref(), + Some("s3://bucket/warehouse/db") + ); assert!(longest_prefix_match(&entries, "s3://other/x").is_none()); let fresh = Some(SystemTime::now() + REFRESH_BUFFER + Duration::from_secs(60)); @@ -876,6 +884,7 @@ mod tests { .load_credential("s3://bucket/warehouse/f") .await .unwrap(); + assert_eq!(credential.prefix.as_deref(), Some("s3://bucket")); match credential.kind { StorageCredentialKind::S3(s3) => { @@ -916,6 +925,7 @@ mod tests { .expect("provider should be built"); let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); + assert_eq!(credential.prefix, None); assert_eq!(s3_access_key_id(&credential), "SEED_AK"); mock.assert_async().await; } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 2771ffd998..c8caceacf1 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1019,6 +1019,7 @@ pub fn iceberg::io::StorageConfig::deserialize<__D>(__deserializer: __D) -> core pub struct iceberg::io::StorageCredential pub iceberg::io::StorageCredential::expires_at: core::option::Option pub iceberg::io::StorageCredential::kind: iceberg::io::StorageCredentialKind +pub iceberg::io::StorageCredential::prefix: core::option::Option impl core::clone::Clone for iceberg::io::StorageCredential pub fn iceberg::io::StorageCredential::clone(&self) -> iceberg::io::StorageCredential impl core::fmt::Debug for iceberg::io::StorageCredential diff --git a/crates/iceberg/src/io/storage/mod.rs b/crates/iceberg/src/io/storage/mod.rs index b10bd380c0..dd812a2d45 100644 --- a/crates/iceberg/src/io/storage/mod.rs +++ b/crates/iceberg/src/io/storage/mod.rs @@ -188,13 +188,17 @@ pub trait StorageCredentialProvider: Debug + Send + Sync { /// `path` is the absolute location being accessed (e.g. /// `s3://bucket/warehouse/db/table/...`). Providers that vend distinct /// credentials per location prefix use it to select the most specific - /// match. + /// match. When the selected credential has a declared + /// [`StorageCredential::prefix`], it must cover `path`. async fn load_credential(&self, path: &str) -> Result; } -/// A vended storage credential together with when it expires. +/// A vended storage credential together with its scope and expiry. #[derive(Clone, Debug)] pub struct StorageCredential { + /// Storage-location prefix this credential is scoped to. `None` represents a + /// credential without a declared scope, sourced from flat storage properties. + pub prefix: Option, /// The backend-specific credential material. pub kind: StorageCredentialKind, /// When the credential expires, if known. `None` means non-expiring and diff --git a/crates/storage/opendal/src/gcs.rs b/crates/storage/opendal/src/gcs.rs index 3bbaba2c5b..311bd1ac88 100644 --- a/crates/storage/opendal/src/gcs.rs +++ b/crates/storage/opendal/src/gcs.rs @@ -30,7 +30,9 @@ use reqsign_core::{Context, Error as ReqsignError, ProvideCredential, Result as use reqsign_google::{Credential as GoogleCredential, Token as GoogleToken}; use url::Url; -use crate::utils::{from_opendal_error, is_truthy, system_time_to_timestamp}; +use crate::utils::{ + from_opendal_error, is_truthy, system_time_to_timestamp, validate_credential_prefix, +}; /// Parse iceberg properties to [`GcsConfig`]. pub(crate) fn gcs_config_parse(mut m: HashMap) -> Result { @@ -172,6 +174,8 @@ impl ProvideCredential for VendedGcsCredentialProvider { .with_source(e) })?; + validate_credential_prefix(&self.path, credential.prefix.as_deref())?; + let expires_at = credential .expires_at .map(system_time_to_timestamp) diff --git a/crates/storage/opendal/src/s3.rs b/crates/storage/opendal/src/s3.rs index 8e57cdfdb6..e6b825347d 100644 --- a/crates/storage/opendal/src/s3.rs +++ b/crates/storage/opendal/src/s3.rs @@ -38,7 +38,9 @@ use reqsign_core::{ }; use url::Url; -use crate::utils::{from_opendal_error, is_truthy, system_time_to_timestamp}; +use crate::utils::{ + from_opendal_error, is_truthy, system_time_to_timestamp, validate_credential_prefix, +}; /// Parse iceberg props to s3 config. pub(crate) fn s3_config_parse(mut m: HashMap) -> Result { @@ -217,6 +219,8 @@ impl ProvideCredential for VendedS3CredentialProvider { .with_source(e) })?; + validate_credential_prefix(&self.path, credential.prefix.as_deref())?; + let expires_in = credential .expires_at .map(system_time_to_timestamp) diff --git a/crates/storage/opendal/src/utils.rs b/crates/storage/opendal/src/utils.rs index 3afd319f22..c168093fc2 100644 --- a/crates/storage/opendal/src/utils.rs +++ b/crates/storage/opendal/src/utils.rs @@ -49,3 +49,21 @@ pub(crate) fn system_time_to_timestamp( reqsign_core::time::Timestamp::from_millisecond(millis) .map_err(|e| reqsign_core::Error::unexpected(format!("invalid credential expiry: {e}"))) } + +/// Validate that a provider's declared credential prefix covers the path for +/// which the backend requested the credential. +#[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] +pub(crate) fn validate_credential_prefix( + path: &str, + prefix: Option<&str>, +) -> reqsign_core::Result<()> { + match prefix { + Some("") => Err(reqsign_core::Error::unexpected( + "vended credential has an empty storage prefix", + )), + Some(prefix) if !path.starts_with(prefix) => Err(reqsign_core::Error::unexpected(format!( + "vended credential prefix {prefix:?} does not cover storage location {path:?}" + ))), + _ => Ok(()), + } +} From f7b065a3fc8a3f4f93320eeafba340ed88cd4e3f Mon Sep 17 00:00:00 2001 From: Zakariya Stasa Date: Fri, 7 Aug 2026 10:38:27 +0100 Subject: [PATCH 4/4] fix(rest): harden vended credential refresh --- crates/catalog/rest/src/catalog.rs | 34 +- crates/catalog/rest/src/client.rs | 285 +++++--- crates/catalog/rest/src/credential.rs | 879 +++++++++++++++++------- crates/iceberg/public-api.txt | 25 +- crates/iceberg/src/io/storage/mod.rs | 116 +++- crates/storage/opendal/src/gcs.rs | 120 +++- crates/storage/opendal/src/lib.rs | 207 +++++- crates/storage/opendal/src/resolving.rs | 77 +++ crates/storage/opendal/src/s3.rs | 132 +++- crates/storage/opendal/src/utils.rs | 18 +- 10 files changed, 1474 insertions(+), 419 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 1a400fd7da..f86e29c0d9 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -41,6 +41,7 @@ use typed_builder::TypedBuilder; use crate::client::{ HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error, }; +use crate::credential::build_vended_credential_provider; use crate::endpoint::{Endpoint, V1_NAMESPACE_EXISTS, V1_TABLE_EXISTS}; use crate::types::{ CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, @@ -315,7 +316,7 @@ impl RestCatalogConfig { HeaderValue::from_str(value).map_err(|e| { Error::new( ErrorKind::DataInvalid, - format!("Invalid header value: {value}"), + format!("Invalid value for header: {key}"), ) .with_source(e) })?, @@ -524,6 +525,7 @@ impl RestCatalog { &self, metadata_location: Option<&str>, extra_config: Option>, + table_auth_config: Option>, ) -> Result { let context = self.context().await?; let mut props = context.config.props.clone(); @@ -558,10 +560,11 @@ impl RestCatalog { // If the catalog vends refreshable credentials for this table's storage, // attach a provider so the backend re-fetches them before they expire. - let credential_provider = crate::credential::build_vended_credential_provider( + let credential_provider = build_vended_credential_provider( context.client.clone(), &context.config.uri, &props, + table_auth_config.as_ref(), )?; let mut builder = FileIOBuilder::new(factory).with_props(props); @@ -872,14 +875,15 @@ impl Catalog for RestCatalog { "Metadata location missing in `create_table` response!", ))?; - let config = response - .config + let table_config = response.config; + let config = table_config + .clone() .into_iter() .chain(self.user_config.props.clone()) .collect(); let file_io = self - .load_file_io(Some(metadata_location), Some(config)) + .load_file_io(Some(metadata_location), Some(config), Some(table_config)) .await?; let mut table_builder = Table::builder() @@ -932,14 +936,19 @@ impl Catalog for RestCatalog { } }; - let config = response - .config + let table_config = response.config; + let config = table_config + .clone() .into_iter() .chain(self.user_config.props.clone()) .collect(); let file_io = self - .load_file_io(response.metadata_location.as_deref(), Some(config)) + .load_file_io( + response.metadata_location.as_deref(), + Some(config), + Some(table_config), + ) .await?; let mut table_builder = Table::builder() @@ -1074,13 +1083,14 @@ impl Catalog for RestCatalog { "Metadata location missing in `register_table` response!", ))?; - let config = response - .config + let table_config = response.config; + let config = table_config + .clone() .into_iter() .chain(self.user_config.props.clone()) .collect(); let file_io = self - .load_file_io(Some(metadata_location), Some(config)) + .load_file_io(Some(metadata_location), Some(config), Some(table_config)) .await?; let mut table_builder = Table::builder() @@ -1156,7 +1166,7 @@ impl Catalog for RestCatalog { }; let file_io = self - .load_file_io(Some(&response.metadata_location), None) + .load_file_io(Some(&response.metadata_location), None, None) .await?; let mut table_builder = Table::builder() diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs index 8ac2bbfb45..bc9157212b 100644 --- a/crates/catalog/rest/src/client.rs +++ b/crates/catalog/rest/src/client.rs @@ -17,6 +17,7 @@ use std::collections::HashMap; use std::fmt::{Debug, Formatter}; +use std::sync::Arc; use http::StatusCode; use iceberg::{Error, ErrorKind, Result}; @@ -34,7 +35,7 @@ pub(crate) struct HttpClient { /// The token to be used for authentication. /// /// It's possible to fetch the token from the server while needed. - token: Mutex>, + token: Arc>>, /// The token endpoint to be used for authentication. token_endpoint: String, /// The credential to be used for authentication. @@ -66,7 +67,7 @@ impl HttpClient { let extra_headers = cfg.extra_headers()?; Ok(HttpClient { client: cfg.client().unwrap_or_default(), - token: Mutex::new(cfg.token()), + token: Arc::new(Mutex::new(cfg.token())), token_endpoint: cfg.get_token_endpoint(), credential: cfg.credential(), extra_headers, @@ -83,14 +84,70 @@ impl HttpClient { pub(crate) fn for_table( &self, catalog_uri: &str, - props: HashMap, + props: &HashMap, ) -> Result { let cfg = RestCatalogConfig::builder() .uri(catalog_uri.to_string()) - .props(props) + .props(props.clone()) .client(Some(self.client.clone())) .build(); - Self::new(&cfg) + let table_token = cfg.token(); + let configured_credential = cfg.credential(); + let has_table_token = props.contains_key("token"); + let has_table_credential = props.contains_key("credential"); + let has_table_auth_header = props.keys().any(|key| { + key.strip_prefix("header.") + .is_some_and(|name| name.eq_ignore_ascii_case("authorization")) + }); + let has_oauth_params = ["scope", "audience", "resource"] + .iter() + .any(|key| props.contains_key(*key)); + let has_oauth_config = props.contains_key("oauth2-server-uri") || has_oauth_params; + let has_oauth_override = + has_table_credential || (has_oauth_config && self.credential.is_some()); + let has_table_auth = has_table_token || has_table_auth_header || has_oauth_override; + let mut extra_headers = self.extra_headers.clone(); + if has_table_auth && !has_table_auth_header { + // Authentication is applied after request headers are initially built, but + // `execute` reapplies `extra_headers`. Do not let an inherited catalog + // Authorization header overwrite table-scoped authentication at that point. + extra_headers.remove(http::header::AUTHORIZATION); + } + extra_headers.extend(cfg.extra_headers()?); + + Ok(Self { + client: self.client.clone(), + token: if has_table_auth { + Arc::new(Mutex::new(table_token)) + } else { + Arc::clone(&self.token) + }, + token_endpoint: if has_oauth_override && props.contains_key("oauth2-server-uri") { + cfg.get_token_endpoint() + } else { + self.token_endpoint.clone() + }, + credential: if has_table_token || has_table_auth_header { + None + } else if has_table_credential { + configured_credential + } else { + self.credential.clone() + }, + extra_headers, + extra_oauth_params: if has_oauth_override && has_oauth_params { + cfg.extra_oauth_params() + } else { + self.extra_oauth_params.clone() + }, + disable_header_redaction: if props + .contains_key(crate::REST_CATALOG_PROP_DISABLE_HEADER_REDACTION) + { + cfg.disable_header_redaction() + } else { + self.disable_header_redaction + }, + }) } /// Update the http client with new configuration. @@ -104,7 +161,10 @@ impl HttpClient { .unwrap_or(self.extra_headers); Ok(HttpClient { client: cfg.client().unwrap_or(self.client), - token: Mutex::new(cfg.token().or_else(|| self.token.into_inner())), + token: match cfg.token() { + Some(token) => Arc::new(Mutex::new(Some(token))), + None => self.token, + }, token_endpoint: if !cfg.get_token_endpoint().is_empty() { cfg.get_token_endpoint() } else { @@ -308,26 +368,10 @@ pub(crate) async fn deserialize_catalog_response( }) } -/// Returns true if the header may carry a secret. -fn is_sensitive_header(name: &str) -> bool { - let name_lower = name.to_lowercase(); - [ - "auth", - "token", - "secret", - "key", - "password", - "cookie", - "credential", - ] - .iter() - .any(|pattern| name_lower.contains(pattern)) -} - -/// Redacts sensitive headers and returns a debug-formatted string. +/// Redacts header values and returns a debug-formatted string. /// /// If `disable_redaction` is true, returns all headers without redaction. -/// Otherwise, replaces sensitive header values with "[REDACTED]". +/// Otherwise, redacts every header value. fn format_headers_redacted(headers: &HeaderMap, disable_redaction: bool) -> String { if disable_redaction { // Return all headers as-is without redaction @@ -338,16 +382,10 @@ fn format_headers_redacted(headers: &HeaderMap, disable_redaction: bool) -> Stri return format!("{all:?}"); } - // Redact sensitive headers by replacing their values with "[REDACTED]" + // Retain names for diagnostics but redact every value let redacted: HashMap<&str, &str> = headers .iter() - .filter_map(|(name, value)| { - if is_sensitive_header(name.as_str()) { - Some((name.as_str(), "[REDACTED]")) - } else { - value.to_str().ok().map(|v| (name.as_str(), v)) - } - }) + .map(|(name, _)| (name.as_str(), "[REDACTED]")) .collect(); format!("{redacted:?}") } @@ -390,96 +428,21 @@ mod tests { } #[test] - fn test_format_headers_redacted_non_sensitive() { + fn test_format_headers_redacts_all_values() { let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer secret-token".parse().unwrap()); headers.insert("content-type", "application/json".parse().unwrap()); headers.insert("x-request-id", "abc123".parse().unwrap()); let result = format_headers_redacted(&headers, false); + assert!(result.contains("authorization")); assert!(result.contains("content-type")); - assert!(result.contains("application/json")); assert!(result.contains("x-request-id")); - assert!(result.contains("abc123")); - } - - #[test] - fn test_format_headers_redacted_filters_sensitive() { - let mut headers = HeaderMap::new(); - headers.insert("authorization", "Bearer secret-token".parse().unwrap()); - headers.insert("x-client-secret", "private-value".parse().unwrap()); - headers.insert("x-client-credential", "credential-value".parse().unwrap()); - headers.insert("content-type", "application/json".parse().unwrap()); - - let result = format_headers_redacted(&headers, false); - - // Sensitive header should be present but with redacted value - assert!(result.contains("authorization")); assert!(result.contains("[REDACTED]")); - // Sensitive value should NOT be present assert!(!result.contains("secret-token")); - assert!(!result.contains("private-value")); - assert!(!result.contains("credential-value")); - // Non-sensitive header should be present with actual value - assert!(result.contains("content-type")); - assert!(result.contains("application/json")); - } - - #[test] - fn test_format_headers_redacted_filters_set_cookie() { - let mut headers = HeaderMap::new(); - headers.insert( - "set-cookie", - "CF_Authorization=sensitive-session-token; Path=/; Secure;" - .parse() - .unwrap(), - ); - headers.insert("server", "cloudflare".parse().unwrap()); - - let result = format_headers_redacted(&headers, false); - - // Sensitive header should be present but with redacted value - assert!(result.contains("set-cookie")); - assert!(result.contains("[REDACTED]")); - // Sensitive value should NOT be present - assert!(!result.contains("sensitive-session-token")); - // Non-sensitive header should be present with actual value - assert!(result.contains("server")); - assert!(result.contains("cloudflare")); - } - - #[test] - fn test_format_headers_redacted_filters_all_sensitive() { - let mut headers = HeaderMap::new(); - headers.insert("authorization", "Bearer token".parse().unwrap()); - headers.insert("proxy-authorization", "Basic creds".parse().unwrap()); - headers.insert("set-cookie", "session=abc".parse().unwrap()); - headers.insert("cookie", "session=abc".parse().unwrap()); - headers.insert("x-api-key", "api-key-123".parse().unwrap()); - headers.insert("x-auth-token", "auth-token-456".parse().unwrap()); - headers.insert("x-request-id", "req-123".parse().unwrap()); - - let result = format_headers_redacted(&headers, false); - - // All sensitive headers should be present but with redacted values - assert!(result.contains("authorization")); - assert!(result.contains("proxy-authorization")); - assert!(result.contains("set-cookie")); - assert!(result.contains("cookie")); - assert!(result.contains("x-api-key")); - assert!(result.contains("x-auth-token")); - assert!(result.contains("[REDACTED]")); - - // Ensure no sensitive values leaked - assert!(!result.contains("Bearer token")); - assert!(!result.contains("Basic creds")); - assert!(!result.contains("session=abc")); - assert!(!result.contains("api-key-123")); - assert!(!result.contains("auth-token-456")); - - // Non-sensitive header should be present with actual value - assert!(result.contains("x-request-id")); - assert!(result.contains("req-123")); + assert!(!result.contains("application/json")); + assert!(!result.contains("abc123")); } #[test] @@ -501,4 +464,102 @@ mod tests { // [REDACTED] should NOT be present when redaction is disabled assert!(!result.contains("[REDACTED]")); } + + #[tokio::test] + async fn table_client_reuses_parent_token_unless_overridden() { + let inherited_props = HashMap::from([( + "credential".to_string(), + "client-id:client-secret".to_string(), + )]); + let config = RestCatalogConfig::builder() + .uri("https://catalog.example".to_string()) + .props(inherited_props.clone()) + .build(); + let client = HttpClient::new(&config).unwrap(); + *client.token.lock().await = Some("catalog-token".to_string()); + + let inherited = client + .for_table("https://catalog.example", &HashMap::new()) + .unwrap(); + assert!(Arc::ptr_eq(&client.token, &inherited.token)); + assert_eq!( + inherited.token.lock().await.as_deref(), + Some("catalog-token") + ); + + let overridden = client + .for_table( + "https://catalog.example", + &HashMap::from([("token".to_string(), "table-token".to_string())]), + ) + .unwrap(); + assert!(!Arc::ptr_eq(&client.token, &overridden.token)); + assert_eq!( + overridden.token.lock().await.as_deref(), + Some("table-token") + ); + } + + #[tokio::test] + async fn table_auth_removes_inherited_authorization_header() { + let config = RestCatalogConfig::builder() + .uri("https://catalog.example".to_string()) + .props(HashMap::from([ + ("token".to_string(), "catalog-token".to_string()), + ( + "header.Authorization".to_string(), + "Bearer catalog-header-token".to_string(), + ), + ])) + .build(); + let client = HttpClient::new(&config).unwrap(); + + let table_client = client + .for_table( + "https://catalog.example", + &HashMap::from([("token".to_string(), "table-token".to_string())]), + ) + .unwrap(); + + assert!( + !table_client + .extra_headers + .contains_key(http::header::AUTHORIZATION) + ); + assert_eq!( + table_client.token.lock().await.as_deref(), + Some("table-token") + ); + } + + #[tokio::test] + async fn table_oauth_options_without_credential_reuse_parent_token() { + let config = RestCatalogConfig::builder() + .uri("https://catalog.example".to_string()) + .props(HashMap::from([( + "token".to_string(), + "catalog-token".to_string(), + )])) + .build(); + let client = HttpClient::new(&config).unwrap(); + let table_props = HashMap::from([ + ("scope".to_string(), "table-scope".to_string()), + ( + "oauth2-server-uri".to_string(), + "https://table-auth.example/token".to_string(), + ), + ]); + + let table_client = client + .for_table("https://catalog.example", &table_props) + .unwrap(); + + assert!(Arc::ptr_eq(&client.token, &table_client.token)); + assert_eq!( + table_client.token.lock().await.as_deref(), + Some("catalog-token") + ); + assert_eq!(table_client.token_endpoint, client.token_endpoint); + assert_eq!(table_client.extra_oauth_params, client.extra_oauth_params); + } } diff --git a/crates/catalog/rest/src/credential.rs b/crates/catalog/rest/src/credential.rs index 457e90299f..77eed648fc 100644 --- a/crates/catalog/rest/src/credential.rs +++ b/crates/catalog/rest/src/credential.rs @@ -28,8 +28,11 @@ //! 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. +//! expected backend-specific type. 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. //! //! # Adding a cloud //! @@ -38,7 +41,7 @@ //! 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::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; @@ -69,7 +72,7 @@ struct CloudRefresh { 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"`. + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. enabled_key: &'static str, /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. jitter_prefetch: bool, @@ -95,7 +98,7 @@ impl CloudRefresh { jitter_prefetch: false, parse_credential: parse_gcs_credential, }; - // TODO: Azure (ADLS) is not yet supported: opendal 0.57's Azdls builder exposes no + // TODO: Azure (ADLS) is not yet supported: opendal'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. @@ -112,7 +115,10 @@ impl CloudRefresh { } fn matches_location(&self, location: &str) -> bool { - scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) } } @@ -127,7 +133,7 @@ const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); /// value through the full value. const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); -/// Maximum failure backoff while a cached credential remains usable. +/// Maximum delay between failed refresh attempts. const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); /// A cached vended credential and its refresh schedule. @@ -142,7 +148,7 @@ struct CachedEntry { impl CachedEntry { fn new(credential: StorageCredential, jitter_prefetch: bool) -> Self { let refresh_at = credential - .expires_at + .expires_at() .map(|expires_at| prefetch_time(expires_at, jitter_prefetch)); Self { credential, @@ -154,7 +160,7 @@ impl CachedEntry { /// immediately due. Otherwise AWS applies the same jitter as it does to a /// freshly fetched value. fn seed(credential: StorageCredential, jitter_prefetch: bool) -> Self { - let due = credential.expires_at.is_some_and(|expires_at| { + 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) @@ -172,11 +178,22 @@ impl CachedEntry { fn is_unexpired(&self, now: SystemTime) -> bool { self.credential - .expires_at + .expires_at() .is_none_or(|expires_at| now < expires_at) } } +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec, + errors: Vec, +} + +struct CredentialError { + prefix: String, + error: Error, +} + /// Cached credentials plus failure-backoff state. struct CacheState { entries: Vec, @@ -184,6 +201,31 @@ struct CacheState { retry_not_before: Option, } +impl CacheState { + fn record_success(&mut self) { + self.consecutive_failures = 0; + self.retry_not_before = None; + } + + fn record_failure(&mut self) { + self.consecutive_failures = self.consecutive_failures.saturating_add(1); + self.retry_not_before = + Instant::now().checked_add(failure_backoff(self.consecutive_failures)); + } + + fn insert_fallback_if_missing(&mut self, fallback: CachedEntry, now: SystemTime) -> bool { + let has_unexpired_entry = self.entries.iter().any(|entry| { + entry.is_unexpired(now) && entry.credential.prefix() == fallback.credential.prefix() + }); + if has_unexpired_entry { + false + } else { + self.entries.push(fallback); + true + } + } +} + struct ConfiguredCloud { cloud: &'static CloudRefresh, endpoint: String, @@ -224,7 +266,7 @@ impl RestVendedCredentialProvider { } /// Fetch fresh credentials from the catalog's credentials endpoint. - async fn fetch(&self, configured: &ConfiguredCloud) -> Result> { + async fn fetch(&self, configured: &ConfiguredCloud) -> Result { let mut request = self.client.request(Method::GET, &configured.endpoint); if let Some(plan_id) = &self.plan_id { request = request.query(&[("planId", plan_id)]); @@ -235,18 +277,39 @@ impl RestVendedCredentialProvider { match response.status() { StatusCode::OK => { let parsed: LoadCredentialsResponse = response.json().await?; - parsed + let mut entries = Vec::new(); + let mut errors = Vec::new(); + let matching_credentials = parsed .storage_credentials .into_iter() - .filter(|sc| configured.cloud.matches_location(&sc.prefix)) - .map(|sc| { - (configured.cloud.parse_credential)(&sc.config, Some(sc.prefix)).map( - |credential| { - CachedEntry::new(credential, configured.cloud.jitter_prefetch) - }, - ) - }) - .collect() + .filter(|credential| configured.cloud.matches_location(&credential.prefix)); + let parse_credential = configured.cloud.parse_credential; + let now = SystemTime::now(); + + for storage_credential in matching_credentials { + let prefix = storage_credential.prefix; + let parsed_credential = + parse_credential(&storage_credential.config, Some(prefix.clone())); + match parsed_credential { + Ok(credential) => { + let entry = + CachedEntry::new(credential, configured.cloud.jitter_prefetch); + if entry.is_unexpired(now) { + entries.push(entry); + } else { + errors.push(CredentialError { + prefix, + error: Error::new( + ErrorKind::DataInvalid, + "invalid vended credential: credential is already expired", + ), + }); + } + } + Err(error) => errors.push(CredentialError { prefix, error }), + } + } + Ok(ParsedCredentials { entries, errors }) } _ => Err(deserialize_unexpected_catalog_error( response, @@ -262,32 +325,101 @@ impl RestVendedCredentialProvider { path: &str, fallback: Option, ) -> Result { - 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 self.fetch(configured).await { + Ok(ParsedCredentials { entries, errors }) => { + let invalid_prefixes = errors + .iter() + .map(|error| error.prefix.clone()) + .collect::>(); + let credential_error = errors + .into_iter() + .filter(|error| path.starts_with(&error.prefix)) + .max_by_key(|error| error.prefix.len()); + let failure = credential_error + .map(|error| error.error) + .unwrap_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!( + "no unexpired vended credential matches storage location: {path}" + ), + ) + }); - match refreshed { - Ok((entries, credential)) => { let mut cache = configured.cache.lock().await; + if entries.is_empty() { + cache.record_failure(); + return fallback + .filter(|entry| entry.is_unexpired(SystemTime::now())) + .map(|entry| entry.credential) + .ok_or(failure); + } + + let now = SystemTime::now(); + let invalid_fallbacks = cache + .entries + .iter() + .filter(|entry| entry.is_unexpired(now)) + .filter(|entry| { + entry + .credential + .prefix() + .is_some_and(|prefix| invalid_prefixes.contains(prefix)) + }) + .cloned() + .collect::>(); cache.entries = entries; - cache.consecutive_failures = 0; - cache.retry_not_before = None; - Ok(credential) + + // An invalid replacement for one prefix must not evict that + // prefix's still-valid cached credential. Iterate in reverse + // because equal-length prefix selection uses the last cached entry. + // This preserves the same credential if a bad response follows a + // response with duplicate prefixes. + let mut restored_fallback_prefixes = HashSet::new(); + for fallback in invalid_fallbacks.into_iter().rev() { + let prefix = fallback.credential.prefix().map(str::to_owned); + if cache.insert_fallback_if_missing(fallback, now) + && let Some(prefix) = prefix + { + restored_fallback_prefixes.insert(prefix); + } + } + + let selected = longest_prefix_match(&cache.entries, path) + .filter(|entry| entry.is_unexpired(now)) + .map(|entry| { + let is_restored_fallback = entry + .credential + .prefix() + .is_some_and(|prefix| restored_fallback_prefixes.contains(prefix)); + (entry.credential.clone(), is_restored_fallback) + }); + + if let Some((credential, is_restored_fallback)) = selected { + if is_restored_fallback { + cache.record_failure(); + } else { + cache.record_success(); + } + return Ok(credential); + } + + // Cache valid credentials for other prefixes while retaining this + // path's still-usable credential when its replacement was absent or + // invalid. Backoff ensures the failed path is retried promptly. + if let Some(fallback) = fallback.filter(|entry| entry.is_unexpired(now)) { + let credential = fallback.credential.clone(); + cache.insert_fallback_if_missing(fallback, now); + cache.record_failure(); + return Ok(credential); + } + + cache.record_failure(); + Err(failure) } 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)); + cache.record_failure(); // Graceful degradation: while the cached credential remains // usable, serve it and retry after jittered backoff. Expired @@ -312,6 +444,14 @@ impl std::fmt::Debug for RestVendedCredentialProvider { enum CacheDecision { Use(StorageCredential), Refresh(Option), + Backoff, +} + +fn refresh_backoff_error(path: &str) -> Error { + Error::new( + ErrorKind::Unexpected, + format!("vended credential refresh is temporarily backed off for storage location: {path}"), + ) } async fn cache_decision(configured: &ConfiguredCloud, path: &str) -> CacheDecision { @@ -326,9 +466,12 @@ async fn cache_decision(configured: &ConfiguredCloud, path: &str) -> CacheDecisi 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()); + return current + .as_ref() + .filter(|entry| entry.is_unexpired(now)) + .map(|entry| CacheDecision::Use(entry.credential.clone())) + .unwrap_or(CacheDecision::Backoff); } CacheDecision::Refresh(current) @@ -357,6 +500,7 @@ impl StorageCredentialProvider for RestVendedCredentialProvider { let current = match cache_decision(configured, path).await { CacheDecision::Use(credential) => return Ok(credential), CacheDecision::Refresh(current) => current, + CacheDecision::Backoff => return Err(refresh_backoff_error(path)), }; // One caller refreshes, while concurrent callers immediately keep using the @@ -379,6 +523,7 @@ impl StorageCredentialProvider for RestVendedCredentialProvider { let current = match cache_decision(configured, path).await { CacheDecision::Use(credential) => return Ok(credential), CacheDecision::Refresh(current) => current, + CacheDecision::Backoff => return Err(refresh_backoff_error(path)), }; self.refresh_credential(configured, path, current).await @@ -392,11 +537,10 @@ fn longest_prefix_match<'a>(entries: &'a [CachedEntry], path: &str) -> Option<&' .filter(|entry| { entry .credential - .prefix - .as_deref() + .prefix() .is_none_or(|prefix| path.starts_with(prefix)) }) - .max_by_key(|entry| entry.credential.prefix.as_deref().map_or(0, str::len)) + .max_by_key(|entry| entry.credential.prefix().map_or(0, str::len)) } /// Compute a successful credential's prefetch time. @@ -429,18 +573,20 @@ fn failure_backoff(consecutive_failures: u32) -> Duration { /// or `None` when no supported cloud advertises an enabled refresh endpoint. /// /// `base_uri` is the catalog URI, used to resolve a relative endpoint. +/// `table_auth_props` is the unmerged config returned by the table endpoint: +/// keeping it separate prevents local FileIO overrides from masking table auth. pub(crate) fn build_vended_credential_provider( client: Arc, base_uri: &str, props: &HashMap, + table_auth_props: Option<&HashMap>, ) -> Result>> { 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)); + .is_none_or(|value| value.eq_ignore_ascii_case("true")); if !enabled { return None; } @@ -473,7 +619,9 @@ pub(crate) fn build_vended_credential_provider( return Ok(None); } - let table_client = Arc::new(client.for_table(base_uri, props.clone())?); + let empty_auth_props = HashMap::new(); + let table_client = + Arc::new(client.for_table(base_uri, table_auth_props.unwrap_or(&empty_auth_props))?); let plan_id = props.get(REST_CATALOG_PROP_SCAN_PLAN_ID).cloned(); Ok(Some(Arc::new(RestVendedCredentialProvider::new( table_client, @@ -509,14 +657,15 @@ fn parse_s3_credential( let secret_access_key = required_nonempty(config, S3_SECRET_ACCESS_KEY)?; let session_token = required_nonempty(config, S3_SESSION_TOKEN)?; let expires_at = required_epoch_millis(config, S3_SESSION_TOKEN_EXPIRES_AT_MS)?; - Ok(StorageCredential { - prefix, - kind: StorageCredentialKind::S3(S3Credential { - access_key_id, - secret_access_key, - session_token: Some(session_token), - }), - expires_at: Some(expires_at), + let credential = StorageCredential::new(StorageCredentialKind::S3(S3Credential::new( + access_key_id, + secret_access_key, + Some(session_token), + ))) + .with_expiration(expires_at); + Ok(match prefix { + Some(prefix) => credential.with_prefix(prefix), + None => credential, }) } @@ -527,10 +676,11 @@ fn parse_gcs_credential( ) -> Result { let token = required_nonempty(config, GCS_TOKEN)?; let expires_at = required_epoch_millis(config, GCS_TOKEN_EXPIRES_AT)?; - Ok(StorageCredential { - prefix, - kind: StorageCredentialKind::Gcs(GcsCredential { token }), - expires_at: Some(expires_at), + let credential = StorageCredential::new(StorageCredentialKind::Gcs(GcsCredential::new(token))) + .with_expiration(expires_at); + Ok(match prefix { + Some(prefix) => credential.with_prefix(prefix), + None => credential, }) } @@ -586,20 +736,23 @@ mod tests { access_key_id: &str, expires_at: Option, ) -> StorageCredential { - StorageCredential { - prefix: prefix.map(str::to_string), - kind: StorageCredentialKind::S3(S3Credential { - access_key_id: access_key_id.to_string(), - secret_access_key: "secret".to_string(), - session_token: None, - }), - expires_at, + let mut credential = StorageCredential::new(StorageCredentialKind::S3(S3Credential::new( + access_key_id, + "secret", + None, + ))); + if let Some(prefix) = prefix { + credential = credential.with_prefix(prefix); } + if let Some(expires_at) = expires_at { + credential = credential.with_expiration(expires_at); + } + credential } fn s3_access_key_id(credential: &StorageCredential) -> &str { - match &credential.kind { - StorageCredentialKind::S3(s3) => &s3.access_key_id, + match credential.kind() { + StorageCredentialKind::S3(s3) => s3.access_key_id(), other => panic!("expected S3 credential, got {other:?}"), } } @@ -608,6 +761,53 @@ mod tests { CachedEntry::new(s3_cred(Some(prefix), access_key_id, expires_at), false) } + fn test_client(base_uri: &str) -> Arc { + let config = RestCatalogConfig::builder() + .uri(base_uri.to_string()) + .build(); + Arc::new(HttpClient::new(&config).unwrap()) + } + + fn aws_refresh_props(endpoint: &str) -> HashMap { + HashMap::from([( + CloudRefresh::AWS.endpoint_key.to_string(), + endpoint.to_string(), + )]) + } + + fn with_s3_seed( + mut props: HashMap, + access_key_id: &str, + expires_at: SystemTime, + ) -> HashMap { + props.extend([ + (S3_ACCESS_KEY_ID.to_string(), access_key_id.to_string()), + (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), + (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), + ( + S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), + epoch_millis(expires_at), + ), + ]); + props + } + + fn test_provider( + base_uri: &str, + props: &HashMap, + ) -> Arc { + build_vended_credential_provider(test_client(base_uri), base_uri, props, None) + .expect("provider construction should succeed") + .expect("provider should be built") + } + + fn s3_response(prefix: &str, access_key_id: &str, expires_at: SystemTime) -> String { + let expires_at = epoch_millis(expires_at); + format!( + r#"{{"storage-credentials":[{{"prefix":"{prefix}","config":{{"s3.access-key-id":"{access_key_id}","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires_at}"}}}}]}}"# + ) + } + #[test] fn resolve_endpoint_matches_java_semantics() { assert_eq!( @@ -640,7 +840,11 @@ mod tests { } #[test] - fn cloud_selection_uses_url_scheme() { + fn cloud_selection_accepts_root_prefixes_and_url_schemes() { + // Java uses these root prefixes for its fallback clients and accepts + // credentials scoped directly to them. + assert!(CloudRefresh::for_location("s3").is_some()); + assert!(CloudRefresh::for_location("gs").is_some()); assert!(CloudRefresh::for_location("s3://b/k").is_some()); assert!(CloudRefresh::for_location("S3://b/k").is_some()); assert!(CloudRefresh::for_location("s3a://b/k").is_some()); @@ -669,13 +873,13 @@ mod tests { "1500".to_string(), ); let credential = parse_s3_credential(&config, Some("s3://bucket".to_string())).unwrap(); - assert_eq!(credential.prefix.as_deref(), Some("s3://bucket")); + assert_eq!(credential.prefix(), Some("s3://bucket")); assert_eq!( - credential.expires_at, + credential.expires_at(), Some(UNIX_EPOCH + Duration::from_millis(1500)) ); - match credential.kind { - StorageCredentialKind::S3(s3) => assert_eq!(s3.session_token.as_deref(), Some("TOK")), + match credential.kind() { + StorageCredentialKind::S3(s3) => assert_eq!(s3.session_token(), Some("TOK")), other => panic!("expected S3, got {other:?}"), } } @@ -689,13 +893,13 @@ mod tests { config.insert(GCS_TOKEN_EXPIRES_AT.to_string(), "2000".to_string()); let credential = parse_gcs_credential(&config, Some("gs://bucket".to_string())).unwrap(); - assert_eq!(credential.prefix.as_deref(), Some("gs://bucket")); - match &credential.kind { - StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token, "ya29.token"), + assert_eq!(credential.prefix(), Some("gs://bucket")); + match credential.kind() { + StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token(), "ya29.token"), other => panic!("expected GCS, got {other:?}"), } assert_eq!( - credential.expires_at, + credential.expires_at(), Some(UNIX_EPOCH + Duration::from_millis(2000)) ); } @@ -730,7 +934,7 @@ mod tests { expires_at - REFRESH_BUFFER ); - for _ in 0..100 { + for _ in 0..16 { let refresh_at = prefetch_time(expires_at, true); assert!(refresh_at >= expires_at - REFRESH_BUFFER); assert!(refresh_at < expires_at - MIN_REFRESH_BUFFER); @@ -752,12 +956,14 @@ mod tests { #[test] fn failure_backoff_is_jittered_and_capped() { - for failures in 1_u32..=40 { - let ceiling = INITIAL_FAILURE_BACKOFF - .checked_mul(1_u32 << failures.saturating_sub(1).min(5)) - .unwrap_or(MAX_FAILURE_BACKOFF) - .min(MAX_FAILURE_BACKOFF); - for _ in 0..100 { + for (failures, ceiling) in [ + (1, Duration::from_secs(1)), + (2, Duration::from_secs(2)), + (5, Duration::from_secs(16)), + (6, MAX_FAILURE_BACKOFF), + (u32::MAX, MAX_FAILURE_BACKOFF), + ] { + for _ in 0..16 { let backoff = failure_backoff(failures); assert!(backoff >= ceiling / 2); assert!(backoff <= ceiling); @@ -774,10 +980,7 @@ mod tests { ]; let got = longest_prefix_match(&entries, "s3://bucket/warehouse/db/t/f").unwrap(); assert_eq!(s3_access_key_id(&got.credential), "narrow"); - assert_eq!( - got.credential.prefix.as_deref(), - Some("s3://bucket/warehouse/db") - ); + assert_eq!(got.credential.prefix(), Some("s3://bucket/warehouse/db")); assert!(longest_prefix_match(&entries, "s3://other/x").is_none()); let fresh = Some(SystemTime::now() + REFRESH_BUFFER + Duration::from_secs(60)); @@ -793,25 +996,34 @@ mod tests { #[test] fn provider_support_tracks_configured_cloud_endpoints() { - let config = RestCatalogConfig::builder() - .uri("http://cat".to_string()) - .build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); - let props = HashMap::from([( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/creds".to_string(), - )]); + let client = test_client("http://cat"); + let props = aws_refresh_props("/v1/creds"); // The provider is configured independently of the table metadata scheme, // but advertises support only for clouds whose endpoint is present. - let provider = build_vended_credential_provider(client.clone(), "http://cat", &props) + let provider = build_vended_credential_provider(client.clone(), "http://cat", &props, None) .unwrap() .unwrap(); assert!(provider.supports_path("s3://b/k")); assert!(!provider.supports_path("abfss://fs@acct.dfs.core.windows.net/k")); + let enabled = HashMap::from([ + ( + CloudRefresh::AWS.endpoint_key.to_string(), + "/v1/creds".to_string(), + ), + ( + CloudRefresh::AWS.enabled_key.to_string(), + "True".to_string(), + ), + ]); + assert!( + build_vended_credential_provider(client.clone(), "http://cat", &enabled, None) + .unwrap() + .is_some() + ); // No endpoint advertised. assert!( - build_vended_credential_provider(client.clone(), "http://cat", &HashMap::new()) + build_vended_credential_provider(client.clone(), "http://cat", &HashMap::new(), None) .unwrap() .is_none() ); @@ -827,29 +1039,27 @@ mod tests { ), ]); assert!( - build_vended_credential_provider(client, "http://cat", &disabled) + build_vended_credential_provider(client, "http://cat", &disabled, None) .unwrap() .is_none() ); // Java's `Strings.isNullOrEmpty` check treats an empty endpoint as absent. let empty = HashMap::from([(CloudRefresh::AWS.endpoint_key.to_string(), String::new())]); - let config = RestCatalogConfig::builder() - .uri("http://cat".to_string()) - .build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); + let client = test_client("http://cat"); assert!( - build_vended_credential_provider(client, "http://cat", &empty) + build_vended_credential_provider(client, "http://cat", &empty, None) .unwrap() .is_none() ); } #[tokio::test] - async fn fetches_and_selects_vended_credential() { + async fn refresh_includes_scan_plan_id() { let mut server = Server::new_async().await; - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); - let body = format!( - r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + let body = s3_response( + "s3://bucket", + "AK", + SystemTime::now() + Duration::from_secs(3600), ); let mock = server .mock("GET", "/v1/credentials") @@ -863,39 +1073,122 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); // No static creds -> no seed -> first load fetches from the endpoint. - let props = HashMap::from([ - ( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - ), - ( - REST_CATALOG_PROP_SCAN_PLAN_ID.to_string(), - "scan-plan-1".to_string(), - ), - ]); + let mut props = aws_refresh_props("/v1/credentials"); + props.insert( + REST_CATALOG_PROP_SCAN_PLAN_ID.to_string(), + "scan-plan-1".to_string(), + ); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .expect("provider construction should succeed") - .expect("provider should be built"); + let provider = test_provider(&server.url(), &props); let credential = provider .load_credential("s3://bucket/warehouse/f") .await .unwrap(); - assert_eq!(credential.prefix.as_deref(), Some("s3://bucket")); + assert_eq!(credential.prefix(), Some("s3://bucket")); - match credential.kind { + match credential.kind() { StorageCredentialKind::S3(s3) => { - assert_eq!(s3.access_key_id, "AK"); - assert_eq!(s3.session_token.as_deref(), Some("TOK")); + assert_eq!(s3.access_key_id(), "AK"); + assert_eq!(s3.session_token(), Some("TOK")); } other => panic!("expected S3, got {other:?}"), } mock.assert_async().await; } + #[tokio::test] + async fn successful_refresh_caches_entries_before_selecting_path() { + let mut server = Server::new_async().await; + let body = s3_response( + "s3://bucket/table-a", + "AK", + SystemTime::now() + Duration::from_secs(3600), + ); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let props = aws_refresh_props("/v1/credentials"); + let provider = test_provider(&server.url(), &props); + + assert!( + provider + .load_credential("s3://bucket/table-b/file") + .await + .is_err() + ); + // A successful response with no matching credential is negatively + // cached instead of immediately hitting the endpoint again. + assert!( + provider + .load_credential("s3://bucket/table-b/file") + .await + .is_err() + ); + assert_eq!( + s3_access_key_id( + &provider + .load_credential("s3://bucket/table-a/file") + .await + .unwrap() + ), + "AK" + ); + mock.assert_async().await; + } + + #[tokio::test] + async fn malformed_entry_does_not_discard_valid_entries() { + let mut server = Server::new_async().await; + let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let body = format!( + r#"{{"storage-credentials":[ + {{"prefix":"s3://bucket/invalid","config":{{"s3.access-key-id":"BAD"}}}}, + {{"prefix":"s3://bucket/valid","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}} + ]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let props = aws_refresh_props("/v1/credentials"); + let provider = test_provider(&server.url(), &props); + + assert!( + provider + .load_credential("s3://bucket/invalid/file") + .await + .is_err() + ); + assert!( + provider + .load_credential("s3://bucket/invalid/file") + .await + .is_err() + ); + assert_eq!( + s3_access_key_id( + &provider + .load_credential("s3://bucket/valid/file") + .await + .unwrap() + ), + "AK" + ); + mock.assert_async().await; + } + #[tokio::test] async fn fresh_seed_is_served_without_fetching() { let mut server = Server::new_async().await; @@ -906,26 +1199,16 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); - let props = HashMap::from([ - ( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - ), - (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), - (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), - (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), - (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expires), - ]); + let props = with_s3_seed( + aws_refresh_props("/v1/credentials"), + "SEED_AK", + SystemTime::now() + Duration::from_secs(3600), + ); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .expect("provider construction should succeed") - .expect("provider should be built"); + let provider = test_provider(&server.url(), &props); let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); - assert_eq!(credential.prefix, None); + assert_eq!(credential.prefix(), None); assert_eq!(s3_access_key_id(&credential), "SEED_AK"); mock.assert_async().await; } @@ -941,27 +1224,17 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); // Seeded credential is within the refresh buffer but not yet expired. - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); - let props = HashMap::from([ - ( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - ), - (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), - (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), - (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), - (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expires), - ]); + let props = with_s3_seed( + aws_refresh_props("/v1/credentials"), + "SEED_AK", + SystemTime::now() + Duration::from_secs(60), + ); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .expect("provider construction should succeed") - .expect("provider should be built"); + let provider = test_provider(&server.url(), &props); // The first refresh fails, but the still-valid seed is served. Immediate // follow-up operations stay inside the first jittered backoff window. - for _ in 0..3 { + for _ in 0..2 { let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); assert_eq!(s3_access_key_id(&credential), "SEED_AK"); } @@ -969,11 +1242,172 @@ mod tests { } #[tokio::test] - async fn concurrent_prefetch_serves_unexpired_credential_without_waiting() { + async fn empty_refresh_preserves_unexpired_seed() { let mut server = Server::new_async().await; - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"storage-credentials":[]}"#) + .create_async() + .await; + + let props = with_s3_seed( + aws_refresh_props("/v1/credentials"), + "SEED_AK", + SystemTime::now() + Duration::from_secs(60), + ); + let provider = test_provider(&server.url(), &props); + + for _ in 0..2 { + let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); + assert_eq!(s3_access_key_id(&credential), "SEED_AK"); + } + mock.assert_async().await; + } + + #[tokio::test] + async fn malformed_specific_refresh_prefers_fallback_over_valid_broader_entry() { + let mut server = Server::new_async().await; + let refreshed_expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let body = format!( + r#"{{"storage-credentials":[ + {{"prefix":"s3://bucket/requested","config":{{"s3.access-key-id":"BAD"}}}}, + {{"prefix":"s3://bucket","config":{{"s3.access-key-id":"NEW_AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{refreshed_expires}"}}}} + ]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let provider = RestVendedCredentialProvider::new(test_client(&server.url()), None, vec![ + ConfiguredCloud { + cloud: &CloudRefresh::AWS, + endpoint: format!("{}/v1/credentials", server.url()), + cache: Mutex::new(CacheState { + entries: vec![cached_s3( + "s3://bucket/requested", + "SEED_AK", + Some(SystemTime::now() + Duration::from_secs(60)), + )], + consecutive_failures: 0, + retry_not_before: None, + }), + refresh: Mutex::new(()), + }, + ]); + + let fallback = provider + .load_credential("s3://bucket/requested/file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&fallback), "SEED_AK"); + + // The malformed specific replacement is a failed refresh for this path. + // Its still-valid fallback wins over the broader fetched credential and + // is backed off, so an immediate retry does not fetch again. + let fallback = provider + .load_credential("s3://bucket/requested/other-file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&fallback), "SEED_AK"); + + let refreshed = provider + .load_credential("s3://bucket/other/file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&refreshed), "NEW_AK"); + mock.assert_async().await; + } + + #[tokio::test] + async fn invalid_refresh_preserves_other_prefix_fallbacks() { + let mut server = Server::new_async().await; + let refreshed_expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); + let expired = epoch_millis(SystemTime::now() - Duration::from_secs(60)); let body = format!( - r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"NEW_AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + r#"{{"storage-credentials":[ + {{"prefix":"s3://bucket/table-a","config":{{"s3.access-key-id":"NEW_AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{refreshed_expires}"}}}}, + {{"prefix":"s3://bucket/table-b","config":{{"s3.access-key-id":"BAD"}}}}, + {{"prefix":"s3://bucket/table-c","config":{{"s3.access-key-id":"EXPIRED_CK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expired}"}}}} + ]}}"# + ); + let mock = server + .mock("GET", "/v1/credentials") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + + let provider = RestVendedCredentialProvider::new(test_client(&server.url()), None, vec![ + ConfiguredCloud { + cloud: &CloudRefresh::AWS, + endpoint: format!("{}/v1/credentials", server.url()), + cache: Mutex::new(CacheState { + entries: vec![ + cached_s3( + "s3://bucket/table-a", + "OLD_AK", + Some(SystemTime::now() + Duration::from_secs(60)), + ), + cached_s3( + "s3://bucket/table-b", + "OLDER_BK", + Some(SystemTime::now() + Duration::from_secs(3600)), + ), + cached_s3( + "s3://bucket/table-b", + "LATEST_BK", + Some(SystemTime::now() + Duration::from_secs(3600)), + ), + cached_s3( + "s3://bucket/table-c", + "VALID_CK", + Some(SystemTime::now() + Duration::from_secs(3600)), + ), + ], + consecutive_failures: 0, + retry_not_before: None, + }), + refresh: Mutex::new(()), + }, + ]); + + let refreshed = provider + .load_credential("s3://bucket/table-a/file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&refreshed), "NEW_AK"); + + let fallback = provider + .load_credential("s3://bucket/table-b/file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&fallback), "LATEST_BK"); + + let fallback = provider + .load_credential("s3://bucket/table-c/file") + .await + .unwrap(); + assert_eq!(s3_access_key_id(&fallback), "VALID_CK"); + mock.assert_async().await; + } + + #[tokio::test] + async fn concurrent_prefetch_serves_unexpired_credential_without_waiting() { + let mut server = Server::new_async().await; + let body = s3_response( + "s3://bucket", + "NEW_AK", + SystemTime::now() + Duration::from_secs(3600), ); let (started_tx, started_rx) = mpsc::channel(); let release = Arc::new(Barrier::new(2)); @@ -991,22 +1425,12 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); - let seed_expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); - let props = HashMap::from([ - ( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - ), - (S3_ACCESS_KEY_ID.to_string(), "SEED_AK".to_string()), - (S3_SECRET_ACCESS_KEY.to_string(), "SEED_SK".to_string()), - (S3_SESSION_TOKEN.to_string(), "SEED_TOK".to_string()), - (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), seed_expires), - ]); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .unwrap() - .unwrap(); + let props = with_s3_seed( + aws_refresh_props("/v1/credentials"), + "SEED_AK", + SystemTime::now() + Duration::from_secs(60), + ); + let provider = test_provider(&server.url(), &props); let first_provider = Arc::clone(&provider); let first = @@ -1041,9 +1465,10 @@ mod tests { #[tokio::test] async fn refresh_uses_table_scoped_token() { let mut server = Server::new_async().await; - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); - let body = format!( - r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + let body = s3_response( + "s3://bucket", + "AK", + SystemTime::now() + Duration::from_secs(3600), ); let mock = server .mock("GET", "/v1/credentials") @@ -1054,18 +1479,28 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); + let config = RestCatalogConfig::builder() + .uri(server.url()) + .props(HashMap::from([( + "token".to_string(), + "catalog-token".to_string(), + )])) + .build(); let client = Arc::new(HttpClient::new(&config).unwrap()); let props = HashMap::from([ ( CloudRefresh::AWS.endpoint_key.to_string(), "/v1/credentials".to_string(), ), - ("token".to_string(), "table-token".to_string()), + // Local properties win in the FileIO configuration merge, but must + // not mask auth returned by the table endpoint. + ("token".to_string(), "user-token".to_string()), ]); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .unwrap() - .unwrap(); + let table_auth = HashMap::from([("token".to_string(), "table-token".to_string())]); + let provider = + build_vended_credential_provider(client, &server.url(), &props, Some(&table_auth)) + .unwrap() + .unwrap(); provider.load_credential("s3://bucket/x/f").await.unwrap(); mock.assert_async().await; @@ -1075,8 +1510,10 @@ mod tests { async fn one_provider_refreshes_multiple_clouds() { let mut server = Server::new_async().await; let expires = epoch_millis(SystemTime::now() + Duration::from_secs(3600)); - let aws_body = format!( - r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + let aws_body = s3_response( + "s3://bucket", + "AK", + SystemTime::now() + Duration::from_secs(3600), ); let gcp_body = format!( r#"{{"storage-credentials":[{{"prefix":"gs://bucket","config":{{"gcs.oauth2.token":"GCS","gcs.oauth2.token-expires-at":"{expires}"}}}}]}}"# @@ -1096,8 +1533,6 @@ mod tests { .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); let props = HashMap::from([ ( CloudRefresh::AWS.endpoint_key.to_string(), @@ -1108,9 +1543,7 @@ mod tests { "/v1/gcp-credentials".to_string(), ), ]); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .unwrap() - .unwrap(); + let provider = test_provider(&server.url(), &props); assert!(provider.supports_path("s3://bucket/x")); assert!(provider.supports_path("gs://bucket/x")); @@ -1122,9 +1555,9 @@ mod tests { .load_credential("gs://bucket/x") .await .unwrap() - .kind + .kind() { - StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token, "GCS"), + StorageCredentialKind::Gcs(gcs) => assert_eq!(gcs.token(), "GCS"), other => panic!("expected GCS credential, got {other:?}"), } aws_mock.assert_async().await; @@ -1136,28 +1569,20 @@ mod tests { let mut server = Server::new_async().await; let mock = server .mock("GET", Matcher::Any) + .expect(1) .with_status(500) .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); - let expired = epoch_millis(SystemTime::now() - Duration::from_secs(60)); - let props = HashMap::from([ - ( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - ), - (S3_ACCESS_KEY_ID.to_string(), "EXPIRED_AK".to_string()), - (S3_SECRET_ACCESS_KEY.to_string(), "EXPIRED_SK".to_string()), - (S3_SESSION_TOKEN.to_string(), "EXPIRED_TOK".to_string()), - (S3_SESSION_TOKEN_EXPIRES_AT_MS.to_string(), expired), - ]); + let props = with_s3_seed( + aws_refresh_props("/v1/credentials"), + "EXPIRED_AK", + SystemTime::now() - Duration::from_secs(60), + ); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .expect("provider construction should succeed") - .expect("provider should be built"); + let provider = test_provider(&server.url(), &props); + assert!(provider.load_credential("s3://bucket/x/f").await.is_err()); assert!(provider.load_credential("s3://bucket/x/f").await.is_err()); mock.assert_async().await; } @@ -1168,33 +1593,27 @@ mod tests { // The vended TTL (60s) is shorter than REFRESH_BUFFER, so every operation // is eligible for refresh. This matches AWS CachedSupplier: successful // results do not receive failure backoff. - let expires = epoch_millis(SystemTime::now() + Duration::from_secs(60)); - let body = format!( - r#"{{"storage-credentials":[{{"prefix":"s3://bucket","config":{{"s3.access-key-id":"AK","s3.secret-access-key":"SK","s3.session-token":"TOK","s3.session-token-expires-at-ms":"{expires}"}}}}]}}"# + let body = s3_response( + "s3://bucket", + "AK", + SystemTime::now() + Duration::from_secs(60), ); let mock = server .mock("GET", "/v1/credentials") - .expect(3) + .expect(2) .with_status(200) .with_header("content-type", "application/json") .with_body(body) .create_async() .await; - let config = RestCatalogConfig::builder().uri(server.url()).build(); - let client = Arc::new(HttpClient::new(&config).unwrap()); // No static creds -> no seed -> each sequential load fetches because the // returned credential is already inside its prefetch window. - let props = HashMap::from([( - CloudRefresh::AWS.endpoint_key.to_string(), - "/v1/credentials".to_string(), - )]); + let props = aws_refresh_props("/v1/credentials"); - let provider = build_vended_credential_provider(client, &server.url(), &props) - .expect("provider construction should succeed") - .expect("provider should be built"); + let provider = test_provider(&server.url(), &props); - for _ in 0..3 { + for _ in 0..2 { let credential = provider.load_credential("s3://bucket/x/f").await.unwrap(); assert_eq!(s3_access_key_id(&credential), "AK"); } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 93eaa414cf..698dcb1a68 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -784,7 +784,10 @@ pub fn iceberg::io::GcsConfig::serialize<__S>(&self, __serializer: __S) -> core: impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::GcsConfig pub fn iceberg::io::GcsConfig::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg::io::GcsCredential -pub iceberg::io::GcsCredential::token: alloc::string::String +impl iceberg::io::GcsCredential +pub fn iceberg::io::GcsCredential::into_token(self) -> alloc::string::String +pub fn iceberg::io::GcsCredential::new(token: impl core::convert::Into) -> Self +pub fn iceberg::io::GcsCredential::token(&self) -> &str impl core::clone::Clone for iceberg::io::GcsCredential pub fn iceberg::io::GcsCredential::clone(&self) -> iceberg::io::GcsCredential impl core::fmt::Debug for iceberg::io::GcsCredential @@ -972,9 +975,12 @@ pub fn iceberg::io::S3Config::serialize<__S>(&self, __serializer: __S) -> core:: impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::S3Config pub fn iceberg::io::S3Config::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg::io::S3Credential -pub iceberg::io::S3Credential::access_key_id: alloc::string::String -pub iceberg::io::S3Credential::secret_access_key: alloc::string::String -pub iceberg::io::S3Credential::session_token: core::option::Option +impl iceberg::io::S3Credential +pub fn iceberg::io::S3Credential::access_key_id(&self) -> &str +pub fn iceberg::io::S3Credential::into_parts(self) -> (alloc::string::String, alloc::string::String, core::option::Option) +pub fn iceberg::io::S3Credential::new(access_key_id: impl core::convert::Into, secret_access_key: impl core::convert::Into, session_token: core::option::Option) -> Self +pub fn iceberg::io::S3Credential::secret_access_key(&self) -> &str +pub fn iceberg::io::S3Credential::session_token(&self) -> core::option::Option<&str> impl core::clone::Clone for iceberg::io::S3Credential pub fn iceberg::io::S3Credential::clone(&self) -> iceberg::io::S3Credential impl core::fmt::Debug for iceberg::io::S3Credential @@ -1017,9 +1023,14 @@ pub fn iceberg::io::StorageConfig::serialize<__S>(&self, __serializer: __S) -> c impl<'de> serde_core::de::Deserialize<'de> for iceberg::io::StorageConfig pub fn iceberg::io::StorageConfig::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg::io::StorageCredential -pub iceberg::io::StorageCredential::expires_at: core::option::Option -pub iceberg::io::StorageCredential::kind: iceberg::io::StorageCredentialKind -pub iceberg::io::StorageCredential::prefix: core::option::Option +impl iceberg::io::StorageCredential +pub fn iceberg::io::StorageCredential::expires_at(&self) -> core::option::Option +pub fn iceberg::io::StorageCredential::into_kind(self) -> iceberg::io::StorageCredentialKind +pub fn iceberg::io::StorageCredential::kind(&self) -> &iceberg::io::StorageCredentialKind +pub fn iceberg::io::StorageCredential::new(kind: iceberg::io::StorageCredentialKind) -> Self +pub fn iceberg::io::StorageCredential::prefix(&self) -> core::option::Option<&str> +pub fn iceberg::io::StorageCredential::with_expiration(self, expires_at: std::time::SystemTime) -> Self +pub fn iceberg::io::StorageCredential::with_prefix(self, prefix: impl core::convert::Into) -> Self impl core::clone::Clone for iceberg::io::StorageCredential pub fn iceberg::io::StorageCredential::clone(&self) -> iceberg::io::StorageCredential impl core::fmt::Debug for iceberg::io::StorageCredential diff --git a/crates/iceberg/src/io/storage/mod.rs b/crates/iceberg/src/io/storage/mod.rs index dd812a2d45..8ce1780c44 100644 --- a/crates/iceberg/src/io/storage/mod.rs +++ b/crates/iceberg/src/io/storage/mod.rs @@ -198,12 +198,55 @@ pub trait StorageCredentialProvider: Debug + Send + Sync { pub struct StorageCredential { /// Storage-location prefix this credential is scoped to. `None` represents a /// credential without a declared scope, sourced from flat storage properties. - pub prefix: Option, + prefix: Option, /// The backend-specific credential material. - pub kind: StorageCredentialKind, + kind: StorageCredentialKind, /// When the credential expires, if known. `None` means non-expiring and /// backends treat such a credential as always valid and never refresh it. - pub expires_at: Option, + expires_at: Option, +} + +impl StorageCredential { + /// Create a storage credential with no declared scope or expiration. + pub fn new(kind: StorageCredentialKind) -> Self { + Self { + prefix: None, + kind, + expires_at: None, + } + } + + /// Set the storage-location prefix this credential is scoped to. + pub fn with_prefix(mut self, prefix: impl Into) -> Self { + self.prefix = Some(prefix.into()); + self + } + + /// Set when this credential expires. + pub fn with_expiration(mut self, expires_at: SystemTime) -> Self { + self.expires_at = Some(expires_at); + self + } + + /// Return the storage-location prefix this credential is scoped to. + pub fn prefix(&self) -> Option<&str> { + self.prefix.as_deref() + } + + /// Return the backend-specific credential material. + pub fn kind(&self) -> &StorageCredentialKind { + &self.kind + } + + /// Consume this credential and return its backend-specific material. + pub fn into_kind(self) -> StorageCredentialKind { + self.kind + } + + /// Return when this credential expires. + pub fn expires_at(&self) -> Option { + self.expires_at + } } /// Backend-specific credential material. @@ -219,11 +262,51 @@ pub enum StorageCredentialKind { #[derive(Clone)] pub struct S3Credential { /// AWS access key ID. - pub access_key_id: String, + access_key_id: String, /// AWS secret access key. - pub secret_access_key: String, + secret_access_key: String, /// AWS session token, set for temporary (STS/vended) credentials. - pub session_token: Option, + session_token: Option, +} + +impl S3Credential { + /// Create temporary Amazon S3 credentials. + pub fn new( + access_key_id: impl Into, + secret_access_key: impl Into, + session_token: Option, + ) -> Self { + Self { + access_key_id: access_key_id.into(), + secret_access_key: secret_access_key.into(), + session_token, + } + } + + /// Return the AWS access key ID. + pub fn access_key_id(&self) -> &str { + &self.access_key_id + } + + /// Return the AWS secret access key. + pub fn secret_access_key(&self) -> &str { + &self.secret_access_key + } + + /// Return the AWS session token, if present. + pub fn session_token(&self) -> Option<&str> { + self.session_token.as_deref() + } + + /// Consume these credentials and return their component values. + pub fn into_parts(self) -> (String, String, Option) { + let Self { + access_key_id, + secret_access_key, + session_token, + } = self; + (access_key_id, secret_access_key, session_token) + } } impl Debug for S3Credential { @@ -236,7 +319,26 @@ impl Debug for S3Credential { #[derive(Clone)] pub struct GcsCredential { /// OAuth2 bearer token used to access GCS. - pub token: String, + token: String, +} + +impl GcsCredential { + /// Create a Google Cloud Storage credential. + pub fn new(token: impl Into) -> Self { + Self { + token: token.into(), + } + } + + /// Return the OAuth2 bearer token used to access GCS. + pub fn token(&self) -> &str { + &self.token + } + + /// Consume this credential and return its OAuth2 bearer token. + pub fn into_token(self) -> String { + self.token + } } impl Debug for GcsCredential { diff --git a/crates/storage/opendal/src/gcs.rs b/crates/storage/opendal/src/gcs.rs index 311bd1ac88..20cf1919a1 100644 --- a/crates/storage/opendal/src/gcs.rs +++ b/crates/storage/opendal/src/gcs.rs @@ -30,6 +30,7 @@ use reqsign_core::{Context, Error as ReqsignError, ProvideCredential, Result as use reqsign_google::{Credential as GoogleCredential, Token as GoogleToken}; use url::Url; +use crate::DynamicCredentialScope; use crate::utils::{ from_opendal_error, is_truthy, system_time_to_timestamp, validate_credential_prefix, }; @@ -82,6 +83,7 @@ pub(crate) fn gcs_config_build( cfg: &GcsConfig, credential_provider: &Option>, path: &str, + credential_scope: Option<&DynamicCredentialScope>, ) -> Result { let url = Url::parse(path)?; if !matches!(url.scheme(), "gs" | "gcs") { @@ -128,6 +130,7 @@ pub(crate) fn gcs_config_build( builder = builder.credential_provider(VendedGcsCredentialProvider::new( Arc::clone(provider), path.to_string(), + credential_scope.cloned(), )); } @@ -142,19 +145,31 @@ struct VendedGcsCredentialProvider { /// Absolute path this operator serves: handed back to the provider so it can /// select the vended credential whose prefix best matches the location. path: String, + /// Exact scope selected when a bulk-delete operator was created. Ordinary + /// operators are unbound so they can follow the provider's best match. + credential_scope: Option, } impl std::fmt::Debug for VendedGcsCredentialProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("VendedGcsCredentialProvider") .field("path", &self.path) + .field("credential_scope", &self.credential_scope) .finish_non_exhaustive() } } impl VendedGcsCredentialProvider { - fn new(provider: Arc, path: String) -> Self { - Self { provider, path } + fn new( + provider: Arc, + path: String, + credential_scope: Option, + ) -> Self { + Self { + provider, + path, + credential_scope, + } } } @@ -174,16 +189,20 @@ impl ProvideCredential for VendedGcsCredentialProvider { .with_source(e) })?; - validate_credential_prefix(&self.path, credential.prefix.as_deref())?; + validate_credential_prefix( + &self.path, + credential.prefix(), + self.credential_scope.as_ref(), + )?; let expires_at = credential - .expires_at + .expires_at() .map(system_time_to_timestamp) .transpose()?; - match credential.kind { + match credential.into_kind() { StorageCredentialKind::Gcs(gcs) => { Ok(Some(GoogleCredential::with_token(GoogleToken { - access_token: gcs.token, + access_token: gcs.into_token(), expires_at, }))) } @@ -193,3 +212,92 @@ impl ProvideCredential for VendedGcsCredentialProvider { } } } + +#[cfg(test)] +mod tests { + use async_trait::async_trait; + use iceberg::io::{GcsCredential, StorageCredential}; + + use super::*; + + #[derive(Debug)] + struct FixedCredentialProvider(StorageCredential); + + #[async_trait] + impl StorageCredentialProvider for FixedCredentialProvider { + async fn load_credential(&self, _path: &str) -> Result { + Ok(self.0.clone()) + } + } + + fn credential(prefix: Option<&str>) -> StorageCredential { + let credential = StorageCredential::new(StorageCredentialKind::Gcs(GcsCredential::new( + "access-token", + ))); + match prefix { + Some(prefix) => credential.with_prefix(prefix), + None => credential, + } + } + + fn provider( + path: &str, + returned_prefix: Option<&str>, + credential_scope: Option, + ) -> VendedGcsCredentialProvider { + VendedGcsCredentialProvider::new( + Arc::new(FixedCredentialProvider(credential(returned_prefix))), + path.to_string(), + credential_scope, + ) + } + + #[tokio::test] + async fn vended_provider_enforces_bound_credential_scope() { + let path = "gs://bucket/table/file.parquet"; + let table_prefix = "gs://bucket/table"; + + assert!( + provider( + path, + Some(table_prefix), + Some(DynamicCredentialScope::Prefix(table_prefix.to_string())), + ) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider( + path, + Some("gs://bucket"), + Some(DynamicCredentialScope::Prefix(table_prefix.to_string())), + ) + .provide_credential(&Context::default()) + .await + .is_err() + ); + assert!( + provider(path, Some("gs://bucket"), None) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider(path, None, Some(DynamicCredentialScope::Unscoped)) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider( + path, + Some(table_prefix), + Some(DynamicCredentialScope::Unscoped), + ) + .provide_credential(&Context::default()) + .await + .is_err() + ); + } +} diff --git a/crates/storage/opendal/src/lib.rs b/crates/storage/opendal/src/lib.rs index 0f301dfc91..e58b94575e 100644 --- a/crates/storage/opendal/src/lib.rs +++ b/crates/storage/opendal/src/lib.rs @@ -145,6 +145,21 @@ impl StorageFactory for OpenDalStorageFactory { config: &StorageConfig, credential_provider: Option>, ) -> Result> { + #[allow(unreachable_patterns)] + let supports_credential_provider = match self { + #[cfg(feature = "opendal-s3")] + OpenDalStorageFactory::S3 { .. } => true, + #[cfg(feature = "opendal-gcs")] + OpenDalStorageFactory::Gcs => true, + _ => false, + }; + if credential_provider.is_some() && !supports_credential_provider { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "OpenDAL storage factory does not support refreshable credentials for this backend", + )); + } + match self { #[cfg(feature = "opendal-memory")] OpenDalStorageFactory::Memory => { @@ -262,6 +277,33 @@ pub enum OpenDalStorage { }, } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum DeleteCredentialScope { + Static, + Dynamic(DynamicCredentialScope), +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum DynamicCredentialScope { + Unscoped, + Prefix(String), +} + +impl DynamicCredentialScope { + fn from_prefix(prefix: Option<&str>) -> Self { + match prefix { + Some(prefix) => Self::Prefix(prefix.to_string()), + None => Self::Unscoped, + } + } +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct DeleteBatchKey { + storage: String, + credential_scope: DeleteCredentialScope, +} + impl OpenDalStorage { /// Creates operator from path. /// @@ -279,6 +321,17 @@ impl OpenDalStorage { pub(crate) fn create_operator<'a>( &self, path: &'a impl AsRef, + ) -> Result<(Operator, &'a str)> { + self.create_operator_with_scope(path, None) + } + + /// Creates an operator, optionally binding dynamic credentials to the exact + /// scope used to group a bulk-delete batch. + #[allow(unreachable_code, unused_variables)] + fn create_operator_with_scope<'a>( + &self, + path: &'a impl AsRef, + credential_scope: Option<&DynamicCredentialScope>, ) -> Result<(Operator, &'a str)> { let path = path.as_ref(); let (operator, relative_path): (Operator, &str) = match self { @@ -310,6 +363,7 @@ impl OpenDalStorage { customized_credential_load, credential_provider, path, + credential_scope, )?; let op_info = op.info(); @@ -336,7 +390,8 @@ impl OpenDalStorage { config, credential_provider, } => { - let operator = gcs_config_build(config, credential_provider, path)?; + let operator = + gcs_config_build(config, credential_provider, path, credential_scope)?; let url = url::Url::parse(path).map_err(|e| { Error::new( ErrorKind::DataInvalid, @@ -415,25 +470,58 @@ impl OpenDalStorage { } } - /// Returns whether `path` uses a path-scoped dynamic credential provider. - /// - /// Such paths cannot share an OpenDAL deleter unless the provider can expose - /// the credential scope that applies to each path. Process them independently - /// so bulk deletion remains bounded without crossing credential boundaries. - fn uses_dynamic_credentials(&self, path: &str) -> bool { - match self { + /// Return the dynamic credential provider that serves `path`, if any. + fn credential_provider_for_path( + &self, + path: &str, + ) -> Option<&Arc> { + let provider: &Arc = (match self { #[cfg(feature = "opendal-s3")] OpenDalStorage::S3 { + customized_credential_load: None, credential_provider: Some(provider), .. - } => provider.supports_path(path), + } => Some(provider), #[cfg(feature = "opendal-gcs")] OpenDalStorage::Gcs { credential_provider: Some(provider), .. - } => provider.supports_path(path), - _ => false, - } + } => Some(provider), + _ => None, + })?; + provider.supports_path(path).then_some(provider) + } + + /// Returns a key that keeps bulk deletes within one operator and credential + /// scope. Loading the credential is normally a cache hit and avoids rebuilding + /// an operator for every path while preventing a batch from crossing prefixes. + async fn delete_batch_key_for_path(&self, path: &str) -> Result { + let credential_scope = match self.credential_provider_for_path(path) { + Some(provider) => { + let credential = provider.load_credential(path).await?; + if credential + .prefix() + .is_some_and(|prefix| prefix.is_empty() || !path.starts_with(prefix)) + { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "vended credential prefix {:?} does not cover storage location {path:?}", + credential.prefix() + ), + )); + } + DeleteCredentialScope::Dynamic(DynamicCredentialScope::from_prefix( + credential.prefix(), + )) + } + None => DeleteCredentialScope::Static, + }; + + Ok(DeleteBatchKey { + storage: self.batch_key_for_path(path), + credential_scope, + }) } /// Extracts the relative path from an absolute path without building an operator. @@ -608,23 +696,21 @@ impl Storage for OpenDalStorage { } async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> { - let mut deleters: HashMap = HashMap::new(); + let mut deleters: HashMap = HashMap::new(); while let Some(path) = paths.next().await { - if self.uses_dynamic_credentials(&path) { - let (op, relative_path) = self.create_operator(&path)?; - op.delete(relative_path).await.map_err(from_opendal_error)?; - continue; - } + let batch_key = self.delete_batch_key_for_path(&path).await?; - let bucket = self.batch_key_for_path(&path); - - let (relative_path, deleter) = match deleters.entry(bucket) { + let (relative_path, deleter) = match deleters.entry(batch_key) { Entry::Occupied(entry) => { (self.relativize_path(&path)?.to_string(), entry.into_mut()) } Entry::Vacant(entry) => { - let (op, rel) = self.create_operator(&path)?; + let credential_scope = match &entry.key().credential_scope { + DeleteCredentialScope::Static => None, + DeleteCredentialScope::Dynamic(scope) => Some(scope), + }; + let (op, rel) = self.create_operator_with_scope(&path, credential_scope)?; let rel = rel.to_string(); let deleter = op.deleter().await.map_err(from_opendal_error)?; (rel, entry.insert(deleter)) @@ -702,8 +788,18 @@ mod tests { #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] #[async_trait] impl StorageCredentialProvider for AlwaysSupportedCredentialProvider { - async fn load_credential(&self, _path: &str) -> Result { - unreachable!("batch key selection must not load a credential") + async fn load_credential(&self, path: &str) -> Result { + let prefix = if path.contains("/table-a/") { + "s3://bucket/table-a" + } else { + "s3://bucket/table-b" + }; + Ok( + iceberg::io::StorageCredential::new(iceberg::io::StorageCredentialKind::S3( + iceberg::io::S3Credential::new("access-key", "secret-key", None), + )) + .with_prefix(prefix), + ) } } @@ -714,6 +810,22 @@ mod tests { assert_eq!(op.info().scheme().to_string(), "memory"); } + #[cfg(all( + feature = "opendal-memory", + any(feature = "opendal-s3", feature = "opendal-gcs") + ))] + #[test] + fn test_factory_rejects_credentials_for_unsupported_backend() { + let error = OpenDalStorageFactory::Memory + .build_with_credentials( + &StorageConfig::new(), + Some(Arc::new(AlwaysSupportedCredentialProvider)), + ) + .expect_err("memory must reject a credential provider"); + + assert_eq!(error.kind(), ErrorKind::FeatureUnsupported); + } + #[cfg(feature = "opendal-memory")] #[test] fn test_relativize_path_memory() { @@ -770,20 +882,53 @@ mod tests { } #[cfg(feature = "opendal-s3")] - #[test] - fn test_dynamic_credentials_use_isolated_deletes() { + #[tokio::test] + async fn test_dynamic_credentials_batch_by_prefix() { let storage = OpenDalStorage::S3 { config: Arc::new(S3Config::default()), customized_credential_load: None, credential_provider: Some(Arc::new(AlwaysSupportedCredentialProvider)), }; let first = "s3://bucket/table-a/file.parquet"; - let second = "s3://bucket/table-b/file.parquet"; + let same_scope = "s3://bucket/table-a/other.parquet"; + let other_scope = "s3://bucket/table-b/file.parquet"; - assert!(storage.uses_dynamic_credentials(first)); - assert!(storage.uses_dynamic_credentials(second)); - assert_eq!(storage.batch_key_for_path(first), "bucket"); - assert_eq!(storage.batch_key_for_path(second), "bucket"); + let first_key = storage.delete_batch_key_for_path(first).await.unwrap(); + assert_eq!( + first_key.credential_scope, + DeleteCredentialScope::Dynamic(DynamicCredentialScope::Prefix( + "s3://bucket/table-a".to_string() + )) + ); + assert_eq!( + first_key, + storage.delete_batch_key_for_path(same_scope).await.unwrap() + ); + assert_ne!( + first_key, + storage + .delete_batch_key_for_path(other_scope) + .await + .unwrap() + ); + } + + #[cfg(feature = "opendal-s3")] + #[tokio::test] + async fn test_custom_s3_credential_loader_ignores_dynamic_provider_for_batching() { + let storage = OpenDalStorage::S3 { + config: Arc::new(S3Config::default()), + customized_credential_load: Some(CustomAwsCredentialLoader::new( + reqsign_aws_v4::StaticCredentialProvider::new("access-key", "secret-key"), + )), + credential_provider: Some(Arc::new(AlwaysSupportedCredentialProvider)), + }; + + let key = storage + .delete_batch_key_for_path("s3://bucket/table-a/file.parquet") + .await + .unwrap(); + assert_eq!(key.credential_scope, DeleteCredentialScope::Static); } #[cfg(feature = "opendal-s3")] diff --git a/crates/storage/opendal/src/resolving.rs b/crates/storage/opendal/src/resolving.rs index 0fad716e52..05194bb292 100644 --- a/crates/storage/opendal/src/resolving.rs +++ b/crates/storage/opendal/src/resolving.rs @@ -80,6 +80,17 @@ fn extract_scheme(path: &str) -> Result<&'static str> { parse_scheme(url.scheme()) } +#[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] +fn supports_dynamic_credentials(scheme: &str) -> bool { + match scheme { + #[cfg(feature = "opendal-s3")] + "s3" => true, + #[cfg(feature = "opendal-gcs")] + "gcs" => true, + _ => false, + } +} + /// Build an [`OpenDalStorage`] variant for the given scheme and config properties. fn build_storage_for_scheme( scheme: &'static str, @@ -200,6 +211,14 @@ impl StorageFactory for OpenDalResolvingStorageFactory { config: &StorageConfig, credential_provider: Option>, ) -> Result> { + #[cfg(not(any(feature = "opendal-s3", feature = "opendal-gcs")))] + if credential_provider.is_some() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "OpenDAL resolving storage does not support refreshable credentials because no compatible backend is enabled", + )); + } + Ok(Arc::new(OpenDalResolvingStorage { props: config.props().clone(), storages: RwLock::new(HashMap::new()), @@ -250,6 +269,21 @@ impl OpenDalResolvingStorage { fn resolve(&self, path: &str) -> Result> { let scheme = extract_scheme(path)?; + #[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] + if self + .credential_provider + .as_ref() + .is_some_and(|provider| provider.supports_path(path)) + && !supports_dynamic_credentials(scheme) + { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!( + "OpenDAL resolving storage does not support refreshable credentials for scheme: {scheme}" + ), + )); + } + // Fast path: check read lock first. { let cache = self @@ -358,8 +392,36 @@ impl Storage for OpenDalResolvingStorage { mod tests { use super::*; + #[derive(Debug)] + struct AllPathsCredentialProvider; + + #[async_trait] + impl StorageCredentialProvider for AllPathsCredentialProvider { + async fn load_credential(&self, _path: &str) -> Result { + unreachable!("unsupported backends must reject the provider before loading") + } + } + + #[cfg(not(any(feature = "opendal-s3", feature = "opendal-gcs")))] + #[test] + fn test_factory_rejects_credentials_without_compatible_backend() { + let error = OpenDalResolvingStorageFactory::new() + .build_with_credentials( + &StorageConfig::new(), + Some(Arc::new(AllPathsCredentialProvider)), + ) + .expect_err("a provider must not be silently discarded"); + + assert_eq!(error.kind(), ErrorKind::FeatureUnsupported); + } + /// Builds a resolving storage with empty props, suitable for `resolve()` /// calls that don't actually hit any backend. + #[cfg(any( + feature = "opendal-s3", + feature = "opendal-gcs", + feature = "opendal-azdls" + ))] fn empty_resolving_storage() -> OpenDalResolvingStorage { OpenDalResolvingStorage { props: HashMap::new(), @@ -387,6 +449,21 @@ mod tests { assert!(Arc::ptr_eq(&a, &c), "s3 and s3n should share one instance"); } + #[cfg(all( + feature = "opendal-memory", + any(feature = "opendal-s3", feature = "opendal-gcs") + ))] + #[test] + fn test_resolver_rejects_credentials_for_unsupported_backend() { + let mut storage = empty_resolving_storage(); + storage.credential_provider = Some(Arc::new(AllPathsCredentialProvider)); + + let error = storage + .resolve("memory:/key") + .expect_err("memory must reject a credential provider"); + assert_eq!(error.kind(), ErrorKind::FeatureUnsupported); + } + #[cfg(feature = "opendal-azdls")] #[test] fn test_resolve_azdls_aliases_share_instance() { diff --git a/crates/storage/opendal/src/s3.rs b/crates/storage/opendal/src/s3.rs index e6b825347d..9d0282e007 100644 --- a/crates/storage/opendal/src/s3.rs +++ b/crates/storage/opendal/src/s3.rs @@ -38,6 +38,7 @@ use reqsign_core::{ }; use url::Url; +use crate::DynamicCredentialScope; use crate::utils::{ from_opendal_error, is_truthy, system_time_to_timestamp, validate_credential_prefix, }; @@ -137,6 +138,7 @@ pub(crate) fn s3_config_build( customized_credential_load: &Option, credential_provider: &Option>, path: &str, + credential_scope: Option<&DynamicCredentialScope>, ) -> Result { let url = Url::parse(path)?; let bucket = url.host_str().ok_or_else(|| { @@ -172,6 +174,7 @@ pub(crate) fn s3_config_build( let chain = ProvideCredentialChain::new().push(VendedS3CredentialProvider::new( Arc::clone(provider), path.to_string(), + credential_scope.cloned(), )); builder = builder.credential_provider_chain(chain); } @@ -187,19 +190,31 @@ struct VendedS3CredentialProvider { /// Absolute path this operator serves; handed back to the provider so it can /// select the vended credential whose prefix best matches the location. path: String, + /// Exact scope selected when a bulk-delete operator was created. Ordinary + /// operators are unbound so they can follow the provider's best match. + credential_scope: Option, } impl std::fmt::Debug for VendedS3CredentialProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("VendedS3CredentialProvider") .field("path", &self.path) + .field("credential_scope", &self.credential_scope) .finish_non_exhaustive() } } impl VendedS3CredentialProvider { - fn new(provider: Arc, path: String) -> Self { - Self { provider, path } + fn new( + provider: Arc, + path: String, + credential_scope: Option, + ) -> Self { + Self { + provider, + path, + credential_scope, + } } } @@ -219,19 +234,26 @@ impl ProvideCredential for VendedS3CredentialProvider { .with_source(e) })?; - validate_credential_prefix(&self.path, credential.prefix.as_deref())?; + validate_credential_prefix( + &self.path, + credential.prefix(), + self.credential_scope.as_ref(), + )?; let expires_in = credential - .expires_at + .expires_at() .map(system_time_to_timestamp) .transpose()?; - match credential.kind { - StorageCredentialKind::S3(s3) => Ok(Some(AwsCredential { - access_key_id: s3.access_key_id, - secret_access_key: s3.secret_access_key, - session_token: s3.session_token, - expires_in, - })), + match credential.into_kind() { + StorageCredentialKind::S3(s3) => { + let (access_key_id, secret_access_key, session_token) = s3.into_parts(); + Ok(Some(AwsCredential { + access_key_id, + secret_access_key, + session_token, + expires_in, + })) + } _ => Err(ReqsignError::unexpected( "S3 storage received a non-S3 credential from the provider", )), @@ -270,9 +292,44 @@ impl CustomAwsCredentialLoader { mod tests { use std::collections::HashMap; - use iceberg::io::S3_PATH_STYLE_ACCESS; + use async_trait::async_trait; + use iceberg::io::{S3_PATH_STYLE_ACCESS, S3Credential, StorageCredential}; + + use super::*; + + #[derive(Debug)] + struct FixedCredentialProvider(StorageCredential); + + #[async_trait] + impl StorageCredentialProvider for FixedCredentialProvider { + async fn load_credential(&self, _path: &str) -> Result { + Ok(self.0.clone()) + } + } - use super::s3_config_parse; + fn credential(prefix: Option<&str>) -> StorageCredential { + let credential = StorageCredential::new(StorageCredentialKind::S3(S3Credential::new( + "access-key", + "secret-key", + Some("session-token".to_string()), + ))); + match prefix { + Some(prefix) => credential.with_prefix(prefix), + None => credential, + } + } + + fn provider( + path: &str, + returned_prefix: Option<&str>, + credential_scope: Option, + ) -> VendedS3CredentialProvider { + VendedS3CredentialProvider::new( + Arc::new(FixedCredentialProvider(credential(returned_prefix))), + path.to_string(), + credential_scope, + ) + } fn parse_with(prop: Option<&str>) -> bool { let mut props = HashMap::new(); @@ -289,4 +346,53 @@ mod tests { assert!(parse_with(Some("false"))); assert!(!parse_with(Some("true"))); } + + #[tokio::test] + async fn vended_provider_enforces_bound_credential_scope() { + let path = "s3://bucket/table/file.parquet"; + let table_prefix = "s3://bucket/table"; + + assert!( + provider( + path, + Some(table_prefix), + Some(DynamicCredentialScope::Prefix(table_prefix.to_string())), + ) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider( + path, + Some("s3://bucket"), + Some(DynamicCredentialScope::Prefix(table_prefix.to_string())), + ) + .provide_credential(&Context::default()) + .await + .is_err() + ); + assert!( + provider(path, Some("s3://bucket"), None) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider(path, None, Some(DynamicCredentialScope::Unscoped)) + .provide_credential(&Context::default()) + .await + .is_ok() + ); + assert!( + provider( + path, + Some(table_prefix), + Some(DynamicCredentialScope::Unscoped), + ) + .provide_credential(&Context::default()) + .await + .is_err() + ); + } } diff --git a/crates/storage/opendal/src/utils.rs b/crates/storage/opendal/src/utils.rs index c168093fc2..ddb4f08b2d 100644 --- a/crates/storage/opendal/src/utils.rs +++ b/crates/storage/opendal/src/utils.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +#[cfg(any(feature = "opendal-s3", feature = "opendal-gcs"))] +use crate::DynamicCredentialScope; + pub(crate) fn is_truthy(value: &str) -> bool { ["true", "t", "1", "on"].contains(&value.to_lowercase().as_str()) } @@ -56,6 +59,7 @@ pub(crate) fn system_time_to_timestamp( pub(crate) fn validate_credential_prefix( path: &str, prefix: Option<&str>, + credential_scope: Option<&DynamicCredentialScope>, ) -> reqsign_core::Result<()> { match prefix { Some("") => Err(reqsign_core::Error::unexpected( @@ -64,6 +68,18 @@ pub(crate) fn validate_credential_prefix( Some(prefix) if !path.starts_with(prefix) => Err(reqsign_core::Error::unexpected(format!( "vended credential prefix {prefix:?} does not cover storage location {path:?}" ))), - _ => Ok(()), + _ => match credential_scope { + Some(DynamicCredentialScope::Unscoped) if prefix.is_some() => { + Err(reqsign_core::Error::unexpected( + "vended credential scope changed while a bulk-delete operator was active", + )) + } + Some(DynamicCredentialScope::Prefix(expected)) if prefix != Some(expected.as_str()) => { + Err(reqsign_core::Error::unexpected( + "vended credential scope changed while a bulk-delete operator was active", + )) + } + _ => Ok(()), + }, } }