From c7e58cb6fcd53016827472587bfc1e93a2092683 Mon Sep 17 00:00:00 2001 From: TheAwiteb Date: Mon, 29 Jan 2024 23:24:35 +0300 Subject: [PATCH] First commit --- .gitignore | 2 + Cargo.toml | 32 ++++ Justfile | 35 ++++ examples/simple_login.rs | 204 ++++++++++++++++++++ src/captcha_gen.rs | 29 +++ src/finder.rs | 395 +++++++++++++++++++++++++++++++++++++++ src/lib.rs | 174 +++++++++++++++++ src/storage.rs | 368 ++++++++++++++++++++++++++++++++++++ 8 files changed, 1239 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 Justfile create mode 100644 examples/simple_login.rs create mode 100644 src/captcha_gen.rs create mode 100644 src/finder.rs create mode 100644 src/lib.rs create mode 100644 src/storage.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..96ef6c0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +/target +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..1828492 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "salvo-captcha" +version = "0.1.0" +rust-version = "1.75.0" +edition = "2021" +authors = ["Awiteb "] +description = "A captcha middleware for Salvo framework." +license = "MIT" +repository = "https://github.com/TheAwiteb/salvo-captcha" +documentation = "https://docs.rs/salvo-captcha" +readme = "README.md" +keywords = ["salvo", "captcha", "middleware"] +categories = ["web-programming", "network-programming"] + + +[dependencies] +async-trait = "0.1.77" +cacache = { version = "12.0.0", default-features = false, features = ["tokio-runtime", "mmap"], optional = true } +captcha = { version = "0.0.9", default-features = false } +easy-ext = "1.0.1" +log = "0.4.20" +salvo_core = "^ 0.65" +uuid = { version = "1.7.0", features = ["v4"], optional = true } + +[features] +cacache-storage = ["dep:cacache", "dep:uuid"] + +[dev-dependencies] +tempfile = ">= 3.9" +tokio = { version = ">= 1.35", features = ["macros", "rt-multi-thread"] } +base64 = ">= 0.21" +salvo = { version = ">= 0.65", features = ["affix"] } \ No newline at end of file diff --git a/Justfile b/Justfile new file mode 100644 index 0000000..dde7453 --- /dev/null +++ b/Justfile @@ -0,0 +1,35 @@ +# This justfile is for the contrbutors of this project, not for the end user. +# +# Requirements for this justfile: +# - Linux distribution, the real programer does not program on garbage OS like Windows or MacOS +# - just (Of course) +# - cargo (For the build and tests) + +set shell := ["/usr/bin/bash", "-c"] + +JUST_EXECUTABLE := "just -u -f " + justfile() +header := "Available tasks:\n" +# Get the MSRV from the Cargo.toml +msrv := `cat Cargo.toml | grep "rust-version" | sed 's/.*"\(.*\)".*/\1/'` + + +_default: + @{{JUST_EXECUTABLE}} --list-heading "{{header}}" --list + +# Run the CI (Local use only) +@ci: + cargo fmt --all --check + cargo clippy --workspace --all-targets --examples --tests --all-features -- -D warnings + cargo nextest run --workspace --all-targets --all-features + @{{JUST_EXECUTABLE}} msrv + +# Install workspace tools +@install: + cargo install cargo-nextest + +# Check that the current MSRV is correct +msrv: + @echo "Current MSRV: {{msrv}}" + @echo "Checking that the current MSRV is correct..." + @cargo +{{msrv}} check --workspace --all-targets --all-features + @echo "The current MSRV is correct" diff --git a/examples/simple_login.rs b/examples/simple_login.rs new file mode 100644 index 0000000..ee1f7d7 --- /dev/null +++ b/examples/simple_login.rs @@ -0,0 +1,204 @@ +// Example of using the `salvo_captcha`, this example is a simple login page with a captcha. +// The page will be in +// You can see a video of this example here: +// +// Run the example with `cargo run --example simple_login --features cacache-storage` +// +// Or set up a crate for it, with the following `Cargo.toml`: +// ```toml +// [dependencies] +// base64 = ">= 0.21" +// salvo = { version = ">= 0.65", features = ["affix"] } +// salvo-captcha = { version = ">= 0.1", features = ["cacache-storage"] } +// tokio = { version = ">= 1.35", features = ["macros", "rt-multi-thread", "time"] } +// ``` + +use std::{sync::Arc, time::Duration}; + +use base64::{engine::GeneralPurpose, Engine}; +use salvo::prelude::*; +use salvo_captcha::*; + +// The type of the captcha +type MyCaptcha = Captcha>; + +// To convert the image to base64, to show it in the browser +const BASE_64_ENGINE: GeneralPurpose = GeneralPurpose::new( + &base64::alphabet::STANDARD, + base64::engine::general_purpose::PAD, +); + +#[handler] +async fn index(_req: &mut Request, res: &mut Response, depot: &mut Depot) { + // Get the captcha from the depot + let captcha_storage = depot.obtain::>().unwrap(); + + // Create a new captcha + let (token, image) = captcha_storage + .as_ref() + .new_captcha(CaptchaName::Mila, CaptchaDifficulty::Medium) + .await + .expect("Failed to save captcha") + .expect("Failed to create captcha"); + + // Convert the image to base64 + let image = BASE_64_ENGINE.encode(image); + + // Set the response content + res.render(Text::Html(index_page(image, token))) +} + +#[handler] +async fn auth(req: &mut Request, res: &mut Response, depot: &mut Depot) { + // Get the captcha state from the depot, where we can know if the captcha is passed + let captcha_state = depot.get_captcha_state().unwrap(); + // Not important, just for demo + let Some(username) = req.form::("username").await else { + res.status_code(StatusCode::BAD_REQUEST); + return res.render(Text::Html("Invalid form submission")); + }; + + // Handle the captcha state, that's all + let content = match captcha_state { + CaptchaState::Passed => { + format!("Welcome, {username}!") + } + CaptchaState::AnswerNotFound => "Captcha answer not found".to_string(), + CaptchaState::TokenNotFound => "Captcha token not found".to_string(), + CaptchaState::WrongAnswer => "Wrong captcha answer".to_string(), + CaptchaState::WrongToken => "Wrong captcha token".to_string(), + CaptchaState::Skipped => "Captcha skipped".to_string(), + CaptchaState::StorageError => "Captcha storage error".to_string(), + }; + + res.render(Text::Html(captcha_result_page(content))) +} + +#[tokio::main] +async fn main() { + let captcha_middleware = MyCaptcha::new( + CacacheStorage::new("./captcha-cache"), + CaptchaFormFinder::new(), + ) + .skipper(|req: &mut Request, _: &Depot| { + // Skip the captcha if the request path is /skipped + req.uri().path() == "/skipped" + }); + let captcha_storage = Arc::new(captcha_middleware.storage().clone()); + + // Set up a task to clean the expired captcha, just call the `clear_expired` function from the storage + let captcha_cleaner = tokio::spawn({ + let cleanner_storage = captcha_storage.clone(); + async move { + let captcha_expired_after = Duration::from_secs(60 * 5); + let clean_interval = Duration::from_secs(60); + + // Just this loop, to clean the expired captcha every 60 seconds, each captcha will be expired after 5 minutes + loop { + cleanner_storage + .clear_expired(captcha_expired_after) + .await + .ok(); // You should log this error + tokio::time::sleep(clean_interval).await; + } + } + }); + + let router = Router::new() + .hoop(affix::inject(captcha_storage.clone())) + .push(Router::with_path("/").get(index)) + .push( + Router::new() + .hoop(captcha_middleware) + .push(Router::with_path("/auth").post(auth)) + .push(Router::with_path("/skipped").post(auth)), + ); + + let acceptor = TcpListener::new(("127.0.0.1", 5800)).bind().await; + Server::new(acceptor).serve(router).await; + captcha_cleaner.await.ok(); +} + +fn index_page(captcha_image: String, captcha_token: String) -> String { + format!( + r#" + + + Salvo Captcha Example + + + +

Salvo Captcha Example

+

Sign In

+ +
+ + + +
+ +
+ +
+ +
+ Or you can skip the captcha +
+ + + + +
+ +
+ +
+ Source Code + + + "# + ) +} + +fn captcha_result_page(captcha_result: String) -> String { + format!( + r#" + + + Salvo Captcha Example + + + +

Salvo Captcha Example

+

Sign In

+ {captcha_result} +
+ Go Back + + + "# + ) +} diff --git a/src/captcha_gen.rs b/src/captcha_gen.rs new file mode 100644 index 0000000..7c6ba4f --- /dev/null +++ b/src/captcha_gen.rs @@ -0,0 +1,29 @@ +use crate::{CaptchaDifficulty, CaptchaName, CaptchaStorage}; + +/// Captcha generator, used to generate a new captcha image. This trait are implemented for all [`CaptchaStorage`]. +pub trait CaptchaGenerator: CaptchaStorage { + /// Create a new captcha image and return the token and the image encoded as png. Will return None if the captcha crate failed to create the captcha. + /// + /// The returned captcha image is 220x110 pixels. + /// + /// For more information about the captcha name and difficulty, see the [`README.md`](https://github.com/TheAwiteb/salvo-captcha/#captcha-name-and-difficulty). + fn new_captcha( + &self, + name: CaptchaName, + difficulty: CaptchaDifficulty, + ) -> impl std::future::Future)>, Self::Error>> + Send + { + async { + let Some((captcha_answer, captcha_image)) = + captcha::by_name(difficulty, name).as_tuple() + else { + return Ok(None); + }; + + let token = self.store_answer(captcha_answer.into()).await?; + Ok(Some((token, captcha_image))) + } + } +} + +impl CaptchaGenerator for T where T: CaptchaStorage {} diff --git a/src/finder.rs b/src/finder.rs new file mode 100644 index 0000000..4bb118a --- /dev/null +++ b/src/finder.rs @@ -0,0 +1,395 @@ +// 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. + +use std::marker::PhantomData; + +use salvo_core::http::header::HeaderName; +use salvo_core::http::Request; + +/// Trait to find the captcha token and answer from the request. +pub trait CaptchaFinder: Send + Sync { + /// The token type + type Token: TryFrom + Sync + Send; + /// The answer type + type Answer: TryFrom + Sync + Send; + + /// The token error type + type TError: std::fmt::Debug + Send; + /// The answer error type + type AError: std::fmt::Debug + Send; + + /// Find the captcha token from the request. + /// + /// Return [`None`] if the request does not contain a captcha token. An error is returned if the token is invalid format. + fn find_token( + &self, + req: &mut Request, + ) -> impl std::future::Future, Self::TError>> + std::marker::Send; + + /// Find the captcha answer from the request. + /// + /// Return [`None`] if the request does not contain a captcha answer. An error is returned if the answer is invalid format. + fn find_answer( + &self, + req: &mut Request, + ) -> impl std::future::Future, Self::AError>> + std::marker::Send; +} + +/// Find the captcha token and answer from the header +#[derive(Debug)] +pub struct CaptchaHeaderFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + phantom: PhantomData<(T, A)>, + + /// The header name of the captcha token + /// + /// Default: "x-captcha-token" + pub token_header: HeaderName, + + /// The header name of the captcha answer + /// + /// Default: "x-captcha-answer" + pub answer_header: HeaderName, +} + +/// Find the captcha token and answer from the form +#[derive(Debug)] +pub struct CaptchaFormFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + phantom: PhantomData<(T, A)>, + + /// The form name of the captcha token + /// + /// Default: "captcha_token" + pub token_name: String, + + /// The form name of the captcha answer + /// + /// Default: "captcha_answer" + pub answer_name: String, +} + +impl CaptchaHeaderFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + /// Create a new CaptchaHeaderFinder + pub fn new() -> Self { + Self::default() + } + + /// Set the token header name + pub fn token_header(mut self, token_header: HeaderName) -> Self { + self.token_header = token_header; + self + } + + /// Set the answer header name + pub fn answer_header(mut self, answer_header: HeaderName) -> Self { + self.answer_header = answer_header; + self + } +} + +impl CaptchaFormFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + /// Create a new CaptchaFormFinder + pub fn new() -> Self { + Self::default() + } + + /// Set the token form name + pub fn token_name(mut self, token_name: String) -> Self { + self.token_name = token_name; + self + } + + /// Set the answer form name + pub fn answer_name(mut self, answer_name: String) -> Self { + self.answer_name = answer_name; + self + } +} + +impl Default for CaptchaHeaderFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + /// Create a default CaptchaHeaderFinder with: + /// - token_header: "x-captcha-token" + /// - answer_header: "x-captcha-answer" + fn default() -> Self { + Self { + phantom: PhantomData, + token_header: HeaderName::from_static("x-captcha-token"), + answer_header: HeaderName::from_static("x-captcha-answer"), + } + } +} + +impl Default for CaptchaFormFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, +{ + /// Create a default CaptchaFormFinder with: + /// - token_name: "captcha_token" + /// - answer_name: "captcha_answer" + fn default() -> Self { + Self { + phantom: PhantomData, + token_name: "captcha_token".to_string(), + answer_name: "captcha_answer".to_string(), + } + } +} + +impl CaptchaFinder for CaptchaHeaderFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, + >::Error: Send, + >::Error: std::fmt::Debug, + >::Error: Send, + >::Error: std::fmt::Debug, +{ + type Token = T; + type Answer = A; + + type TError = >::Error; + type AError = >::Error; + + async fn find_token(&self, req: &mut Request) -> Result, Self::TError> { + req.headers() + .get(&self.token_header) + .and_then(|t| t.to_str().ok()) + .map(|t| Self::Token::try_from(t.to_string())) + .transpose() + } + + async fn find_answer(&self, req: &mut Request) -> Result, Self::AError> { + req.headers() + .get(&self.answer_header) + .and_then(|a| a.to_str().ok()) + .map(|a| Self::Answer::try_from(a.to_string())) + .transpose() + } +} + +impl CaptchaFinder for CaptchaFormFinder +where + T: TryFrom + Sync + Send, + A: TryFrom + Sync + Send, + >::Error: Send, + >::Error: std::fmt::Debug, + >::Error: Send, + >::Error: std::fmt::Debug, +{ + type Token = T; + type Answer = A; + + type TError = >::Error; + type AError = >::Error; + + async fn find_token(&self, req: &mut Request) -> Result, Self::TError> { + req.form::(&self.token_name) + .await + .map(|t| Self::Token::try_from(t.to_string())) + .transpose() + } + + async fn find_answer(&self, req: &mut Request) -> Result, Self::AError> { + req.form::(&self.answer_name) + .await + .map(|a| Self::Answer::try_from(a.to_string())) + .transpose() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use salvo_core::http::headers::ContentType; + use salvo_core::http::Request; + use salvo_core::http::{header::*, ReqBody}; + + #[tokio::test] + async fn test_captcha_header_finder() { + let finder = CaptchaHeaderFinder::::new(); + let mut req = Request::default(); + let headers = req.headers_mut(); + let token = uuid::Uuid::new_v4(); + headers.insert( + HeaderName::from_static("x-captcha-token"), + HeaderValue::from_str(&token.to_string()).unwrap(), + ); + headers.insert( + HeaderName::from_static("x-captcha-answer"), + HeaderValue::from_static("answer"), + ); + assert_eq!( + finder.find_token(&mut req).await, + Ok(Some(token.to_string())) + ); + assert!(matches!( + finder.find_answer(&mut req).await, + Ok(Some(a)) if a == *"answer" + )); + } + + #[tokio::test] + async fn test_captcha_header_finder_customized() { + let finder = CaptchaHeaderFinder::::new() + .token_header(HeaderName::from_static("token")) + .answer_header(HeaderName::from_static("answer")); + let mut req = Request::default(); + let headers = req.headers_mut(); + let token = uuid::Uuid::new_v4(); + headers.insert( + HeaderName::from_static("token"), + HeaderValue::from_str(&token.to_string()).unwrap(), + ); + headers.insert( + HeaderName::from_static("answer"), + HeaderValue::from_static("answer"), + ); + assert_eq!( + finder.find_token(&mut req).await, + Ok(Some(token.to_string())) + ); + assert!(matches!( + finder.find_answer(&mut req).await, + Ok(Some(a)) if a == *"answer" + )); + } + + #[tokio::test] + async fn test_captcha_header_finder_none() { + let finder = CaptchaHeaderFinder::::new(); + let mut req = Request::default(); + + assert_eq!(finder.find_token(&mut req).await, Ok(None)); + assert_eq!(finder.find_answer(&mut req).await, Ok(None)); + } + + #[tokio::test] + async fn test_captcha_header_finder_customized_none() { + let finder = CaptchaHeaderFinder::::new() + .token_header(HeaderName::from_static("token")) + .answer_header(HeaderName::from_static("answer")); + let mut req = Request::default(); + + assert_eq!(finder.find_token(&mut req).await, Ok(None)); + assert_eq!(finder.find_answer(&mut req).await, Ok(None)); + } + + #[tokio::test] + async fn test_captcha_form_finder() { + let finder = CaptchaFormFinder::::new(); + let mut req = Request::default(); + *req.body_mut() = ReqBody::Once("captcha_token=token&captcha_answer=answer".into()); + let headers = req.headers_mut(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_str(&ContentType::form_url_encoded().to_string()).unwrap(), + ); + + assert_eq!( + finder.find_token(&mut req).await, + Ok(Some("token".to_string())) + ); + assert!(matches!( + finder.find_answer(&mut req).await, + Ok(Some(a)) if a == *"answer" + )); + } + + #[tokio::test] + async fn test_captcha_form_finder_customized() { + let finder = CaptchaFormFinder::::new() + .token_name("token".to_string()) + .answer_name("answer".to_string()); + let mut req = Request::default(); + *req.body_mut() = ReqBody::Once("token=token&answer=answer".into()); + let headers = req.headers_mut(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_str(&ContentType::form_url_encoded().to_string()).unwrap(), + ); + + assert_eq!( + finder.find_token(&mut req).await, + Ok(Some("token".to_string())) + ); + assert!(matches!( + finder.find_answer(&mut req).await, + Ok(Some(a)) if a == *"answer" + )); + } + + #[tokio::test] + async fn test_captcha_form_finder_none() { + let finder = CaptchaFormFinder::::new(); + let mut req = Request::default(); + *req.body_mut() = ReqBody::Once("".into()); + let headers = req.headers_mut(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_str(&ContentType::form_url_encoded().to_string()).unwrap(), + ); + + assert_eq!(finder.find_token(&mut req).await, Ok(None)); + assert_eq!(finder.find_answer(&mut req).await, Ok(None)); + } + + #[tokio::test] + async fn test_captcha_form_finder_customized_none() { + let finder = CaptchaFormFinder::::new() + .token_name("token".to_string()) + .answer_name("answer".to_string()); + let mut req = Request::default(); + *req.body_mut() = ReqBody::Once("".into()); + let headers = req.headers_mut(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_str(&ContentType::form_url_encoded().to_string()).unwrap(), + ); + + assert_eq!(finder.find_token(&mut req).await, Ok(None)); + assert_eq!(finder.find_answer(&mut req).await, Ok(None)); + } + + #[tokio::test] + async fn test_captcha_form_finder_invalid() { + let finder = CaptchaFormFinder::::new(); + let mut req = Request::default(); + *req.body_mut() = ReqBody::Once("captcha_token=token&captcha_answer=answer".into()); + let headers = req.headers_mut(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_str(&ContentType::json().to_string()).unwrap(), + ); + + assert_eq!(finder.find_token(&mut req).await, Ok(None)); + assert_eq!(finder.find_answer(&mut req).await, Ok(None)); + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..0e6a7b6 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,174 @@ +// 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. + +#![doc = include_str!("../README.md")] +#![deny(warnings)] +#![deny(missing_docs)] +#![deny(clippy::print_stdout)] + +mod captcha_gen; +mod finder; +mod storage; + +use salvo_core::{ + handler::{none_skipper, Skipper}, + Depot, FlowCtrl, Handler, Request, Response, +}; +pub use {captcha_gen::*, finder::*, storage::*}; + +// Exports from other crates +pub use captcha::{CaptchaName, Difficulty as CaptchaDifficulty}; + +/// Key used to insert the captcha state into the depot +pub const CAPTCHA_STATE_KEY: &str = "::salvo_captcha::captcha_state"; + +/// Captcha struct, contains the token and answer. +#[non_exhaustive] +#[allow(clippy::type_complexity)] +pub struct Captcha +where + S: CaptchaStorage, + F: CaptchaFinder, +{ + /// The captcha finder, used to find the captcha token and answer from the request. + finder: F, + /// The storage of the captcha, used to store and get the captcha token and answer. + storage: S, + /// The skipper of the captcha, used to skip the captcha check. + skipper: Box, +} + +/// The captcha states of the request +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CaptchaState { + /// The captcha is checked and passed. If the captcha is passed, it will be cleared from the storage. + Passed, + /// The captcha check is skipped. This depends on the skipper. + Skipped, + /// Can't find the captcha token in the request + TokenNotFound, + /// Can't find the captcha answer in the request + AnswerNotFound, + /// The captcha token is wrong, can't find the captcha in the storage. + /// Maybe the captcha token entered by the user is wrong, or the captcha is expired, because the storage has been cleared. + WrongToken, + /// The captcha answer is wrong. This will not clear the captcha from the storage. + WrongAnswer, + /// Storage error + StorageError, +} + +impl Captcha +where + S: CaptchaStorage, + F: CaptchaFinder, +{ + /// Create a new Captcha + pub fn new(storage: impl Into, finder: impl Into) -> Self { + Self { + finder: finder.into(), + storage: storage.into(), + skipper: Box::new(none_skipper), + } + } + + /// Returns the captcha storage + pub fn storage(&self) -> &S { + &self.storage + } + + /// Set the captcha skipper, the skipper will be used to check if the captcha check should be skipped. + pub fn skipper(mut self, skipper: impl Skipper) -> Self { + self.skipper = Box::new(skipper); + self + } +} + +/// The captcha extension of the depot. +/// Used to get the captcha info from the depot. +#[easy_ext::ext(CaptchaDepotExt)] +impl Depot { + /// Get the captcha state from the depot + pub fn get_captcha_state(&self) -> Option<&CaptchaState> { + self.get(CAPTCHA_STATE_KEY).ok() + } +} + +#[async_trait::async_trait] +impl Handler for Captcha +where + S: CaptchaStorage, + F: CaptchaFinder + 'static, +{ + async fn handle( + &self, + req: &mut Request, + depot: &mut Depot, + _: &mut Response, + _: &mut FlowCtrl, + ) { + if self.skipper.skipped(req, depot) { + log::info!("Captcha check is skipped"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::Skipped); + return; + } + + let token = match self.finder.find_token(req).await { + Ok(Some(token)) => token, + Ok(None) => { + log::info!("Captcha token is not found in request"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::TokenNotFound); + return; + } + Err(err) => { + log::error!("Failed to find captcha token from request: {err:?}"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::WrongToken); + return; + } + }; + + let answer = match self.finder.find_answer(req).await { + Ok(Some(answer)) => answer, + Ok(None) => { + log::info!("Captcha answer is not found in request"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::AnswerNotFound); + return; + } + Err(err) => { + log::error!("Failed to find captcha answer from request: {err:?}"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::WrongAnswer); + return; + } + }; + + match self.storage.get_answer(&token).await { + Ok(Some(captch_answer)) => { + log::info!("Captcha answer is exist in storage for token: {token}"); + if captch_answer == answer { + log::info!("Captcha answer is correct for token: {token}"); + self.storage.clear_by_token(&token).await.ok(); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::Passed); + } else { + log::info!("Captcha answer is wrong for token: {token}"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::WrongAnswer); + } + } + Ok(None) => { + log::info!("Captcha answer is not exist in storage for token: {token}"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::WrongToken); + } + Err(err) => { + log::error!("Failed to get captcha answer from storage: {err}"); + depot.insert(CAPTCHA_STATE_KEY, CaptchaState::StorageError); + } + }; + } +} diff --git a/src/storage.rs b/src/storage.rs new file mode 100644 index 0000000..8b2e06c --- /dev/null +++ b/src/storage.rs @@ -0,0 +1,368 @@ +// 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. + +use std::{sync::Arc, time::Duration}; + +#[cfg(feature = "cacache-storage")] +use std::{ + path::{Path, PathBuf}, + time::SystemTime, +}; + +/// Trait to store the captcha token and answer. is also clear the expired captcha. +/// +/// The trait will be implemented for `Arc` if `T` implements the trait. +/// +/// The trait is thread safe, so the storage can be shared between threads. +pub trait CaptchaStorage: Send + Sync + 'static +where + Self: Clone + std::fmt::Debug, +{ + /// The token type + type Token: TryFrom + std::fmt::Display + Send + Sync; + /// The answer type + type Answer: std::cmp::PartialEq + From + Send + Sync; + /// The error type of the storage. + type Error: std::fmt::Display + std::fmt::Debug + Send; + + /// Store the captcha token and answer. + fn store_answer( + &self, + answer: Self::Answer, + ) -> impl std::future::Future> + Send; + + /// Returns the answer of the captcha token. This method will return None if the token is not exist. + fn get_answer( + &self, + token: &Self::Token, + ) -> impl std::future::Future, Self::Error>> + Send; + + /// Clear the expired captcha. + fn clear_expired( + &self, + expired_after: Duration, + ) -> impl std::future::Future> + Send; + + /// Clear the captcha by token. + fn clear_by_token( + &self, + token: &Self::Token, + ) -> impl std::future::Future> + Send; +} + +/// The [`cacache`] storage. +/// +/// [`cacache`]: https://github.com/zkat/cacache-rs +#[cfg(feature = "cacache-storage")] +#[derive(Debug, Clone)] +pub struct CacacheStorage { + /// The cacache cache directory. + cache_dir: PathBuf, +} + +#[cfg(feature = "cacache-storage")] +impl CacacheStorage { + /// Create a new CacacheStorage + pub fn new(cache_dir: impl Into) -> Self { + Self { + cache_dir: cache_dir.into(), + } + } + + /// Get the cacache cache directory. + pub fn cache_dir(&self) -> &Path { + &self.cache_dir + } +} + +#[cfg(feature = "cacache-storage")] +impl CaptchaStorage for CacacheStorage { + type Error = cacache::Error; + type Token = String; + type Answer = String; + + async fn store_answer(&self, answer: Self::Answer) -> Result { + let token = uuid::Uuid::new_v4(); + log::info!("Storing captcha answer to cacache for token: {token}"); + cacache::write(&self.cache_dir, token.to_string(), answer.as_bytes()).await?; + Ok(token.to_string()) + } + + async fn get_answer(&self, token: &Self::Token) -> Result, Self::Error> { + log::info!("Getting captcha answer from cacache for token: {token}"); + match cacache::read(&self.cache_dir, token.to_string()).await { + Ok(answer) => { + log::info!("Captcha answer is exist in cacache for token: {token}"); + Ok(Some( + String::from_utf8(answer) + .expect("All the stored captcha answer should be utf8"), + )) + } + Err(cacache::Error::EntryNotFound(_, _)) => { + log::info!("Captcha answer is not exist in cacache for token: {token}"); + Ok(None) + } + Err(err) => { + log::error!("Failed to get captcha answer from cacache for token: {token}"); + Err(err) + } + } + } + + async fn clear_expired(&self, expired_after: Duration) -> Result<(), Self::Error> { + let now = SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("SystemTime must be later than UNIX_EPOCH") + .as_millis(); + let expired_after = expired_after.as_millis(); + + let expr_keys = cacache::index::ls(&self.cache_dir).filter_map(|meta| { + if let Ok(meta) = meta { + if now >= (meta.time + expired_after) { + return Some(meta.key); + } + } + None + }); + + for key in expr_keys { + cacache::RemoveOpts::new() + .remove_fully(true) + .remove(&self.cache_dir, &key) + .await + .ok(); + } + Ok(()) + } + + async fn clear_by_token(&self, token: &Self::Token) -> Result<(), Self::Error> { + log::info!("Clearing captcha token from cacache: {token}"); + let remove_opts = cacache::RemoveOpts::new().remove_fully(true); + remove_opts.remove(&self.cache_dir, token.to_string()).await + } +} + +impl CaptchaStorage for Arc +where + T: CaptchaStorage, +{ + type Error = T::Error; + type Token = T::Token; + type Answer = T::Answer; + + fn store_answer( + &self, + answer: Self::Answer, + ) -> impl std::future::Future> + Send { + self.as_ref().store_answer(answer) + } + + fn get_answer( + &self, + token: &Self::Token, + ) -> impl std::future::Future, Self::Error>> + Send { + self.as_ref().get_answer(token) + } + + fn clear_expired( + &self, + expired_after: Duration, + ) -> impl std::future::Future> + Send { + self.as_ref().clear_expired(expired_after) + } + + fn clear_by_token( + &self, + token: &Self::Token, + ) -> impl std::future::Future> + Send { + self.as_ref().clear_by_token(token) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + #[cfg(feature = "cacache-storage")] + async fn cacache_store_captcha() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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] + #[cfg(feature = "cacache-storage")] + async fn cacache_clear_expired() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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] + #[cfg(feature = "cacache-storage")] + async fn cacache_clear_by_token() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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] + #[cfg(feature = "cacache-storage")] + async fn cacache_is_token_exist() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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".to_owned()) + .await + .expect("failed to check if token is exist") + .is_none()); + } + + #[tokio::test] + #[cfg(feature = "cacache-storage")] + async fn cacache_get_answer() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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()) + ); + assert!(storage + .get_answer(&"token".to_owned()) + .await + .expect("failed to get captcha answer") + .is_none()); + } + + #[tokio::test] + #[cfg(feature = "cacache-storage")] + async fn cacache_cache_dir() { + let cache_dir = tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(); + let storage = CacacheStorage::new(cache_dir.clone()); + assert_eq!(storage.cache_dir(), &cache_dir); + } + + #[tokio::test] + #[cfg(feature = "cacache-storage")] + async fn cacache_clear_expired_with_expired_after() { + let storage = CacacheStorage::new( + tempfile::tempdir() + .expect("failed to create temp file") + .path() + .to_owned(), + ); + + 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()); + } +}