From 110a8728e1d41c46c6e9f2d3f889a08bc2ffb4bf Mon Sep 17 00:00:00 2001 From: Awiteb Date: Sun, 11 Aug 2024 12:09:28 +0000 Subject: [PATCH] feat: Add query finder Signed-off-by: Awiteb --- src/finder/mod.rs | 2 + src/finder/query_finder.rs | 151 +++++++++++++++++++++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100644 src/finder/query_finder.rs diff --git a/src/finder/mod.rs b/src/finder/mod.rs index e0568c7..b75c5b4 100644 --- a/src/finder/mod.rs +++ b/src/finder/mod.rs @@ -13,9 +13,11 @@ use salvo_core::http::Request; mod form_finder; mod header_finder; +mod query_finder; pub use form_finder::*; pub use header_finder::*; +pub use query_finder::*; /// Trait to find the captcha token and answer from the request. pub trait CaptchaFinder: Send + Sync { diff --git a/src/finder/query_finder.rs b/src/finder/query_finder.rs new file mode 100644 index 0000000..7f98629 --- /dev/null +++ b/src/finder/query_finder.rs @@ -0,0 +1,151 @@ +// 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 salvo_core::http::Request; + +use crate::CaptchaFinder; + +/// Find the captcha token and answer from the url query +#[derive(Debug)] +pub struct CaptchaQueryFinder { + /// The query name of the captcha token + /// + /// Default: "c_t" + pub token_name: String, + + /// The query name of the captcha answer + /// + /// Default: "c_a" + pub answer_name: String, +} + +impl CaptchaQueryFinder { + /// Create a new [`CaptchaQueryFinder`] + pub fn new() -> Self { + Self::default() + } + + /// Set the token query name + pub fn token_name(mut self, token_name: String) -> Self { + self.token_name = token_name; + self + } + + /// Set the answer query name + pub fn answer_name(mut self, answer_name: String) -> Self { + self.answer_name = answer_name; + self + } +} + +impl Default for CaptchaQueryFinder { + /// Create a default [`CaptchaQueryFinder`] with: + /// - token_name: "c_t" + /// - answer_name: "c_a" + fn default() -> Self { + Self { + token_name: "c_t".to_string(), + answer_name: "c_a".to_string(), + } + } +} + +impl CaptchaFinder for CaptchaQueryFinder { + async fn find_token(&self, req: &mut Request) -> Option> { + req.queries() + .get(&self.token_name) + .map(|o| Some(o.to_owned())) + } + + async fn find_answer(&self, req: &mut Request) -> Option> { + req.queries() + .get(&self.answer_name) + .map(|o| Some(o.to_owned())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + #[rstest::rstest] + #[case::not_found(None, None, None, None, None, None)] + #[case::normal( + None, + None, + Some(("c_t", "token")), + Some(("c_a", "answer")), + Some(Some("token")), + Some(Some("answer")) + )] + #[case::custom_keys( + Some("cc_t"), + Some("cc_a"), + Some(("cc_t", "token")), + Some(("cc_a", "answer")), + Some(Some("token")), + Some(Some("answer")) + )] + #[case::only_token( + None, + None, + Some(("c_t", "token")), + None, + Some(Some("token")), + None + )] + #[case::only_answer(None, None, None, Some(("c_a", "ans")), None, Some(Some("ans")))] + #[case::custom_not_found(Some("cc_t"), Some("cc_a"), None, None, None, None)] + #[case::custom_not_found_with_query( + Some("cc_t"), + Some("cc_a"), + Some(("c_t", "token")), + Some(("c_a", "answer")), + None, + None + )] + async fn test_query_finder( + #[case] custom_token_key: Option<&'static str>, + #[case] custom_answer_key: Option<&'static str>, + #[case] token_key_val: Option<(&'static str, &'static str)>, + #[case] answer_key_val: Option<(&'static str, &'static str)>, + #[case] excepted_token: Option>, + #[case] excepted_answer: Option>, + ) { + let mut req = Request::default(); + let mut finder = CaptchaQueryFinder::new(); + if let Some(token_key) = custom_token_key { + finder = finder.token_name(token_key.to_string()) + } + if let Some(answer_key) = custom_answer_key { + finder = finder.answer_name(answer_key.to_string()) + } + + let queries = req.queries_mut(); + + if let Some((k, v)) = token_key_val { + queries.insert(k.to_owned(), v.to_owned()); + } + if let Some((k, v)) = answer_key_val { + queries.insert(k.to_owned(), v.to_owned()); + } + + assert_eq!( + finder.find_token(&mut req).await, + excepted_token.map(|o| o.map(ToOwned::to_owned)) + ); + assert_eq!( + finder.find_answer(&mut req).await, + excepted_answer.map(|o| o.map(ToOwned::to_owned)) + ); + } +}