From a9ad7e2c0104d08535f7cc19cc8d5785c1a67ba6 Mon Sep 17 00:00:00 2001 From: Kynan Rilee Date: Mon, 8 Jun 2026 13:35:57 -0400 Subject: [PATCH] [SS-69] support custom RequestAuthenticator for REST catalog client (#1) * support custom RequestAuthenticator for REST catalog client * safer Debug impls and error ctx * mutex-guard the entire OAuth2 token-fetch process --- crates/catalog/rest/src/catalog.rs | 376 ++++++++++++++++++++++++++--- crates/catalog/rest/src/client.rs | 204 ++++------------ crates/catalog/rest/src/lib.rs | 2 + crates/catalog/rest/src/token.rs | 306 +++++++++++++++++++++++ 4 files changed, 686 insertions(+), 202 deletions(-) create mode 100644 crates/catalog/rest/src/token.rs diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 0626ce5061..03fda00bea 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -40,6 +40,9 @@ use typed_builder::TypedBuilder; use crate::client::{ HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error, }; +use crate::token::{ + BearerTokenAuthenticator, OAuth2TokenProvider, RequestAuthenticator, StaticTokenProvider, +}; use crate::types::{ CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult, @@ -74,6 +77,7 @@ impl Default for RestCatalogBuilder { warehouse: None, props: HashMap::new(), client: None, + authenticator: None, }, storage_factory: None, runtime: None, @@ -129,6 +133,8 @@ impl CatalogBuilder for RestCatalogBuilder { ErrorKind::DataInvalid, "Catalog uri is required", )) + } else if let Err(error) = self.config.authenticator_cfg().strict_validate() { + Err(error) } else { let runtime = self.runtime.unwrap_or_else(Runtime::current); Ok(RestCatalog::new(self.config, self.storage_factory, runtime)) @@ -145,6 +151,129 @@ impl RestCatalogBuilder { self.config.client = Some(client); self } + + /// Configures the catalog with a custom request authenticator, in place + /// of the `token` / `credential` props. + /// + /// For bearer-token flows, wrap a `TokenProvider` in `BearerTokenAuthenticator`. + /// For auth that needs to see the outgoing request (e.g. AWS SigV4), + /// implement `RequestAuthenticator` directly. + pub fn with_authenticator(mut self, authenticator: Arc) -> Self { + self.config.authenticator = Some(authenticator); + self + } +} + +/// There are multiple ways to configure the auth mechanism used by the catalog client. +/// On our side, we attempt to enforce that only one mechanism is used (see `strict_validate`). +/// But the catalog server can supply arbitrary overrides, so we must implement semantics +/// for multiple options being specified at once. +/// +/// Option 0: No auth. +pub struct AuthenticatorConfig { + /// Option 1: Custom request authenticator. + pub custom_authenticator: Option>, + + /// Option 2: Default OAuth flow. + /// OAuth endpoint. + pub token_endpoint: String, + /// OAuth credentials. + pub credential: Option<(Option, String)>, + /// OAuth params. + pub extra_oauth_params: HashMap, + + /// Option 3: Static auth token. + pub token: Option, +} + +impl AuthenticatorConfig { + /// Someone (e.g. the catalog server's "config" endpoint) gave us overrides to apply to our config. + /// N.B. `Some(_)` means override existing value. `None` or empty means no-op. + /// + /// Returns `true` if we should generate a new authenticator from the merged config. + pub(crate) fn merge(mut self, overrides: Self) -> (Self, bool) { + let mut should_regenerate = false; + + if overrides.custom_authenticator.is_some() { + self.custom_authenticator = overrides.custom_authenticator; + should_regenerate = true; + } + + if !overrides.token_endpoint.is_empty() && self.token_endpoint != overrides.token_endpoint { + self.token_endpoint = overrides.token_endpoint; + should_regenerate = true; + } + + if overrides.credential.is_some() && self.credential != overrides.credential { + self.credential = overrides.credential; + should_regenerate = true; + } + + if !overrides.extra_oauth_params.is_empty() + && self.extra_oauth_params != overrides.extra_oauth_params + { + self.extra_oauth_params = overrides.extra_oauth_params; + should_regenerate = true; + } + + if overrides.token.is_some() && self.token != overrides.token { + self.token = overrides.token; + should_regenerate = true; + } + + (self, should_regenerate) + } + + /// Build a new authenticator. + pub(crate) fn authenticator( + &self, + client: Client, + extra_headers: HeaderMap, + ) -> Option> { + // If the user specified a custom authenticator, never override it. + if self.custom_authenticator.is_some() { + return self.custom_authenticator.clone(); + } + + // If there are OAuth credentials, use them. + if let Some((client_id, client_secret)) = &self.credential { + let provider = Arc::new(OAuth2TokenProvider::new( + client, + client_id.clone(), + client_secret.clone(), + self.token_endpoint.clone(), + extra_headers, + self.extra_oauth_params.clone(), + // If there's a preexisting token, seed the cache with that token. + self.token.clone(), + )); + return Some(Arc::new(BearerTokenAuthenticator::new(provider))); + } + + // If there's only a token, use that. + if let Some(token) = &self.token { + let provider = Arc::new(StaticTokenProvider::new(token)); + return Some(Arc::new(BearerTokenAuthenticator::new(provider))); + } + + // There's no auth configured. + None + } + + /// We expect our users to configure only one auth mechanism. + /// (These rules do not apply to server-supplied overrides, hence the "strict".) + fn strict_validate(&self) -> Result<()> { + if self.custom_authenticator.is_some() + && (self.token.is_some() || self.credential.is_some()) + { + Err(Error::new( + ErrorKind::DataInvalid, + "Custom authenticator is not compatible with `token` or `credential` props.", + )) + } else { + Ok(()) + } + } } /// Rest catalog configuration. @@ -163,6 +292,10 @@ pub(crate) struct RestCatalogConfig { #[builder(default)] client: Option, + + /// A custom request authenticator, as an alternative to the auth configuration props. + #[builder(default, setter(strip_option))] + authenticator: Option>, } impl RestCatalogConfig { @@ -178,7 +311,8 @@ impl RestCatalogConfig { [&self.uri, PATH_V1, "config"].join("/") } - pub(crate) fn get_token_endpoint(&self) -> String { + /// Helper for [`authenticator_cfg`]. + fn get_token_endpoint(&self) -> String { if let Some(oauth2_uri) = self.props.get("oauth2-server-uri") { oauth2_uri.to_string() } else { @@ -220,13 +354,30 @@ impl RestCatalogConfig { self.client.clone() } + // Depends on some combination of: + // - explicit `authenticator` field + // - `credential` + `get_token_endpoint` + `extra_oauth_params` + // - `token` + // Or no auth. + pub(crate) fn authenticator_cfg(&self) -> AuthenticatorConfig { + AuthenticatorConfig { + custom_authenticator: self.authenticator.clone(), + token_endpoint: self.get_token_endpoint(), + credential: self.credential(), + extra_oauth_params: self.extra_oauth_params(), + token: self.token(), + } + } + + /// Helper for [`authenticator_cfg`]: /// Get the token from the config. /// /// The client can use this token to send requests. - pub(crate) fn token(&self) -> Option { + fn token(&self) -> Option { self.props.get("token").cloned() } + /// Helper for [`authenticator_cfg`]: /// Get the credentials from the config. The client can use these credentials to fetch a new /// token. /// @@ -235,7 +386,7 @@ impl RestCatalogConfig { /// - `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)> { + fn credential(&self) -> Option<(Option, String)> { let cred = self.props.get("credential")?; match cred.split_once(':') { @@ -294,8 +445,9 @@ impl RestCatalogConfig { Ok(headers) } + /// Helper for [`authenticator_cfg`]: /// Get the optional OAuth headers from the config. - pub(crate) fn extra_oauth_params(&self) -> HashMap { + fn extra_oauth_params(&self) -> HashMap { let mut params = HashMap::new(); if let Some(scope) = self.props.get("scope") { @@ -490,6 +642,7 @@ impl RestCatalog { /// Invalidate the current token without generating a new one. On the next request, the client /// will attempt to generate a new token. + /// (For alternative auth mechanisms, perform the analogous operation.) pub async fn invalidate_token(&self) -> Result<()> { self.context().await?.client.invalidate_token().await } @@ -500,6 +653,8 @@ impl RestCatalog { /// /// If credential is invalid, or the request fails, this method will return an error and leave /// the current token unchanged. + /// + /// (For alternative auth mechanisms, perform the analogous operation.) pub async fn regenerate_token(&self) -> Result<()> { self.context().await?.client.regenerate_token().await } @@ -1178,11 +1333,24 @@ mod tests { .await } + /// We use this mock to inspect the bearer tokens we get from the `TokenProvider`. + async fn create_authed_list_ns_mock(server: &mut ServerGuard, token: &str) -> Mock { + server + .mock("GET", "/v1/namespaces") + .match_header("authorization", format!("Bearer {token}").as_str()) + .with_status(200) + .with_body(r#"{"namespaces": []}"#) + .expect(1) + .create_async() + .await + } + #[tokio::test] async fn test_oauth() { 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 = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1196,10 +1364,11 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; - oauth_mock.assert_async().await; + // Check that "list namespaces" actually retrieves the token and attaches it to our request. + catalog.list_namespaces(None).await.unwrap(); + oauth_mock.assert_async().await; // Fetched the token. config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; // Made a request with the token attached. } #[tokio::test] @@ -1234,6 +1403,7 @@ mod tests { .await; let config_mock = create_config_mock(&mut server).await; + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let catalog = RestCatalog::new( RestCatalogConfig::builder() @@ -1244,11 +1414,11 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; - - oauth_mock.assert_async().await; + // Check that "list namespaces" actually retrieves the token and attaches it to our request. + catalog.list_namespaces(None).await.unwrap(); + oauth_mock.assert_async().await; // Fetched the token. config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; // Made a request with the token attached. } #[tokio::test] @@ -1256,6 +1426,7 @@ mod tests { 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 = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1269,18 +1440,24 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; + // Make a request and check that we have an auth token. + catalog.list_namespaces(None).await.unwrap(); oauth_mock.assert_async().await; config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; + + // Invalidate any cached token. + catalog.invalidate_token().await.unwrap(); + // Make a request and check that we retrieve a new token and attach it to our request. + // Note: The new token ends with '1' instead of '0'. 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())); + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000001").await; + catalog.list_namespaces(None).await.unwrap(); + oauth_mock.assert_async().await; // Fetched a new token. + list_ns_mock.assert_async().await; // Made a request with the new token attached. } #[tokio::test] @@ -1288,6 +1465,7 @@ mod tests { 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 = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1301,18 +1479,32 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; + // Make a request and check that we have an auth token. + catalog.list_namespaces(None).await.unwrap(); oauth_mock.assert_async().await; config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; + + // Invalidate any cached token. + catalog.invalidate_token().await.unwrap(); + // The OAuth server is broken, so we can't get a new token. + // Check that we can't make authenticated "list namespaces" requests. let oauth_mock = - create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 500) + create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000000", 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); + assert!(catalog.list_namespaces(None).await.is_err()); // Request failed. + oauth_mock.assert_async().await; // Tried fetching a new token. + + // The OAuth server works again. + // Check that we can make authenticated "list namespaces" requests. + let oauth_mock = + create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000000", 200) + .await; + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000000").await; + catalog.list_namespaces(None).await.unwrap(); + oauth_mock.assert_async().await; // Fetched a new token. + list_ns_mock.assert_async().await; // Made a request with the new token attached. } #[tokio::test] @@ -1320,6 +1512,7 @@ mod tests { 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 = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1333,18 +1526,23 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; + // Make a request and check that we have an auth token. + catalog.list_namespaces(None).await.unwrap(); oauth_mock.assert_async().await; config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; + // Regenerate the token. 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())); + oauth_mock.assert_async().await; // Fetched a new token. + + // Check that we use the new token. + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000001").await; + catalog.list_namespaces(None).await.unwrap(); + list_ns_mock.assert_async().await; // Made a request with the new token. } #[tokio::test] @@ -1352,6 +1550,7 @@ mod tests { 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 = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1365,21 +1564,116 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; + // Make a request and check that we have an auth token. + catalog.list_namespaces(None).await.unwrap(); oauth_mock.assert_async().await; config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; + // Fail to generate a new token. 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()); + let regenerate_result = catalog.regenerate_token().await; + assert!(regenerate_result.is_err()); oauth_mock.assert_async().await; - let token = catalog.context().await.unwrap().client.token().await; - // original token is left intact - assert_eq!(token, Some("ey000000000000".to_string())); + // Make a request and check that we're reusing the old token. + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000000").await; + catalog.list_namespaces(None).await.unwrap(); + list_ns_mock.assert_async().await; + } + + // Demonstrate that a `RequestAuthenticator` can add whatever headers it wants. + // (For example, AWS SigV4 needs `Authorization: AWS4-HMAC-SHA256 ...` and `X-Amz-Date`.) + #[tokio::test] + async fn test_custom_authenticator() { + use async_trait::async_trait; + use reqwest::Request; + + #[derive(Debug)] + struct TwoHeaderAuth; + + #[async_trait] + impl RequestAuthenticator for TwoHeaderAuth { + async fn authenticate_request(&self, req: &mut Request) -> Result<()> { + req.headers_mut() + .insert("x-custom-auth", "sig-abc".parse().unwrap()); + req.headers_mut() + .insert("x-auth-date", "20260512T000000Z".parse().unwrap()); + Ok(()) + } + + async fn invalidate_cache(&self) -> Result<()> { + Ok(()) + } + + async fn regenerate_cache(&self) -> Result<()> { + Ok(()) + } + } + + let mut server = Server::new_async().await; + + let config_mock = server + .mock("GET", "/v1/config") + .match_header("x-custom-auth", "sig-abc") + .match_header("x-auth-date", "20260512T000000Z") + .with_status(200) + .with_body(r#"{"overrides": {}, "defaults": {}}"#) + .create_async() + .await; + let list_ns_mock = server + .mock("GET", "/v1/namespaces") + .match_header("x-custom-auth", "sig-abc") + .match_header("x-auth-date", "20260512T000000Z") + .with_status(200) + .with_body(r#"{"namespaces": []}"#) + .create_async() + .await; + + let catalog = RestCatalog::new( + RestCatalogConfig::builder() + .uri(server.url()) + .authenticator(Arc::new(TwoHeaderAuth)) + .build(), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + ); + + catalog.list_namespaces(None).await.unwrap(); + config_mock.assert_async().await; + list_ns_mock.assert_async().await; + } + + #[tokio::test] + async fn test_strict_validate_rejects_authenticator_with_token() { + let cfg = AuthenticatorConfig { + custom_authenticator: Some(Arc::new(BearerTokenAuthenticator::new(Arc::new( + StaticTokenProvider::new("t"), + )))), + token_endpoint: String::new(), + credential: None, + extra_oauth_params: HashMap::new(), + token: Some("t".into()), + }; + let err = cfg.strict_validate().unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[tokio::test] + async fn test_strict_validate_rejects_authenticator_with_credential() { + let cfg = AuthenticatorConfig { + custom_authenticator: Some(Arc::new(BearerTokenAuthenticator::new(Arc::new( + StaticTokenProvider::new("t"), + )))), + token_endpoint: String::new(), + credential: Some((Some("id".into()), "sec".into())), + extra_oauth_params: HashMap::new(), + token: None, + }; + let err = cfg.strict_validate().unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); } #[tokio::test] @@ -1452,10 +1746,12 @@ mod tests { assert_eq!(headers, expected_headers); } + // The `oauth2-server-uri` prop overrides the default OAuth endpoint. #[tokio::test] async fn test_oauth_with_oauth2_server_uri() { let mut server = Server::new_async().await; let config_mock = create_config_mock(&mut server).await; + let list_ns_mock = create_authed_list_ns_mock(&mut server, "ey000000000000").await; let mut auth_server = Server::new_async().await; let auth_server_path = "/some/path"; @@ -1479,11 +1775,11 @@ mod tests { Runtime::current(), ); - let token = catalog.context().await.unwrap().client.token().await; - - oauth_mock.assert_async().await; + // Make a request and check that we fetch and use the OAuth2 server's token. + catalog.list_namespaces(None).await.unwrap(); + oauth_mock.assert_async().await; // Fetched a token from the OAuth2 server. config_mock.assert_async().await; - assert_eq!(token, Some("ey000000000000".to_string())); + list_ns_mock.assert_async().await; // Made a request with the token attached. } #[tokio::test] diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs index 07dc0620da..b61144369e 100644 --- a/crates/catalog/rest/src/client.rs +++ b/crates/catalog/rest/src/client.rs @@ -17,32 +17,29 @@ 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 serde::de::DeserializeOwned; -use tokio::sync::Mutex; -use crate::RestCatalogConfig; -use crate::types::{ErrorResponse, TokenResponse}; +use crate::token::RequestAuthenticator; +use crate::{AuthenticatorConfig, RestCatalogConfig}; pub(crate) 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)>, + /// Invoked before each request. Mutates the request to add authentication + /// (e.g. inserting an Authorization header, signing with SigV4, etc). + authenticator: Option>, + + /// This generated the current `authenticator`. When we receive overrides, + /// we apply them to this config and generate a new `authenticator`. + token_provider_cfg: AuthenticatorConfig, + /// 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, } @@ -60,13 +57,13 @@ impl HttpClient { /// Create a new http client. pub fn new(cfg: &RestCatalogConfig) -> Result { let extra_headers = cfg.extra_headers()?; + let token_provider_cfg = cfg.authenticator_cfg(); + let client = cfg.client().unwrap_or_default(); Ok(HttpClient { - client: cfg.client().unwrap_or_default(), - token: Mutex::new(cfg.token()), - token_endpoint: cfg.get_token_endpoint(), - credential: cfg.credential(), + client: client.clone(), + authenticator: token_provider_cfg.authenticator(client, extra_headers.clone()), + token_provider_cfg, extra_headers, - extra_oauth_params: cfg.extra_oauth_params(), disable_header_redaction: cfg.disable_header_redaction(), }) } @@ -80,109 +77,33 @@ impl HttpClient { .then(|| cfg.extra_headers()) .transpose()? .unwrap_or(self.extra_headers); + let client = cfg.client().unwrap_or(self.client); + + let (token_provider_cfg, should_regenerate) = + self.token_provider_cfg.merge(cfg.authenticator_cfg()); + let authenticator = if should_regenerate { + token_provider_cfg.authenticator(client.clone(), extra_headers.clone()) + } else { + self.authenticator + }; + 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), + client, + authenticator, + token_provider_cfg, 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(), }) } - /// This API is testing only to assert the token. - #[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() - } - - 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. + /// (For alternative auth mechanisms, perform the analogous operation.) pub(crate) async fn invalidate_token(&self) -> Result<()> { - *self.token.lock().await = None; - Ok(()) + match self.authenticator.as_ref() { + None => Ok(()), + Some(authenticator) => authenticator.invalidate_cache().await, + } } /// Invalidate the current token and set a new one. Generates a new token before invalidating @@ -191,56 +112,13 @@ impl HttpClient { /// /// 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(()); + /// (For alternative auth mechanisms, perform the analogous operation.) + pub(crate) async fn regenerate_token(&self) -> Result<()> { + match self.authenticator.as_ref() { + None => Ok(()), + Some(authenticator) => authenticator.regenerate_cache().await, } - - // 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(()) } #[inline] @@ -259,7 +137,9 @@ impl HttpClient { // 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?; + if let Some(auth) = &self.authenticator { + auth.authenticate_request(&mut request).await?; + } self.execute(request).await } diff --git a/crates/catalog/rest/src/lib.rs b/crates/catalog/rest/src/lib.rs index 6bee950970..a7c663598c 100644 --- a/crates/catalog/rest/src/lib.rs +++ b/crates/catalog/rest/src/lib.rs @@ -53,7 +53,9 @@ mod catalog; mod client; +mod token; mod types; pub use catalog::*; +pub use token::*; pub use types::*; diff --git a/crates/catalog/rest/src/token.rs b/crates/catalog/rest/src/token.rs new file mode 100644 index 0000000000..c1a8685728 --- /dev/null +++ b/crates/catalog/rest/src/token.rs @@ -0,0 +1,306 @@ +// 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 client. +//! +//! Most generic trait: [`RequestAuthenticator`] +//! - Allows arbitrary mutations to the outgoing [`reqwest::Request`]. +//! +//! Less generic trait: [`TokenProvider`] +//! - Produces a bearer token. e.g. OAuth +//! - Use [`BearerTokenAuthenticator`] to wrap a [`TokenProvider`] into a [`RequestAuthenticator`]. + +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use http::{Method, StatusCode}; +use iceberg::{Error, ErrorKind, Result}; +use reqwest::header::HeaderMap; +use reqwest::{Client, Request}; +use tokio::sync::Mutex; + +use crate::{ErrorResponse, TokenResponse}; + +/// Authenticates outgoing REST-catalog requests by mutating them in place. +/// +/// Implementations may add headers, rewrite the URL, or do anything else they +/// need before the request is sent. This is the lowest-level auth hook; +/// callers that just need to attach a bearer token should typically implement +/// [`TokenProvider`] and rely on the [`BearerTokenAuthenticator`] adapter. +#[async_trait] +pub trait RequestAuthenticator: Send + Sync + Debug { + /// Apply authentication to the outgoing request. + async fn authenticate_request(&self, req: &mut Request) -> Result<()>; + + /// Discard any cached auth state (tokens, signatures, credentials, etc). + async fn invalidate_cache(&self) -> Result<()>; + + /// Refresh the cached auth state (e.g. a bearer token). + /// + /// Expected flow: + /// 1. Attempt to regenerate the auth state (e.g. fetch a new bearer token). + /// 2a. If regeneration succeeded, use the new state and discard the old state. (e.g. replace the token) + /// 2b. If regeneration failed, keep the old state but surface the error. (e.g. keep the old token) + async fn regenerate_cache(&self) -> Result<()>; +} + +/// Produces bearer tokens for authenticated REST catalog requests. +/// Caching, expiry, and refresh are up to the implementer. +#[async_trait] +pub trait TokenProvider: Send + Sync + Debug { + /// Get a token for the next request. + async fn token(&self) -> Result; + + /// Discard any cached token if it exists. + async fn invalidate(&self) -> Result<()>; + + /// Try generating a new cached token. + async fn regenerate(&self) -> Result<()>; +} + +/// Adapts a [`TokenProvider`] into a [`RequestAuthenticator`] by inserting +/// `Authorization: Bearer ` on the outgoing request. +#[derive(Debug, Clone)] +pub struct BearerTokenAuthenticator { + provider: Arc, +} + +impl BearerTokenAuthenticator { + /// Wrap the given token provider so it can be used as a `RequestAuthenticator`. + pub fn new(provider: Arc) -> Self { + Self { provider } + } +} + +#[async_trait] +impl RequestAuthenticator for BearerTokenAuthenticator { + async fn authenticate_request(&self, req: &mut Request) -> Result<()> { + let token = self.provider.token().await?; + req.headers_mut().insert( + http::header::AUTHORIZATION, + format!("Bearer {token}").parse().map_err(|e| { + Error::new(ErrorKind::DataInvalid, "Invalid authorization token!").with_source(e) + })?, + ); + Ok(()) + } + + async fn invalidate_cache(&self) -> Result<()> { + self.provider.invalidate().await + } + + async fn regenerate_cache(&self) -> Result<()> { + self.provider.regenerate().await + } +} + +/// Always produce the same bearer token. +pub struct StaticTokenProvider { + token: String, +} + +impl Debug for StaticTokenProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("StaticTokenProvider") + .finish_non_exhaustive() + } +} + +impl StaticTokenProvider { + /// We'll provide the same bearer token for every request. + pub fn new(token: impl Into) -> Self { + Self { + token: token.into(), + } + } +} + +#[async_trait] +impl TokenProvider for StaticTokenProvider { + async fn token(&self) -> Result { + Ok(self.token.clone()) + } + + async fn invalidate(&self) -> Result<()> { + // No-op since we only have the one token. + Ok(()) + } + + async fn regenerate(&self) -> Result<()> { + // No-op since we only have the one token. + Ok(()) + } +} + +/// Use the standard OAuth2 flow to obtain bearer tokens. +pub struct OAuth2TokenProvider { + client: Client, + + /// Authenticate as this entity. + client_id: Option, + /// Authenticate using this secret. + client_secret: String, + + /// Authenticate here. + token_endpoint: String, + + /// Request headers. + extra_headers: HeaderMap, + + /// Form params. + extra_oauth_params: HashMap, + + /// Most recently fetched token. + cached_token: Mutex>, +} + +impl Debug for OAuth2TokenProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Omit `client_secret` and `cached_token`: they're secrets. For + // `extra_headers`/`extra_oauth_params` show only the keys, since the + // values may carry secrets (e.g. a Basic `Authorization` header). + f.debug_struct("OAuth2TokenProvider") + .field("client_id", &self.client_id) + .field("token_endpoint", &self.token_endpoint) + .field( + "extra_headers", + &self.extra_headers.keys().collect::>(), + ) + .field( + "extra_oauth_params", + &self.extra_oauth_params.keys().collect::>(), + ) + .finish_non_exhaustive() + } +} + +impl OAuth2TokenProvider { + /// Construct a new token provider. + pub fn new( + client: Client, + client_id: Option, + client_secret: String, + token_endpoint: String, + extra_headers: HeaderMap, + extra_oauth_params: HashMap, + cached_token: Option, + ) -> Self { + Self { + client, + client_id, + client_secret, + token_endpoint, + extra_headers, + extra_oauth_params, + cached_token: Mutex::new(cached_token), + } + } + + /// Just fetch a token. Don't store it or do anything else. + async fn exchange_credential_for_token(&self) -> Result { + let mut params = HashMap::with_capacity(4); + params.insert("grant_type", "client_credentials"); + if let Some(client_id) = &self.client_id { + params.insert("client_id", client_id); + } + params.insert("client_secret", &self.client_secret); + params.extend( + self.extra_oauth_params + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())), + ); + + let mut auth_req = self + .client + .request(Method::POST, &self.token_endpoint) + .headers(self.extra_headers.clone()) + .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| { + // We omit the response text from the error context + // because an OK response could include an auth token (even if we can't parse it). + Error::new( + ErrorKind::Unexpected, + "Failed to parse response from OAuth2 token provider!", + ) + .with_context("operation", "auth") + .with_context("url", auth_url.to_string()) + .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) + } +} + +#[async_trait] +impl TokenProvider for OAuth2TokenProvider { + /// Fetch a new token if we don't already have one cached. + async fn token(&self) -> Result { + let mut cached = self.cached_token.lock().await; + + if let Some(token) = cached.clone() { + return Ok(token); + } + + let token = self.exchange_credential_for_token().await?; + *cached = Some(token.clone()); + Ok(token) + } + + async fn invalidate(&self) -> Result<()> { + *self.cached_token.lock().await = None; + Ok(()) + } + + /// Try fetching and caching a new token. If that fails, keep the old token. + async fn regenerate(&self) -> Result<()> { + let mut cached = self.cached_token.lock().await; + *cached = Some(self.exchange_credential_for_token().await?); + Ok(()) + } +}