iceberg_catalog_rest/auth/
oauth2.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::collections::HashMap;
19use std::fmt::{Debug, Formatter};
20use std::sync::Arc;
21
22use async_trait::async_trait;
23use http::StatusCode;
24use iceberg::sensitive::SensitiveString;
25use iceberg::{Error, ErrorKind, Result};
26use reqwest::header::HeaderMap;
27use tokio::sync::Mutex;
28
29use super::{AuthManager, AuthSession, HttpRequest};
30use crate::catalog::{
31    REST_CATALOG_PROP_URI, RestCatalogConfig, credential_from_props, default_token_endpoint,
32    explicit_headers_from_props,
33};
34use crate::client::HttpClient;
35use crate::types::{ErrorResponse, TokenResponse};
36
37/// The manager's own OAuth2 options, which properties are merged onto.
38struct OAuth2Params {
39    extra_headers: HeaderMap,
40    token_endpoint: String,
41    credential: Option<(Option<String>, SensitiveString)>,
42    extra_oauth_params: HashMap<String, String>,
43}
44
45/// [`AuthManager`] implementing the OAuth2 client-credentials flow used by
46/// Iceberg REST catalogs.
47///
48/// A configured `token` is used directly; otherwise `credential` is exchanged
49/// for a token at the token endpoint and cached. The cached token is shared
50/// across sessions so it survives the config handshake.
51pub struct OAuth2Manager {
52    token: Arc<Mutex<Option<SensitiveString>>>,
53    init_params: OAuth2Params,
54    /// True when the token endpoint was derived from the catalog URI (not
55    /// explicitly configured): it is then recomputed from the merged URI in
56    /// [`Self::catalog_session`], since `/v1/config` may override the URI.
57    endpoint_is_default: bool,
58}
59
60impl OAuth2Manager {
61    /// Creates a manager exchanging credentials at `token_endpoint`, with no
62    /// token or credential configured. Combine with the `with_*` methods:
63    ///
64    /// ```rust,ignore
65    /// let manager = OAuth2Manager::new("https://auth.example.com/v1/oauth/tokens")
66    ///     .with_credential(Some("client-id".into()), "client-secret".into());
67    /// ```
68    pub fn new(token_endpoint: impl Into<String>) -> Self {
69        Self {
70            token: Arc::new(Mutex::new(None)),
71            init_params: OAuth2Params {
72                extra_headers: HeaderMap::new(),
73                token_endpoint: token_endpoint.into(),
74                credential: None,
75                // Same default as the configuration path: the catalog scope.
76                extra_oauth_params: HashMap::from([("scope".to_string(), "catalog".to_string())]),
77            },
78            endpoint_is_default: false,
79        }
80    }
81
82    /// Sets a bearer token used directly (takes precedence over `credential`).
83    pub fn with_token(mut self, token: impl Into<String>) -> Self {
84        self.token = Arc::new(Mutex::new(Some(SensitiveString::from(token.into()))));
85        self
86    }
87
88    /// Sets the client credential exchanged for a token at the token endpoint.
89    pub fn with_credential(mut self, client_id: Option<String>, client_secret: String) -> Self {
90        self.init_params.credential = Some((client_id, client_secret.into()));
91        self
92    }
93
94    /// Sets extra headers sent with token requests.
95    pub fn with_extra_headers(mut self, headers: HeaderMap) -> Self {
96        self.init_params.extra_headers = headers;
97        self
98    }
99
100    /// Adds extra OAuth2 form parameters (e.g. `scope`, `audience`), merged
101    /// onto the defaults: provide a `scope` entry to replace the default
102    /// `catalog` scope.
103    pub fn with_extra_oauth_params(mut self, params: HashMap<String, String>) -> Self {
104        self.init_params.extra_oauth_params.extend(params);
105        self
106    }
107
108    pub(crate) fn from_config(cfg: &RestCatalogConfig) -> Result<Self> {
109        Ok(Self {
110            token: Arc::new(Mutex::new(cfg.token().map(SensitiveString::from))),
111            init_params: OAuth2Params {
112                extra_headers: cfg.extra_headers()?,
113                token_endpoint: cfg.get_token_endpoint(),
114                credential: cfg.credential().map(|(id, secret)| (id, secret.into())),
115                extra_oauth_params: cfg.extra_oauth_params(),
116            },
117            endpoint_is_default: cfg.explicit_oauth2_server_uri().is_none(),
118        })
119    }
120}
121
122impl Debug for OAuth2Manager {
123    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("OAuth2Manager")
125            .field("token_endpoint", &self.init_params.token_endpoint)
126            .finish_non_exhaustive()
127    }
128}
129
130#[async_trait]
131impl AuthManager for OAuth2Manager {
132    async fn init_session(
133        &self,
134        client: &HttpClient,
135        props: &HashMap<String, String>,
136    ) -> Result<Box<dyn AuthSession>> {
137        Ok(Box::new(self.session_from(client, props).await?))
138    }
139
140    async fn catalog_session(
141        &self,
142        client: &HttpClient,
143        props: &HashMap<String, String>,
144    ) -> Result<Arc<dyn AuthSession>> {
145        Ok(Arc::new(self.session_from(client, props).await?))
146    }
147}
148
149impl OAuth2Manager {
150    /// Builds a session from the manager's options with `props` merged onto
151    /// them, so an injected manager keeps whatever a property doesn't
152    /// override. The manager's token cell is shared with every session it
153    /// builds, so a token cached during the handshake survives it.
154    async fn session_from(
155        &self,
156        client: &HttpClient,
157        props: &HashMap<String, String>,
158    ) -> Result<OAuth2Session> {
159        // The properties may carry a new token (or restate the user's).
160        if let Some(token) = props.get("token") {
161            *self.token.lock().await = Some(SensitiveString::from(token.clone()));
162        }
163
164        let mut extra_headers = self.init_params.extra_headers.clone();
165        extra_headers.extend(explicit_headers_from_props(props)?);
166
167        let mut extra_oauth_params = self.init_params.extra_oauth_params.clone();
168        for key in ["scope", "audience", "resource"] {
169            if let Some(value) = props.get(key) {
170                extra_oauth_params.insert(key.to_string(), value.to_string());
171            }
172        }
173
174        let token_endpoint = match props.get("oauth2-server-uri") {
175            Some(uri) if !uri.is_empty() => uri.clone(),
176            // A default endpoint follows the merged catalog URI (which
177            // `/v1/config` may have overridden); explicit ones are kept.
178            _ if self.endpoint_is_default => props
179                .get(REST_CATALOG_PROP_URI)
180                .map(|uri| default_token_endpoint(uri))
181                .unwrap_or_else(|| self.init_params.token_endpoint.clone()),
182            _ => self.init_params.token_endpoint.clone(),
183        };
184
185        let credential = credential_from_props(props)
186            .map(|(id, secret)| (id, secret.into()))
187            .or_else(|| self.init_params.credential.clone());
188
189        Ok(OAuth2Session {
190            token: self.token.clone(),
191            // A configured token takes precedence over the credential: the
192            // token cell is pre-seeded, and the credential only comes into
193            // play once that token is gone.
194            token_source: match credential {
195                Some(credential) => {
196                    TokenSource::ClientCredentials(Box::new(ClientCredentialsConfig {
197                        client: client.clone(),
198                        credential,
199                        token_endpoint,
200                        extra_headers,
201                        extra_oauth_params,
202                    }))
203                }
204                None => TokenSource::StaticToken,
205            },
206        })
207    }
208}
209
210/// Attaches `token` as a `Authorization: Bearer <token>` header, marked
211/// sensitive so `Debug`-formatted requests redact it.
212fn attach_bearer(req: &mut HttpRequest, token: &SensitiveString) -> Result<()> {
213    let mut value: http::HeaderValue =
214        format!("Bearer {}", token.expose()).parse().map_err(|e| {
215            Error::new(
216                ErrorKind::DataInvalid,
217                "Invalid token received from catalog server!",
218            )
219            .with_source(e)
220        })?;
221    value.set_sensitive(true);
222    req.headers_mut().insert(http::header::AUTHORIZATION, value);
223    Ok(())
224}
225
226/// [`AuthSession`] attaching an OAuth2 bearer token.
227///
228/// The token is a configured one (which replaces whatever the cell holds), a
229/// token cached by an earlier session (the cell is shared with the owning
230/// [`OAuth2Manager`]), or — with [`TokenSource::ClientCredentials`] — one
231/// exchanged for the credential on demand.
232///
233/// # TODO: Support automatic token refreshing.
234struct OAuth2Session {
235    token: Arc<Mutex<Option<SensitiveString>>>,
236    token_source: TokenSource,
237}
238
239/// How an [`OAuth2Session`] obtains a token once none is cached.
240enum TokenSource {
241    /// Nothing to obtain: the session attaches the configured token, or no
242    /// authentication at all when there is none.
243    StaticToken,
244    /// The credential is exchanged for a token at the token endpoint.
245    ClientCredentials(Box<ClientCredentialsConfig>),
246}
247
248struct ClientCredentialsConfig {
249    client: HttpClient,
250    credential: (Option<String>, SensitiveString),
251    token_endpoint: String,
252    extra_headers: HeaderMap,
253    extra_oauth_params: HashMap<String, String>,
254}
255
256impl Debug for OAuth2Session {
257    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
258        let mut out = f.debug_struct("OAuth2Session");
259        if let TokenSource::ClientCredentials(config) = &self.token_source {
260            out.field("token_endpoint", &config.token_endpoint);
261        }
262        out.finish_non_exhaustive()
263    }
264}
265
266impl ClientCredentialsConfig {
267    async fn exchange_credential_for_token(&self) -> Result<String> {
268        let (client_id, client_secret) = &self.credential;
269
270        let mut params = HashMap::with_capacity(4);
271        params.insert("grant_type", "client_credentials");
272        if let Some(client_id) = client_id {
273            params.insert("client_id", client_id);
274        }
275        params.insert("client_secret", client_secret.expose());
276        params.extend(
277            self.extra_oauth_params
278                .iter()
279                .map(|(k, v)| (k.as_str(), v.as_str())),
280        );
281
282        let response = self
283            .client
284            .post_form(&self.token_endpoint, &self.extra_headers, &params)
285            .await?;
286        let status = response.status();
287        let body = response.body();
288
289        let auth_res: TokenResponse = if status == StatusCode::OK {
290            Ok(serde_json::from_slice(body).map_err(|e| {
291                Error::new(
292                    ErrorKind::Unexpected,
293                    "Failed to parse response from rest catalog server!",
294                )
295                .with_context("operation", "auth")
296                .with_context("url", self.token_endpoint.clone())
297                .with_context("json", String::from_utf8_lossy(body))
298                .with_source(e)
299            })?)
300        } else {
301            let e: ErrorResponse = serde_json::from_slice(body).map_err(|e| {
302                Error::new(ErrorKind::Unexpected, "Received unexpected response")
303                    .with_context("code", status.to_string())
304                    .with_context("operation", "auth")
305                    .with_context("url", self.token_endpoint.clone())
306                    .with_context("json", String::from_utf8_lossy(body))
307                    .with_source(e)
308            })?;
309            Err(Error::from(e))
310        }?;
311        Ok(auth_res.access_token)
312    }
313}
314
315#[async_trait]
316impl AuthSession for OAuth2Session {
317    /// Uses the cached token when present; otherwise exchanges the credential
318    /// for one, caches it, then uses it. Without a credential and without a
319    /// token, no authentication is attached.
320    async fn authenticate(&self, req: &mut HttpRequest) -> Result<()> {
321        // The lock is held across the exchange: waiters reuse a successful
322        // result, and retry themselves after a failure.
323        let token = {
324            let mut token = self.token.lock().await;
325            match (&*token, &self.token_source) {
326                (Some(token), _) => Some(token.clone()),
327                (None, TokenSource::StaticToken) => None,
328                (None, TokenSource::ClientCredentials(config)) => {
329                    let new_token =
330                        SensitiveString::from(config.exchange_credential_for_token().await?);
331                    *token = Some(new_token.clone());
332                    Some(new_token)
333                }
334            }
335        };
336
337        match token {
338            Some(token) => attach_bearer(req, &token),
339            None => Ok(()),
340        }
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::collections::HashMap;
347
348    use reqwest::Client;
349
350    use super::*;
351    use crate::RestCatalogConfig;
352
353    fn test_client() -> HttpClient {
354        HttpClient::new(
355            &RestCatalogConfig::builder()
356                .uri("http://localhost".to_string())
357                .build(),
358        )
359        .unwrap()
360    }
361
362    #[tokio::test]
363    async fn test_static_token_session_attaches_token() {
364        // Token-only config: the token is attached as-is.
365        let manager = OAuth2Manager::new("http://localhost/unused").with_token("tok-static");
366        let session = manager
367            .init_session(&test_client(), &HashMap::new())
368            .await
369            .unwrap();
370
371        let mut req = HttpRequest::new(
372            Client::new()
373                .get("https://rest.example.com/v1/config")
374                .build()
375                .unwrap(),
376        );
377        session.authenticate(&mut req).await.unwrap();
378        assert_eq!(
379            req.headers().get("authorization").unwrap(),
380            "Bearer tok-static"
381        );
382    }
383}
iceberg_catalog_rest/auth/oauth2.rs

iceberg_catalog_rest/auth/
oauth2.rs
