bitwarden_auth/token_management/
secrets_manager_token_handler.rs

1//! Token handler implementation for Bitwarden Secrets Manager authentication.
2
3use std::sync::{Arc, RwLock};
4
5use bitwarden_core::{
6    NotAuthenticatedError, OrganizationId,
7    auth::{TokenHandler, login::LoginError},
8    client::login_method::{LoginMethod, ServiceAccountLoginMethod},
9    key_management::KeyIds,
10};
11use bitwarden_crypto::KeyStore;
12use chrono::Utc;
13
14use super::middleware::{MiddlewareExt, MiddlewareWrapper};
15use crate::token_management::middleware::TOKEN_RENEW_MARGIN_SECONDS;
16
17/// Token handler for Bitwarden authentication.
18#[derive(Clone, Default)]
19pub struct SecretsManagerTokenHandler {
20    inner: Arc<RwLock<SecretsManagerTokenHandlerInner>>,
21}
22
23#[derive(Clone, Default)]
24struct SecretsManagerTokenHandlerInner {
25    access_token: Option<String>,
26    expires_on: Option<i64>,
27
28    // The following are passed as optional as they are filled in when instantiating the
29    // middleware.
30    login_method: Option<Arc<RwLock<Option<Arc<LoginMethod>>>>>,
31    identity_config: Option<bitwarden_api_api::Configuration>,
32    key_store: Option<KeyStore<KeyIds>>,
33}
34
35impl TokenHandler for SecretsManagerTokenHandler {
36    fn initialize_middleware(
37        &self,
38        login_method: Arc<RwLock<Option<Arc<LoginMethod>>>>,
39        identity_config: bitwarden_api_api::Configuration,
40        key_store: KeyStore<KeyIds>,
41    ) -> Arc<dyn reqwest_middleware::Middleware> {
42        {
43            let mut inner = self.inner.write().expect("RwLock is not poisoned");
44            inner.login_method = Some(login_method);
45            inner.identity_config = Some(identity_config);
46            inner.key_store = Some(key_store);
47        }
48        Arc::new(MiddlewareWrapper(self.clone()))
49    }
50
51    fn set_tokens(&self, access_token: String, _refresh_token: Option<String>, expires_in: u64) {
52        let mut inner = self.inner.write().expect("RwLock is not poisoned");
53        inner.access_token = Some(access_token);
54        inner.expires_on = Some(Utc::now().timestamp() + expires_in as i64);
55    }
56}
57
58impl SecretsManagerTokenHandler {
59    /// Get the organization ID associated with the current access token, if available.
60    pub fn get_access_token_organization(&self) -> Option<OrganizationId> {
61        let guard = self.inner.read().ok()?;
62        let login_method = guard.login_method.as_ref()?.read().ok()?;
63
64        match login_method.as_ref()?.as_ref() {
65            LoginMethod::User(_) => None,
66            LoginMethod::ServiceAccount(ServiceAccountLoginMethod::AccessToken {
67                organization_id,
68                ..
69            }) => Some(*organization_id),
70        }
71    }
72}
73
74#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
75#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
76impl MiddlewareExt for SecretsManagerTokenHandler {
77    async fn get_token(&self) -> Result<Option<String>, LoginError> {
78        // We're not holding on to a lock for the duration of the token renewal, so if multiple
79        // requests come in at the same time when the token is expired, we may end up renewing the
80        // token multiple times. This is not ideal, but it's the behavior of the previous
81        // implementation. We should be able to introduce an async semaphore or something
82        // similar to prevent this if it becomes an issue in practice.
83        let inner = self.inner.read().expect("RwLock is not poisoned").clone();
84
85        // Validate the token, returning early if it's still valid.
86        if let Some(expires) = inner.expires_on
87            && Utc::now().timestamp() < expires - TOKEN_RENEW_MARGIN_SECONDS
88        {
89            return Ok(inner.access_token.clone());
90        }
91
92        // These should always be set by initialize_middleware before we get here, but we return an
93        // error if not.
94        let login_method = inner.login_method.ok_or(NotAuthenticatedError)?;
95        let identity_config = inner.identity_config.ok_or(NotAuthenticatedError)?;
96        let key_store = inner.key_store.ok_or(NotAuthenticatedError)?;
97
98        let login_method = login_method
99            .read()
100            .expect("RwLock is not poisoned")
101            .clone()
102            .ok_or(NotAuthenticatedError)?;
103
104        let LoginMethod::ServiceAccount(service_account_login_method) = login_method.as_ref()
105        else {
106            return Err(NotAuthenticatedError.into());
107        };
108
109        let (access_token, refresh_token, expires_in) =
110            bitwarden_core::auth::renew::renew_sm_token_sdk_managed(
111                service_account_login_method,
112                identity_config,
113                key_store,
114            )
115            .await?;
116
117        self.set_tokens(access_token.clone(), refresh_token, expires_in);
118        Ok(Some(access_token))
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use std::{
125        str::FromStr,
126        sync::{Arc, RwLock},
127    };
128
129    use bitwarden_core::{
130        auth::{AccessToken, TokenHandler},
131        client::login_method::{LoginMethod, ServiceAccountLoginMethod},
132        key_management::KeyIds,
133    };
134    use bitwarden_crypto::KeyStore;
135    use wiremock::MockServer;
136
137    use super::*;
138    use crate::token_management::test_utils::*;
139
140    fn service_account_login_method() -> Arc<RwLock<Option<Arc<LoginMethod>>>> {
141        let access_token = AccessToken::from_str(
142            "0.ec2c1d46-6a4b-4751-a310-af9601317f2d.C2IgxjjLF7qSshsbwe8JGcbM075YXw:X8vbvA0bduihIDe/qrzIQQ==",
143        )
144        .unwrap();
145
146        Arc::new(RwLock::new(Some(Arc::new(LoginMethod::ServiceAccount(
147            ServiceAccountLoginMethod::AccessToken {
148                access_token,
149                organization_id: "00000000-0000-0000-0000-000000000001".parse().unwrap(),
150                state_file: None,
151            },
152        )))))
153    }
154
155    #[tokio::test]
156    async fn attaches_existing_token_when_not_expired() {
157        let app_server = start_app_server().await;
158        let identity_server = MockServer::start().await;
159
160        let handler = SecretsManagerTokenHandler::default();
161        handler.set_tokens("original-token".to_string(), None, 3600);
162
163        let client = build_client(handler.initialize_middleware(
164            service_account_login_method(),
165            identity_config(&identity_server.uri()),
166            KeyStore::<KeyIds>::default(),
167        ));
168
169        let auth = send_auth_request(&client, &app_server).await;
170        assert_eq!(auth.as_deref(), Some("Bearer original-token"));
171        assert_eq!(identity_server.received_requests().await.unwrap().len(), 0);
172        assert_eq!(app_server.received_requests().await.unwrap().len(), 1);
173    }
174
175    #[tokio::test]
176    async fn renews_expired_token() {
177        let app_server = start_app_server().await;
178        let identity_server = start_renewal_server("renewed-token").await;
179
180        let handler = SecretsManagerTokenHandler::default();
181        // expires_in=0 means the token is immediately considered expired
182        handler.set_tokens("expired-token".to_string(), None, 0);
183
184        let client = build_client(handler.initialize_middleware(
185            service_account_login_method(),
186            identity_config(&identity_server.uri()),
187            KeyStore::<KeyIds>::default(),
188        ));
189
190        let auth = send_auth_request(&client, &app_server).await;
191        assert_eq!(auth.as_deref(), Some("Bearer renewed-token"));
192        assert_eq!(identity_server.received_requests().await.unwrap().len(), 1);
193        assert_eq!(app_server.received_requests().await.unwrap().len(), 1);
194    }
195}
bitwarden_auth/token_management/secrets_manager_token_handler.rs

bitwarden_auth/token_management/
secrets_manager_token_handler.rs
