mcp_oauth/store/
json_file.rs

1//! JSON-file-backed implementations of the storage traits.
2//!
3//! This is the default backend, replicating the original behaviour:
4//! in-memory `HashMap`s with atomic JSON file persistence on every
5//! mutation that affects durable state.
6
7use std::collections::HashMap;
8use std::io::Write as _;
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use serde::{Deserialize, Serialize};
13use tokio::sync::Mutex;
14use webauthn_rs::prelude::{AuthenticationResult, Passkey};
15
16use super::{
17    AccessTokenEntry, AuthCode, ClientStore, PasskeyStore, RefreshTokenEntry, RegisteredClient,
18    StoreError, TokenStore,
19};
20
21// ---------------------------------------------------------------------------
22// Capacity limits (caller-provided via StoreCaps)
23// ---------------------------------------------------------------------------
24
25/// Capacity limits injected into the JSON-file stores at construction.
26///
27/// Mirrors the relevant fields from [`crate::CapacityConfig`]; kept as a
28/// separate crate-private type so `create_json_file_stores` has a stable
29/// signature independent of future additions to the public config.
30#[derive(Debug, Clone, Copy)]
31#[expect(
32    clippy::struct_field_names,
33    reason = "`max_` prefix carries meaning (cap semantics) and mirrors field names on the public CapacityConfig"
34)]
35pub(crate) struct StoreCaps {
36    pub(crate) max_access_tokens: usize,
37    pub(crate) max_refresh_tokens: usize,
38    pub(crate) max_auth_codes: usize,
39    pub(crate) max_registered_clients: Option<usize>,
40}
41
42use super::TRANSIENT_STATE_TTL_SECS;
43
44// ---------------------------------------------------------------------------
45// Shared persistence state (tokens + clients live in the same file)
46// ---------------------------------------------------------------------------
47
48#[derive(Serialize, Deserialize, Default)]
49struct PersistedTokens {
50    access_tokens: HashMap<String, AccessTokenEntry>,
51    refresh_tokens: HashMap<String, RefreshTokenEntry>,
52    registered_clients: HashMap<String, RegisteredClient>,
53}
54
55struct SharedState {
56    access_tokens: HashMap<String, AccessTokenEntry>,
57    refresh_tokens: HashMap<String, RefreshTokenEntry>,
58    registered_clients: HashMap<String, RegisteredClient>,
59    /// Auth codes are *not* persisted (short TTL, in-memory only).
60    auth_codes: HashMap<String, AuthCode>,
61    tokens_path: PathBuf,
62    caps: StoreCaps,
63}
64
65impl SharedState {
66    /// Persist the three durable maps atomically.
67    fn persist(&self) -> Result<(), StoreError> {
68        let persisted = PersistedTokens {
69            access_tokens: self.access_tokens.clone(),
70            refresh_tokens: self.refresh_tokens.clone(),
71            registered_clients: self.registered_clients.clone(),
72        };
73        save_tokens(&self.tokens_path, &persisted).map_err(StoreError::Backend)
74    }
75}
76
77// ---------------------------------------------------------------------------
78// Constructors
79// ---------------------------------------------------------------------------
80
81/// Create the default JSON-file-backed token and client stores.
82///
83/// Both stores share the same underlying state so that persistence
84/// (to `tokens.json`) remains atomic. Returns `(token_store, client_store)`
85/// plus a summary of what was loaded for logging.
86#[must_use]
87pub(crate) fn create_json_file_stores(
88    passkey_store_path: &Path,
89    caps: StoreCaps,
90) -> (impl TokenStore, impl ClientStore, StoreSummary) {
91    let tp = tokens_path(passkey_store_path);
92    let persisted = load_tokens(&tp);
93
94    let summary = StoreSummary {
95        access_tokens: persisted.access_tokens.len(),
96        refresh_tokens: persisted.refresh_tokens.len(),
97        registered_clients: persisted.registered_clients.len(),
98        tokens_path: tp.clone(),
99    };
100
101    let shared = Arc::new(Mutex::new(SharedState {
102        access_tokens: persisted.access_tokens,
103        refresh_tokens: persisted.refresh_tokens,
104        registered_clients: persisted.registered_clients,
105        auth_codes: HashMap::new(),
106        tokens_path: tp,
107        caps,
108    }));
109
110    let token_store = JsonFileTokenStore {
111        state: Arc::clone(&shared),
112    };
113    let client_store = JsonFileClientStore { state: shared };
114
115    (token_store, client_store, summary)
116}
117
118/// Summary of data loaded from disk (for startup logging).
119pub(crate) struct StoreSummary {
120    pub access_tokens: usize,
121    pub refresh_tokens: usize,
122    pub registered_clients: usize,
123    pub tokens_path: PathBuf,
124}
125
126// ---------------------------------------------------------------------------
127// JsonFileTokenStore
128// ---------------------------------------------------------------------------
129
130/// JSON-file-backed [`TokenStore`].
131pub struct JsonFileTokenStore {
132    state: Arc<Mutex<SharedState>>,
133}
134
135impl TokenStore for JsonFileTokenStore {
136    async fn store_auth_code(&self, code: String, entry: AuthCode) -> Result<(), StoreError> {
137        let mut s = self.state.lock().await;
138        // Clean up expired codes before inserting
139        let now = crate::now_epoch();
140        s.auth_codes
141            .retain(|_, v| now.saturating_sub(v.created_at) <= TRANSIENT_STATE_TTL_SECS);
142        if s.auth_codes.len() >= s.caps.max_auth_codes {
143            return Err(StoreError::CapacityExceeded);
144        }
145        s.auth_codes.insert(code, entry);
146        drop(s);
147        Ok(())
148    }
149
150    async fn consume_auth_code(&self, code: &str) -> Result<Option<AuthCode>, StoreError> {
151        let mut s = self.state.lock().await;
152        Ok(s.auth_codes.remove(code))
153    }
154
155    async fn store_access_token(
156        &self,
157        token: String,
158        entry: AccessTokenEntry,
159    ) -> Result<(), StoreError> {
160        let mut s = self.state.lock().await;
161        // Clean up expired tokens
162        let now = crate::now_epoch();
163        s.access_tokens
164            .retain(|_, v| now.saturating_sub(v.created_at) < v.expires_in_secs);
165        if s.access_tokens.len() >= s.caps.max_access_tokens {
166            return Err(StoreError::CapacityExceeded);
167        }
168        s.access_tokens.insert(token, entry);
169        s.persist()
170    }
171
172    async fn get_access_token(&self, token: &str) -> Result<Option<AccessTokenEntry>, StoreError> {
173        let s = self.state.lock().await;
174        Ok(s.access_tokens.get(token).cloned())
175    }
176
177    async fn revoke_access_tokens_by_refresh(&self, refresh_token: &str) -> Result<(), StoreError> {
178        let mut s = self.state.lock().await;
179        s.access_tokens
180            .retain(|_, v| v.refresh_token != refresh_token);
181        s.persist()
182    }
183
184    async fn store_refresh_token(
185        &self,
186        token: String,
187        entry: RefreshTokenEntry,
188    ) -> Result<(), StoreError> {
189        let mut s = self.state.lock().await;
190        if s.refresh_tokens.len() >= s.caps.max_refresh_tokens {
191            return Err(StoreError::CapacityExceeded);
192        }
193        s.refresh_tokens.insert(token, entry);
194        s.persist()
195    }
196
197    async fn get_refresh_token(
198        &self,
199        token: &str,
200    ) -> Result<Option<RefreshTokenEntry>, StoreError> {
201        let s = self.state.lock().await;
202        Ok(s.refresh_tokens.get(token).cloned())
203    }
204
205    async fn consume_refresh_token(
206        &self,
207        token: &str,
208    ) -> Result<Option<RefreshTokenEntry>, StoreError> {
209        let mut s = self.state.lock().await;
210        let entry = s.refresh_tokens.remove(token);
211        if entry.is_some() {
212            s.persist()?;
213            drop(s);
214        }
215        Ok(entry)
216    }
217
218    async fn cleanup_expired_tokens(&self, now: u64) -> Result<(), StoreError> {
219        let mut s = self.state.lock().await;
220        let before = s.access_tokens.len();
221        s.access_tokens
222            .retain(|_, v| now.saturating_sub(v.created_at) < v.expires_in_secs);
223        if s.access_tokens.len() != before {
224            s.persist()?;
225            drop(s);
226        }
227        Ok(())
228    }
229}
230
231// ---------------------------------------------------------------------------
232// JsonFileClientStore
233// ---------------------------------------------------------------------------
234
235/// JSON-file-backed [`ClientStore`].
236pub struct JsonFileClientStore {
237    state: Arc<Mutex<SharedState>>,
238}
239
240impl ClientStore for JsonFileClientStore {
241    async fn register_client(
242        &self,
243        id: String,
244        client: RegisteredClient,
245    ) -> Result<(), StoreError> {
246        let mut s = self.state.lock().await;
247        s.registered_clients.insert(id, client);
248        s.persist()
249    }
250
251    async fn try_register_client(
252        &self,
253        id: String,
254        client: RegisteredClient,
255    ) -> Result<bool, StoreError> {
256        use std::collections::hash_map::Entry;
257
258        let mut s = self.state.lock().await;
259        if let Some(cap) = s.caps.max_registered_clients
260            && s.registered_clients.len() >= cap
261        {
262            return Ok(false);
263        }
264        // Defensive: never clobber an existing client. Callers that supply
265        // their own id (custom stores, future handlers) must not be able to
266        // silently overwrite a registered client's secret or redirect URIs.
267        match s.registered_clients.entry(id) {
268            Entry::Occupied(_) => Ok(false),
269            Entry::Vacant(slot) => {
270                slot.insert(client);
271                s.persist()?;
272                drop(s);
273                Ok(true)
274            }
275        }
276    }
277
278    async fn get_client(&self, id: &str) -> Result<Option<RegisteredClient>, StoreError> {
279        let s = self.state.lock().await;
280        Ok(s.registered_clients.get(id).cloned())
281    }
282
283    async fn client_count(&self) -> Result<usize, StoreError> {
284        let s = self.state.lock().await;
285        Ok(s.registered_clients.len())
286    }
287}
288
289// ---------------------------------------------------------------------------
290// JsonFilePasskeyStore
291// ---------------------------------------------------------------------------
292
293/// JSON-file-backed [`PasskeyStore`].
294pub struct JsonFilePasskeyStore {
295    passkeys: Mutex<Vec<Passkey>>,
296    path: PathBuf,
297}
298
299impl JsonFilePasskeyStore {
300    /// Create a new passkey store, loading existing passkeys from `path`.
301    #[must_use]
302    pub fn new(path: PathBuf) -> Self {
303        let passkeys = load_passkeys(&path);
304        Self {
305            passkeys: Mutex::new(passkeys),
306            path,
307        }
308    }
309
310    /// Return the number of passkeys loaded at construction time.
311    ///
312    /// This is intended for startup logging only.
313    pub async fn passkey_count(&self) -> usize {
314        self.passkeys.lock().await.len()
315    }
316}
317
318impl PasskeyStore for JsonFilePasskeyStore {
319    async fn list_passkeys(&self) -> Result<Vec<Passkey>, StoreError> {
320        Ok(self.passkeys.lock().await.clone())
321    }
322
323    async fn add_passkey_if_none(&self, passkey: Passkey) -> Result<bool, StoreError> {
324        let mut pks = self.passkeys.lock().await;
325        if !pks.is_empty() {
326            return Ok(false);
327        }
328        pks.push(passkey);
329        save_passkeys(&self.path, &pks).map_err(StoreError::Backend)?;
330        drop(pks);
331        Ok(true)
332    }
333
334    async fn add_passkey(&self, passkey: Passkey) -> Result<(), StoreError> {
335        let mut pks = self.passkeys.lock().await;
336        pks.push(passkey);
337        let result = save_passkeys(&self.path, &pks).map_err(StoreError::Backend);
338        drop(pks);
339        result
340    }
341
342    async fn update_passkey(&self, auth_result: &AuthenticationResult) -> Result<(), StoreError> {
343        let mut pks = self.passkeys.lock().await;
344        for pk in pks.iter_mut() {
345            pk.update_credential(auth_result);
346        }
347        let result = save_passkeys(&self.path, &pks).map_err(StoreError::Backend);
348        drop(pks);
349        result
350    }
351
352    async fn has_passkeys(&self) -> Result<bool, StoreError> {
353        Ok(!self.passkeys.lock().await.is_empty())
354    }
355}
356
357// ---------------------------------------------------------------------------
358// File I/O helpers (moved from lib.rs)
359// ---------------------------------------------------------------------------
360
361// M3: Atomic file write with restrictive permissions (0o600 on Unix).
362//
363// SECURITY NOTE: Persisted token files (tokens.json, passkeys.json) contain
364// plaintext secrets. Ensure the data directory is owned by the service user
365// and not world-readable. On a public-facing deployment, consider mounting
366// the data directory on a tmpfs or encrypted filesystem.
367pub(crate) fn atomic_write(
368    path: &Path,
369    data: &[u8],
370) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
371    if let Some(parent) = path.parent() {
372        std::fs::create_dir_all(parent)?;
373    }
374    let temp_path = path.with_extension("tmp");
375    {
376        let mut opts = std::fs::OpenOptions::new();
377        opts.write(true).create(true).truncate(true);
378        #[cfg(unix)]
379        {
380            use std::os::unix::fs::OpenOptionsExt;
381            opts.mode(0o600);
382        }
383        let mut file = opts.open(&temp_path)?;
384        file.write_all(data)?;
385        file.flush()?;
386    }
387    std::fs::rename(&temp_path, path)?;
388    Ok(())
389}
390
391pub(crate) fn load_passkeys(path: &Path) -> Vec<Passkey> {
392    std::fs::read_to_string(path)
393        .ok()
394        .and_then(|s| serde_json::from_str(&s).ok())
395        .unwrap_or_default()
396}
397
398fn save_passkeys(
399    path: &Path,
400    passkeys: &[Passkey],
401) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
402    atomic_write(path, serde_json::to_string_pretty(passkeys)?.as_bytes())
403}
404
405fn tokens_path(passkey_path: &Path) -> PathBuf {
406    passkey_path.with_file_name("tokens.json")
407}
408
409fn load_tokens(path: &Path) -> PersistedTokens {
410    std::fs::read_to_string(path)
411        .ok()
412        .and_then(|s| serde_json::from_str(&s).ok())
413        .unwrap_or_default()
414}
415
416fn save_tokens(
417    path: &Path,
418    tokens: &PersistedTokens,
419) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
420    atomic_write(path, serde_json::to_string_pretty(tokens)?.as_bytes())
421}
422
423#[cfg(test)]
424#[expect(
425    clippy::unwrap_used,
426    reason = "test module: invariants are established by the test fixtures themselves, so .unwrap() is idiomatic and a panic on violation is the desired test failure mode"
427)]
428mod tests {
429    use super::*;
430    use crate::store::{ClientStore, PasskeyStore, TokenStore};
431
432    /// Generous caps for tests that don't care about capacity enforcement.
433    fn large_caps() -> StoreCaps {
434        StoreCaps {
435            max_access_tokens: 10_000,
436            max_refresh_tokens: 10_000,
437            max_auth_codes: 10_000,
438            max_registered_clients: Some(1),
439        }
440    }
441
442    // -- TokenStore tests --
443
444    #[tokio::test]
445    async fn test_store_and_consume_auth_code() {
446        let dir = tempfile::tempdir().unwrap();
447        let (store, _, _) =
448            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
449
450        let code = AuthCode::new("cid".into(), "uri".into(), "ch".into(), 1000);
451        store.store_auth_code("code1".into(), code).await.unwrap();
452
453        // First consume returns the entry
454        let entry = store.consume_auth_code("code1").await.unwrap();
455        assert!(entry.is_some());
456        assert_eq!(entry.unwrap().client_id, "cid");
457
458        // Second consume returns None (single-use)
459        let entry = store.consume_auth_code("code1").await.unwrap();
460        assert!(entry.is_none());
461    }
462
463    #[tokio::test]
464    async fn test_store_and_get_access_token() {
465        let dir = tempfile::tempdir().unwrap();
466        let (store, _, _) =
467            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
468
469        let entry = AccessTokenEntry::new("cid".into(), 1000, 3600, "rt1".into());
470        store.store_access_token("at1".into(), entry).await.unwrap();
471
472        let got = store.get_access_token("at1").await.unwrap();
473        assert!(got.is_some());
474        let got = got.unwrap();
475        assert_eq!(got.client_id, "cid");
476        assert_eq!(got.expires_in_secs, 3600);
477        assert_eq!(got.refresh_token, "rt1");
478
479        // Non-existent token returns None
480        assert!(store.get_access_token("nope").await.unwrap().is_none());
481    }
482
483    #[tokio::test]
484    async fn test_store_and_consume_refresh_token() {
485        let dir = tempfile::tempdir().unwrap();
486        let (store, _, _) =
487            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
488
489        let entry = RefreshTokenEntry::new("cid".into());
490        store
491            .store_refresh_token("rt1".into(), entry)
492            .await
493            .unwrap();
494
495        // Peek (non-destructive)
496        let got = store.get_refresh_token("rt1").await.unwrap();
497        assert!(got.is_some());
498
499        // Consume
500        let got = store.consume_refresh_token("rt1").await.unwrap();
501        assert!(got.is_some());
502        assert_eq!(got.unwrap().client_id, "cid");
503
504        // Second consume returns None
505        assert!(store.consume_refresh_token("rt1").await.unwrap().is_none());
506    }
507
508    #[tokio::test]
509    async fn test_revoke_access_tokens_by_refresh() {
510        let dir = tempfile::tempdir().unwrap();
511        let (store, _, _) =
512            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
513
514        // Store two access tokens with different refresh tokens
515        let now = crate::now_epoch();
516        store
517            .store_access_token(
518                "at1".into(),
519                AccessTokenEntry::new("cid".into(), now, 3600, "rt-a".into()),
520            )
521            .await
522            .unwrap();
523        store
524            .store_access_token(
525                "at2".into(),
526                AccessTokenEntry::new("cid".into(), now, 3600, "rt-b".into()),
527            )
528            .await
529            .unwrap();
530
531        // Revoke only tokens associated with rt-a
532        store.revoke_access_tokens_by_refresh("rt-a").await.unwrap();
533
534        // at1 should be gone, at2 should remain
535        assert!(store.get_access_token("at1").await.unwrap().is_none());
536        assert!(store.get_access_token("at2").await.unwrap().is_some());
537    }
538
539    #[tokio::test]
540    async fn test_cleanup_expired_tokens() {
541        let dir = tempfile::tempdir().unwrap();
542        let (store, _, _) =
543            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
544
545        let now = crate::now_epoch();
546        // Expired token (created 10000s ago, expires in 3600s)
547        store
548            .store_access_token(
549                "expired".into(),
550                AccessTokenEntry::new("cid".into(), now - 10000, 3600, "rt1".into()),
551            )
552            .await
553            .unwrap();
554        // Fresh token
555        store
556            .store_access_token(
557                "fresh".into(),
558                AccessTokenEntry::new("cid".into(), now, 3600, "rt2".into()),
559            )
560            .await
561            .unwrap();
562
563        store.cleanup_expired_tokens(now).await.unwrap();
564
565        assert!(store.get_access_token("expired").await.unwrap().is_none());
566        assert!(store.get_access_token("fresh").await.unwrap().is_some());
567    }
568
569    #[tokio::test]
570    async fn test_auth_code_capacity_exceeded() {
571        let dir = tempfile::tempdir().unwrap();
572        let caps = StoreCaps {
573            max_auth_codes: 3,
574            ..large_caps()
575        };
576        let (store, _, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
577
578        let now = crate::now_epoch();
579        for i in 0..3 {
580            store
581                .store_auth_code(
582                    format!("code-{i}"),
583                    AuthCode::new("cid".into(), "uri".into(), "ch".into(), now),
584                )
585                .await
586                .unwrap();
587        }
588
589        let result = store
590            .store_auth_code(
591                "overflow".into(),
592                AuthCode::new("cid".into(), "uri".into(), "ch".into(), now),
593            )
594            .await;
595        assert!(matches!(result, Err(StoreError::CapacityExceeded)));
596    }
597
598    #[tokio::test]
599    async fn test_access_token_capacity_exceeded() {
600        let dir = tempfile::tempdir().unwrap();
601        let caps = StoreCaps {
602            max_access_tokens: 3,
603            ..large_caps()
604        };
605        let (store, _, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
606
607        let now = crate::now_epoch();
608        for i in 0..3 {
609            store
610                .store_access_token(
611                    format!("at-{i}"),
612                    AccessTokenEntry::new("cid".into(), now, 3600, format!("rt-{i}")),
613                )
614                .await
615                .unwrap();
616        }
617
618        let result = store
619            .store_access_token(
620                "overflow".into(),
621                AccessTokenEntry::new("cid".into(), now, 3600, "rt-overflow".into()),
622            )
623            .await;
624        assert!(matches!(result, Err(StoreError::CapacityExceeded)));
625    }
626
627    #[tokio::test]
628    async fn test_refresh_token_capacity_exceeded() {
629        let dir = tempfile::tempdir().unwrap();
630        let caps = StoreCaps {
631            max_refresh_tokens: 3,
632            ..large_caps()
633        };
634        let (store, _, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
635
636        for i in 0..3 {
637            store
638                .store_refresh_token(format!("rt-{i}"), RefreshTokenEntry::new("cid".into()))
639                .await
640                .unwrap();
641        }
642
643        let result = store
644            .store_refresh_token("overflow".into(), RefreshTokenEntry::new("cid".into()))
645            .await;
646        assert!(matches!(result, Err(StoreError::CapacityExceeded)));
647    }
648
649    #[tokio::test]
650    async fn test_token_persistence_across_reload() {
651        let dir = tempfile::tempdir().unwrap();
652        let passkey_path = dir.path().join("passkeys.json");
653
654        // Store tokens and a client
655        {
656            let (token_store, client_store, _) =
657                create_json_file_stores(&passkey_path, large_caps());
658            let now = crate::now_epoch();
659            token_store
660                .store_access_token(
661                    "at1".into(),
662                    AccessTokenEntry::new("cid".into(), now, 86400, "rt1".into()),
663                )
664                .await
665                .unwrap();
666            token_store
667                .store_refresh_token("rt1".into(), RefreshTokenEntry::new("cid".into()))
668                .await
669                .unwrap();
670            client_store
671                .register_client(
672                    "client1".into(),
673                    RegisteredClient::new("secret".into(), vec!["uri".into()]),
674                )
675                .await
676                .unwrap();
677        }
678
679        // Reload from same path
680        let (token_store, client_store, summary) =
681            create_json_file_stores(&passkey_path, large_caps());
682        assert_eq!(summary.access_tokens, 1);
683        assert_eq!(summary.refresh_tokens, 1);
684        assert_eq!(summary.registered_clients, 1);
685
686        // Verify data survived
687        assert!(token_store.get_access_token("at1").await.unwrap().is_some());
688        assert!(
689            token_store
690                .get_refresh_token("rt1")
691                .await
692                .unwrap()
693                .is_some()
694        );
695        assert!(client_store.get_client("client1").await.unwrap().is_some());
696    }
697
698    // -- ClientStore tests --
699
700    #[tokio::test]
701    async fn test_try_register_client_default_single_client_cap() {
702        // Default cap is Some(1) — the historical single-client lock.
703        let dir = tempfile::tempdir().unwrap();
704        let (_, client_store, _) =
705            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
706
707        let ok = client_store
708            .try_register_client(
709                "c1".into(),
710                RegisteredClient::new("secret1".into(), vec!["uri".into()]),
711            )
712            .await
713            .unwrap();
714        assert!(ok);
715
716        let ok = client_store
717            .try_register_client(
718                "c2".into(),
719                RegisteredClient::new("secret2".into(), vec!["uri".into()]),
720            )
721            .await
722            .unwrap();
723        assert!(!ok);
724
725        assert!(client_store.get_client("c1").await.unwrap().is_some());
726        assert!(client_store.get_client("c2").await.unwrap().is_none());
727        assert_eq!(client_store.client_count().await.unwrap(), 1);
728    }
729
730    #[tokio::test]
731    async fn test_try_register_client_respects_configurable_cap() {
732        let dir = tempfile::tempdir().unwrap();
733        let caps = StoreCaps {
734            max_registered_clients: Some(2),
735            ..large_caps()
736        };
737        let (_, client_store, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
738
739        for i in 0..2 {
740            let ok = client_store
741                .try_register_client(
742                    format!("c{i}"),
743                    RegisteredClient::new(format!("s{i}"), vec!["uri".into()]),
744                )
745                .await
746                .unwrap();
747            assert!(ok, "registration {i} should succeed");
748        }
749
750        let ok = client_store
751            .try_register_client(
752                "c2".into(),
753                RegisteredClient::new("s2".into(), vec!["uri".into()]),
754            )
755            .await
756            .unwrap();
757        assert!(!ok, "third registration should be rejected by cap");
758        assert_eq!(client_store.client_count().await.unwrap(), 2);
759    }
760
761    #[tokio::test]
762    async fn test_try_register_client_does_not_clobber_existing_id() {
763        // Defence-in-depth: even when under the cap, re-registering with an
764        // existing id must return Ok(false) and leave the original client
765        // intact. The real handler uses CSPRNG ids so collision is
766        // astronomically unlikely, but custom stores or future callers may
767        // pass known ids.
768        let dir = tempfile::tempdir().unwrap();
769        let caps = StoreCaps {
770            max_registered_clients: Some(5),
771            ..large_caps()
772        };
773        let (_, client_store, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
774
775        let ok = client_store
776            .try_register_client(
777                "dup".into(),
778                RegisteredClient::new("original-secret".into(), vec!["uri-1".into()]),
779            )
780            .await
781            .unwrap();
782        assert!(ok);
783
784        let ok = client_store
785            .try_register_client(
786                "dup".into(),
787                RegisteredClient::new("attacker-secret".into(), vec!["uri-evil".into()]),
788            )
789            .await
790            .unwrap();
791        assert!(!ok, "duplicate id should be refused");
792
793        let existing = client_store.get_client("dup").await.unwrap().unwrap();
794        assert_eq!(existing.client_secret, "original-secret");
795        assert_eq!(existing.redirect_uris, vec!["uri-1"]);
796        assert_eq!(client_store.client_count().await.unwrap(), 1);
797    }
798
799    #[tokio::test]
800    async fn test_try_register_client_unlimited_when_cap_is_none() {
801        let dir = tempfile::tempdir().unwrap();
802        let caps = StoreCaps {
803            max_registered_clients: None,
804            ..large_caps()
805        };
806        let (_, client_store, _) = create_json_file_stores(&dir.path().join("passkeys.json"), caps);
807
808        for i in 0..5 {
809            let ok = client_store
810                .try_register_client(
811                    format!("c{i}"),
812                    RegisteredClient::new(format!("s{i}"), vec!["uri".into()]),
813                )
814                .await
815                .unwrap();
816            assert!(ok, "registration {i} should succeed with unlimited cap");
817        }
818        assert_eq!(client_store.client_count().await.unwrap(), 5);
819    }
820
821    #[tokio::test]
822    async fn test_client_persistence_across_reload() {
823        let dir = tempfile::tempdir().unwrap();
824        let passkey_path = dir.path().join("passkeys.json");
825
826        {
827            let (_, client_store, _) = create_json_file_stores(&passkey_path, large_caps());
828            client_store
829                .register_client(
830                    "c1".into(),
831                    RegisteredClient::new("secret".into(), vec!["u".into()]),
832                )
833                .await
834                .unwrap();
835        }
836
837        let (_, client_store, _) = create_json_file_stores(&passkey_path, large_caps());
838        let client = client_store.get_client("c1").await.unwrap();
839        assert!(client.is_some());
840        assert_eq!(client.unwrap().redirect_uris, vec!["u"]);
841    }
842
843    // -- Regression: get_refresh_token is non-destructive --
844
845    #[tokio::test]
846    async fn test_get_refresh_token_does_not_consume() {
847        let dir = tempfile::tempdir().unwrap();
848        let (store, _, _) =
849            create_json_file_stores(&dir.path().join("passkeys.json"), large_caps());
850
851        store
852            .store_refresh_token("rt1".into(), RefreshTokenEntry::new("cid".into()))
853            .await
854            .unwrap();
855
856        // get_refresh_token should not remove it
857        let entry = store.get_refresh_token("rt1").await.unwrap();
858        assert!(entry.is_some());
859
860        // Should still be there
861        let entry = store.get_refresh_token("rt1").await.unwrap();
862        assert!(entry.is_some());
863
864        // consume_refresh_token removes it
865        let entry = store.consume_refresh_token("rt1").await.unwrap();
866        assert!(entry.is_some());
867
868        // Now it's gone
869        assert!(store.get_refresh_token("rt1").await.unwrap().is_none());
870    }
871
872    // -- PasskeyStore tests --
873
874    #[tokio::test]
875    async fn test_passkey_store_empty() {
876        let dir = tempfile::tempdir().unwrap();
877        let store = JsonFilePasskeyStore::new(dir.path().join("passkeys.json"));
878
879        assert!(!store.has_passkeys().await.unwrap());
880        assert!(store.list_passkeys().await.unwrap().is_empty());
881    }
882}
mcp_oauth/store/json_file.rs

mcp_oauth/store/
json_file.rs
