salvo_captcha/storage/
memory_storage.rs

1// Copyright (c) 2024-2025, Awiteb <a@4rs.nl>
2//     A captcha middleware for Salvo framework.
3//
4// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
5// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
6// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
7// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
8// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
9// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
10// THE SOFTWARE.
11
12#![allow(warnings)]
13
14use std::{
15    collections::HashMap,
16    convert::Infallible,
17    time::{Duration, SystemTime},
18};
19use tokio::sync::RwLock;
20
21use crate::CaptchaStorage;
22
23/// Captcha storage implementation using an in-memory [HashMap].
24#[derive(Debug)]
25pub struct MemoryStorage(RwLock<HashMap<String, (u64, String)>>);
26
27impl MemoryStorage {
28    /// Create a new instance of [`MemoryStorage`].
29    pub fn new() -> Self {
30        Self(RwLock::new(HashMap::new()))
31    }
32}
33
34impl CaptchaStorage for MemoryStorage {
35    /// This storage does not return any error.
36    type Error = Infallible;
37
38    async fn store_answer(&self, answer: String) -> Result<String, Self::Error> {
39        let token = uuid::Uuid::new_v4().to_string();
40        let mut write_lock = self.0.write().await;
41        write_lock.insert(token.clone(), (now(), answer));
42
43        Ok(token)
44    }
45
46    async fn get_answer(&self, token: &str) -> Result<Option<String>, Self::Error> {
47        let reader = self.0.read().await;
48        Ok(reader.get(token).map(|(_, answer)| answer.to_owned()))
49    }
50
51    async fn clear_expired(&self, expired_after: Duration) -> Result<(), Self::Error> {
52        let expired_after = now() - expired_after.as_secs();
53
54        let mut write_lock = self.0.write().await;
55        write_lock.retain(|_, (timestamp, _)| *timestamp > expired_after);
56
57        Ok(())
58    }
59
60    async fn clear_by_token(&self, token: &str) -> Result<(), Self::Error> {
61        let mut write_lock = self.0.write().await;
62        write_lock.retain(|c_token, (_, _)| c_token != token);
63        Ok(())
64    }
65}
66
67fn now() -> u64 {
68    SystemTime::now()
69        .duration_since(std::time::UNIX_EPOCH)
70        .expect("SystemTime before UNIX EPOCH!")
71        .as_secs()
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[tokio::test]
79    async fn memory_store_captcha() {
80        let storage = MemoryStorage::new();
81
82        let token = storage
83            .store_answer("answer".to_owned())
84            .await
85            .expect("failed to store captcha");
86        assert_eq!(
87            storage
88                .get_answer(&token)
89                .await
90                .expect("failed to get captcha answer"),
91            Some("answer".to_owned())
92        );
93    }
94
95    #[tokio::test]
96    async fn memory_clear_expired() {
97        let storage = MemoryStorage::new();
98
99        let token = storage
100            .store_answer("answer".to_owned())
101            .await
102            .expect("failed to store captcha");
103        storage
104            .clear_expired(Duration::from_secs(0))
105            .await
106            .expect("failed to clear expired captcha");
107        assert!(storage
108            .get_answer(&token)
109            .await
110            .expect("failed to get captcha answer")
111            .is_none());
112    }
113
114    #[tokio::test]
115    async fn memory_clear_by_token() {
116        let storage = MemoryStorage::new();
117
118        let token = storage
119            .store_answer("answer".to_owned())
120            .await
121            .expect("failed to store captcha");
122        storage
123            .clear_by_token(&token)
124            .await
125            .expect("failed to clear captcha by token");
126        assert!(storage
127            .get_answer(&token)
128            .await
129            .expect("failed to get captcha answer")
130            .is_none());
131    }
132
133    #[tokio::test]
134    async fn memory_is_token_exist() {
135        let storage = MemoryStorage::new();
136
137        let token = storage
138            .store_answer("answer".to_owned())
139            .await
140            .expect("failed to store captcha");
141        assert!(storage
142            .get_answer(&token)
143            .await
144            .expect("failed to check if token is exist")
145            .is_some());
146        assert!(storage
147            .get_answer("token")
148            .await
149            .expect("failed to check if token is exist")
150            .is_none());
151    }
152
153    #[tokio::test]
154    async fn memory_clear_expired_with_expired_after() {
155        let storage = MemoryStorage::new();
156
157        let token = storage
158            .store_answer("answer".to_owned())
159            .await
160            .expect("failed to store captcha");
161        storage
162            .clear_expired(Duration::from_secs(1))
163            .await
164            .expect("failed to clear expired captcha");
165        assert_eq!(
166            storage
167                .get_answer(&token)
168                .await
169                .expect("failed to get captcha answer"),
170            Some("answer".to_owned())
171        );
172        tokio::time::sleep(Duration::from_secs(1)).await;
173        storage
174            .clear_expired(Duration::from_secs(1))
175            .await
176            .expect("failed to clear expired captcha");
177        assert!(storage
178            .get_answer(&token)
179            .await
180            .expect("failed to get captcha answer")
181            .is_none());
182    }
183}
salvo_captcha/storage/memory_storage.rs

salvo_captcha/storage/
memory_storage.rs
