1use shared_state::ttl::TtlCache;
7use std::time::Duration;
8
9#[async_trait::async_trait]
11pub trait CsrfStore: Send + Sync {
12 async fn store(&self, key: &str, nonce: &str, redirect_url: Option<String>) -> bool;
15
16 async fn retrieve_and_consume(&self, key: &str) -> Option<(String, Option<String>)>;
19
20 async fn cleanup_expired(&self);
22}
23
24pub 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 pub fn new() -> Self {
38 Self::with_ttl(Duration::from_secs(600))
39 }
40
41 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 self.entries.take_fresh(&key.to_string())
61 }
62
63 async fn cleanup_expired(&self) {
64 self.entries.cleanup();
65 }
66}
67
68pub enum CsrfStoreType {
70 Memory(MemoryCsrfStore),
72}
73
74impl Default for CsrfStoreType {
75 fn default() -> Self {
76 Self::Memory(MemoryCsrfStore::new())
77 }
78}
79
80impl CsrfStoreType {
81 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#[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}