feat: Add new memory storage
Signed-off-by: Awiteb <a@4rs.nl>
This commit is contained in:
parent
e21ce0d524
commit
27d985114a
2 changed files with 184 additions and 0 deletions
182
src/storage/memory_storage.rs
Normal file
182
src/storage/memory_storage.rs
Normal file
|
@ -0,0 +1,182 @@
|
||||||
|
// Copyright (c) 2024, Awiteb <a@4rs.nl>
|
||||||
|
// 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<HashMap<String, (u64, String)>>);
|
||||||
|
|
||||||
|
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<String, Self::Error> {
|
||||||
|
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<Option<String>, 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());
|
||||||
|
}
|
||||||
|
}
|
|
@ -13,9 +13,11 @@ use std::{sync::Arc, time::Duration};
|
||||||
|
|
||||||
#[cfg(feature = "cacache-storage")]
|
#[cfg(feature = "cacache-storage")]
|
||||||
mod cacache_storage;
|
mod cacache_storage;
|
||||||
|
mod memory_storage;
|
||||||
|
|
||||||
#[cfg(feature = "cacache-storage")]
|
#[cfg(feature = "cacache-storage")]
|
||||||
pub use cacache_storage::*;
|
pub use cacache_storage::*;
|
||||||
|
pub use memory_storage::*;
|
||||||
|
|
||||||
/// Trait to store the captcha token and answer. is also clear the expired captcha.
|
/// Trait to store the captcha token and answer. is also clear the expired captcha.
|
||||||
///
|
///
|
||||||
|
|
Loading…
Reference in a new issue