zakariya-s commented on code in PR #2932: URL: https://github.com/apache/iceberg-rust/pull/2932#discussion_r4378781608
########## crates/catalog/rest/src/credential.rs: ########## @@ -0,0 +1,3675 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Vended storage credentials from a REST catalog. +//! +//! A REST catalog can vend short-lived storage credentials whose lifetime the +//! client does not control. [`RestVendedCredentialProvider`] implements the +//! core [`StorageCredentialProvider`] trait so storage backends re-fetch those +//! credentials from the catalog's table credentials endpoint before they +//! expire, keeping long-running jobs authenticated instead of failing with a +//! `403` once the initial token's TTL elapses. +//! +//! Unlike the Java client, which has one provider per cloud SDK, this is a +//! single backend-agnostic provider with an independent endpoint and cache for +//! each configured cloud. The path being accessed selects the cloud cache. Like +//! Java's `StorageCredential`, the returned [`StorageCredential`] holds the +//! backend's storage properties, which the storage adapter turns into its own +//! credential. This preserves Java's per-cloud credential selection and +//! prefetch policies while supporting mixed-cloud tables through a resolving +//! FileIO. Unlike Java's scheduled refresh, which permanently stops +//! after a failed fetch, transient failures are retried here with jittered +//! exponential backoff while an unexpired credential remains available. +//! +//! Like Java's `VendedCredentialsProvider`, a provider is rebuilt from the +//! FileIO properties after [`FileIO`](iceberg::io::FileIO) serialization and +//! connects to the catalog lazily in the receiving process. This requires +//! catalog authentication that can be rebuilt from `rest.auth.type`. +//! +//! # Adding a cloud +//! +//! The refresh policy for each cloud lives in one [`CloudRefresh`] constant. To +//! add a backend, teach its storage adapter to read its credential properties +//! from a [`StorageCredential`]. Then write its `parse_*` function, add a +//! `CloudRefresh` constant, and list it in [`CloudRefresh::SUPPORTED`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use iceberg::io::{ + ADLS_REFRESH_CREDENTIALS_ENABLED, ADLS_REFRESH_CREDENTIALS_ENDPOINT, + ADLS_SAS_TOKEN_EXPIRES_AT_MS_PREFIX, ADLS_SAS_TOKEN_PREFIX, AWS_REFRESH_CREDENTIALS_ENABLED, + AWS_REFRESH_CREDENTIALS_ENDPOINT, GCS_REFRESH_CREDENTIALS_ENABLED, + GCS_REFRESH_CREDENTIALS_ENDPOINT, GCS_TOKEN, GCS_TOKEN_EXPIRES_AT, S3_ACCESS_KEY_ID, + S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SESSION_TOKEN_EXPIRES_AT_MS, StorageConfig, + StorageCredential, StorageCredentialProvider, StorageCredentialProviderFactory, +}; +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Url}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, OnceCell}; + +use crate::auth::{AUTH_TYPE_NONE, AuthManager, load_auth_manager, static_token_session}; +use crate::catalog::{ + REST_CATALOG_PROP_AUTH_TYPE, REST_CATALOG_PROP_REFERENCED_BY, REST_CATALOG_PROP_SCAN_PLAN_ID, + RestCatalogConfig, +}; +use crate::client::{HttpClient, unexpected_catalog_error_without_body}; +use crate::request::HttpRequest; +use crate::types::LoadCredentialsResponse; + +type CredentialParser = + fn(config: &HashMap<String, String>, prefix: String) -> Result<VendedCredential>; +type KeyedSeedCredentialParser = + fn(config: &HashMap<String, String>) -> HashMap<String, VendedCredential>; +type KeyedSeedPathResolver = fn(path: &str) -> Result<KeyedSeedPath>; + +/// A credential vended by the catalog, with the expiry its config declares. +#[derive(Clone)] +struct VendedCredential { + credential: StorageCredential, + expires_at: SystemTime, +} + +enum SeedStrategy { + /// One credential stored in the backend's flat properties, for every + /// location with one of the scheme prefixes, like Java's root clients. + Flat { prefixes: &'static [&'static str] }, + /// Credentials selected by a backend-specific key derived from each path. + Keyed(KeyedSeedStrategy), +} + +struct KeyedSeedStrategy { + parse_credentials: KeyedSeedCredentialParser, + resolve_path: KeyedSeedPathResolver, +} + +struct KeyedSeedPath { + /// Seed keys for the path, most specific first. + keys: Vec<String>, + scope: String, +} + +/// Cloud-specific details regarding vended-credential refresh. +/// +/// It contains the location schemes it backs, the property keys it is configured +/// with, and how to parse its credential. The generic provider stays free of +/// any per-cloud knowledge. +struct CloudRefresh { + /// Location URL schemes this backend serves. + schemes: &'static [&'static str], + /// Table property naming the refresh endpoint (absolute or catalog-relative). + endpoint_key: &'static str, + /// Table property controlling refresh; only missing or case-insensitive `"true"` enables it. + enabled_key: &'static str, + /// Table property that must be present for refresh to apply, as in Java, + /// where it decides whether the backend uses vended credentials at all. + required_key: Option<&'static str>, + /// Whether to jitter successful prefetch times like AWS `CachedSupplier`. + jitter_prefetch: bool, + /// Parse a complete credential from catalog-supplied properties. + parse_credential: CredentialParser, + /// How initial credentials in the table properties are selected. + seed_strategy: SeedStrategy, +} + +impl CloudRefresh { + /// S3 / AWS + const AWS: Self = Self { + schemes: &["s3", "s3a", "s3n"], + endpoint_key: AWS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: AWS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: true, + parse_credential: parse_s3_credential, + // `s3` also covers `s3a` and `s3n` locations. + seed_strategy: SeedStrategy::Flat { prefixes: &["s3"] }, + }; + /// Google Cloud Storage + const GCP: Self = Self { + schemes: &["gs", "gcs"], + endpoint_key: GCS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: GCS_REFRESH_CREDENTIALS_ENABLED, + // Java's GCS client only refreshes a vended OAuth2 token; without + // one it uses Google's default credentials. + required_key: Some(GCS_TOKEN), + jitter_prefetch: false, + parse_credential: parse_gcs_credential, + seed_strategy: SeedStrategy::Flat { + prefixes: &["gs", "gcs"], + }, + }; + /// Azure Data Lake Storage + const AZURE: Self = Self { + schemes: &["abfs", "abfss", "wasb", "wasbs"], + endpoint_key: ADLS_REFRESH_CREDENTIALS_ENDPOINT, + enabled_key: ADLS_REFRESH_CREDENTIALS_ENABLED, + required_key: None, + jitter_prefetch: false, + parse_credential: parse_azdls_credential, + seed_strategy: SeedStrategy::Keyed(KeyedSeedStrategy { + parse_credentials: parse_azdls_account_seeds, + resolve_path: resolve_azdls_seed_path, + }), + }; + + /// Backends with refresh support + const SUPPORTED: &[Self] = &[Self::AWS, Self::GCP, Self::AZURE]; + + fn matches_location(&self, location: &str) -> bool { + self.schemes + .iter() + .any(|scheme| location.eq_ignore_ascii_case(scheme)) + || scheme_of(location).is_some_and(|scheme| self.schemes.contains(&scheme.as_str())) + } +} + +/// Re-fetch a credential once it is within this window of expiry, so a fresh +/// token is in hand before the object store would reject the old one. +const REFRESH_BUFFER: Duration = Duration::from_mins(5); + +/// AWS keeps at least one minute between its jittered prefetch time and expiry. +const MIN_REFRESH_BUFFER: Duration = Duration::from_mins(1); + +/// Initial ceiling for failure backoff. Equal jitter chooses from half this +/// value through the full value. +const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1); + +/// Maximum delay between failed refresh attempts. +const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(30); + +/// A cached vended credential and its refresh schedule. +#[derive(Clone)] +struct CachedEntry { + credential: StorageCredential, + expires_at: SystemTime, + /// When this entry becomes eligible for prefetch. + refresh_at: SystemTime, +} + +impl CachedEntry { + fn new(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let VendedCredential { + credential, + expires_at, + } = vended; + Self { + credential, + expires_at, + refresh_at: prefetch_time(SystemTime::now(), expires_at, jitter_prefetch), + } + } + + /// Seed entries that are already inside the nominal five-minute window are + /// immediately due. Otherwise AWS applies the same jitter as it does to a + /// freshly fetched value. + fn seed(vended: VendedCredential, jitter_prefetch: bool) -> Self { + let due = SystemTime::now() + .checked_add(REFRESH_BUFFER) + .is_none_or(|refresh_boundary| refresh_boundary >= vended.expires_at); + let mut entry = Self::new(vended, jitter_prefetch); + if due { + entry.refresh_at = UNIX_EPOCH; + } + entry + } + + fn is_fresh(&self, now: SystemTime) -> bool { + now < self.refresh_at + } + + fn is_unexpired(&self, now: SystemTime) -> bool { + now < self.expires_at + } +} + +/// Parsed entries from one successful credentials response. +struct ParsedCredentials { + entries: Vec<CachedEntry>, + errors: Vec<CredentialError>, +} + +struct CredentialError { + prefix: String, + error: Error, +} + +/// Cached credentials plus failure-backoff state. +/// +/// A fetch returns the credentials for every prefix of the table, so a path +/// the last response did not cover would not be covered by an immediate +/// re-fetch either. Backoff is therefore shared by the whole cloud. +struct CacheState { + entries: Vec<CachedEntry>, + /// Prefixes whose credential in the last successful response was invalid. + /// Paths under them must not use a broader credential: the catalog scoped + /// them more tightly. + failed_prefixes: HashSet<String>, + consecutive_failures: u32, + retry_not_before: Option<Instant>, + /// Why the last refresh failed, reported while refresh is backed off. + last_failure: Option<String>, +} + +impl CacheState { + fn new(entries: Vec<CachedEntry>) -> Self { + Self { + entries, + failed_prefixes: HashSet::new(), + consecutive_failures: 0, + retry_not_before: None, + last_failure: None, + } + } + + /// The index of the unexpired credential with the longest prefix covering + /// `path`. + fn unexpired_match(&self, path: &str, now: SystemTime) -> Option<usize> { + self.entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_unexpired(now) && entry.credential.covers(path)) Review Comment: it now picks the longest covering credential even when it has expired and refreshes it instead of serving a broader one using `StorageCredential::select` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
