rango_core/
forgery.rs

1//! Cross-site request forgery protection.
2
3use axum::{
4    body::to_bytes,
5    http::{
6        Method,
7        header::{COOKIE, HeaderMap, HeaderValue, SET_COOKIE},
8    },
9    middleware::Next,
10};
11
12use crate::{
13    error::Error,
14    view::{Request, Response},
15};
16
17/// Maximum request body size the guard reads (2 MiB).
18pub const LIMIT: usize = 2 * 1024 * 1024;
19
20const NAME: &str = "forgery";
21
22/// Request-local token the guard attaches to processed requests.
23#[derive(Clone)]
24pub struct Token(
25    /// Token value.
26    pub String,
27);
28
29fn generate() -> String {
30    uuid::Uuid::new_v4().to_string()
31}
32
33fn safe(method: &Method) -> bool {
34    matches!(
35        method,
36        &Method::GET | &Method::HEAD | &Method::OPTIONS | &Method::TRACE
37    )
38}
39
40fn parse_field(body: &[u8]) -> Option<String> {
41    let map: std::collections::HashMap<String, String> = serde_urlencoded::from_bytes(body).ok()?;
42    map.get(NAME).cloned()
43}
44
45fn signed(token: &str) -> HeaderValue {
46    HeaderValue::from_str(&format!("{NAME}={token}; Path=/; HttpOnly; SameSite=Lax")).unwrap()
47}
48
49/// Forgery token value from the `forgery` cookie, if present.
50pub fn cookie(headers: &HeaderMap) -> Option<String> {
51    named(headers, NAME)
52}
53
54/// Cookie value of a single named cookie.
55pub fn named(headers: &HeaderMap, name: &str) -> Option<String> {
56    let value = headers.get(COOKIE)?.to_str().ok()?;
57    value.split(';').map(str::trim).find_map(|part| {
58        let (key, value) = part.split_once('=')?;
59        (key == name).then(|| value.to_string())
60    })
61}
62
63/// Token for the request: the extension value, the cookie, or a fresh one.
64pub fn token(req: &Request) -> String {
65    if let Some(token) = req.extensions().get::<Token>() {
66        return token.0.clone();
67    }
68    cookie(req.headers()).unwrap_or_else(generate)
69}
70
71/// Middleware that checks the token field on unsafe methods and issues a cookie.
72pub async fn guard(req: Request, next: Next) -> Result<Response, Error> {
73    let safe = safe(req.method());
74    let (mut parts, body) = req.into_parts();
75    let bytes = to_bytes(body, LIMIT)
76        .await
77        .map_err(|_| Error::BadRequest("body too large".into()))?;
78    let existing = cookie(&parts.headers);
79
80    if !safe {
81        let Some(token) = existing.clone() else {
82            return Err(Error::Forbidden);
83        };
84        if parse_field(&bytes).as_deref() != Some(token.as_str()) {
85            return Err(Error::Forbidden);
86        }
87    }
88
89    let token = existing.clone().unwrap_or_else(generate);
90    parts.extensions.insert(Token(token.clone()));
91    let mut response = next.run(Request::from_parts(parts, bytes.into())).await;
92
93    if existing.is_none() {
94        response.headers_mut().insert(SET_COOKIE, signed(&token));
95    }
96    Ok(response)
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use axum::{Router, body::Body, http::StatusCode, middleware, routing::get};
103    use tower::ServiceExt;
104
105    #[test]
106    fn cookies() {
107        let mut headers = HeaderMap::new();
108        headers.insert(COOKIE, "a=1; forgery=xyz; b=2".parse().unwrap());
109        assert_eq!(cookie(&headers), Some("xyz".into()));
110        assert_eq!(named(&headers, "b"), Some("2".into()));
111        assert_eq!(named(&headers, "none"), None);
112    }
113
114    #[test]
115    fn tokens() {
116        let mut with_ext = axum::http::Request::<()>::builder()
117            .body(Body::empty())
118            .unwrap();
119        with_ext.extensions_mut().insert(Token("abc".into()));
120        assert_eq!(token(&with_ext), "abc");
121        let req = axum::http::Request::<()>::builder()
122            .body(Body::empty())
123            .unwrap();
124        assert_ne!(token(&req), token(&req));
125    }
126
127    async fn probe(req: Request) -> String {
128        req.extensions()
129            .get::<Token>()
130            .map(|t| t.0.clone())
131            .unwrap_or_default()
132    }
133
134    fn build(method: &str, set: Option<&str>, body: &'static str) -> axum::extract::Request {
135        let mut request = axum::http::Request::<()>::builder()
136            .method(method)
137            .uri("/")
138            .header("content-type", "application/x-www-form-urlencoded");
139        if let Some(cookie) = set {
140            request = request.header(COOKIE, cookie);
141        }
142        request.body(Body::from(body)).unwrap()
143    }
144
145    async fn body_of(response: Response) -> String {
146        let bytes = axum::body::to_bytes(response.into_body(), 1024)
147            .await
148            .unwrap();
149        String::from_utf8(bytes.to_vec()).unwrap()
150    }
151
152    fn app() -> Router {
153        Router::new()
154            .route("/", get(probe))
155            .layer(middleware::from_fn(guard))
156    }
157
158    #[tokio::test]
159    async fn issues_token() {
160        let res = app().oneshot(build("GET", None, "")).await.unwrap();
161        assert_eq!(res.status(), StatusCode::OK);
162        let issued = res
163            .headers()
164            .get(SET_COOKIE)
165            .unwrap()
166            .to_str()
167            .unwrap()
168            .split(';')
169            .next()
170            .unwrap()
171            .strip_prefix("forgery=")
172            .unwrap()
173            .to_string();
174        assert_eq!(issued, body_of(res).await);
175    }
176
177    #[tokio::test]
178    async fn keeps_cookie() {
179        let res = app()
180            .oneshot(build("GET", Some("forgery=stated"), ""))
181            .await
182            .unwrap();
183        assert_eq!(res.status(), StatusCode::OK);
184        assert!(res.headers().get(SET_COOKIE).is_none());
185        assert_eq!(body_of(res).await, "stated");
186    }
187
188    #[tokio::test]
189    async fn guards_post() {
190        let app = Router::new()
191            .route("/", get(probe).post(probe))
192            .layer(middleware::from_fn(guard));
193        let denied = app
194            .clone()
195            .oneshot(build("POST", None, "forgery=nope"))
196            .await
197            .unwrap();
198        assert_eq!(denied.status(), StatusCode::FORBIDDEN);
199        let wrong = app
200            .clone()
201            .oneshot(build("POST", Some("forgery=said"), "forgery=other"))
202            .await
203            .unwrap();
204        assert_eq!(wrong.status(), StatusCode::FORBIDDEN);
205        let ok = app
206            .oneshot(build("POST", Some("forgery=said"), "forgery=said"))
207            .await
208            .unwrap();
209        assert_eq!(ok.status(), StatusCode::OK);
210        assert_eq!(body_of(ok).await, "said");
211    }
212
213    #[tokio::test]
214    async fn too_large() {
215        let mut req = build("POST", Some("forgery=x"), "forgery=x");
216        *req.body_mut() = Body::from(vec![b'x'; LIMIT + 1]);
217        let res = app().oneshot(req).await.unwrap();
218        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
219    }
220}
rango_core/forgery.rs

rango_core/
forgery.rs
