bitwarden_auth/token_management/
password_manager_token_handler.rs

1//! Token handler implementation for Bitwarden Password Manager authentication.
2
3use std::sync::{Arc, RwLock};
4
5use bitwarden_core::{
6    NotAuthenticatedError,
7    auth::{TokenHandler, login::LoginError},
8    client::login_method::LoginMethod,
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 PasswordManagerTokenHandler {
20    inner: Arc<RwLock<PasswordManagerTokenHandlerInner>>,
21}
22
23#[derive(Clone, Default)]
24struct PasswordManagerTokenHandlerInner {
25    access_token: Option<String>,
26    expires_on: Option<i64>,
27
28    refresh_token: Option<String>,
29
30    // The following are passed as optional as they are filled in when instantiating the
31    // middleware.
32    login_method: Option<Arc<RwLock<Option<Arc<LoginMethod>>>>>,
33    identity_config: Option<bitwarden_api_api::Configuration>,
34}
35
36impl TokenHandler for PasswordManagerTokenHandler {
37    fn initialize_middleware(
38        &self,
39        login_method: Arc<RwLock<Option<Arc<LoginMethod>>>>,
40        identity_config: bitwarden_api_api::Configuration,
41        _key_store: KeyStore<KeyIds>,
42    ) -> Arc<dyn reqwest_middleware::Middleware> {
43        {
44            let mut inner = self.inner.write().expect("RwLock is not poisoned");
45            inner.login_method = Some(login_method);
46            inner.identity_config = Some(identity_config);
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.refresh_token = refresh_token;
55        inner.expires_on = Some(Utc::now().timestamp() + expires_in as i64);
56    }
57}
58
59#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
60#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
61impl MiddlewareExt for PasswordManagerTokenHandler {
62    async fn get_token(&self) -> Result<Option<String>, LoginError> {
63        // We're not holding on to a lock for the duration of the token renewal, so if multiple
64        // requests come in at the same time when the token is expired, we may end up renewing the
65        // token multiple times. This is not ideal, but it's the behavior of the previous
66        // implementation. We should be able to introduce an async semaphore or something
67        // similar to prevent this if it becomes an issue in practice.
68        let inner = self.inner.read().expect("RwLock is not poisoned").clone();
69
70        // Validate the token, returning early if it's still valid.
71        if let Some(expires) = inner.expires_on
72            && Utc::now().timestamp() < expires - TOKEN_RENEW_MARGIN_SECONDS
73        {
74            return Ok(inner.access_token.clone());
75        }
76
77        // These should always be set by initialize_middleware before we get here, but we return an
78        // error if not.
79        let login_method = inner.login_method.ok_or(NotAuthenticatedError)?;
80        let identity_config = inner.identity_config.ok_or(NotAuthenticatedError)?;
81
82        let login_method = login_method
83            .read()
84            .expect("RwLock is not poisoned")
85            .clone()
86            .ok_or(NotAuthenticatedError)?;
87
88        let LoginMethod::User(user_login_method) = login_method.as_ref() else {
89            return Err(NotAuthenticatedError.into());
90        };
91
92        let (access_token, refresh_token, expires_in) =
93            bitwarden_core::auth::renew::renew_pm_token_sdk_managed(
94                inner.refresh_token,
95                user_login_method,
96                identity_config,
97            )
98            .await?;
99
100        self.set_tokens(access_token.clone(), refresh_token, expires_in);
101        Ok(Some(access_token))
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use std::sync::{Arc, RwLock};
108
109    use bitwarden_core::{
110        auth::TokenHandler,
111        client::login_method::{LoginMethod, UserLoginMethod},
112        key_management::KeyIds,
113    };
114    use bitwarden_crypto::{Kdf, KeyStore};
115    use wiremock::MockServer;
116
117    use super::*;
118    use crate::token_management::test_utils::*;
119
120    fn api_key_login_method() -> Arc<RwLock<Option<Arc<LoginMethod>>>> {
121        Arc::new(RwLock::new(Some(Arc::new(LoginMethod::User(
122            UserLoginMethod::ApiKey {
123                client_id: "test-client".to_string(),
124                client_secret: "test-secret".to_string(),
125                email: "[email protected]".to_string(),
126                kdf: Kdf::default_pbkdf2(),
127            },
128        )))))
129    }
130
131    #[tokio::test]
132    async fn attaches_existing_token_when_not_expired() {
133        let app_server = start_app_server().await;
134        let identity_server = MockServer::start().await;
135
136        let handler = PasswordManagerTokenHandler::default();
137        handler.set_tokens(
138            "original-token".to_string(),
139            Some("refresh".to_string()),
140            5000,
141        );
142        let client = build_client(handler.initialize_middleware(
143            api_key_login_method(),
144            identity_config(&identity_server.uri()),
145            KeyStore::<KeyIds>::default(),
146        ));
147
148        let auth = send_auth_request(&client, &app_server).await;
149        assert_eq!(auth.as_deref(), Some("Bearer original-token"));
150        assert_eq!(identity_server.received_requests().await.unwrap().len(), 0);
151        assert_eq!(app_server.received_requests().await.unwrap().len(), 1);
152    }
153
154    #[tokio::test]
155    async fn renews_expired_token() {
156        let app_server = start_app_server().await;
157        let identity_server = start_renewal_server("renewed-token").await;
158
159        let handler = PasswordManagerTokenHandler::default();
160        // expires_in=0 means the token is considered expired as it's less than the margin
161        handler.set_tokens(
162            "expired-token".to_string(),
163            Some("old-refresh".to_string()),
164            0,
165        );
166
167        let client = build_client(handler.initialize_middleware(
168            api_key_login_method(),
169            identity_config(&identity_server.uri()),
170            KeyStore::<KeyIds>::default(),
171        ));
172
173        let auth = send_auth_request(&client, &app_server).await;
174        assert_eq!(auth.as_deref(), Some("Bearer renewed-token"));
175        assert_eq!(identity_server.received_requests().await.unwrap().len(), 1);
176        assert_eq!(app_server.received_requests().await.unwrap().len(), 1);
177    }
178}
bitwarden_auth/token_management/password_manager_token_handler.rs

bitwarden_auth/token_management/
password_manager_token_handler.rs
