Skip to main content

oauth_toolkit/
csrf.rs

1//! CSRF state management with pluggable backends.
2//!
3//! The in-memory store delegates expiry handling to
4//! [`shared_state::ttl::TtlCache`] instead of hand-rolling a TTL map.
5
6use shared_state::ttl::TtlCache;
7use std::time::Duration;
8
9/// CSRF state store trait.
10#[async_trait::async_trait]
11pub trait CsrfStore: Send + Sync {
12    /// Store a state nonce with optional redirect URL.
13    /// Returns true if stored successfully.
14    async fn store(&self, key: &str, nonce: &str, redirect_url: Option<String>) -> bool;
15
16    /// Retrieve and consume a state nonce (one-time use).
17    /// Returns (nonce, redirect_url) if found and not expired.
18    async fn retrieve_and_consume(&self, key: &str) -> Option<(String, Option<String>)>;
19
20    /// Remove expired entries.
21    async fn cleanup_expired(&self);
22}
23
24/// In-memory CSRF store with TTL-backed one-time semantics.
25pub struct MemoryCsrfStore {
26    entries: TtlCache<String, (String, Option<String>)>,
27}
28
29impl Default for MemoryCsrfStore {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl MemoryCsrfStore {
36    /// Create a new in-memory store with 10-minute TTL.
37    pub fn new() -> Self {
38        Self::with_ttl(Duration::from_secs(600))
39    }
40
41    /// Create with custom TTL.
42    pub fn with_ttl(ttl: Duration) -> Self {
43        Self {
44            entries: TtlCache::new(ttl),
45        }
46    }
47}
48
49#[async_trait::async_trait]
50impl CsrfStore for MemoryCsrfStore {
51    async fn store(&self, key: &str, nonce: &str, redirect_url: Option<String>) -> bool {
52        self.entries
53            .insert(key.to_string(), (nonce.to_string(), redirect_url));
54        true
55    }
56
57    async fn retrieve_and_consume(&self, key: &str) -> Option<(String, Option<String>)> {
58        // One-time use: `take_fresh` removes the entry and refuses
59        // expired ones, so a stale state can never be replayed.
60        self.entries.take_fresh(&key.to_string())
61    }
62
63    async fn cleanup_expired(&self) {
64        self.entries.cleanup();
65    }
66}
67
68/// Multi-backend CSRF store.
69pub enum CsrfStoreType {
70    /// In-memory backend (single-instance only).
71    Memory(MemoryCsrfStore),
72}
73
74impl Default for CsrfStoreType {
75    fn default() -> Self {
76        Self::Memory(MemoryCsrfStore::new())
77    }
78}
79
80impl CsrfStoreType {
81    /// Create a new store. Currently only in-memory is supported.
82    pub fn new() -> Self {
83        Self::Memory(MemoryCsrfStore::new())
84    }
85}
86
87#[async_trait::async_trait]
88impl CsrfStore for CsrfStoreType {
89    async fn store(&self, key: &str, nonce: &str, redirect_url: Option<String>) -> bool {
90        match self {
91            Self::Memory(s) => s.store(key, nonce, redirect_url).await,
92        }
93    }
94
95    async fn retrieve_and_consume(&self, key: &str) -> Option<(String, Option<String>)> {
96        match self {
97            Self::Memory(s) => s.retrieve_and_consume(key).await,
98        }
99    }
100
101    async fn cleanup_expired(&self) {
102        match self {
103            Self::Memory(s) => s.cleanup_expired().await,
104        }
105    }
106}
107
108#[cfg(test)]
109// Test code: unwrap/expect are the idiomatic way to assert setup success.
110#[allow(clippy::unwrap_used, clippy::expect_used)]
111mod tests {
112    use super::*;
113
114    #[tokio::test]
115    async fn memory_store_roundtrip() {
116        let store = MemoryCsrfStore::new();
117        store
118            .store("key1", "nonce1", Some("https://app.com".into()))
119            .await;
120        let result = store.retrieve_and_consume("key1").await;
121        assert_eq!(
122            result,
123            Some(("nonce1".into(), Some("https://app.com".into())))
124        );
125        assert!(store.retrieve_and_consume("key1").await.is_none());
126    }
127
128    #[tokio::test]
129    async fn memory_store_missing_key() {
130        let store = MemoryCsrfStore::new();
131        assert!(store.retrieve_and_consume("nonexistent").await.is_none());
132    }
133
134    #[tokio::test]
135    async fn memory_store_cleanup() {
136        let store = MemoryCsrfStore::with_ttl(Duration::from_millis(1));
137        store.store("key1", "nonce1", None).await;
138        tokio::time::sleep(Duration::from_millis(5)).await;
139        store.cleanup_expired().await;
140        assert!(store.retrieve_and_consume("key1").await.is_none());
141    }
142
143    #[tokio::test]
144    async fn csrf_store_type_delegates() {
145        let store = CsrfStoreType::new();
146        store.store("key1", "nonce1", None).await;
147        assert!(store.retrieve_and_consume("key1").await.is_some());
148    }
149}