cognite/api/
authenticator.rs

1use crate::dto::utils::MaybeStringU64;
2use async_trait::async_trait;
3use futures_locks::RwLock;
4use reqwest::{
5    header::{HeaderMap, HeaderValue},
6    StatusCode,
7};
8use reqwest_middleware::ClientWithMiddleware;
9use serde::{Deserialize, Serialize};
10use std::time::{Duration, Instant};
11use std::{fmt::Display, sync::Arc};
12use thiserror::Error;
13
14/// Type of closure for a synchronous auth callback.
15type CustomAuthCallback =
16    dyn Fn(&mut HeaderMap, &ClientWithMiddleware) -> Result<(), AuthenticatorError> + Send + Sync;
17
18#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
19#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
20/// Trait for a custom authenticator. This should set the necessary headers in `headers` before each
21/// request. Note that this may be called from multiple places in parallel.
22pub trait CustomAuthenticator {
23    /// Set the required headers for authentication. This may use the provided
24    /// `client` to perform a request, if necessary. This will be called frequently, so
25    /// make sure it only makes external requests when needed.
26    ///
27    /// # Arguments
28    ///
29    /// * `headers` - Header map to modify.
30    /// * `client` - Client used to perform any external authentication requests.
31    async fn set_headers(
32        &self,
33        headers: &mut HeaderMap,
34        client: &ClientWithMiddleware,
35    ) -> Result<(), AuthenticatorError>;
36}
37
38/// Enumeration of the possible authentication methods available.
39#[derive(Clone)]
40pub enum AuthHeaderManager {
41    /// Authenticator that makes OIDC requests to obtain tokens.
42    OIDCToken(Arc<Authenticator>),
43    /// A fixed OIDC token
44    FixedToken(String),
45    /// An internal auth ticket.
46    AuthTicket(String),
47    /// A synchronous authentication method.
48    Custom(Arc<CustomAuthCallback>),
49    /// An async authentication method.
50    CustomAsync(Arc<dyn CustomAuthenticator + Send + Sync>),
51}
52
53impl AuthHeaderManager {
54    /// Set necessary headers in `headers`. This will sometimes request tokens from
55    /// the identity provider.
56    ///
57    /// # Arguments
58    ///
59    /// * `headers` - Request header collection.
60    /// * `client` - Reqwest client used to send authentication requests, if necessary.
61    pub async fn set_headers(
62        &self,
63        headers: &mut HeaderMap,
64        client: &ClientWithMiddleware,
65    ) -> Result<(), AuthenticatorError> {
66        match self {
67            AuthHeaderManager::OIDCToken(a) => {
68                let token = a.get_token(client).await?;
69                let auth_header_value =
70                    HeaderValue::from_str(&format!("Bearer {token}")).map_err(|e| {
71                        AuthenticatorError::internal_error(
72                            "Failed to set authorization bearer token".to_string(),
73                            Some(e.to_string()),
74                        )
75                    })?;
76                headers.insert("Authorization", auth_header_value);
77            }
78            AuthHeaderManager::FixedToken(token) => {
79                let auth_header_value =
80                    HeaderValue::from_str(&format!("Bearer {token}")).map_err(|e| {
81                        AuthenticatorError::internal_error(
82                            "Failed to set authorization bearer token".to_string(),
83                            Some(e.to_string()),
84                        )
85                    })?;
86                headers.insert("Authorization", auth_header_value);
87            }
88            AuthHeaderManager::AuthTicket(t) => {
89                let auth_ticket_header_value = HeaderValue::from_str(t).map_err(|e| {
90                    AuthenticatorError::internal_error(
91                        "Failed to set auth ticket".to_string(),
92                        Some(e.to_string()),
93                    )
94                })?;
95                headers.insert("auth-ticket", auth_ticket_header_value);
96            }
97            AuthHeaderManager::Custom(c) => c(headers, client)?,
98            AuthHeaderManager::CustomAsync(c) => c.set_headers(headers, client).await?,
99        }
100        Ok(())
101    }
102}
103
104/// Configuration for authentication using the OIDC authenticator
105pub struct AuthenticatorConfig {
106    /// Service principal client ID.
107    pub client_id: String,
108    /// IdP token URL.
109    pub token_url: String,
110    /// Service principal client secret.
111    pub secret: String,
112    /// Optional resource.
113    pub resource: Option<String>,
114    /// Optional audience.
115    pub audience: Option<String>,
116    /// Optional space separate list of scopes.
117    pub scopes: Option<String>,
118    /// Optional default token expiry time, in seconds.
119    /// If this is set, the authenticator will fall back on this if
120    /// the identity provider returns a token response without `expires_in`.
121    /// If this is not set, and `expires_in` is missing, the authenticator will return an error.
122    pub default_expires_in: Option<u64>,
123}
124
125#[derive(Serialize, Deserialize, Debug)]
126struct AuthenticatorRequest {
127    client_id: String,
128    client_secret: String,
129    resource: Option<String>,
130    audience: Option<String>,
131    scope: Option<String>,
132    grant_type: String,
133}
134
135impl AuthenticatorRequest {
136    fn new(config: AuthenticatorConfig) -> AuthenticatorRequest {
137        AuthenticatorRequest {
138            client_id: config.client_id,
139            client_secret: config.secret,
140            grant_type: "client_credentials".to_string(),
141            resource: config.resource,
142            audience: config.audience,
143            scope: config.scopes,
144        }
145    }
146}
147
148#[derive(Serialize, Deserialize, Debug)]
149struct AuthenticatorResponse {
150    access_token: String,
151    expires_in: Option<MaybeStringU64>,
152}
153
154#[derive(Serialize, Deserialize, Debug, Error)]
155/// Error from an authenticator request.
156pub struct AuthenticatorError {
157    /// Error message
158    pub error: String,
159    /// Detailed error description.
160    pub error_description: Option<String>,
161    /// Error URI.
162    pub error_uri: Option<String>,
163}
164
165impl AuthenticatorError {
166    /// Create an authenticator error from message and description.
167    ///
168    /// # Arguments
169    ///
170    /// * `error` - Short error message
171    /// * `error_description` - Detailed error description.
172    pub fn internal_error(error: String, error_description: Option<String>) -> Self {
173        Self {
174            error,
175            error_description,
176            error_uri: None,
177        }
178    }
179}
180
181impl Display for AuthenticatorError {
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        write!(f, "{}", self.error,)?;
184        if let Some(error_description) = &self.error_description {
185            write!(f, ": {error_description}")?;
186        }
187        if let Some(error_uri) = &self.error_uri {
188            write!(f, " ({error_uri})")?;
189        }
190        Ok(())
191    }
192}
193
194struct AuthenticatorState {
195    last_token: Option<String>,
196    current_token_expiry: Instant,
197}
198
199/// Result from getting a token, including expiry time.
200pub struct AuthenticatorResult {
201    /// The token string.
202    token: String,
203    /// The time when the token will expire.
204    expiry: Instant,
205}
206
207/// Simple OIDC authenticator.
208pub struct Authenticator {
209    req: AuthenticatorRequest,
210    state: RwLock<AuthenticatorState>,
211    token_url: String,
212    default_expires_in: Option<Duration>,
213}
214
215impl AuthenticatorResult {
216    /// Get the token string.
217    pub fn token(&self) -> &str {
218        &self.token
219    }
220
221    /// Consume self and get the token string.
222    pub fn into_token(self) -> String {
223        self.token
224    }
225
226    /// Get the expiry time.
227    pub fn expiry(&self) -> Instant {
228        self.expiry
229    }
230}
231
232impl Authenticator {
233    /// Create a new authenticator with given config.
234    ///
235    /// # Arguments
236    ///
237    /// * `config` - Authenticator configuration.
238    pub fn new(config: AuthenticatorConfig) -> Authenticator {
239        Authenticator {
240            token_url: config.token_url.clone(),
241            default_expires_in: config.default_expires_in.map(Duration::from_secs),
242            req: AuthenticatorRequest::new(config),
243            state: RwLock::new(AuthenticatorState {
244                last_token: None,
245                current_token_expiry: Instant::now(),
246            }),
247        }
248    }
249
250    async fn request_token(
251        &self,
252        client: &ClientWithMiddleware,
253    ) -> Result<AuthenticatorResult, AuthenticatorError> {
254        let response = client
255            .post(&self.token_url)
256            .form(&self.req)
257            .send()
258            .await
259            .map_err(|e| {
260                AuthenticatorError::internal_error(
261                    "Something went wrong when sending the request".to_string(),
262                    Some(e.to_string()),
263                )
264            })?;
265
266        let status = response.status();
267
268        let start = Instant::now();
269
270        let response = response.text().await.map_err(|e| {
271            AuthenticatorError::internal_error(
272                "Failed to receive response contents".to_owned(),
273                Some(e.to_string()),
274            )
275        })?;
276
277        if status != StatusCode::OK {
278            return match serde_json::from_str(&response) {
279                Ok(e) => Err(e),
280                Err(e) => Err(AuthenticatorError::internal_error(
281                    format!("Something went wrong (status: {status}), but the response error couldn't be deserialized. Raw response: {response}")
282                    , Some(e.to_string())))
283            };
284        }
285
286        let response: AuthenticatorResponse = serde_json::from_str(&response).map_err(|e| {
287            AuthenticatorError::internal_error(
288                "Failed to deserialize response from OAuth endpoint".to_string(),
289                Some(e.to_string()),
290            )
291        })?;
292
293        let token = response.access_token;
294        let Some(expires_in) = response
295            .expires_in
296            // Subtract 60 as a buffer. We do retry on 401s, but it's best to renew the
297            // token before it expires. If for whatever reason expires_in is less than 60,
298            // we will just always renew before sending a request. We won't (hopefully)
299            // get an infinite loop.
300            .map(|m| Duration::from_secs(m.0.saturating_sub(60)))
301            .or(self.default_expires_in)
302        else {
303            return Err(AuthenticatorError::internal_error(
304                "Missing expires_in in response, and no default expiration configured".to_owned(),
305                None,
306            ));
307        };
308
309        Ok(AuthenticatorResult {
310            token,
311            expiry: start + expires_in,
312        })
313    }
314
315    /// Get a token. This will only fetch a new token if it is about
316    /// to expire (will expire in the next 60 seconds). This also
317    /// returns when the next token will be requested. This is the time
318    /// when the authenticator will refresh the token, so the actual
319    /// expiry time minus 60 seconds.
320    ///
321    /// # Arguments
322    ///
323    /// * `client` - Reqwest client to use for requests to the IdP.
324    pub async fn get_token_with_expiry(
325        &self,
326        client: &ClientWithMiddleware,
327    ) -> Result<AuthenticatorResult, AuthenticatorError> {
328        let now = Instant::now();
329        {
330            let state = &*self.state.read().await;
331            if let Some(last) = &state.last_token {
332                if state.current_token_expiry > now {
333                    return Ok(AuthenticatorResult {
334                        token: last.clone(),
335                        expiry: state.current_token_expiry,
336                    });
337                }
338            }
339        }
340
341        // If the token is expired, release the read lock and try to acquire a write lock.
342        let mut write = self.state.write().await;
343
344        // Need to check here too, in case we were blocked in this write lock by another thread
345        // fetching the token.
346        if let Some(last) = &write.last_token {
347            if write.current_token_expiry > now {
348                return Ok(AuthenticatorResult {
349                    token: last.clone(),
350                    expiry: write.current_token_expiry,
351                });
352            }
353        }
354
355        match self.request_token(client).await {
356            Ok(response) => {
357                write.current_token_expiry = response.expiry;
358                write.last_token = Some(response.token.clone());
359                Ok(response)
360            }
361            Err(e) => Err(e),
362        }
363    }
364
365    /// Get a token. This will only fetch a new token if it is about
366    /// to expire (will expire in the next 60 seconds).
367    ///
368    /// # Arguments
369    ///
370    /// * `client` - Reqwest client to use for requests to the IdP.
371    pub async fn get_token(
372        &self,
373        client: &ClientWithMiddleware,
374    ) -> Result<String, AuthenticatorError> {
375        Ok(self.get_token_with_expiry(client).await?.token)
376    }
377}
cognite/api/authenticator.rs

cognite/api/
authenticator.rs
