salvo_captcha/finder/
header_finder.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
12use salvo_core::http::{HeaderName, Request};
13
14use crate::CaptchaFinder;
15
16/// Find the captcha token and answer from the header
17#[derive(Debug)]
18pub struct CaptchaHeaderFinder {
19    /// The header name of the captcha token
20    ///
21    /// Default: "x-captcha-token"
22    pub token_header: HeaderName,
23
24    /// The header name of the captcha answer
25    ///
26    /// Default: "x-captcha-answer"
27    pub answer_header: HeaderName,
28}
29
30impl CaptchaHeaderFinder {
31    /// Create a new CaptchaHeaderFinder
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Set the token header name
37    pub fn token_header(mut self, token_header: HeaderName) -> Self {
38        self.token_header = token_header;
39        self
40    }
41
42    /// Set the answer header name
43    pub fn answer_header(mut self, answer_header: HeaderName) -> Self {
44        self.answer_header = answer_header;
45        self
46    }
47}
48
49impl Default for CaptchaHeaderFinder {
50    /// Create a default CaptchaHeaderFinder with:
51    /// - token_header: "x-captcha-token"
52    /// - answer_header: "x-captcha-answer"
53    fn default() -> Self {
54        Self {
55            token_header: HeaderName::from_static("x-captcha-token"),
56            answer_header: HeaderName::from_static("x-captcha-answer"),
57        }
58    }
59}
60
61impl CaptchaFinder for CaptchaHeaderFinder {
62    async fn find_token(&self, req: &mut Request) -> Option<Option<String>> {
63        req.headers()
64            .get(&self.token_header)
65            .map(|t| t.to_str().map(ToString::to_string).ok())
66    }
67
68    async fn find_answer(&self, req: &mut Request) -> Option<Option<String>> {
69        req.headers()
70            .get(&self.answer_header)
71            .map(|a| a.to_str().map(ToString::to_string).ok())
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use salvo_core::http::HeaderValue;
79
80    #[tokio::test]
81    #[rstest::rstest]
82    #[case::not_found(None, None, None, None, None, None)]
83    #[case::normal(
84         None,
85         None,
86         Some(("x-captcha-token", "token")),
87         Some(("x-captcha-answer", "answer")),
88         Some(Some("token")),
89         Some(Some("answer"))
90    )]
91    #[case::custom_headers(
92         Some("custom-token"),
93         Some("custom-answer"),
94         Some(("custom-token", "token")),
95         Some(("custom-answer", "answer")),
96         Some(Some("token")),
97         Some(Some("answer"))
98    )]
99    #[case::only_token(
100         None,
101         None,
102         Some(("x-captcha-token", "token")),
103         None,
104         Some(Some("token")),
105         None
106    )]
107    #[case::only_answer(
108         None,
109         None,
110         None,
111         Some(("x-captcha-answer", "answer")),
112         None,
113         Some(Some("answer"))
114    )]
115    #[case::custom_not_found(Some("custom-token"), Some("custom-answer"), None, None, None, None)]
116    #[case::custom_not_found_with_headers(
117         Some("custom-token"),
118         Some("custom-answer"),
119         Some(("x-captcha-token", "token")),
120         Some(("x-captcha-answer", "answer")),
121         None,
122         None
123    )]
124    async fn test_header_finder(
125        #[case] custom_token_header: Option<&'static str>,
126        #[case] custom_answer_header: Option<&'static str>,
127        #[case] token_header_name_value: Option<(&'static str, &'static str)>,
128        #[case] answer_header_name_value: Option<(&'static str, &'static str)>,
129        #[case] excepted_token: Option<Option<&'static str>>,
130        #[case] excepted_answer: Option<Option<&'static str>>,
131    ) {
132        let mut finder = CaptchaHeaderFinder::new();
133        if let Some(custom_token) = custom_token_header {
134            finder = finder.token_header(HeaderName::from_static(custom_token));
135        }
136        if let Some(custom_answer) = custom_answer_header {
137            finder = finder.answer_header(HeaderName::from_static(custom_answer));
138        }
139
140        let mut req = Request::default();
141        let headers = req.headers_mut();
142        if let Some((token_header_name, token_header_value)) = token_header_name_value {
143            headers.insert(
144                HeaderName::from_static(token_header_name),
145                HeaderValue::from_static(token_header_value),
146            );
147        }
148        if let Some((answer_header_name, answer_header_value)) = answer_header_name_value {
149            headers.insert(
150                HeaderName::from_static(answer_header_name),
151                HeaderValue::from_static(answer_header_value),
152            );
153        }
154
155        assert_eq!(
156            finder.find_token(&mut req).await,
157            excepted_token.map(|o| o.map(ToOwned::to_owned))
158        );
159        assert_eq!(
160            finder.find_answer(&mut req).await,
161            excepted_answer.map(|o| o.map(ToOwned::to_owned))
162        );
163    }
164}
salvo_captcha/finder/header_finder.rs

salvo_captcha/finder/
header_finder.rs
