oauth_toolkit/
csrf.rs

1//! CSRF state management with pluggable backends.
2
3use dashmap::DashMap;
4use std::time::{Duration, Instant};
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
82            .retain(|_, entry| now.duration_since(entry.created_at) <= self.ttl);
83    }
84}
85
86/// Multi-backend CSRF store.
87pub enum CsrfStoreType {
88    /// In-memory backend (single-instance only).
89    Memory(MemoryCsrfStore),
90}
91
92impl Default for CsrfStoreType {
93    fn default() -> Self {
94        Self::Memory(MemoryCsrfStore::new())
95    }
96}
97
98impl CsrfStoreType {
99    /// Create a new store. Currently only in-memory is supported.
100    pub fn new() -> Self {
101        Self::Memory(MemoryCsrfStore::new())
102    }
103}
104
105#[async_trait::async_trait]
106impl CsrfStore for CsrfStoreType {
107    async fn store(&self, key: &str, nonce: &str, redirect_url: Option<String>) -> bool {
108        match self {
109            Self::Memory(s) => s.store(key, nonce, redirect_url).await,
110        }
111    }
112
113    async fn retrieve_and_consume(&self, key: &str) -> Option<(String, Option<String>)> {
114        match self {
115            Self::Memory(s) => s.retrieve_and_consume(key).await,
116        }
117    }
118
119    async fn cleanup_expired(&self) {
120        match self {
121            Self::Memory(s) => s.cleanup_expired().await,
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[tokio::test]
131    async fn memory_store_roundtrip() {
132        let store = MemoryCsrfStore::new();
133        store
134            .store("key1", "nonce1", Some("https://app.com".into()))
135            .await;
136        let result = store.retrieve_and_consume("key1").await;
137        assert_eq!(
138            result,
139            Some(("nonce1".into(), Some("https://app.com".into())))
140        );
141        assert!(store.retrieve_and_consume("key1").await.is_none());
142    }
143
144    #[tokio::test]
145    async fn memory_store_missing_key() {
146        let store = MemoryCsrfStore::new();
147        assert!(store.retrieve_and_consume("nonexistent").await.is_none());
148    }
149
150    #[tokio::test]
151    async fn memory_store_cleanup() {
152        let store = MemoryCsrfStore::with_ttl(Duration::from_millis(1));
153        store.store("key1", "nonce1", None).await;
154        tokio::time::sleep(Duration::from_millis(5)).await;
155        store.cleanup_expired().await;
156        assert!(store.retrieve_and_consume("key1").await.is_none());
157    }
158
159    #[tokio::test]
160    async fn csrf_store_type_delegates() {
161        let store = CsrfStoreType::new();
162        store.store("key1", "nonce1", None).await;
163        assert!(store.retrieve_and_consume("key1").await.is_some());
164    }
165}
oauth_toolkit/csrf.rs

oauth_toolkit/
csrf.rs
