From 27d985114a3825ea1f9f4a284d56badf30dfb990 Mon Sep 17 00:00:00 2001 From: Awiteb Date: Mon, 12 Aug 2024 13:15:47 +0000 Subject: [PATCH] feat: Add new memory storage Signed-off-by: Awiteb --- src/storage/memory_storage.rs | 182 ++++++++++++++++++++++++++++++++++ src/storage/mod.rs | 2 + 2 files changed, 184 insertions(+) create mode 100644 src/storage/memory_storage.rs diff --git a/src/storage/memory_storage.rs b/src/storage/memory_storage.rs new file mode 100644 index 0000000..42e6b12 --- /dev/null +++ b/src/storage/memory_storage.rs @@ -0,0 +1,182 @@ +// Copyright (c) 2024, Awiteb +// A captcha middleware for Salvo framework. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +#![allow(warnings)] + +use std::{ + collections::HashMap, + convert::Infallible, + time::{Duration, SystemTime}, +}; +use tokio::sync::RwLock; + +use crate::CaptchaStorage; + +/// Captcha storage implementation using an in-memory HashMap. +#[derive(Debug)] +pub struct MemoryStorage(RwLock>); + +impl MemoryStorage { + /// Create a new instance of [`MemoryStorage`]. + pub fn new() -> Self { + Self(RwLock::new(HashMap::new())) + } +} + +impl CaptchaStorage for MemoryStorage { + type Error = Infallible; + + async fn store_answer(&self, answer: String) -> Result { + let token = uuid::Uuid::new_v4().to_string(); + let mut write_lock = self.0.write().await; + write_lock.insert(token.clone(), (now(), answer)); + + Ok(token) + } + + async fn get_answer(&self, token: &str) -> Result, Self::Error> { + let reader = self.0.read().await; + Ok(reader.get(token).map(|(_, answer)| answer.to_owned())) + } + + async fn clear_expired(&self, expired_after: Duration) -> Result<(), Self::Error> { + let expired_after = now() - expired_after.as_secs(); + + let mut write_lock = self.0.write().await; + write_lock.retain(|_, (timestamp, _)| *timestamp > expired_after); + + Ok(()) + } + + async fn clear_by_token(&self, token: &str) -> Result<(), Self::Error> { + let mut write_lock = self.0.write().await; + write_lock.retain(|c_token, (_, _)| c_token != token); + Ok(()) + } +} + +fn now() -> u64 { + SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("SystemTime before UNIX EPOCH!") + .as_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn memory_store_captcha() { + let storage = MemoryStorage::new(); + + let token = storage + .store_answer("answer".to_owned()) + .await + .expect("failed to store captcha"); + assert_eq!( + storage + .get_answer(&token) + .await + .expect("failed to get captcha answer"), + Some("answer".to_owned()) + ); + } + + #[tokio::test] + async fn memory_clear_expired() { + let storage = MemoryStorage::new(); + + let token = storage + .store_answer("answer".to_owned()) + .await + .expect("failed to store captcha"); + storage + .clear_expired(Duration::from_secs(0)) + .await + .expect("failed to clear expired captcha"); + assert!(storage + .get_answer(&token) + .await + .expect("failed to get captcha answer") + .is_none()); + } + + #[tokio::test] + async fn memory_clear_by_token() { + let storage = MemoryStorage::new(); + + let token = storage + .store_answer("answer".to_owned()) + .await + .expect("failed to store captcha"); + storage + .clear_by_token(&token) + .await + .expect("failed to clear captcha by token"); + assert!(storage + .get_answer(&token) + .await + .expect("failed to get captcha answer") + .is_none()); + } + + #[tokio::test] + async fn memory_is_token_exist() { + let storage = MemoryStorage::new(); + + let token = storage + .store_answer("answer".to_owned()) + .await + .expect("failed to store captcha"); + assert!(storage + .get_answer(&token) + .await + .expect("failed to check if token is exist") + .is_some()); + assert!(storage + .get_answer("token") + .await + .expect("failed to check if token is exist") + .is_none()); + } + + #[tokio::test] + async fn memory_clear_expired_with_expired_after() { + let storage = MemoryStorage::new(); + + let token = storage + .store_answer("answer".to_owned()) + .await + .expect("failed to store captcha"); + storage + .clear_expired(Duration::from_secs(1)) + .await + .expect("failed to clear expired captcha"); + assert_eq!( + storage + .get_answer(&token) + .await + .expect("failed to get captcha answer"), + Some("answer".to_owned()) + ); + tokio::time::sleep(Duration::from_secs(1)).await; + storage + .clear_expired(Duration::from_secs(1)) + .await + .expect("failed to clear expired captcha"); + assert!(storage + .get_answer(&token) + .await + .expect("failed to get captcha answer") + .is_none()); + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 4f9f917..a76c609 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -13,9 +13,11 @@ use std::{sync::Arc, time::Duration}; #[cfg(feature = "cacache-storage")] mod cacache_storage; +mod memory_storage; #[cfg(feature = "cacache-storage")] pub use cacache_storage::*; +pub use memory_storage::*; /// Trait to store the captcha token and answer. is also clear the expired captcha. ///