diff --git a/Cargo.lock b/Cargo.lock index 633534233e..902f213648 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3891,7 +3891,9 @@ name = "iceberg-catalog-rest" version = "0.10.0" dependencies = [ "async-trait", + "bytes", "chrono", + "futures", "http 1.5.0", "iceberg", "iceberg_test_utils", @@ -3902,6 +3904,7 @@ dependencies = [ "serde_derive", "serde_json", "tokio", + "tracing", "typed-builder", "uuid", ] diff --git a/crates/catalog/rest/Cargo.toml b/crates/catalog/rest/Cargo.toml index 247709efd4..8dc9a86d7f 100644 --- a/crates/catalog/rest/Cargo.toml +++ b/crates/catalog/rest/Cargo.toml @@ -40,12 +40,17 @@ serde = { workspace = true } serde_derive = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } +tracing = { workspace = true } typed-builder = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] +bytes = { workspace = true } +futures = { workspace = true } iceberg_test_utils = { path = "../../test_utils", features = ["tests"] } mockito = { workspace = true } +# `stream` lets tests build a streaming body to exercise HttpRequestBody::Streaming. +reqwest = { workspace = true, features = ["stream"] } tokio = { workspace = true } [lints] diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 776b11c40a..61c6b795ba 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -1,4 +1,19 @@ pub mod iceberg_catalog_rest +pub enum iceberg_catalog_rest::HttpRequestBody<'a> +pub iceberg_catalog_rest::HttpRequestBody::Buffered(&'a [u8]) +pub iceberg_catalog_rest::HttpRequestBody::Empty +pub iceberg_catalog_rest::HttpRequestBody::Streaming +impl<'a> iceberg_catalog_rest::HttpRequestBody<'a> +pub fn iceberg_catalog_rest::HttpRequestBody<'a>::as_bytes(&self) -> core::option::Option<&'a [u8]> +impl<'a> core::clone::Clone for iceberg_catalog_rest::HttpRequestBody<'a> +pub fn iceberg_catalog_rest::HttpRequestBody<'a>::clone(&self) -> iceberg_catalog_rest::HttpRequestBody<'a> +impl<'a> core::cmp::Eq for iceberg_catalog_rest::HttpRequestBody<'a> +impl<'a> core::cmp::PartialEq for iceberg_catalog_rest::HttpRequestBody<'a> +pub fn iceberg_catalog_rest::HttpRequestBody<'a>::eq(&self, other: &iceberg_catalog_rest::HttpRequestBody<'a>) -> bool +impl<'a> core::fmt::Debug for iceberg_catalog_rest::HttpRequestBody<'a> +pub fn iceberg_catalog_rest::HttpRequestBody<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl<'a> core::marker::Copy for iceberg_catalog_rest::HttpRequestBody<'a> +impl<'a> core::marker::StructuralPartialEq for iceberg_catalog_rest::HttpRequestBody<'a> pub struct iceberg_catalog_rest::CommitTableRequest pub iceberg_catalog_rest::CommitTableRequest::identifier: core::option::Option pub iceberg_catalog_rest::CommitTableRequest::requirements: alloc::vec::Vec @@ -109,6 +124,23 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::ErrorResponse pub fn iceberg_catalog_rest::ErrorResponse::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::ErrorResponse pub fn iceberg_catalog_rest::ErrorResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::HttpClient +impl iceberg_catalog_rest::HttpClient +pub async fn iceberg_catalog_rest::HttpClient::post_form(&self, url: &str, headers: &http::header::map::HeaderMap, form: &std::collections::hash::map::HashMap<&str, &str>) -> iceberg::error::Result<(http::status::StatusCode, alloc::vec::Vec)> +pub fn iceberg_catalog_rest::HttpClient::with_session(&self, session: alloc::sync::Arc) -> Self +pub fn iceberg_catalog_rest::HttpClient::without_session(&self) -> Self +impl core::clone::Clone for iceberg_catalog_rest::HttpClient +pub fn iceberg_catalog_rest::HttpClient::clone(&self) -> iceberg_catalog_rest::HttpClient +impl core::fmt::Debug for iceberg_catalog_rest::HttpClient +pub fn iceberg_catalog_rest::HttpClient::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg_catalog_rest::HttpRequest +impl iceberg_catalog_rest::HttpRequest +pub fn iceberg_catalog_rest::HttpRequest::body(&self) -> iceberg_catalog_rest::HttpRequestBody<'_> +pub fn iceberg_catalog_rest::HttpRequest::headers(&self) -> &http::header::map::HeaderMap +pub fn iceberg_catalog_rest::HttpRequest::headers_mut(&mut self) -> &mut http::header::map::HeaderMap +pub fn iceberg_catalog_rest::HttpRequest::method(&self) -> &http::method::Method +pub fn iceberg_catalog_rest::HttpRequest::new(inner: reqwest::async_impl::request::Request) -> Self +pub fn iceberg_catalog_rest::HttpRequest::url_str(&self) -> &str pub struct iceberg_catalog_rest::ListNamespaceResponse pub iceberg_catalog_rest::ListNamespaceResponse::namespaces: alloc::vec::Vec pub iceberg_catalog_rest::ListNamespaceResponse::next_page_token: core::option::Option @@ -175,6 +207,24 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::NamespaceResponse pub fn iceberg_catalog_rest::NamespaceResponse::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::NamespaceResponse pub fn iceberg_catalog_rest::NamespaceResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::NoopAuthManager +impl core::fmt::Debug for iceberg_catalog_rest::NoopAuthManager +pub fn iceberg_catalog_rest::NoopAuthManager::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg_catalog_rest::AuthManager for iceberg_catalog_rest::NoopAuthManager +pub fn iceberg_catalog_rest::NoopAuthManager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _client: &'life1 iceberg_catalog_rest::HttpClient, _props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg_catalog_rest::NoopAuthManager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _client: &'life1 iceberg_catalog_rest::HttpClient, _props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub struct iceberg_catalog_rest::OAuth2Manager +impl iceberg_catalog_rest::OAuth2Manager +pub fn iceberg_catalog_rest::OAuth2Manager::new(token_endpoint: impl core::convert::Into) -> Self +pub fn iceberg_catalog_rest::OAuth2Manager::with_credential(self, client_id: core::option::Option, client_secret: alloc::string::String) -> Self +pub fn iceberg_catalog_rest::OAuth2Manager::with_extra_headers(self, headers: http::header::map::HeaderMap) -> Self +pub fn iceberg_catalog_rest::OAuth2Manager::with_extra_oauth_params(self, params: std::collections::hash::map::HashMap) -> Self +pub fn iceberg_catalog_rest::OAuth2Manager::with_token(self, token: impl core::convert::Into) -> Self +impl core::fmt::Debug for iceberg_catalog_rest::OAuth2Manager +pub fn iceberg_catalog_rest::OAuth2Manager::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg_catalog_rest::AuthManager for iceberg_catalog_rest::OAuth2Manager +pub fn iceberg_catalog_rest::OAuth2Manager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg_catalog_rest::OAuth2Manager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub struct iceberg_catalog_rest::RegisterTableRequest pub iceberg_catalog_rest::RegisterTableRequest::metadata_location: alloc::string::String pub iceberg_catalog_rest::RegisterTableRequest::name: alloc::string::String @@ -207,9 +257,6 @@ pub fn iceberg_catalog_rest::RenameTableRequest::serialize<__S>(&self, __seriali impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::RenameTableRequest pub fn iceberg_catalog_rest::RenameTableRequest::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::RestCatalog -impl iceberg_catalog_rest::RestCatalog -pub async fn iceberg_catalog_rest::RestCatalog::invalidate_token(&self) -> iceberg::error::Result<()> -pub async fn iceberg_catalog_rest::RestCatalog::regenerate_token(&self) -> iceberg::error::Result<()> impl core::fmt::Debug for iceberg_catalog_rest::RestCatalog pub fn iceberg_catalog_rest::RestCatalog::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::catalog::Catalog for iceberg_catalog_rest::RestCatalog @@ -230,6 +277,7 @@ pub fn iceberg_catalog_rest::RestCatalog::update_namespace<'life0, 'life1, 'asyn pub fn iceberg_catalog_rest::RestCatalog::update_table<'life0, 'async_trait>(&'life0 self, commit: iceberg::catalog::TableCommit) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait pub struct iceberg_catalog_rest::RestCatalogBuilder impl iceberg_catalog_rest::RestCatalogBuilder +pub fn iceberg_catalog_rest::RestCatalogBuilder::with_auth_manager(self, auth_manager: alloc::sync::Arc) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_client(self, client: reqwest::async_impl::client::Client) -> Self impl core::default::Default for iceberg_catalog_rest::RestCatalogBuilder pub fn iceberg_catalog_rest::RestCatalogBuilder::default() -> Self @@ -287,6 +335,20 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::UpdateNamespacePropert pub fn iceberg_catalog_rest::UpdateNamespacePropertiesResponse::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::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::AUTH_TYPE_NONE: &str +pub const iceberg_catalog_rest::AUTH_TYPE_OAUTH2: &str +pub const iceberg_catalog_rest::REST_CATALOG_PROP_AUTH_TYPE: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_URI: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_WAREHOUSE: &str +pub trait iceberg_catalog_rest::AuthManager: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg_catalog_rest::AuthManager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg_catalog_rest::AuthManager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +impl iceberg_catalog_rest::AuthManager for iceberg_catalog_rest::NoopAuthManager +pub fn iceberg_catalog_rest::NoopAuthManager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _client: &'life1 iceberg_catalog_rest::HttpClient, _props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg_catalog_rest::NoopAuthManager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _client: &'life1 iceberg_catalog_rest::HttpClient, _props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +impl iceberg_catalog_rest::AuthManager for iceberg_catalog_rest::OAuth2Manager +pub fn iceberg_catalog_rest::OAuth2Manager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg_catalog_rest::OAuth2Manager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub trait iceberg_catalog_rest::AuthSession: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg_catalog_rest::AuthSession::authenticate<'life0, 'life1, 'async_trait>(&'life0 self, request: &'life1 mut iceberg_catalog_rest::HttpRequest) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait diff --git a/crates/catalog/rest/src/auth/mod.rs b/crates/catalog/rest/src/auth/mod.rs new file mode 100644 index 0000000000..347b14bf64 --- /dev/null +++ b/crates/catalog/rest/src/auth/mod.rs @@ -0,0 +1,152 @@ +// 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. + +//! Pluggable authentication for the REST catalog, mirroring Iceberg Java's +//! `AuthManager`/`AuthSession` API. + +mod oauth2; + +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use iceberg::Result; +pub use oauth2::OAuth2Manager; + +use crate::client::HttpClient; +use crate::request::HttpRequest; + +/// `rest.auth.type` value disabling authentication. +pub const AUTH_TYPE_NONE: &str = "none"; +/// `rest.auth.type` value selecting OAuth2 token authentication. +pub const AUTH_TYPE_OAUTH2: &str = "oauth2"; + +/// Creates the [`AuthSession`]s used to authenticate REST catalog requests. +/// +/// A manager is created once per catalog, either from the `rest.auth.type` +/// property or injected through `RestCatalogBuilder::with_auth_manager`. It +/// builds the sessions the catalog then keeps. +/// +/// Both methods are handed the catalog's [`HttpClient`], which an +/// implementation may reuse for its own requests (e.g. a token exchange) so +/// that they share the catalog's connection pool and configuration. +#[async_trait] +pub trait AuthManager: Debug + Send + Sync { + /// Session used for the initial `/v1/config` handshake, given the + /// user-supplied properties. + /// + /// Returns a [`Box`]: an init session is used once and released, unlike + /// the shared [`AuthManager::catalog_session`]. + async fn init_session( + &self, + client: &HttpClient, + props: &HashMap, + ) -> Result>; + + /// Session used for all subsequent catalog requests, given the properties + /// merged from the user configuration and the server's config response. + /// + /// Returns an [`Arc`]: this session is shared by concurrent requests for + /// the rest of the catalog's lifetime. Implementations may carry state + /// (e.g. a cached token) over from the init session. + async fn catalog_session( + &self, + client: &HttpClient, + props: &HashMap, + ) -> Result>; +} + +/// Authenticates outgoing REST catalog requests. +#[async_trait] +pub trait AuthSession: Debug + Send + Sync { + /// Applies authentication to the request (adds headers, signs, ...). + async fn authenticate(&self, request: &mut HttpRequest) -> Result<()>; +} + +/// [`AuthManager`] that performs no authentication. +#[derive(Debug)] +pub struct NoopAuthManager; + +/// [`AuthSession`] that performs no authentication. +#[derive(Debug)] +pub(crate) struct NoopSession; + +#[async_trait] +impl AuthManager for NoopAuthManager { + async fn init_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + Ok(Box::new(NoopSession)) + } + + async fn catalog_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + Ok(Arc::new(NoopSession)) + } +} + +#[async_trait] +impl AuthSession for NoopSession { + async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use reqwest::Client; + + use super::*; + use crate::RestCatalogConfig; + + fn test_client() -> HttpClient { + HttpClient::new( + &RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .build(), + ) + .unwrap() + } + + #[tokio::test] + async fn test_static_token_session_attaches_token() { + // Token-only config: the token is attached as-is. + let manager = OAuth2Manager::new("http://localhost/unused").with_token("tok-static"); + let session = manager + .init_session(&test_client(), &HashMap::new()) + .await + .unwrap(); + + let mut req = HttpRequest::new( + Client::new() + .get("https://rest.example.com/v1/config") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + assert_eq!( + req.headers().get("authorization").unwrap(), + "Bearer tok-static" + ); + } +} diff --git a/crates/catalog/rest/src/auth/oauth2.rs b/crates/catalog/rest/src/auth/oauth2.rs new file mode 100644 index 0000000000..3747d853d9 --- /dev/null +++ b/crates/catalog/rest/src/auth/oauth2.rs @@ -0,0 +1,338 @@ +// 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. + +use std::collections::HashMap; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use async_trait::async_trait; +use http::StatusCode; +use iceberg::{Credential, Error, ErrorKind, Result}; +use reqwest::header::HeaderMap; +use tokio::sync::Mutex; + +use super::{AuthManager, AuthSession, HttpRequest}; +use crate::catalog::{ + REST_CATALOG_PROP_URI, RestCatalogConfig, credential_from_props, default_token_endpoint, + explicit_headers_from_props, +}; +use crate::client::HttpClient; +use crate::types::{ErrorResponse, TokenResponse}; + +/// The manager's own OAuth2 options, which properties are merged onto. +struct OAuth2Params { + extra_headers: HeaderMap, + token_endpoint: String, + credential: Option<(Option, Credential)>, + extra_oauth_params: HashMap, +} + +/// [`AuthManager`] implementing the OAuth2 client-credentials flow used by +/// Iceberg REST catalogs. +/// +/// A configured `token` is used directly; otherwise `credential` is exchanged +/// for a token at the token endpoint and cached. The cached token is shared +/// across sessions so it survives the config handshake. +pub struct OAuth2Manager { + token: Arc>>, + init_params: OAuth2Params, + /// True when the token endpoint was derived from the catalog URI (not + /// explicitly configured): it is then recomputed from the merged URI in + /// [`Self::catalog_session`], since `/v1/config` may override the URI. + endpoint_is_default: bool, +} + +impl OAuth2Manager { + /// Creates a manager exchanging credentials at `token_endpoint`, with no + /// token or credential configured. Combine with the `with_*` methods: + /// + /// ```rust,ignore + /// let manager = OAuth2Manager::new("https://auth.example.com/v1/oauth/tokens") + /// .with_credential(Some("client-id".into()), "client-secret".into()); + /// ``` + pub fn new(token_endpoint: impl Into) -> Self { + Self { + token: Arc::new(Mutex::new(None)), + init_params: OAuth2Params { + extra_headers: HeaderMap::new(), + token_endpoint: token_endpoint.into(), + credential: None, + // Same default as the configuration path: the catalog scope. + extra_oauth_params: HashMap::from([("scope".to_string(), "catalog".to_string())]), + }, + endpoint_is_default: false, + } + } + + /// Sets a bearer token used directly (takes precedence over `credential`). + pub fn with_token(mut self, token: impl Into) -> Self { + self.token = Arc::new(Mutex::new(Some(Credential::from(token.into())))); + self + } + + /// Sets the client credential exchanged for a token at the token endpoint. + pub fn with_credential(mut self, client_id: Option, client_secret: String) -> Self { + self.init_params.credential = Some((client_id, client_secret.into())); + self + } + + /// Sets extra headers sent with token requests. + pub fn with_extra_headers(mut self, headers: HeaderMap) -> Self { + self.init_params.extra_headers = headers; + self + } + + /// Adds extra OAuth2 form parameters (e.g. `scope`, `audience`), merged + /// onto the defaults: provide a `scope` entry to replace the default + /// `catalog` scope. + pub fn with_extra_oauth_params(mut self, params: HashMap) -> Self { + self.init_params.extra_oauth_params.extend(params); + self + } + + pub(crate) fn from_config(cfg: &RestCatalogConfig) -> Result { + Ok(Self { + token: Arc::new(Mutex::new(cfg.token().map(Credential::from))), + init_params: OAuth2Params { + extra_headers: cfg.extra_headers()?, + token_endpoint: cfg.get_token_endpoint(), + credential: cfg.credential().map(|(id, secret)| (id, secret.into())), + extra_oauth_params: cfg.extra_oauth_params(), + }, + endpoint_is_default: cfg.explicit_oauth2_server_uri().is_none(), + }) + } +} + +impl Debug for OAuth2Manager { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OAuth2Manager") + .field("token_endpoint", &self.init_params.token_endpoint) + .finish_non_exhaustive() + } +} + +#[async_trait] +impl AuthManager for OAuth2Manager { + async fn init_session( + &self, + client: &HttpClient, + props: &HashMap, + ) -> Result> { + Ok(Box::new(self.session_from(client, props).await?)) + } + + async fn catalog_session( + &self, + client: &HttpClient, + props: &HashMap, + ) -> Result> { + Ok(Arc::new(self.session_from(client, props).await?)) + } +} + +impl OAuth2Manager { + /// Builds a session from the manager's options with `props` merged onto + /// them, so an injected manager keeps whatever a property doesn't + /// override. The manager's token cell is shared with every session it + /// builds, so a token cached during the handshake survives it. + async fn session_from( + &self, + client: &HttpClient, + props: &HashMap, + ) -> Result { + // The properties may carry a new token (or restate the user's). + if let Some(token) = props.get("token") { + *self.token.lock().await = Some(Credential::from(token.clone())); + } + + let mut extra_headers = self.init_params.extra_headers.clone(); + extra_headers.extend(explicit_headers_from_props(props)?); + + let mut extra_oauth_params = self.init_params.extra_oauth_params.clone(); + for key in ["scope", "audience", "resource"] { + if let Some(value) = props.get(key) { + extra_oauth_params.insert(key.to_string(), value.to_string()); + } + } + + let token_endpoint = match props.get("oauth2-server-uri") { + Some(uri) if !uri.is_empty() => uri.clone(), + // A default endpoint follows the merged catalog URI (which + // `/v1/config` may have overridden); explicit ones are kept. + _ if self.endpoint_is_default => props + .get(REST_CATALOG_PROP_URI) + .map(|uri| default_token_endpoint(uri)) + .unwrap_or_else(|| self.init_params.token_endpoint.clone()), + _ => self.init_params.token_endpoint.clone(), + }; + + let credential = credential_from_props(props) + .map(|(id, secret)| (id, secret.into())) + .or_else(|| self.init_params.credential.clone()); + + Ok(OAuth2Session { + token: self.token.clone(), + // A configured token takes precedence over the credential: the + // token cell is pre-seeded, and the credential only comes into + // play once that token is gone. + token_source: match credential { + Some(credential) => { + TokenSource::ClientCredentials(Box::new(ClientCredentialsConfig { + client: client.clone(), + credential, + token_endpoint, + extra_headers, + extra_oauth_params, + })) + } + None => TokenSource::StaticToken, + }, + }) + } +} + +/// Attaches `token` as a `Authorization: Bearer ` header, marked +/// sensitive so `Debug`-formatted requests redact it. +fn attach_bearer(req: &mut HttpRequest, token: &Credential) -> Result<()> { + let mut value: http::HeaderValue = + format!("Bearer {}", token.expose()).parse().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + "Invalid token received from catalog server!", + ) + .with_source(e) + })?; + value.set_sensitive(true); + req.headers_mut().insert(http::header::AUTHORIZATION, value); + Ok(()) +} + +/// [`AuthSession`] attaching an OAuth2 bearer token. +/// +/// The token is a configured one (which replaces whatever the cell holds), a +/// token cached by an earlier session (the cell is shared with the owning +/// [`OAuth2Manager`]), or — with [`TokenSource::ClientCredentials`] — one +/// exchanged for the credential on demand. +/// +/// # TODO: Support automatic token refreshing. +struct OAuth2Session { + token: Arc>>, + token_source: TokenSource, +} + +/// How an [`OAuth2Session`] obtains a token once none is cached. +enum TokenSource { + /// Nothing to obtain: the session attaches the configured token, or no + /// authentication at all when there is none. + StaticToken, + /// The credential is exchanged for a token at the token endpoint. + ClientCredentials(Box), +} + +struct ClientCredentialsConfig { + client: HttpClient, + credential: (Option, Credential), + token_endpoint: String, + extra_headers: HeaderMap, + extra_oauth_params: HashMap, +} + +impl Debug for OAuth2Session { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mut out = f.debug_struct("OAuth2Session"); + if let TokenSource::ClientCredentials(config) = &self.token_source { + out.field("token_endpoint", &config.token_endpoint); + } + out.finish_non_exhaustive() + } +} + +impl ClientCredentialsConfig { + async fn exchange_credential_for_token(&self) -> Result { + let (client_id, client_secret) = &self.credential; + + let mut params = HashMap::with_capacity(4); + params.insert("grant_type", "client_credentials"); + if let Some(client_id) = client_id { + params.insert("client_id", client_id); + } + params.insert("client_secret", client_secret.expose()); + params.extend( + self.extra_oauth_params + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())), + ); + + let (status, body) = self + .client + .post_form(&self.token_endpoint, &self.extra_headers, ¶ms) + .await?; + + let auth_res: TokenResponse = if status == StatusCode::OK { + Ok(serde_json::from_slice(&body).map_err(|e| { + Error::new( + ErrorKind::Unexpected, + "Failed to parse response from rest catalog server!", + ) + .with_context("operation", "auth") + .with_context("url", self.token_endpoint.clone()) + .with_context("json", String::from_utf8_lossy(&body)) + .with_source(e) + })?) + } else { + let e: ErrorResponse = serde_json::from_slice(&body).map_err(|e| { + Error::new(ErrorKind::Unexpected, "Received unexpected response") + .with_context("code", status.to_string()) + .with_context("operation", "auth") + .with_context("url", self.token_endpoint.clone()) + .with_context("json", String::from_utf8_lossy(&body)) + .with_source(e) + })?; + Err(Error::from(e)) + }?; + Ok(auth_res.access_token) + } +} + +#[async_trait] +impl AuthSession for OAuth2Session { + /// Uses the cached token when present; otherwise exchanges the credential + /// for one, caches it, then uses it. Without a credential and without a + /// token, no authentication is attached. + async fn authenticate(&self, req: &mut HttpRequest) -> Result<()> { + // The lock is held across the exchange: waiters reuse a successful + // result, and retry themselves after a failure. + let token = { + let mut token = self.token.lock().await; + match (&*token, &self.token_source) { + (Some(token), _) => Some(token.clone()), + (None, TokenSource::StaticToken) => None, + (None, TokenSource::ClientCredentials(config)) => { + let new_token = Credential::from(config.exchange_credential_for_token().await?); + *token = Some(new_token.clone()); + Some(new_token) + } + } + }; + + match token { + Some(token) => attach_bearer(req, &token), + None => Ok(()), + } + } +} diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 8642b32d22..24f19faf8f 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -18,9 +18,10 @@ //! This module contains the iceberg REST catalog implementation. use std::collections::{HashMap, HashSet}; +use std::fmt::{Debug, Formatter}; use std::future::Future; use std::str::FromStr; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use async_trait::async_trait; use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; @@ -34,14 +35,16 @@ use itertools::Itertools; use reqwest::header::{ HeaderMap, HeaderName, HeaderValue, {self}, }; -use reqwest::{Client, Method, StatusCode, Url}; +use reqwest::{Client, Method, Response, StatusCode, Url}; use tokio::sync::OnceCell; use typed_builder::TypedBuilder; +use crate::auth::{AUTH_TYPE_NONE, AUTH_TYPE_OAUTH2, AuthManager, NoopAuthManager, OAuth2Manager}; use crate::client::{ HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error, }; use crate::endpoint::{Endpoint, V1_NAMESPACE_EXISTS, V1_TABLE_EXISTS}; +use crate::request::HttpRequest; use crate::types::{ CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult, @@ -52,8 +55,13 @@ use crate::types::{ pub const REST_CATALOG_PROP_URI: &str = "uri"; /// REST catalog warehouse location pub const REST_CATALOG_PROP_WAREHOUSE: &str = "warehouse"; -/// Disable header redaction in error logs (defaults to false for security) +/// Disable header redaction in error logs and `Debug` output (defaults to +/// false for security) pub const REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str = "disable-header-redaction"; +/// Authentication scheme: `none` or `oauth2`. When unset, `oauth2` is used +/// if a `token`, `credential` or `oauth2-server-uri` is configured, `none` +/// otherwise. +pub const REST_CATALOG_PROP_AUTH_TYPE: &str = "rest.auth.type"; const ICEBERG_REST_SPEC_VERSION: &str = "0.14.1"; const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -63,6 +71,7 @@ const PATH_V1: &str = "v1"; #[derive(Debug)] pub struct RestCatalogBuilder { config: RestCatalogConfig, + auth_manager: Option>, storage_factory: Option>, kms_client_factory: Option>, runtime: Option, @@ -77,7 +86,9 @@ impl Default for RestCatalogBuilder { warehouse: None, props: HashMap::new(), client: None, + default_client: Arc::new(OnceLock::new()), }, + auth_manager: None, storage_factory: None, kms_client_factory: None, runtime: None, @@ -146,6 +157,7 @@ impl CatalogBuilder for RestCatalogBuilder { }; Ok(RestCatalog::new( self.config, + self.auth_manager, self.storage_factory, runtime, kms_client, @@ -161,10 +173,16 @@ impl RestCatalogBuilder { self.config.client = Some(client); self } + + /// Injects a custom auth manager, overriding the `rest.auth.type` configuration. + pub fn with_auth_manager(mut self, auth_manager: Arc) -> Self { + self.auth_manager = Some(auth_manager); + self + } } /// Rest catalog configuration. -#[derive(Clone, Debug, TypedBuilder)] +#[derive(Clone, TypedBuilder)] pub(crate) struct RestCatalogConfig { #[builder(default, setter(strip_option))] name: Option, @@ -179,6 +197,49 @@ pub(crate) struct RestCatalogConfig { #[builder(default)] client: Option, + + /// Lazily-created default HTTP client, shared through clones of this + /// config so OAuth and catalog traffic reuse one connection pool + /// (matching the single-client behavior before the AuthManager refactor). + #[builder(default)] + default_client: Arc>, +} + +/// Property keys whose values are secrets, or may embed them (headers, +/// connection strings, keys like `adls.account-key` or `s3.sse.key`). +fn is_sensitive_prop(key: &str) -> bool { + key.contains("token") + || key.contains("credential") + || key.contains("secret") + || key.contains("password") + || key.contains("key") + || key.contains("connection-string") + || key.starts_with("header.") +} + +/// Redacts secret property values: this config is printed by +/// [`RestCatalog`]'s derived `Debug`. +impl Debug for RestCatalogConfig { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let props: HashMap<&str, &str> = self + .props + .iter() + .map(|(key, value)| { + let value = if is_sensitive_prop(key) { + "[REDACTED]" + } else { + value.as_str() + }; + (key.as_str(), value) + }) + .collect(); + f.debug_struct("RestCatalogConfig") + .field("name", &self.name) + .field("uri", &self.uri) + .field("warehouse", &self.warehouse) + .field("props", &props) + .finish_non_exhaustive() + } } impl RestCatalogConfig { @@ -195,11 +256,13 @@ impl RestCatalogConfig { } pub(crate) fn get_token_endpoint(&self) -> String { - if let Some(oauth2_uri) = self.props.get("oauth2-server-uri") { - oauth2_uri.to_string() - } else { - [&self.uri, PATH_V1, "oauth", "tokens"].join("/") - } + self.explicit_oauth2_server_uri() + .unwrap_or_else(|| default_token_endpoint(&self.uri)) + } + + /// The `oauth2-server-uri` property, only when explicitly configured. + pub(crate) fn explicit_oauth2_server_uri(&self) -> Option { + self.props.get("oauth2-server-uri").cloned() } fn namespaces_endpoint(&self) -> String { @@ -231,9 +294,13 @@ impl RestCatalogConfig { ]) } - /// Get the client from the config. - pub(crate) fn client(&self) -> Option { - self.client.clone() + /// The HTTP client: the configured one, or a lazily-created default that + /// is shared across every user of this config (and its clones), so token + /// and catalog requests keep sharing one connection pool. + pub(crate) fn client(&self) -> Client { + self.client + .clone() + .unwrap_or_else(|| self.default_client.get_or_init(Client::default).clone()) } /// Get the token from the config. @@ -245,89 +312,18 @@ impl RestCatalogConfig { /// Get the credentials from the config. The client can use these credentials to fetch a new /// token. - /// - /// ## Output - /// - /// - `None`: No credential is set. - /// - `Some(None, client_secret)`: No client_id is set, use client_secret directly. - /// - `Some(Some(client_id), client_secret)`: Both client_id and client_secret are set. pub(crate) fn credential(&self) -> Option<(Option, String)> { - let cred = self.props.get("credential")?; - - match cred.split_once(':') { - Some((client_id, client_secret)) => { - Some((Some(client_id.to_string()), client_secret.to_string())) - } - None => Some((None, cred.to_string())), - } + credential_from_props(&self.props) } - /// Get the extra headers from config, which includes: - /// - /// - `content-type` - /// - `x-client-version` - /// - `user-agent` - /// - All headers specified by `header.xxx` in props. + /// Get the extra headers from config, see [`extra_headers_from_props`]. pub(crate) fn extra_headers(&self) -> Result { - let mut headers = HeaderMap::from_iter([ - ( - header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ), - ( - HeaderName::from_static("x-client-version"), - HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), - ), - ( - header::USER_AGENT, - HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), - ), - ]); - - for (key, value) in self - .props - .iter() - .filter_map(|(k, v)| k.strip_prefix("header.").map(|k| (k, v))) - { - headers.insert( - HeaderName::from_str(key).map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!("Invalid header name: {key}"), - ) - .with_source(e) - })?, - HeaderValue::from_str(value).map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!("Invalid header value: {value}"), - ) - .with_source(e) - })?, - ); - } - - Ok(headers) + extra_headers_from_props(&self.props) } /// Get the optional OAuth headers from the config. pub(crate) fn extra_oauth_params(&self) -> HashMap { - let mut params = HashMap::new(); - - if let Some(scope) = self.props.get("scope") { - params.insert("scope".to_string(), scope.to_string()); - } else { - params.insert("scope".to_string(), "catalog".to_string()); - } - - let optional_params = ["audience", "resource"]; - for param_name in optional_params { - if let Some(value) = self.props.get(param_name) { - params.insert(param_name.to_string(), value.to_string()); - } - } - - params + oauth_params_from_props(&self.props) } /// Check if header redaction is disabled in error logs. @@ -349,6 +345,11 @@ impl RestCatalogConfig { let mut props = config.defaults; props.extend(self.props); + // The builder moved the client warehouse off the props; restore it + // between defaults and overrides (default < client < override). + if let Some(warehouse) = &self.warehouse { + props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone()); + } props.extend(config.overrides); self.props = props; @@ -356,8 +357,111 @@ impl RestCatalogConfig { } } +/// Parses the `credential` property. +/// +/// ## Output +/// +/// - `None`: No credential is set. +/// - `Some(None, client_secret)`: No client_id is set, use client_secret directly. +/// - `Some(Some(client_id), client_secret)`: Both client_id and client_secret are set. +pub(crate) fn credential_from_props( + props: &HashMap, +) -> Option<(Option, String)> { + let cred = props.get("credential")?; + + match cred.split_once(':') { + Some((client_id, client_secret)) => { + Some((Some(client_id.to_string()), client_secret.to_string())) + } + None => Some((None, cred.to_string())), + } +} + +/// The extra headers added to each request, which include: +/// +/// - `content-type` +/// - `x-client-version` +/// - `user-agent` +/// - All headers specified by `header.xxx` in props. +pub(crate) fn extra_headers_from_props(props: &HashMap) -> Result { + let mut headers = HeaderMap::from_iter([ + ( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ), + ( + HeaderName::from_static("x-client-version"), + HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), + ), + ( + header::USER_AGENT, + HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), + ), + ]); + + headers.extend(explicit_headers_from_props(props)?); + + Ok(headers) +} + +/// The default OAuth2 token endpoint for a catalog `uri`. +pub(crate) fn default_token_endpoint(uri: &str) -> String { + [uri, PATH_V1, "oauth", "tokens"].join("/") +} + +/// Only the headers explicitly configured via `header.xxx` props (no defaults). +pub(crate) fn explicit_headers_from_props(props: &HashMap) -> Result { + let mut headers = HeaderMap::new(); + for (key, value) in props + .iter() + .filter_map(|(k, v)| k.strip_prefix("header.").map(|k| (k, v))) + { + headers.insert( + HeaderName::from_str(key).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid header name: {key}"), + ) + .with_source(e) + })?, + HeaderValue::from_str(value).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + // The value itself is omitted: it may be a secret. + format!("Invalid value for header: {key}"), + ) + .with_source(e) + })?, + ); + } + + Ok(headers) +} + +/// The optional OAuth parameters added to each authentication request. +pub(crate) fn oauth_params_from_props(props: &HashMap) -> HashMap { + let mut params = HashMap::new(); + + if let Some(scope) = props.get("scope") { + params.insert("scope".to_string(), scope.to_string()); + } else { + params.insert("scope".to_string(), "catalog".to_string()); + } + + let optional_params = ["audience", "resource"]; + for param_name in optional_params { + if let Some(value) = props.get(param_name) { + params.insert(param_name.to_string(), value.to_string()); + } + } + + params +} + #[derive(Debug)] struct RestContext { + /// Carries the session the auth manager derived from the merged + /// configuration, so every request below is authenticated. client: HttpClient, /// Runtime config is fetched from rest server and stored here. /// @@ -367,9 +471,25 @@ struct RestContext { endpoints: HashSet, } +impl RestContext { + /// Testing only: the bearer token the catalog session would attach. + #[cfg(test)] + async fn token(&self) -> Option { + self.client.token().await + } + + /// Sends `request`, authenticated by the client's session. + async fn query_catalog(&self, request: HttpRequest) -> Result { + self.client.query_catalog(request).await + } +} + /// Rest catalog implementation. #[derive(Debug)] pub struct RestCatalog { + /// Injected through [`RestCatalogBuilder::with_auth_manager`]; otherwise + /// one is resolved from `rest.auth.type` when the context is built. + auth_manager: Option>, /// User config is stored as-is and never be changed. /// /// It could be different from the config fetched from the server and used at runtime. @@ -386,11 +506,13 @@ impl RestCatalog { /// Creates a `RestCatalog` from a [`RestCatalogConfig`]. fn new( config: RestCatalogConfig, + auth_manager: Option>, storage_factory: Option>, runtime: Runtime, kms_client: Option>, ) -> Self { Self { + auth_manager, user_config: config, ctx: OnceCell::new(), storage_factory, @@ -411,8 +533,8 @@ impl RestCatalog { request_builder = request_builder.query(&[("purgeRequested", "true")]); } - let request = request_builder.build()?; - let http_response = context.client.query_catalog(request).await?; + let request = HttpRequest::build(request_builder)?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::NO_CONTENT | StatusCode::OK => Ok(()), @@ -428,12 +550,100 @@ impl RestCatalog { } } + /// The properties handed to the [`AuthManager`], with the catalog `uri` + /// and `warehouse` made explicit. + fn auth_props(config: &RestCatalogConfig) -> HashMap { + // `oauth2-server-uri` stays absent unless explicitly configured, so an + // injected manager keeps its own endpoint. The resolved `uri` and + // `warehouse` ARE passed: the builder moved them off the props, and + // the built-in manager recomputes its token endpoint from the URI. + let mut props = config.props.clone(); + props.insert(REST_CATALOG_PROP_URI.to_string(), config.uri.clone()); + if let Some(warehouse) = &config.warehouse { + // A fallback only: after the handshake the merged props hold + // the resolved warehouse, server override included. + props + .entry(REST_CATALOG_PROP_WAREHOUSE.to_string()) + .or_insert_with(|| warehouse.clone()); + } + props + } + + /// The configured auth scheme: explicit `rest.auth.type` (matched + /// case-insensitively) when set; otherwise `oauth2` when a `token`, + /// `credential` or `oauth2-server-uri` is configured (preserving + /// pre-`rest.auth.type` setups), `none` when none is. + fn auth_type(config: &RestCatalogConfig) -> String { + config + .props + .get(REST_CATALOG_PROP_AUTH_TYPE) + // Matched case-insensitively, as the other flag properties are. + .map(|auth_type| auth_type.to_ascii_lowercase()) + .unwrap_or_else(|| { + if config.token().is_some() + || config.credential().is_some() + || config.explicit_oauth2_server_uri().is_some() + { + AUTH_TYPE_OAUTH2.to_string() + } else { + AUTH_TYPE_NONE.to_string() + } + }) + } + + /// Resolves the auth manager: a `with_auth_manager` override wins, + /// otherwise one is built from the `rest.auth.type` configuration. + fn resolve_auth_manager(&self) -> Result> { + if let Some(auth_manager) = &self.auth_manager { + return Ok(auth_manager.clone()); + } + let config = &self.user_config; + let auth_type = Self::auth_type(config); + // Java parity (`AuthManagers`): make the inference visible so users + // configure the type explicitly. + if auth_type == AUTH_TYPE_OAUTH2 && !config.props.contains_key(REST_CATALOG_PROP_AUTH_TYPE) + { + tracing::warn!( + "Inferring {REST_CATALOG_PROP_AUTH_TYPE}={AUTH_TYPE_OAUTH2} from the configured \ + OAuth properties; set it explicitly to avoid this warning" + ); + } + match auth_type.as_str() { + AUTH_TYPE_NONE => Ok(Arc::new(NoopAuthManager)), + AUTH_TYPE_OAUTH2 => Ok(Arc::new(OAuth2Manager::from_config(config)?)), + other => Err(Error::new( + ErrorKind::DataInvalid, + format!( + "unknown '{REST_CATALOG_PROP_AUTH_TYPE}': {other}; use \ + `RestCatalogBuilder::with_auth_manager` to inject a \ + custom auth manager" + ), + )), + } + } + /// Gets the [`RestContext`] from the catalog. async fn context(&self) -> Result<&RestContext> { self.ctx .get_or_try_init(|| async { let client = HttpClient::new(&self.user_config)?; - let catalog_config = RestCatalog::load_config(&client, &self.user_config).await?; + let auth_manager = self.resolve_auth_manager()?; + // The init session lives only for the config handshake, so a + // manager whose session guards a one-shot resource can release + // it before deriving the catalog session. + let catalog_config = { + let init_session = auth_manager + .init_session( + &client.without_session(), + &Self::auth_props(&self.user_config), + ) + .await?; + RestCatalog::load_config( + &client.with_session(Arc::from(init_session)), + &self.user_config, + ) + .await? + }; // Use the advertised endpoints as-is, falling back to // `DEFAULT_ENDPOINTS` when absent or empty. let endpoints = match &catalog_config.endpoints { @@ -444,10 +654,15 @@ impl RestCatalog { }; let config = self.user_config.clone().merge_with_config(catalog_config); let client = client.update_with(&config)?; + // The manager is handed an unauthenticated client: its own + // requests must not be signed by the session it is deriving. + let session = auth_manager + .catalog_session(&client.without_session(), &Self::auth_props(&config)) + .await?; Ok(RestContext { config, - client, + client: client.with_session(session), endpoints, }) }) @@ -464,8 +679,8 @@ impl RestCatalog { /// Issue a `HEAD` request to `url` and interpret it as an existence check: /// `2xx` means it exists, `404` means it doesn't. async fn check_exists_via_head(&self, context: &RestContext, url: String) -> Result { - let request = context.client.request(Method::HEAD, url).build()?; - let http_response = context.client.query_catalog(request).await?; + let request = HttpRequest::build(context.client.request(Method::HEAD, url))?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::NO_CONTENT | StatusCode::OK => Ok(true), @@ -491,7 +706,7 @@ impl RestCatalog { request_builder = request_builder.query(&[("warehouse", warehouse_location)]); } - let request = request_builder.build()?; + let request = HttpRequest::build(request_builder)?; let http_response = client.query_catalog(request).await?; @@ -545,22 +760,6 @@ impl RestCatalog { Ok(file_io) } - - /// Invalidate the current token without generating a new one. On the next request, the client - /// will attempt to generate a new token. - pub async fn invalidate_token(&self) -> Result<()> { - self.context().await?.client.invalidate_token().await - } - - /// Invalidate the current token and set a new one. Generates a new token before invalidating - /// the current token, meaning the old token will be used until this function acquires the lock - /// and overwrites the token. - /// - /// If credential is invalid, or the request fails, this method will return an error and leave - /// the current token unchanged. - pub async fn regenerate_token(&self) -> Result<()> { - self.context().await?.client.regenerate_token().await - } } /// All requests and expected responses are derived from the REST catalog API spec: @@ -588,7 +787,7 @@ impl Catalog for RestCatalog { request = request.query(&[("pageToken", token)]); } - let http_response = context.client.query_catalog(request.build()?).await?; + let http_response = context.query_catalog(HttpRequest::build(request)?).await?; match http_response.status() { StatusCode::OK => { @@ -629,16 +828,17 @@ impl Catalog for RestCatalog { ) -> Result { let context = self.context().await?; - let request = context - .client - .request(Method::POST, context.config.namespaces_endpoint()) - .json(&CreateNamespaceRequest { - namespace: namespace.clone(), - properties, - }) - .build()?; + let request = HttpRequest::build( + context + .client + .request(Method::POST, context.config.namespaces_endpoint()) + .json(&CreateNamespaceRequest { + namespace: namespace.clone(), + properties, + }), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::OK => { @@ -661,12 +861,13 @@ impl Catalog for RestCatalog { async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result { let context = self.context().await?; - let request = context - .client - .request(Method::GET, context.config.namespace_endpoint(namespace)) - .build()?; + let request = HttpRequest::build( + context + .client + .request(Method::GET, context.config.namespace_endpoint(namespace)), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::OK => { @@ -718,12 +919,13 @@ impl Catalog for RestCatalog { async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> { let context = self.context().await?; - let request = context - .client - .request(Method::DELETE, context.config.namespace_endpoint(namespace)) - .build()?; + let request = HttpRequest::build( + context + .client + .request(Method::DELETE, context.config.namespace_endpoint(namespace)), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::NO_CONTENT | StatusCode::OK => Ok(()), @@ -752,7 +954,7 @@ impl Catalog for RestCatalog { request = request.query(&[("pageToken", token)]); } - let http_response = context.client.query_catalog(request.build()?).await?; + let http_response = context.query_catalog(HttpRequest::build(request)?).await?; match http_response.status() { StatusCode::OK => { @@ -800,21 +1002,22 @@ impl Catalog for RestCatalog { let table_ident = TableIdent::new(namespace.clone(), creation.name.clone()); - let request = context - .client - .request(Method::POST, context.config.tables_endpoint(namespace)) - .json(&CreateTableRequest { - name: creation.name, - location: creation.location, - schema: creation.schema, - partition_spec: creation.partition_spec, - write_order: creation.sort_order, - stage_create: Some(false), - properties: creation.properties, - }) - .build()?; - - let http_response = context.client.query_catalog(request).await?; + let request = HttpRequest::build( + context + .client + .request(Method::POST, context.config.tables_endpoint(namespace)) + .json(&CreateTableRequest { + name: creation.name, + location: creation.location, + schema: creation.schema, + partition_spec: creation.partition_spec, + write_order: creation.sort_order, + stage_create: Some(false), + properties: creation.properties, + }), + )?; + + let http_response = context.query_catalog(request).await?; let response = match http_response.status() { StatusCode::OK => { @@ -880,12 +1083,13 @@ impl Catalog for RestCatalog { async fn load_table(&self, table_ident: &TableIdent) -> Result { let context = self.context().await?; - let request = context - .client - .request(Method::GET, context.config.table_endpoint(table_ident)) - .build()?; + let request = HttpRequest::build( + context + .client + .request(Method::GET, context.config.table_endpoint(table_ident)), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; let response = match http_response.status() { StatusCode::OK | StatusCode::NOT_MODIFIED => { @@ -965,16 +1169,17 @@ impl Catalog for RestCatalog { async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { let context = self.context().await?; - let request = context - .client - .request(Method::POST, context.config.rename_table_endpoint()) - .json(&RenameTableRequest { - source: src.clone(), - destination: dest.clone(), - }) - .build()?; + let request = HttpRequest::build( + context + .client + .request(Method::POST, context.config.rename_table_endpoint()) + .json(&RenameTableRequest { + source: src.clone(), + destination: dest.clone(), + }), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; match http_response.status() { StatusCode::NO_CONTENT | StatusCode::OK => Ok(()), @@ -1001,22 +1206,23 @@ impl Catalog for RestCatalog { ) -> Result
{ let context = self.context().await?; - let request = context - .client - .request( - Method::POST, - context - .config - .register_table_endpoint(table_ident.namespace()), - ) - .json(&RegisterTableRequest { - name: table_ident.name.clone(), - metadata_location: metadata_location.clone(), - overwrite: Some(false), - }) - .build()?; + let request = HttpRequest::build( + context + .client + .request( + Method::POST, + context + .config + .register_table_endpoint(table_ident.namespace()), + ) + .json(&RegisterTableRequest { + name: table_ident.name.clone(), + metadata_location: metadata_location.clone(), + overwrite: Some(false), + }), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; let response: LoadTableResult = match http_response.status() { StatusCode::OK => { @@ -1065,20 +1271,21 @@ impl Catalog for RestCatalog { async fn update_table(&self, mut commit: TableCommit) -> Result
{ let context = self.context().await?; - let request = context - .client - .request( - Method::POST, - context.config.table_endpoint(commit.identifier()), - ) - .json(&CommitTableRequest { - identifier: Some(commit.identifier().clone()), - requirements: commit.take_requirements(), - updates: commit.take_updates(), - }) - .build()?; + let request = HttpRequest::build( + context + .client + .request( + Method::POST, + context.config.table_endpoint(commit.identifier()), + ) + .json(&CommitTableRequest { + identifier: Some(commit.identifier().clone()), + requirements: commit.take_requirements(), + updates: commit.take_updates(), + }), + )?; - let http_response = context.client.query_catalog(request).await?; + let http_response = context.query_catalog(request).await?; let response: CommitTableResponse = match http_response.status() { StatusCode::OK => deserialize_catalog_response(http_response).await?, @@ -1159,6 +1366,28 @@ mod tests { use uuid::uuid; use super::*; + use crate::auth::AuthSession; + use crate::request::HttpRequest; + + fn test_catalog(config: RestCatalogConfig) -> RestCatalog { + test_catalog_with(config, None) + } + + fn test_catalog_with( + config: RestCatalogConfig, + auth_manager: Option>, + ) -> RestCatalog { + RestCatalog::new(config, auth_manager, None, Runtime::current(), None) + } + + fn test_client() -> HttpClient { + HttpClient::new( + &RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .build(), + ) + .unwrap() + } #[tokio::test] async fn test_update_config() { @@ -1180,6 +1409,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1257,6 +1487,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1288,6 +1519,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1323,6 +1555,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1377,12 +1610,13 @@ mod tests { .uri(server.url()) .props(props) .build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; + let token = catalog.context().await.unwrap().token().await; oauth_mock.assert_async().await; config_mock.assert_async().await; assert_eq!(token, Some("ey000000000000".to_string())); @@ -1426,12 +1660,13 @@ mod tests { .uri(server.url()) .props(props) .build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; + let token = catalog.context().await.unwrap().token().await; oauth_mock.assert_async().await; config_mock.assert_async().await; @@ -1439,330 +1674,942 @@ mod tests { } #[tokio::test] - async fn test_invalidate_token() { + async fn test_http_headers() { + let server = Server::new_async().await; + let mut props = HashMap::new(); + props.insert("credential".to_string(), "client1:secret1".to_string()); + + let config = RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(); + let headers: HeaderMap = config.extra_headers().unwrap(); + + let expected_headers = HeaderMap::from_iter([ + ( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ), + ( + HeaderName::from_static("x-client-version"), + HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), + ), + ( + header::USER_AGENT, + HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), + ), + ]); + assert_eq!(headers, expected_headers); + } + + #[tokio::test] + async fn test_http_headers_with_custom_headers() { + let server = Server::new_async().await; + let mut props = HashMap::new(); + props.insert("credential".to_string(), "client1:secret1".to_string()); + props.insert( + "header.content-type".to_string(), + "application/yaml".to_string(), + ); + props.insert( + "header.customized-header".to_string(), + "some/value".to_string(), + ); + + let config = RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(); + let headers: HeaderMap = config.extra_headers().unwrap(); + + let expected_headers = HeaderMap::from_iter([ + ( + header::CONTENT_TYPE, + HeaderValue::from_static("application/yaml"), + ), + ( + HeaderName::from_static("x-client-version"), + HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), + ), + ( + header::USER_AGENT, + HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), + ), + ( + HeaderName::from_static("customized-header"), + HeaderValue::from_static("some/value"), + ), + ]); + assert_eq!(headers, expected_headers); + } + + #[tokio::test] + async fn test_oauth_with_oauth2_server_uri() { let mut server = Server::new_async().await; - let oauth_mock = create_oauth_mock(&mut server).await; let config_mock = create_config_mock(&mut server).await; + let mut auth_server = Server::new_async().await; + let auth_server_path = "/some/path"; + let oauth_mock = + create_oauth_mock_with_path(&mut auth_server, auth_server_path, "ey000000000000", 200) + .await; + let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); + props.insert( + "oauth2-server-uri".to_string(), + format!("{}{}", auth_server.url(), auth_server_path).to_string(), + ); let catalog = RestCatalog::new( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; + let token = catalog.context().await.unwrap().token().await; + oauth_mock.assert_async().await; config_mock.assert_async().await; assert_eq!(token, Some("ey000000000000".to_string())); + } - let oauth_mock = - create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 200) - .await; - catalog.invalidate_token().await.unwrap(); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; - assert_eq!(token, Some("ey000000000001".to_string())); + #[tokio::test] + async fn test_config_override() { + let mut server = Server::new_async().await; + let mut redirect_server = Server::new_async().await; + let new_uri = redirect_server.url(); + + let config_mock = server + .mock("GET", "/v1/config") + .with_status(200) + .with_body( + json!( + { + "overrides": { + "uri": new_uri, + "warehouse": "s3://iceberg-catalog", + "prefix": "ice/warehouses/my" + }, + "defaults": {}, + } + ) + .to_string(), + ) + .create_async() + .await; + + let list_ns_mock = redirect_server + .mock("GET", "/v1/ice/warehouses/my/namespaces") + .with_body( + r#"{ + "namespaces": [] + }"#, + ) + .create_async() + .await; + + let catalog = RestCatalog::new( + RestCatalogConfig::builder().uri(server.url()).build(), + None, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + let _namespaces = catalog.list_namespaces(None).await.unwrap(); + + config_mock.assert_async().await; + list_ns_mock.assert_async().await; } #[tokio::test] - async fn test_invalidate_token_failing_request() { + async fn test_list_namespace() { let mut server = Server::new_async().await; - let oauth_mock = create_oauth_mock(&mut server).await; + let config_mock = create_config_mock(&mut server).await; - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); + let list_ns_mock = server + .mock("GET", "/v1/namespaces") + .with_body( + r#"{ + "namespaces": [ + ["ns1", "ns11"], + ["ns2"] + ] + }"#, + ) + .create_async() + .await; let catalog = RestCatalog::new( - RestCatalogConfig::builder() - .uri(server.url()) - .props(props) - .build(), + RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; + let namespaces = catalog.list_namespaces(None).await.unwrap(); + + let expected_ns = vec![ + NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), + NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(), + ]; + + assert_eq!(expected_ns, namespaces); + config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; + } - let oauth_mock = - create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 500) - .await; - catalog.invalidate_token().await.unwrap(); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; - assert_eq!(token, None); + #[tokio::test] + async fn test_auth_type_none_disables_auth() { + // An explicit `rest.auth.type=none` wins over a configured token. + let props = HashMap::from([ + (REST_CATALOG_PROP_AUTH_TYPE.to_string(), "none".to_string()), + ("token".to_string(), "some-oauth-token".to_string()), + ]); + let config = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(props) + .build(); + + let session = test_catalog(config) + .resolve_auth_manager() + .unwrap() + .init_session(&test_client(), &HashMap::new()) + .await + .unwrap(); + let mut req = HttpRequest::new( + Client::new() + .get("https://rest.example.com/v1/config") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + assert!(req.headers().get("authorization").is_none()); } #[tokio::test] - async fn test_regenerate_token() { + async fn test_header_prop_overrides_token_on_the_wire() { + // Pre-AuthManager behavior, preserved: extra headers are applied after + // authentication, so a user-configured `header.authorization` wins + // over a configured token. let mut server = Server::new_async().await; - let oauth_mock = create_oauth_mock(&mut server).await; let config_mock = create_config_mock(&mut server).await; + let list_ns_mock = server + .mock("GET", "/v1/namespaces") + .match_header("authorization", "Basic xyz") + .with_body(r#"{"namespaces": []}"#) + .create_async() + .await; - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); - + let props = HashMap::from([ + ("token".to_string(), "some-oauth-token".to_string()), + ("header.authorization".to_string(), "Basic xyz".to_string()), + ]); let catalog = RestCatalog::new( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; + catalog.list_namespaces(None).await.unwrap(); config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); - - let oauth_mock = - create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 200) - .await; - catalog.regenerate_token().await.unwrap(); - oauth_mock.assert_async().await; - let token = catalog.context().await.unwrap().client.token().await; - assert_eq!(token, Some("ey000000000001".to_string())); + list_ns_mock.assert_async().await; } #[tokio::test] - async fn test_regenerate_token_failing_request() { - let mut server = Server::new_async().await; - let oauth_mock = create_oauth_mock(&mut server).await; - let config_mock = create_config_mock(&mut server).await; - - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); + async fn test_builtin_oauth_endpoint_follows_uri_override() { + // When `/v1/config` overrides `uri` (and no explicit `oauth2-server-uri` + // is set), the built-in manager's default token endpoint must follow + // the merged URI. + let mut bootstrap = Server::new_async().await; + let overridden = Server::new_async().await; + + let config_mock = bootstrap + .mock("GET", "/v1/config") + .with_status(200) + .with_body(format!( + r#"{{"overrides": {{"uri": "{}"}}, "defaults": {{}}}}"#, + overridden.url() + )) + .create_async() + .await; + // Handshake exchange still uses the bootstrap-derived default. + let bootstrap_oauth_mock = + create_oauth_mock_with_path(&mut bootstrap, "/v1/oauth/tokens", "tok-boot", 200).await; + let props = HashMap::from([("credential".to_string(), "client1:secret1".to_string())]); let catalog = RestCatalog::new( RestCatalogConfig::builder() - .uri(server.url()) + .uri(bootstrap.url()) .props(props) .build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; + let context = catalog.context().await.unwrap(); config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + bootstrap_oauth_mock.assert_async().await; + // The catalog session's endpoint follows the overridden URI (visible + // via the session's Debug, which prints its token endpoint). + let session_debug = format!("{:?}", context.client.session()); + assert!(session_debug.contains(&format!("{}/v1/oauth/tokens", overridden.url()))); + } + #[tokio::test] + async fn test_concurrent_authenticate_single_token_exchange() { + // Concurrent requests that all find no cached token must trigger ONE + // credential exchange (the lock is held across it), not one each. + let mut server = Server::new_async().await; + // create_oauth_mock_with_path expects exactly 1 hit. let oauth_mock = - create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 500) - .await; - let invalidate_result = catalog.regenerate_token().await; - assert!(invalidate_result.is_err()); + create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "tok-once", 200).await; + + let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url())) + .with_credential(Some("client1".to_string()), "secret1".to_string()); + let session: Arc = Arc::from( + manager + .init_session(&test_client(), &HashMap::new()) + .await + .unwrap(), + ); + + let client = Client::new(); + let attempts = (0..8).map(|_| { + let session = session.clone(); + let client = client.clone(); + async move { + let mut req = HttpRequest::new( + client + .get("https://rest.example.com/v1/config") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + req.headers() + .get("authorization") + .unwrap() + .to_str() + .unwrap() + .to_string() + } + }); + let bearers = futures::future::join_all(attempts).await; + oauth_mock.assert_async().await; - let token = catalog.context().await.unwrap().client.token().await; + assert!(bearers.iter().all(|b| b == "Bearer tok-once")); + } - // original token is left intact - assert_eq!(token, Some("ey000000000000".to_string())); + #[tokio::test] + async fn test_seeded_token_takes_precedence_over_credential() { + // token + credential: the seeded token is attached without any + // credential exchange. + let mut server = Server::new_async().await; + let oauth_mock = server + .mock("POST", "/v1/oauth/tokens") + .expect(0) + .create_async() + .await; + + let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url())) + .with_token("tok-seed") + .with_credential(Some("client1".to_string()), "secret1".to_string()); + let session = manager + .init_session(&test_client(), &HashMap::new()) + .await + .unwrap(); + + let mut req = HttpRequest::new( + Client::new() + .get("https://rest.example.com/v1/config") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + assert_eq!( + req.headers().get("authorization").unwrap(), + "Bearer tok-seed" + ); + + oauth_mock.assert_async().await; } #[tokio::test] - async fn test_http_headers() { - let server = Server::new_async().await; - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); + async fn test_injected_oauth_manager_keeps_endpoint_and_options() { + // An injected OAuth2Manager must keep its own token endpoint, extra + // headers and OAuth params across the config handshake: only explicit + // properties may override them, never synthesized defaults. + let mut server = Server::new_async().await; + // The server vends the credential, so the exchange runs through the + // post-handshake catalog session (exercising its property merging). + let config_mock = server + .mock("GET", "/v1/config") + .with_status(200) + .with_body(r#"{"defaults": {"credential": "client1:secret1"}, "overrides": {}}"#) + .create_async() + .await; - let config = RestCatalogConfig::builder() - .uri(server.url()) - .props(props) - .build(); - let headers: HeaderMap = config.extra_headers().unwrap(); + // The catalog-host default endpoint must never see the credential. + let default_endpoint_mock = server + .mock("POST", "/v1/oauth/tokens") + .expect(0) + .create_async() + .await; + // The exchange hits the injected endpoint, carrying the injected + // header and OAuth param. + let custom_endpoint_mock = server + .mock("POST", "/custom/oauth/tokens") + .match_header("x-tenant", "t1") + // The default catalog scope must survive alongside the injected + // audience (with_extra_oauth_params merges onto the defaults). + .match_body(mockito::Matcher::AllOf(vec![ + mockito::Matcher::Regex("scope=catalog".to_string()), + mockito::Matcher::Regex("audience=aud-1".to_string()), + ])) + .with_status(200) + .with_body( + r#"{ + "access_token": "ey000000000000", + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "expires_in": 86400 + }"#, + ) + .create_async() + .await; + let ns_mock = server + .mock("GET", "/v1/namespaces") + .match_header("authorization", "Bearer ey000000000000") + .with_body(r#"{"namespaces": []}"#) + .create_async() + .await; - let expected_headers = HeaderMap::from_iter([ - ( - header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ), - ( - HeaderName::from_static("x-client-version"), - HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), - ), - ( - header::USER_AGENT, - HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), - ), - ]); - assert_eq!(headers, expected_headers); + let manager = OAuth2Manager::new(format!("{}/custom/oauth/tokens", server.url())) + .with_extra_headers(HeaderMap::from_iter([( + HeaderName::from_static("x-tenant"), + HeaderValue::from_static("t1"), + )])) + .with_extra_oauth_params(HashMap::from([( + "audience".to_string(), + "aud-1".to_string(), + )])); + let catalog = RestCatalog::new( + RestCatalogConfig::builder().uri(server.url()).build(), + Some(Arc::new(manager)), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + catalog.list_namespaces(None).await.unwrap(); + + config_mock.assert_async().await; + custom_endpoint_mock.assert_async().await; + default_endpoint_mock.assert_async().await; + ns_mock.assert_async().await; + } + + #[tokio::test] + async fn test_props_token_takes_precedence_over_props_credential() { + // Both arriving through the properties rather than the builder. + let mut server = Server::new_async().await; + let oauth_mock = server + .mock("POST", "/v1/oauth/tokens") + .expect(0) + .create_async() + .await; + + let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url())); + let session = manager + .init_session( + &test_client(), + &HashMap::from([ + ("token".to_string(), "tok-props".to_string()), + ("credential".to_string(), "client1:secret1".to_string()), + ]), + ) + .await + .unwrap(); + + let mut req = HttpRequest::new( + Client::new() + .get("https://rest.example.com/v1/config") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + assert_eq!( + req.headers().get("authorization").unwrap(), + "Bearer tok-props" + ); + oauth_mock.assert_async().await; } #[tokio::test] - async fn test_http_headers_with_custom_headers() { - let server = Server::new_async().await; - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); - props.insert( - "header.content-type".to_string(), - "application/yaml".to_string(), - ); - props.insert( - "header.customized-header".to_string(), - "some/value".to_string(), - ); + async fn test_manager_exchanges_over_the_catalog_client() { + // The manager exchanges the credential over the client it is handed, + // sharing the catalog's connection pool. Its own extra headers go on + // the token request; the catalog's do not. + let mut server = Server::new_async().await; + let token_mock = server + .mock("POST", "/v1/oauth/tokens") + .match_header("x-from", "manager") + .match_header("x-catalog-only", mockito::Matcher::Missing) + .with_status(200) + .with_body(r#"{"access_token": "tok", "token_type": "Bearer"}"#) + .create_async() + .await; - let config = RestCatalogConfig::builder() - .uri(server.url()) - .props(props) - .build(); - let headers: HeaderMap = config.extra_headers().unwrap(); + let catalog_client = HttpClient::new( + &RestCatalogConfig::builder() + .uri(server.url()) + .props(HashMap::from([( + "header.x-catalog-only".to_string(), + "not-on-token-requests".to_string(), + )])) + .build(), + ) + .unwrap(); + + let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url())) + .with_credential(Some("client1".to_string()), "secret1".to_string()) + .with_extra_headers(HeaderMap::from_iter([( + HeaderName::from_static("x-from"), + HeaderValue::from_static("manager"), + )])); + let session = manager + .init_session(&catalog_client, &HashMap::new()) + .await + .unwrap(); - let expected_headers = HeaderMap::from_iter([ - ( - header::CONTENT_TYPE, - HeaderValue::from_static("application/yaml"), - ), - ( - HeaderName::from_static("x-client-version"), - HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), - ), - ( - header::USER_AGENT, - HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(), - ), - ( - HeaderName::from_static("customized-header"), - HeaderValue::from_static("some/value"), - ), - ]); - assert_eq!(headers, expected_headers); + let mut req = HttpRequest::new( + Client::new() + .get("https://rest.example.com/v1/namespaces") + .build() + .unwrap(), + ); + session.authenticate(&mut req).await.unwrap(); + token_mock.assert_async().await; } #[tokio::test] - async fn test_oauth_with_oauth2_server_uri() { + async fn test_handshake_is_authenticated_by_the_init_session() { + // `/v1/config` goes out with the init session's authentication, not + // unauthenticated and not with a later one. let mut server = Server::new_async().await; - let config_mock = create_config_mock(&mut server).await; - - let mut auth_server = Server::new_async().await; - let auth_server_path = "/some/path"; - let oauth_mock = - create_oauth_mock_with_path(&mut auth_server, auth_server_path, "ey000000000000", 200) - .await; + let config_mock = server + .mock("GET", "/v1/config") + .match_header("authorization", "Bearer tok-init") + .with_status(200) + .with_body(r#"{"defaults": {}, "overrides": {}}"#) + .create_async() + .await; - let mut props = HashMap::new(); - props.insert("credential".to_string(), "client1:secret1".to_string()); - props.insert( - "oauth2-server-uri".to_string(), - format!("{}{}", auth_server.url(), auth_server_path).to_string(), + let catalog = RestCatalog::new( + RestCatalogConfig::builder() + .uri(server.url()) + .props(HashMap::from([( + "token".to_string(), + "tok-init".to_string(), + )])) + .build(), + None, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, ); + catalog.context().await.unwrap(); + config_mock.assert_async().await; + } + + #[tokio::test] + async fn test_init_session_receives_user_props() { + use tokio::sync::Mutex as AsyncMutex; + + // A custom manager initializes from the user configuration: the + // props carry the catalog `uri` and the user's credentials. + #[derive(Debug)] + struct PlainSession; + #[async_trait] + impl AuthSession for PlainSession { + async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> { + Ok(()) + } + } + + #[derive(Debug)] + struct CapturingManager(Arc>>>); + #[async_trait] + impl AuthManager for CapturingManager { + async fn init_session( + &self, + _client: &HttpClient, + props: &HashMap, + ) -> Result> { + *self.0.lock().await = Some(props.clone()); + Ok(Box::new(PlainSession)) + } + async fn catalog_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + Ok(Arc::new(PlainSession)) + } + } + + let mut server = Server::new_async().await; + let config_mock = create_config_mock(&mut server).await; + let captured = Arc::new(AsyncMutex::new(None)); let catalog = RestCatalog::new( RestCatalogConfig::builder() .uri(server.url()) - .props(props) + .props(HashMap::from([( + "token".to_string(), + "tok-user".to_string(), + )])) .build(), + Some(Arc::new(CapturingManager(captured.clone()))), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let token = catalog.context().await.unwrap().client.token().await; - - oauth_mock.assert_async().await; + catalog.context().await.unwrap(); config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + let props = captured.lock().await.clone().unwrap(); + assert_eq!(props.get("token").map(String::as_str), Some("tok-user")); + assert_eq!( + props.get(REST_CATALOG_PROP_URI).map(String::as_str), + Some(server.url().as_str()) + ); } #[tokio::test] - async fn test_config_override() { - let mut server = Server::new_async().await; - let mut redirect_server = Server::new_async().await; - let new_uri = redirect_server.url(); + async fn test_catalog_session_receives_resolved_warehouse() { + use tokio::sync::Mutex as AsyncMutex; + + // A custom manager must receive the resolved warehouse in the props + // handed to `catalog_session`, with the standard precedence: + // server default < client-side warehouse < server override. + #[derive(Debug)] + struct PlainSession; + #[async_trait] + impl AuthSession for PlainSession { + async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> { + Ok(()) + } + } + + #[derive(Debug)] + struct CapturingManager(Arc>>>); + #[async_trait] + impl AuthManager for CapturingManager { + async fn init_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + Ok(Box::new(PlainSession)) + } + async fn catalog_session( + &self, + _client: &HttpClient, + props: &HashMap, + ) -> Result> { + *self.0.lock().await = Some(props.clone()); + Ok(Arc::new(PlainSession)) + } + } + // Client warehouse wins over a server default. + let mut server = Server::new_async().await; let config_mock = server .mock("GET", "/v1/config") + .match_query(mockito::Matcher::UrlEncoded( + "warehouse".to_string(), + "client-wh".to_string(), + )) .with_status(200) - .with_body( - json!( - { - "overrides": { - "uri": new_uri, - "warehouse": "s3://iceberg-catalog", - "prefix": "ice/warehouses/my" - }, - "defaults": {}, - } - ) - .to_string(), - ) + .with_body(r#"{"defaults": {"warehouse": "default-wh"}, "overrides": {}}"#) .create_async() .await; + let captured = Arc::new(AsyncMutex::new(None)); + let catalog = RestCatalog::new( + RestCatalogConfig::builder() + .uri(server.url()) + .warehouse("client-wh".to_string()) + .build(), + Some(Arc::new(CapturingManager(captured.clone()))), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + catalog.context().await.unwrap(); + config_mock.assert_async().await; + let props = captured.lock().await.clone().unwrap(); + assert_eq!( + props.get("warehouse").map(String::as_str), + Some("client-wh") + ); - let list_ns_mock = redirect_server - .mock("GET", "/v1/ice/warehouses/my/namespaces") - .with_body( - r#"{ - "namespaces": [] - }"#, - ) + // A server override wins over the client warehouse. + let mut server = Server::new_async().await; + let config_mock = server + .mock("GET", "/v1/config") + .match_query(mockito::Matcher::UrlEncoded( + "warehouse".to_string(), + "client-wh".to_string(), + )) + .with_status(200) + .with_body(r#"{"defaults": {}, "overrides": {"warehouse": "override-wh"}}"#) .create_async() .await; - + let captured = Arc::new(AsyncMutex::new(None)); let catalog = RestCatalog::new( - RestCatalogConfig::builder().uri(server.url()).build(), + RestCatalogConfig::builder() + .uri(server.url()) + .warehouse("client-wh".to_string()) + .build(), + Some(Arc::new(CapturingManager(captured.clone()))), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - - let _namespaces = catalog.list_namespaces(None).await.unwrap(); - + catalog.context().await.unwrap(); config_mock.assert_async().await; - list_ns_mock.assert_async().await; + let props = captured.lock().await.clone().unwrap(); + assert_eq!( + props.get("warehouse").map(String::as_str), + Some("override-wh") + ); } #[tokio::test] - async fn test_list_namespace() { - let mut server = Server::new_async().await; + async fn test_init_session_dropped_before_catalog_session() { + use std::sync::atomic::{AtomicBool, Ordering}; + + // A manager whose init session guards a one-shot resource (released on + // drop) must see it released before `catalog_session` is invoked. + #[derive(Debug)] + struct GuardSession(Arc); + impl Drop for GuardSession { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + #[async_trait] + impl AuthSession for GuardSession { + async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> { + Ok(()) + } + } - let config_mock = create_config_mock(&mut server).await; + #[derive(Debug)] + struct PlainSession; + #[async_trait] + impl AuthSession for PlainSession { + async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> { + Ok(()) + } + } - let list_ns_mock = server - .mock("GET", "/v1/namespaces") - .with_body( - r#"{ - "namespaces": [ - ["ns1", "ns11"], - ["ns2"] - ] - }"#, - ) - .create_async() - .await; + #[derive(Debug)] + struct GuardManager(Arc); + #[async_trait] + impl AuthManager for GuardManager { + async fn init_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + Ok(Box::new(GuardSession(self.0.clone()))) + } + async fn catalog_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + if !self.0.load(Ordering::SeqCst) { + return Err(Error::new( + ErrorKind::Unexpected, + "init session must be dropped before catalog_session", + )); + } + Ok(Arc::new(PlainSession)) + } + } + let mut server = Server::new_async().await; + let config_mock = create_config_mock(&mut server).await; + + let dropped = Arc::new(AtomicBool::new(false)); let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + Some(Arc::new(GuardManager(dropped.clone()))), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, ); - let namespaces = catalog.list_namespaces(None).await.unwrap(); + catalog.context().await.unwrap(); + config_mock.assert_async().await; + assert!(dropped.load(Ordering::SeqCst)); + } - let expected_ns = vec![ - NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), - NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(), - ]; + #[test] + fn test_config_debug_redacts_secrets() { + let config = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([ + ("token".to_string(), "tok-secret".to_string()), + ("credential".to_string(), "id:cred-secret".to_string()), + ("header.authorization".to_string(), "Basic xyz".to_string()), + ("adls.account-key".to_string(), "adls-secret".to_string()), + ("s3.sse.key".to_string(), "sse-secret".to_string()), + ( + "adls.connection-string".to_string(), + "cs-secret".to_string(), + ), + ("warehouse".to_string(), "wh1".to_string()), + ])) + .build(); - assert_eq!(expected_ns, namespaces); + let out = format!("{config:?}"); + assert!(!out.contains("tok-secret")); + assert!(!out.contains("cred-secret")); + assert!(!out.contains("Basic xyz")); + assert!(!out.contains("adls-secret")); + assert!(!out.contains("sse-secret")); + assert!(!out.contains("cs-secret")); + assert!(out.contains("[REDACTED]")); + assert!(out.contains("wh1")); + } - config_mock.assert_async().await; - list_ns_mock.assert_async().await; + #[tokio::test] + async fn test_auth_type_defaults() { + // Unset `rest.auth.type`: `oauth2` when any OAuth material is + // configured (existing setups keep working), `none` otherwise. + let bare = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .build(); + assert!( + format!("{:?}", test_catalog(bare).resolve_auth_manager().unwrap()) + .contains("NoopAuthManager") + ); + + let with_token = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([("token".to_string(), "tok".to_string())])) + .build(); + assert!( + format!( + "{:?}", + test_catalog(with_token).resolve_auth_manager().unwrap() + ) + .contains("OAuth2Manager") + ); + + // An explicit type is matched case-insensitively. + let mixed_case = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + "OAuth2".to_string(), + )])) + .build(); + assert!( + format!( + "{:?}", + test_catalog(mixed_case).resolve_auth_manager().unwrap() + ) + .contains("OAuth2Manager") + ); + + // An explicit OAuth endpoint is oauth2 intent too: the manager can + // still pick up a server-supplied token from `/v1/config`. + let with_endpoint = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([( + "oauth2-server-uri".to_string(), + "http://auth.example.com/tokens".to_string(), + )])) + .build(); + assert!( + format!( + "{:?}", + test_catalog(with_endpoint).resolve_auth_manager().unwrap() + ) + .contains("OAuth2Manager") + ); + } + + #[tokio::test] + async fn test_unknown_auth_type_is_rejected() { + let props = HashMap::from([( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + "kerberos".to_string(), + )]); + let config = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(props) + .build(); + + let err = test_catalog(config).resolve_auth_manager().unwrap_err(); + assert!(err.message().contains(REST_CATALOG_PROP_AUTH_TYPE)); + } + + #[tokio::test] + async fn test_with_auth_manager_overrides_config() { + // A custom auth manager takes precedence over `rest.auth.type`. + #[derive(Debug)] + struct StubAuthManager; + #[async_trait] + impl AuthManager for StubAuthManager { + async fn init_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + unimplemented!() + } + async fn catalog_session( + &self, + _client: &HttpClient, + _props: &HashMap, + ) -> Result> { + unimplemented!() + } + } + + let config = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([( + REST_CATALOG_PROP_AUTH_TYPE.to_string(), + "kerberos".to_string(), + )])) + .build(); + + // The unknown auth type is never consulted. + let catalog = test_catalog_with(config, Some(Arc::new(StubAuthManager))); + assert!(catalog.resolve_auth_manager().is_ok()); } #[tokio::test] @@ -1800,6 +2647,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1901,6 +2749,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -1956,6 +2805,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2001,6 +2851,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2036,6 +2887,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2073,6 +2925,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2103,6 +2956,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2145,6 +2999,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2215,6 +3070,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2348,6 +3204,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2394,6 +3251,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2425,6 +3283,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2464,6 +3323,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2497,6 +3357,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2533,6 +3394,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2652,6 +3514,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2690,6 +3553,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2841,6 +3705,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -2912,6 +3777,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -3058,6 +3924,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -3125,6 +3992,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, @@ -3178,6 +4046,7 @@ mod tests { let catalog = RestCatalog::new( RestCatalogConfig::builder().uri(server.url()).build(), + None, Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), None, diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs index 07dc0620da..7baeca40d9 100644 --- a/crates/catalog/rest/src/client.rs +++ b/crates/catalog/rest/src/client.rs @@ -17,57 +17,113 @@ use std::collections::HashMap; use std::fmt::{Debug, Formatter}; +use std::sync::Arc; -use http::StatusCode; use iceberg::{Error, ErrorKind, Result}; use reqwest::header::HeaderMap; -use reqwest::{Client, IntoUrl, Method, Request, RequestBuilder, Response}; +use reqwest::{Client, IntoUrl, Method, RequestBuilder, Response, StatusCode}; use serde::de::DeserializeOwned; -use tokio::sync::Mutex; use crate::RestCatalogConfig; -use crate::types::{ErrorResponse, TokenResponse}; +use crate::auth::{AuthSession, NoopSession}; +use crate::request::HttpRequest; -pub(crate) struct HttpClient { +/// The catalog's HTTP client, handed to an [`AuthManager`] so its own +/// requests share the catalog's connection pool and configuration. +/// +/// [`AuthManager`]: crate::auth::AuthManager +#[derive(Clone)] +pub struct HttpClient { client: Client, - /// The token to be used for authentication. - /// - /// It's possible to fetch the token from the server while needed. - token: Mutex>, - /// The token endpoint to be used for authentication. - token_endpoint: String, - /// The credential to be used for authentication. - credential: Option<(Option, String)>, /// Extra headers to be added to each request. extra_headers: HeaderMap, - /// Extra oauth parameters to be added to each authentication request. - extra_oauth_params: HashMap, /// Whether to disable header redaction in error logs (defaults to false for security). disable_header_redaction: bool, + /// Authenticates everything this client sends. A client handed to an + /// [`AuthManager`] carries no authentication, since that is what the + /// manager is about to create. + /// + /// [`AuthManager`]: crate::auth::AuthManager + session: Arc, } impl Debug for HttpClient { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // The inner client is omitted: an injected one may carry secrets in + // its default headers. f.debug_struct("HttpClient") - .field("client", &self.client) - .field("extra_headers", &self.extra_headers) + .field( + // `header.*` values may hold secrets (e.g. `authorization`). + "extra_headers", + &format_headers_redacted(&self.extra_headers, self.disable_header_redaction), + ) .finish_non_exhaustive() } } impl HttpClient { + /// The same client authenticating with `session` instead: a derived + /// session reuses the connection pool and headers of the client it came + /// from, carrying its own authentication. + pub fn with_session(&self, session: Arc) -> Self { + Self { + session, + ..self.clone() + } + } + + /// The same client with no authentication, for the requests that must not + /// carry it — a session refreshing its own token over the client it + /// authenticates would otherwise recurse. + pub fn without_session(&self) -> Self { + self.with_session(Arc::new(NoopSession)) + } + + /// Sends a form-encoded POST and returns the response status and body, + /// which is what an [`AuthManager`] needs to exchange a credential for a + /// token. Only `headers` are sent; the catalog's own extra headers are + /// not merged in. + /// + /// Like every request, it carries this client's session; call + /// [`Self::without_session`] first to send it unauthenticated. + /// + /// [`AuthManager`]: crate::auth::AuthManager + pub async fn post_form( + &self, + url: &str, + headers: &HeaderMap, + form: &HashMap<&str, &str>, + ) -> Result<(StatusCode, Vec)> { + let mut request = HttpRequest::build( + self.client + .request(Method::POST, url) + .headers(headers.clone()) + .form(form), + )?; + // `headers` may carry a `content-type: application/json` that `form` + // leaves in place. + request.headers_mut().insert( + reqwest::header::CONTENT_TYPE, + reqwest::header::HeaderValue::from_static("application/x-www-form-urlencoded"), + ); + self.session.authenticate(&mut request).await?; + let response = self.client.execute(request.into_inner()).await?; + let status = response.status(); + let body = response + .bytes() + .await + .map_err(|err| err.with_url(url.parse().unwrap_or_else(|_| "/".parse().unwrap())))?; + Ok((status, body.to_vec())) + } + /// Create a new http client. - pub fn new(cfg: &RestCatalogConfig) -> Result { - let extra_headers = cfg.extra_headers()?; + pub(crate) fn new(cfg: &RestCatalogConfig) -> Result { Ok(HttpClient { - client: cfg.client().unwrap_or_default(), - token: Mutex::new(cfg.token()), - token_endpoint: cfg.get_token_endpoint(), - credential: cfg.credential(), - extra_headers, - extra_oauth_params: cfg.extra_oauth_params(), + client: cfg.client(), + extra_headers: cfg.extra_headers()?, disable_header_redaction: cfg.disable_header_redaction(), + session: Arc::new(NoopSession), }) } @@ -75,192 +131,64 @@ impl HttpClient { /// /// If cfg carries new value, we will use cfg instead. /// Otherwise, we will keep the old value. - pub fn update_with(self, cfg: &RestCatalogConfig) -> Result { + pub(crate) fn update_with(self, cfg: &RestCatalogConfig) -> Result { let extra_headers = (!cfg.extra_headers()?.is_empty()) .then(|| cfg.extra_headers()) .transpose()? .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_endpoint: if !cfg.get_token_endpoint().is_empty() { - cfg.get_token_endpoint() - } else { - self.token_endpoint - }, - credential: cfg.credential().or(self.credential), + // `cfg.client()` returns the same shared client. + client: cfg.client(), extra_headers, - extra_oauth_params: if !cfg.extra_oauth_params().is_empty() { - cfg.extra_oauth_params() - } else { - self.extra_oauth_params - }, disable_header_redaction: cfg.disable_header_redaction(), + session: self.session, }) } - /// This API is testing only to assert the token. + /// Testing only: the session authenticating this client's requests. #[cfg(test)] - pub(crate) async fn token(&self) -> Option { - let mut req = self - .request(Method::GET, &self.token_endpoint) - .build() - .unwrap(); - self.authenticate(&mut req).await.ok(); - self.token.lock().await.clone() + pub(crate) fn session(&self) -> &Arc { + &self.session } - async fn exchange_credential_for_token(&self) -> Result { - // Credential must exist here. - let (client_id, client_secret) = self.credential.as_ref().ok_or_else(|| { - Error::new( - ErrorKind::DataInvalid, - "Credential must be provided for authentication", - ) - })?; - - let mut params = HashMap::with_capacity(4); - params.insert("grant_type", "client_credentials"); - if let Some(client_id) = client_id { - params.insert("client_id", client_id); - } - params.insert("client_secret", client_secret); - params.extend( - self.extra_oauth_params - .iter() - .map(|(k, v)| (k.as_str(), v.as_str())), - ); - - let mut auth_req = self - .request(Method::POST, &self.token_endpoint) - .form(¶ms) - .build()?; - // extra headers add content-type application/json header it's necessary to override it with proper type - // note that form call doesn't add content-type header if already present - auth_req.headers_mut().insert( - http::header::CONTENT_TYPE, - http::HeaderValue::from_static("application/x-www-form-urlencoded"), - ); - let auth_url = auth_req.url().clone(); - let auth_resp = self.client.execute(auth_req).await?; - - let auth_res: TokenResponse = if auth_resp.status() == StatusCode::OK { - let text = auth_resp - .bytes() - .await - .map_err(|err| err.with_url(auth_url.clone()))?; - Ok(serde_json::from_slice(&text).map_err(|e| { - Error::new( - ErrorKind::Unexpected, - "Failed to parse response from rest catalog server!", - ) - .with_context("operation", "auth") - .with_context("url", auth_url.to_string()) - .with_context("json", String::from_utf8_lossy(&text)) - .with_source(e) - })?) - } else { - let code = auth_resp.status(); - let text = auth_resp - .bytes() - .await - .map_err(|err| err.with_url(auth_url.clone()))?; - let e: ErrorResponse = serde_json::from_slice(&text).map_err(|e| { - Error::new(ErrorKind::Unexpected, "Received unexpected response") - .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)) - }?; - Ok(auth_res.access_token) - } - - /// Invalidate the current token without generating a new one. On the next request, the client - /// will attempt to generate a new token. - pub(crate) async fn invalidate_token(&self) -> Result<()> { - *self.token.lock().await = None; - Ok(()) - } - - /// Invalidate the current token and set a new one. Generates a new token before invalidating - /// the current token, meaning the old token will be used until this function acquires the lock - /// and overwrites the token. + /// Testing only: the bearer token `session` would attach. /// - /// If credential is invalid, or the request fails, this method will return an error and leave - /// the current token unchanged. - pub(crate) async fn regenerate_token(&self) -> Result<()> { - let new_token = self.exchange_credential_for_token().await?; - *self.token.lock().await = Some(new_token.clone()); - Ok(()) - } - - /// Authenticates the request by adding a bearer token to the authorization header. - /// - /// This method supports three authentication modes: - /// - /// 1. **No authentication** - Skip authentication when both `credential` and `token` are missing. - /// 2. **Token authentication** - Use the provided `token` directly for authentication. - /// 3. **OAuth authentication** - Exchange `credential` for a token, cache it, then use it for authentication. - /// - /// When both `credential` and `token` are present, `token` takes precedence. - /// - /// # TODO: Support automatic token refreshing. - async fn authenticate(&self, req: &mut Request) -> Result<()> { - // Clone the token from lock without holding the lock for entire function. - let token = self.token.lock().await.clone(); - - if self.credential.is_none() && token.is_none() { - return Ok(()); - } - - // Either use the provided token or exchange credential for token, cache and use that - let token = match token { - Some(token) => token, - None => { - let token = self.exchange_credential_for_token().await?; - // Update token so that we use it for next request instead of - // exchanging credential for token from the server again - *self.token.lock().await = Some(token.clone()); - token - } - }; - - // Insert token in request. - req.headers_mut().insert( - http::header::AUTHORIZATION, - format!("Bearer {token}").parse().map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - "Invalid token received from catalog server!", - ) - .with_source(e) - })?, - ); - - Ok(()) + /// Authenticates a throwaway request (never sent) and reads the header + /// back, so it works for any [`AuthSession`]. + #[cfg(test)] + pub(crate) async fn token(&self) -> Option { + let mut request = HttpRequest::build( + self.client + .request(Method::GET, "http://localhost/token-probe"), + ) + .ok()?; + self.session.authenticate(&mut request).await.ok()?; + let request = request.into_inner(); + request + .headers() + .get(reqwest::header::AUTHORIZATION)? + .to_str() + .ok()? + .strip_prefix("Bearer ") + .map(str::to_string) } #[inline] - pub fn request(&self, method: Method, url: U) -> RequestBuilder { + pub(crate) fn request(&self, method: Method, url: U) -> RequestBuilder { self.client .request(method, url) .headers(self.extra_headers.clone()) } - /// Executes the given `Request` and returns a `Response`. - pub async fn execute(&self, mut request: Request) -> Result { - request.headers_mut().extend(self.extra_headers.clone()); - Ok(self.client.execute(request).await?) - } - // Queries the Iceberg REST catalog after authentication with the given `Request` and // returns a `Response`. - pub async fn query_catalog(&self, mut request: Request) -> Result { - self.authenticate(&mut request).await?; - self.execute(request).await + pub(crate) async fn query_catalog(&self, mut request: HttpRequest) -> Result { + // Authenticate first, then apply extra headers, so a configured + // `header.authorization` keeps overriding a token (unchanged behavior). + self.session.authenticate(&mut request).await?; + let mut request = request.into_inner(); + request.headers_mut().extend(self.extra_headers.clone()); + Ok(self.client.execute(request).await?) } /// Returns whether header redaction is disabled for this client. @@ -287,20 +215,21 @@ pub(crate) async fn deserialize_catalog_response( }) } -/// 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 (matched by substring, so +/// e.g. `x-client-secret` is covered along with `authorization`). 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. @@ -361,6 +290,56 @@ pub(crate) async fn deserialize_unexpected_catalog_error( mod tests { use super::*; + #[derive(Debug)] + struct StaticSession; + + #[async_trait::async_trait] + impl AuthSession for StaticSession { + async fn authenticate(&self, request: &mut HttpRequest) -> Result<()> { + request + .headers_mut() + .insert("authorization", "Bearer tok".parse().unwrap()); + Ok(()) + } + } + + #[tokio::test] + async fn test_post_form_carries_the_session_until_it_is_removed() { + // Every request a client sends carries its session; a caller that + // needs an unauthenticated one removes the session first. + let mut server = mockito::Server::new_async().await; + let signed = server + .mock("POST", "/token") + .match_header("authorization", "Bearer tok") + .with_status(200) + .create_async() + .await; + let unsigned = server + .mock("POST", "/token") + .match_header("authorization", mockito::Matcher::Missing) + .with_status(200) + .create_async() + .await; + + let client = HttpClient::new(&RestCatalogConfig::builder().uri(server.url()).build()) + .unwrap() + .with_session(Arc::new(StaticSession)); + let url = format!("{}/token", server.url()); + + client + .post_form(&url, &HeaderMap::new(), &HashMap::new()) + .await + .unwrap(); + signed.assert_async().await; + + client + .without_session() + .post_form(&url, &HeaderMap::new(), &HashMap::new()) + .await + .unwrap(); + unsigned.assert_async().await; + } + #[test] fn test_format_headers_redacted_empty() { let headers = HeaderMap::new(); @@ -382,6 +361,31 @@ mod tests { assert!(result.contains("abc123")); } + #[tokio::test] + async fn test_http_client_debug_redacts_headers() { + let config = RestCatalogConfig::builder() + .uri("http://localhost".to_string()) + .props(HashMap::from([ + ("header.authorization".to_string(), "Basic xyz".to_string()), + ( + "header.x-client-secret".to_string(), + "shh-secret".to_string(), + ), + ( + "header.x-client-credential".to_string(), + "cred-value".to_string(), + ), + ])) + .build(); + let client = HttpClient::new(&config).unwrap(); + + let out = format!("{client:?}"); + assert!(!out.contains("Basic xyz")); + assert!(!out.contains("shh-secret")); + assert!(!out.contains("cred-value")); + assert!(out.contains("[REDACTED]")); + } + #[test] fn test_format_headers_redacted_filters_sensitive() { let mut headers = HeaderMap::new(); diff --git a/crates/catalog/rest/src/lib.rs b/crates/catalog/rest/src/lib.rs index 383728401f..5670ef3c9c 100644 --- a/crates/catalog/rest/src/lib.rs +++ b/crates/catalog/rest/src/lib.rs @@ -51,11 +51,16 @@ #![deny(missing_docs)] +mod auth; mod catalog; mod client; +pub use client::HttpClient; +mod request; +pub use request::{HttpRequest, HttpRequestBody}; mod endpoint; mod types; +pub use auth::*; pub use catalog::*; pub use endpoint::Endpoint; pub use types::*; diff --git a/crates/catalog/rest/src/request.rs b/crates/catalog/rest/src/request.rs new file mode 100644 index 0000000000..7a3c8e084c --- /dev/null +++ b/crates/catalog/rest/src/request.rs @@ -0,0 +1,152 @@ +// 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. + +//! The request type REST catalog authentication works with. + +use http::{HeaderMap, Method}; +use iceberg::Result; +use reqwest::{Request, RequestBuilder}; + +/// An outgoing REST request being authenticated by an +/// [`AuthSession`](crate::AuthSession). +/// +/// Wraps the request so a session mutates it through the stable +/// `http` crate types rather than the concrete request type the REST catalog +/// uses internally. +pub struct HttpRequest { + inner: Request, +} + +impl HttpRequest { + /// Wraps a request, e.g. to unit-test a custom + /// [`AuthSession`](crate::AuthSession). + pub fn new(inner: Request) -> Self { + Self { inner } + } + + /// Builds the request `builder` describes. + pub(crate) fn build(builder: RequestBuilder) -> Result { + Ok(Self::new(builder.build()?)) + } + + /// The wrapped request, for the client that sends it. + pub(crate) fn into_inner(self) -> Request { + self.inner + } + + /// The request method. + pub fn method(&self) -> &Method { + self.inner.method() + } + + /// The request URL, as a string (scheme, host, path and query). + pub fn url_str(&self) -> &str { + self.inner.url().as_str() + } + + /// The request headers. + pub fn headers(&self) -> &HeaderMap { + self.inner.headers() + } + + /// The mutable request headers, e.g. to add an `Authorization` header. + pub fn headers_mut(&mut self) -> &mut HeaderMap { + self.inner.headers_mut() + } + + /// The request body, distinguishing an absent body from a streaming one: + /// signers can sign [`HttpRequestBody::Empty`] (empty-payload hash) and + /// [`HttpRequestBody::Buffered`], but not [`HttpRequestBody::Streaming`]. + pub fn body(&self) -> HttpRequestBody<'_> { + match self.inner.body() { + None => HttpRequestBody::Empty, + Some(body) => match body.as_bytes() { + Some(bytes) => HttpRequestBody::Buffered(bytes), + None => HttpRequestBody::Streaming, + }, + } + } +} + +/// The body of an [`HttpRequest`], as seen by authentication. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HttpRequestBody<'a> { + /// No body is set. + Empty, + /// An in-memory body. + Buffered(&'a [u8]), + /// A streaming body, whose bytes are not available for e.g. signing. + Streaming, +} + +impl<'a> HttpRequestBody<'a> { + /// The signable bytes: empty for [`Self::Empty`], the buffer for + /// [`Self::Buffered`], and `None` for [`Self::Streaming`]. + pub fn as_bytes(&self) -> Option<&'a [u8]> { + match self { + HttpRequestBody::Empty => Some(&[]), + HttpRequestBody::Buffered(bytes) => Some(bytes), + HttpRequestBody::Streaming => None, + } + } +} + +#[cfg(test)] +mod tests { + use reqwest::Client; + + use super::*; + + #[test] + fn test_http_request_body_states() { + let client = Client::new(); + + // No body at all. + let req = client + .get("https://rest.example.com/v1/config") + .build() + .unwrap(); + let http_req = HttpRequest::new(req); + let body = http_req.body(); + assert_eq!(body, HttpRequestBody::Empty); + assert_eq!(body.as_bytes(), Some(&[] as &[u8])); + + // An in-memory body. + let req = client + .post("https://rest.example.com/v1/namespaces") + .body("{}") + .build() + .unwrap(); + let http_req = HttpRequest::new(req); + let body = http_req.body(); + assert_eq!(body, HttpRequestBody::Buffered(b"{}")); + assert_eq!(body.as_bytes(), Some(b"{}" as &[u8])); + + // A streaming body: bytes are unavailable, so it must not sign as empty. + let req = client + .post("https://rest.example.com/v1/namespaces") + .body(reqwest::Body::wrap_stream(futures::stream::once(async { + Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"chunk")) + }))) + .build() + .unwrap(); + let http_req = HttpRequest::new(req); + let body = http_req.body(); + assert_eq!(body, HttpRequestBody::Streaming); + assert_eq!(body.as_bytes(), None); + } +}