CTTY commented on code in PR #2838: URL: https://github.com/apache/iceberg-rust/pull/2838#discussion_r3738982282
########## crates/catalog/rest/src/auth/oauth2.rs: ########## @@ -0,0 +1,362 @@ +// 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 reqwest::{Client, Method}; +use tokio::sync::Mutex; + +use super::{AuthManager, AuthRequest, AuthSession}; +use crate::catalog::{ + REST_CATALOG_PROP_URI, RestCatalogConfig, credential_from_props, default_token_endpoint, + explicit_headers_from_props, +}; +use crate::types::{ErrorResponse, TokenResponse}; + +/// Per-phase OAuth2 parameters (init vs. post-handshake catalog phase). +#[derive(Clone)] +struct OAuth2Params { + extra_headers: HeaderMap, + token_endpoint: String, + credential: Option<(Option<String>, Credential)>, + extra_oauth_params: HashMap<String, String>, +} + +/// [`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 { + client: Client, + token: Arc<Mutex<Option<Credential>>>, + 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<String>) -> Self { + Self { + client: Client::default(), + 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<String>) -> 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<String>, client_secret: String) -> Self { + self.init_params.credential = Some((client_id, client_secret.into())); + self + } + + /// Sets the HTTP client used for token requests. + pub fn with_client(mut self, client: Client) -> Self { + self.client = client; + 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<String, String>) -> Self { + self.init_params.extra_oauth_params.extend(params); + self + } + + pub(crate) fn from_config(cfg: &RestCatalogConfig) -> Result<Self> { + Ok(Self { + client: cfg.client(), + 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) -> Result<Box<dyn AuthSession>> { + Ok(self.build_session(self.init_params.clone())) + } + + async fn catalog_session( + &self, + props: &HashMap<String, String>, + ) -> Result<Arc<dyn AuthSession>> { + // The server config 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())); + } + + // Explicit property overrides merge ONTO the manager's options, so an + // injected manager keeps whatever a property doesn't override. + 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(), + }; + + Ok(Arc::from( + self.build_session(OAuth2Params { + extra_headers, + token_endpoint, + credential: credential_from_props(props) + .map(|(id, secret)| (id, secret.into())) + .or_else(|| self.init_params.credential.clone()), + extra_oauth_params, + }), + )) + } +} + +impl OAuth2Manager { + /// Builds the session matching the configured mode: + /// + /// - a `credential` yields a [`ClientCredentialsSession`] (its token cache + /// pre-seeded when a `token` is also set, and the token then takes + /// precedence over the credential); + /// - otherwise a [`StaticTokenSession`], which attaches the configured + /// token as-is — or nothing when none is set. + /// + /// Both share the manager's token cell, so a cached token survives the + /// config handshake. + fn build_session(&self, params: OAuth2Params) -> Box<dyn AuthSession> { Review Comment: This can be inlined to init_session if init_session takes in a prop map ########## crates/catalog/rest/src/catalog.rs: ########## @@ -341,6 +348,55 @@ impl RestCatalogConfig { .unwrap_or(false) } + /// The configured auth scheme: explicit `rest.auth.type` 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(&self) -> String { Review Comment: we should make this case-insensitive ########## crates/catalog/rest/src/auth/mod.rs: ########## @@ -0,0 +1,236 @@ +// 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 http::{HeaderMap, Method}; +use iceberg::Result; +pub use oauth2::OAuth2Manager; +use reqwest::Request; + +/// `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`, and +/// lives for the lifetime of the catalog. +#[async_trait] +pub trait AuthManager: Debug + Send + Sync { + /// Session used for the initial `/v1/config` handshake, built from the + /// user-supplied configuration. + /// + /// Returns a [`Box`]: an init session is used once and released, unlike + /// the shared [`AuthManager::catalog_session`]. + async fn init_session(&self) -> Result<Box<dyn AuthSession>>; + + /// 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, + props: &HashMap<String, String>, + ) -> Result<Arc<dyn AuthSession>>; +} + +/// An outgoing REST request being authenticated by an [`AuthSession`]. +/// +/// Wraps the request so authentication implementations depend only on the +/// stable `http` crate and standard types, not on the concrete HTTP client the +/// REST catalog uses internally. +pub struct AuthRequest<'a> { + inner: &'a mut Request, +} + +impl<'a> AuthRequest<'a> { + /// Wraps a request, e.g. to unit-test a custom [`AuthSession`]. + pub fn new(inner: &'a mut Request) -> Self { + 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 [`AuthRequestBody::Empty`] (empty-payload hash) and + /// [`AuthRequestBody::Buffered`], but not [`AuthRequestBody::Streaming`]. + pub fn body(&self) -> AuthRequestBody<'_> { + match self.inner.body() { + None => AuthRequestBody::Empty, + Some(body) => match body.as_bytes() { + Some(bytes) => AuthRequestBody::Buffered(bytes), + None => AuthRequestBody::Streaming, + }, + } + } +} + +/// The body of an [`AuthRequest`], as seen by authentication. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthRequestBody<'a> { Review Comment: nit: same here, I think `HttpRequestBody` will be a better name ########## crates/catalog/rest/src/auth/mod.rs: ########## @@ -0,0 +1,313 @@ +// 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 http::{HeaderMap, Method}; +use iceberg::Result; +pub use oauth2::OAuth2Manager; +use reqwest::Request; + +/// `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`, and +/// lives for the lifetime of the catalog. +#[async_trait] +pub trait AuthManager: Debug + Send + Sync { + /// Session used for the initial `/v1/config` handshake, built from the + /// user-supplied configuration. + /// + /// Returns a [`Box`]: an init session is used once and released, unlike + /// the shared [`AuthManager::catalog_session`]. + async fn init_session(&self) -> Result<Box<dyn AuthSession>>; Review Comment: I think Jannik's direction above is correct, eventually we will need to move auth manager outside of client. Should we just do it in this PR? My concern is mainly around custom auth manager. currently init_session doesn't take any properties, and a custom auth manager may need to rely on props from `RestCatalogConfig` to get token or credential to even initialize. Outh2Manager in this PR uses `from_config` but I don't think it's generalizable ########## crates/catalog/rest/src/auth/mod.rs: ########## @@ -0,0 +1,313 @@ +// 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 http::{HeaderMap, Method}; +use iceberg::Result; +pub use oauth2::OAuth2Manager; +use reqwest::Request; + +/// `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`, and +/// lives for the lifetime of the catalog. +#[async_trait] +pub trait AuthManager: Debug + Send + Sync { + /// Session used for the initial `/v1/config` handshake, built from the + /// user-supplied configuration. + /// + /// Returns a [`Box`]: an init session is used once and released, unlike + /// the shared [`AuthManager::catalog_session`]. + async fn init_session(&self) -> Result<Box<dyn AuthSession>>; Review Comment: Took another look at Java's API again, [initSession](https://github.com/apache/iceberg/blob/e938683ad436b83570414877d44228420a164457/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Manager.java#L75) and catalogSession both take in a client and a properties map. The main use case from what I understand is AuthManager can inherit the existing client from the catalog and modify necessary fields like AuthSession carried by the client, so this client can be used to do something else like refreshing token. I think supporting refreshing token can be a different PR, but can we investigate this more to make sure the existing design does not block that? ########## crates/catalog/rest/src/auth/mod.rs: ########## @@ -0,0 +1,236 @@ +// 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 http::{HeaderMap, Method}; +use iceberg::Result; +pub use oauth2::OAuth2Manager; +use reqwest::Request; + +/// `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`, and +/// lives for the lifetime of the catalog. +#[async_trait] +pub trait AuthManager: Debug + Send + Sync { + /// Session used for the initial `/v1/config` handshake, built from the + /// user-supplied configuration. + /// + /// Returns a [`Box`]: an init session is used once and released, unlike + /// the shared [`AuthManager::catalog_session`]. + async fn init_session(&self) -> Result<Box<dyn AuthSession>>; + + /// 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, + props: &HashMap<String, String>, + ) -> Result<Arc<dyn AuthSession>>; +} + +/// An outgoing REST request being authenticated by an [`AuthSession`]. +/// +/// Wraps the request so authentication implementations depend only on the +/// stable `http` crate and standard types, not on the concrete HTTP client the +/// REST catalog uses internally. +pub struct AuthRequest<'a> { Review Comment: nit: I think `HttpRequest` is a better name. it will be more consistent to `HttpClient`. and `HttpClient` will only work with `HttpRequest` not directly with the underlying `Request` ########## crates/catalog/rest/src/client.rs: ########## @@ -250,17 +140,16 @@ impl HttpClient { .headers(self.extra_headers.clone()) } - /// Executes the given `Request` and returns a `Response`. - pub async fn execute(&self, mut request: Request) -> Result<Response> { Review Comment: I think we should change the signature to take in `AuthRequest`/`HttpRequest`, `HttpClient` should be interacting with the wrapper `HttpRequest`. and `HttpClient` will use the underlying `Client` to work with the nested `Request` ########## crates/catalog/rest/src/auth/oauth2.rs: ########## @@ -0,0 +1,362 @@ +// 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 reqwest::{Client, Method}; +use tokio::sync::Mutex; + +use super::{AuthManager, AuthRequest, AuthSession}; +use crate::catalog::{ + REST_CATALOG_PROP_URI, RestCatalogConfig, credential_from_props, default_token_endpoint, + explicit_headers_from_props, +}; +use crate::types::{ErrorResponse, TokenResponse}; + +/// Per-phase OAuth2 parameters (init vs. post-handshake catalog phase). +#[derive(Clone)] +struct OAuth2Params { + extra_headers: HeaderMap, + token_endpoint: String, + credential: Option<(Option<String>, Credential)>, + extra_oauth_params: HashMap<String, String>, +} + +/// [`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 { + client: Client, + token: Arc<Mutex<Option<Credential>>>, + 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<String>) -> Self { + Self { + client: Client::default(), + 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<String>) -> 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<String>, client_secret: String) -> Self { + self.init_params.credential = Some((client_id, client_secret.into())); + self + } + + /// Sets the HTTP client used for token requests. + pub fn with_client(mut self, client: Client) -> Self { + self.client = client; + 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<String, String>) -> Self { + self.init_params.extra_oauth_params.extend(params); + self + } + + pub(crate) fn from_config(cfg: &RestCatalogConfig) -> Result<Self> { + Ok(Self { + client: cfg.client(), + 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) -> Result<Box<dyn AuthSession>> { + Ok(self.build_session(self.init_params.clone())) + } + + async fn catalog_session( + &self, + props: &HashMap<String, String>, + ) -> Result<Arc<dyn AuthSession>> { + // The server config 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())); + } + + // Explicit property overrides merge ONTO the manager's options, so an + // injected manager keeps whatever a property doesn't override. + 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(), + }; + + Ok(Arc::from( + self.build_session(OAuth2Params { + extra_headers, + token_endpoint, + credential: credential_from_props(props) + .map(|(id, secret)| (id, secret.into())) + .or_else(|| self.init_params.credential.clone()), + extra_oauth_params, + }), + )) + } +} + +impl OAuth2Manager { + /// Builds the session matching the configured mode: + /// + /// - a `credential` yields a [`ClientCredentialsSession`] (its token cache + /// pre-seeded when a `token` is also set, and the token then takes + /// precedence over the credential); + /// - otherwise a [`StaticTokenSession`], which attaches the configured + /// token as-is — or nothing when none is set. + /// + /// Both share the manager's token cell, so a cached token survives the + /// config handshake. + fn build_session(&self, params: OAuth2Params) -> Box<dyn AuthSession> { + match params.credential { + Some(credential) => Box::new(ClientCredentialsSession { + client: self.client.clone(), + token: self.token.clone(), + credential, + token_endpoint: params.token_endpoint, + extra_headers: params.extra_headers, + extra_oauth_params: params.extra_oauth_params, + }), + None => Box::new(StaticTokenSession { + token: self.token.clone(), + }), + } + } +} + +/// Attaches `token` as a `Authorization: Bearer <token>` header, marked +/// sensitive so `Debug`-formatted requests redact it. +fn attach_bearer(req: &mut AuthRequest<'_>, 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`] for a pre-configured bearer token: attaches it as-is and +/// cannot obtain a new one (there is no credential to exchange). +#[derive(Debug)] +struct StaticTokenSession { Review Comment: I think the only difference between StaticTokenSession and ClientCredentialsSession is how they behave when the existing token is None. Can we just have one Oauth2Session? Some thing like below ```rust struct OAuth2Session { token: Arc<Mutex<Option<Credential>>>, token_source: TokenSource, } enum TokenSource { StaticToken ClientCredentials(ClientCredentialsConfig), } struct ClientCredentialsConfig { client: Client, credential: (Option<String>, Credential), token_endpoint: String, extra_headers: HeaderMap, extra_oauth_params: HashMap<String, String>, } ``` TokenSource can be refactored to `TokenProvider` if needed in the future to further reduce duplicate code -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
