pidge_client/auth/
backend.rs

1//! Pluggable token persistence for [`super::AuthClient`].
2//!
3//! The CLI resolves each account's backend (OS keychain or plaintext file)
4//! from `config.yaml`; that behaviour lives in [`LocalBackend`] and stays the
5//! default. Hosted consumers (e.g. the remote MCP server) implement
6//! [`TokenBackend`] themselves so tokens can live in a per-user secret store
7//! instead; `AuthClient` and `GraphClient` don't care which.
8
9use async_trait::async_trait;
10use pidge_core::TokenStorage;
11
12use crate::auth::jwt;
13use crate::auth::token_store::TokenStore;
14use crate::auth::tokens::TokenSet;
15use crate::error::ClientError;
16
17/// Where an account's [`TokenSet`] is loaded from and saved to.
18///
19/// Implementations must be safe to share across tasks; `AuthClient` holds one
20/// behind an `Arc`.
21#[async_trait]
22pub trait TokenBackend: Send + Sync {
23    /// Load the stored tokens for `email`. `Ok(None)` means "no session".
24    async fn load(&self, email: &str) -> Result<Option<TokenSet>, ClientError>;
25
26    /// Persist `tokens` for `email`, replacing whatever was stored before.
27    async fn save(&self, email: &str, tokens: &TokenSet) -> Result<(), ClientError>;
28
29    /// Hook invoked with every access token handed out by
30    /// [`super::AuthClient::get_valid_token`]. The default is a no-op; the
31    /// CLI backend uses it to backfill account metadata.
32    fn on_access_token(&self, _email: &str, _access_token: &str) {}
33}
34
35/// The CLI's backend: consults `config.yaml` for the account's
36/// [`TokenStorage`] and dispatches to the keychain or file store.
37#[derive(Debug, Default, Clone, Copy)]
38pub struct LocalBackend;
39
40impl LocalBackend {
41    /// Resolve the token storage backend for an email by consulting `config.yaml`.
42    /// Falls back to [`TokenStorage::Keychain`] (the default) if the config can't
43    /// be read or the account isn't listed yet.
44    fn storage_for(email: &str) -> TokenStorage {
45        pidge_core::Config::load()
46            .ok()
47            .and_then(|c| c.find(email).map(|a| a.storage))
48            .unwrap_or_default()
49    }
50}
51
52#[async_trait]
53impl TokenBackend for LocalBackend {
54    async fn load(&self, email: &str) -> Result<Option<TokenSet>, ClientError> {
55        TokenStore::load(email, Self::storage_for(email))
56    }
57
58    async fn save(&self, email: &str, tokens: &TokenSet) -> Result<(), ClientError> {
59        TokenStore::save(email, tokens, Self::storage_for(email))
60    }
61
62    /// Opportunistic backfill: accounts added before pidge requested the
63    /// `openid` scope have an empty tenant_id in config. Microsoft Graph
64    /// access tokens are JWTs that carry the `tid` claim, so we can fix
65    /// this once per such account on the next Graph call without any
66    /// user action. Silent on any failure: cosmetic, not a correctness
67    /// requirement.
68    fn on_access_token(&self, email: &str, access_token: &str) {
69        let Ok(mut config) = pidge_core::Config::load() else {
70            return;
71        };
72        let Some(existing) = config.find(email).cloned() else {
73            return;
74        };
75        if !existing.tenant_id.is_empty() {
76            return;
77        }
78        let Some(tid) = jwt::extract_tenant_id(access_token) else {
79            return;
80        };
81        if tid.is_empty() {
82            return;
83        }
84        let mut updated = existing;
85        updated.tenant_id = tid;
86        config.add_account(updated);
87        let _ = config.save();
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use std::collections::HashMap;
94    use std::sync::{Arc, Mutex};
95
96    use chrono::{Duration, Utc};
97    use wiremock::matchers::{method, path};
98    use wiremock::{Mock, MockServer, ResponseTemplate};
99
100    use super::*;
101    use crate::auth::AuthClient;
102
103    /// The kind of backend a hosted consumer would write: tokens live in a
104    /// map instead of the keychain, and the refresh path must round-trip
105    /// through it.
106    #[derive(Default)]
107    struct MemoryBackend {
108        tokens: Mutex<HashMap<String, TokenSet>>,
109        seen_access_tokens: Mutex<Vec<String>>,
110    }
111
112    #[async_trait]
113    impl TokenBackend for MemoryBackend {
114        async fn load(&self, email: &str) -> Result<Option<TokenSet>, ClientError> {
115            Ok(self.tokens.lock().unwrap().get(email).cloned())
116        }
117
118        async fn save(&self, email: &str, tokens: &TokenSet) -> Result<(), ClientError> {
119            self.tokens
120                .lock()
121                .unwrap()
122                .insert(email.to_string(), tokens.clone());
123            Ok(())
124        }
125
126        fn on_access_token(&self, _email: &str, access_token: &str) {
127            self.seen_access_tokens
128                .lock()
129                .unwrap()
130                .push(access_token.to_string());
131        }
132    }
133
134    #[tokio::test]
135    async fn custom_backend_receives_refreshed_tokens() {
136        let server = MockServer::start().await;
137        Mock::given(method("POST"))
138            .and(path("/oauth2/v2.0/token"))
139            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
140                "access_token": "NEW_AT",
141                "refresh_token": "NEW_RT",
142                "expires_in": 3600
143            })))
144            .expect(1)
145            .mount(&server)
146            .await;
147
148        let backend = Arc::new(MemoryBackend::default());
149        backend
150            .save(
151                "jane@example.com",
152                &TokenSet {
153                    access_token: "OLD_AT".into(),
154                    refresh_token: "OLD_RT".into(),
155                    expires_at: Utc::now() - Duration::seconds(60),
156                    scope: None,
157                },
158            )
159            .await
160            .unwrap();
161
162        let auth = AuthClient::for_test("cid", server.uri()).with_backend(backend.clone());
163        let token = auth.get_valid_token("jane@example.com").await.unwrap();
164
165        assert_eq!(token, "NEW_AT");
166        let stored = backend.load("jane@example.com").await.unwrap().unwrap();
167        assert_eq!(stored.refresh_token, "NEW_RT");
168        assert_eq!(*backend.seen_access_tokens.lock().unwrap(), vec!["NEW_AT"]);
169    }
170
171    /// Microsoft's token endpoint as it treats an account consented before
172    /// `MailboxSettings.ReadWrite` was added: a refresh asking for it is
173    /// refused (`invalid_grant`), anything else succeeds.
174    struct OldConsent;
175
176    impl wiremock::Respond for OldConsent {
177        fn respond(&self, req: &wiremock::Request) -> ResponseTemplate {
178            let form: HashMap<String, String> = url::form_urlencoded::parse(&req.body)
179                .into_owned()
180                .collect();
181            if form["scope"].contains("MailboxSettings") {
182                return ResponseTemplate::new(400).set_body_json(serde_json::json!({
183                    "error": "invalid_grant",
184                    "error_description": "AADSTS65001: consent required"
185                }));
186            }
187            ResponseTemplate::new(200).set_body_json(serde_json::json!({
188                "access_token": "NEW_AT",
189                "refresh_token": "NEW_RT",
190                "expires_in": 3600,
191                "scope": "https://graph.microsoft.com/Mail.ReadWrite https://graph.microsoft.com/User.Read"
192            }))
193        }
194    }
195
196    #[tokio::test]
197    async fn a_token_consented_before_a_new_scope_keeps_refreshing() {
198        let server = MockServer::start().await;
199        Mock::given(method("POST"))
200            .and(path("/oauth2/v2.0/token"))
201            .respond_with(OldConsent)
202            .mount(&server)
203            .await;
204        let backend = Arc::new(MemoryBackend::default());
205        let expired = |scope: Option<&str>| TokenSet {
206            access_token: "OLD_AT".into(),
207            refresh_token: "OLD_RT".into(),
208            expires_at: Utc::now() - Duration::seconds(60),
209            scope: scope.map(str::to_string),
210        };
211        let auth = AuthClient::for_test("cid", server.uri()).with_backend(backend.clone());
212        assert!(auth.scope().contains("MailboxSettings.ReadWrite"));
213
214        // Stored before grants were recorded, then once more with a grant
215        // that lacks the scope: both refresh with the original set.
216        for scope in [None, Some("Mail.ReadWrite User.Read")] {
217            backend
218                .save("jane@example.com", &expired(scope))
219                .await
220                .unwrap();
221            assert_eq!(
222                auth.get_valid_token("jane@example.com").await.unwrap(),
223                "NEW_AT"
224            );
225            let stored = backend.load("jane@example.com").await.unwrap().unwrap();
226            assert!(stored.scope.unwrap().contains("Mail.ReadWrite"));
227        }
228
229        // A sign-in that granted the new scope refreshes with it.
230        backend
231            .save(
232                "jane@example.com",
233                &expired(Some("Mail.ReadWrite MailboxSettings.ReadWrite")),
234            )
235            .await
236            .unwrap();
237        let err = auth.get_valid_token("jane@example.com").await.unwrap_err();
238        assert!(matches!(err, ClientError::SessionExpired { .. }), "{err:?}");
239        let asked: Vec<String> = server
240            .received_requests()
241            .await
242            .unwrap()
243            .iter()
244            .map(|r| String::from_utf8_lossy(&r.body).into_owned())
245            .collect();
246        assert_eq!(asked.len(), 3);
247        assert!(!asked[0].contains("MailboxSettings") && !asked[1].contains("MailboxSettings"));
248        assert!(asked[2].contains("MailboxSettings"));
249    }
250
251    #[tokio::test]
252    async fn fresh_token_is_returned_without_touching_microsoft() {
253        let server = MockServer::start().await;
254        Mock::given(method("POST"))
255            .and(path("/oauth2/v2.0/token"))
256            .respond_with(ResponseTemplate::new(500))
257            .expect(0)
258            .mount(&server)
259            .await;
260
261        let backend = Arc::new(MemoryBackend::default());
262        backend
263            .save(
264                "jane@example.com",
265                &TokenSet {
266                    access_token: "FRESH".into(),
267                    refresh_token: "RT".into(),
268                    expires_at: Utc::now() + Duration::seconds(3600),
269                    scope: None,
270                },
271            )
272            .await
273            .unwrap();
274
275        let auth = AuthClient::for_test("cid", server.uri()).with_backend(backend);
276        assert_eq!(
277            auth.get_valid_token("jane@example.com").await.unwrap(),
278            "FRESH"
279        );
280    }
281
282    #[tokio::test]
283    async fn missing_session_maps_to_session_expired() {
284        let server = MockServer::start().await;
285        let auth = AuthClient::for_test("cid", server.uri())
286            .with_backend(Arc::new(MemoryBackend::default()));
287        let err = auth
288            .get_valid_token("nobody@example.com")
289            .await
290            .unwrap_err();
291        assert!(matches!(err, ClientError::SessionExpired { .. }));
292    }
293}
pidge_client/auth/backend.rs

pidge_client/auth/
backend.rs
