oauth_toolkit/
csrf.rs

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

oauth_toolkit/
csrf.rs
