salvo_csrf/
finder.rs

1use std::collections::HashMap;
2
3use salvo_core::{Request, async_trait};
4use serde_json::Value;
5
6/// Used to find csrf token from request.
7#[async_trait]
8pub trait CsrfTokenFinder: Send + Sync + 'static {
9    /// Find token from request.
10    async fn find_token(&self, req: &mut Request) -> Option<String>;
11}
12
13/// Find token from http request header.
14#[derive(Clone, Debug)]
15pub struct HeaderFinder {
16    header_name: String,
17}
18impl HeaderFinder {
19    /// Creates a new `HeaderFinder`; you can use a value like `x-csrf-token`.
20    #[inline]
21    pub fn new(header_name: impl Into<String>) -> Self {
22        Self {
23            header_name: header_name.into(),
24        }
25    }
26}
27#[async_trait]
28impl CsrfTokenFinder for HeaderFinder {
29    #[inline]
30    async fn find_token(&self, req: &mut Request) -> Option<String> {
31        req.header(&self.header_name)
32    }
33}
34
35/// Find token from request form body.
36#[derive(Clone, Debug)]
37pub struct FormFinder {
38    field_name: String,
39}
40impl FormFinder {
41    /// Creates a new `FormFinder`.
42    #[inline]
43    pub fn new(field_name: impl Into<String>) -> Self {
44        Self {
45            field_name: field_name.into(),
46        }
47    }
48}
49#[async_trait]
50impl CsrfTokenFinder for FormFinder {
51    #[inline]
52    async fn find_token(&self, req: &mut Request) -> Option<String> {
53        req.form(&self.field_name).await
54    }
55}
56
57/// Find token from request json body.
58#[derive(Clone, Debug)]
59pub struct JsonFinder {
60    field_name: String,
61}
62impl JsonFinder {
63    /// Creates a new `JsonFinder`.
64    #[inline]
65    pub fn new(field_name: impl Into<String>) -> Self {
66        Self {
67            field_name: field_name.into(),
68        }
69    }
70}
71#[async_trait]
72impl CsrfTokenFinder for JsonFinder {
73    async fn find_token(&self, req: &mut Request) -> Option<String> {
74        let data = req.parse_json::<HashMap<String, Value>>().await;
75        if let Ok(data) = data
76            && let Some(value) = data.get(&self.field_name)
77            && let Some(token) = value.as_str()
78        {
79            Some(token.to_owned())
80        } else {
81            None
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use salvo_core::test::TestClient;
89
90    use super::*;
91
92    #[tokio::test]
93    async fn test_header_finder() {
94        let header_finder = HeaderFinder::new("x-csrf-token");
95        let mut req = TestClient::get("http://test.com")
96            .add_header("x-csrf-token", "test_token", true)
97            .build();
98        let token = header_finder.find_token(&mut req).await;
99        assert_eq!(token, Some("test_token".to_owned()));
100    }
101
102    #[tokio::test]
103    async fn test_form_finder() {
104        let form_finder = FormFinder::new("csrf-token");
105        let mut req = TestClient::get("http://test.com")
106            .raw_form("csrf-token=test_token")
107            .build();
108        let token = form_finder.find_token(&mut req).await;
109        assert_eq!(token, Some("test_token".to_owned()));
110    }
111
112    #[tokio::test]
113    async fn test_json_finder() {
114        let json_finder = JsonFinder::new("csrf-token");
115        let mut req = TestClient::get("http://test.com")
116            .raw_json(r#"{"csrf-token":"test_token"}"#)
117            .build();
118        let token = json_finder.find_token(&mut req).await;
119        assert_eq!(token, Some("test_token".to_owned()));
120    }
121
122    #[tokio::test]
123    async fn test_json_finder_not_string() {
124        let json_finder = JsonFinder::new("csrf-token");
125        let mut req = TestClient::get("http://test.com")
126            .raw_json(r#"{"csrf-token":123}"#)
127            .build();
128        let token = json_finder.find_token(&mut req).await;
129        assert_eq!(token, None);
130    }
131}
salvo_csrf/finder.rs

salvo_csrf/
finder.rs
