limiting_factor_axum/api/
guards.rs

1/*  -------------------------------------------------------------
2    Limiting Factor :: axum :: API :: Guards
3    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
4    Project:        Nasqueron
5    License:        BSD-2-Clause
6    -------------------------------------------------------------    */
7
8//! # API extractors
9//!
10//! This module provides reusable extractors to use with axum.
11
12use axum::{
13    extract::{FromRequest, Request},
14    http::StatusCode,
15    response::{IntoResponse, Response},
16};
17use http_body_util::BodyExt;
18
19use limiting_factor_core::api::guards::{RequestBody, REQUEST_BODY_LIMIT};
20
21// New-type wrapper for Axum-specific implementations
22#[derive(Debug, Clone)]
23pub struct AxumRequestBody(pub RequestBody);
24
25impl AxumRequestBody {
26    pub fn new() -> Self {
27        Self(RequestBody::new())
28    }
29
30    // Delegate methods
31    pub fn into_string(self) -> String {
32        self.0.into_string()
33    }
34
35    pub fn into_optional_string(self) -> Option<String> {
36        self.0.into_optional_string()
37    }
38}
39
40impl From<RequestBody> for AxumRequestBody {
41    fn from(data: RequestBody) -> Self {
42        Self(data)
43    }
44}
45
46impl From<AxumRequestBody> for RequestBody {
47    fn from(body: AxumRequestBody) -> Self {
48        body.0
49    }
50}
51
52/// Error type during a request body extraction
53#[derive(Debug)]
54pub enum RequestBodyError {
55    /// Body size is greater than REQUEST_BODY_LIMIT (DoS risk)
56    TooLarge,
57
58    /// Not in UTF-8 encoding
59    InvalidEncoding,
60
61    /// I/O error
62    ReadError(String),
63}
64
65impl IntoResponse for RequestBodyError {
66    fn into_response(self) -> Response {
67        let (status, message) = match self {
68            RequestBodyError::TooLarge => (
69                StatusCode::PAYLOAD_TOO_LARGE,
70                "Request body too large".to_string(),
71            ),
72
73            RequestBodyError::InvalidEncoding => (
74                StatusCode::BAD_REQUEST,
75                "Request body contains invalid characters when trying to decode as UTF-8".to_string(),
76            ),
77
78            RequestBodyError::ReadError(err) => (
79                StatusCode::INTERNAL_SERVER_ERROR,
80                format!("Failed to read request body: {}", err),
81            ),
82        };
83
84        (status, message).into_response()
85    }
86}
87
88impl<S> FromRequest<S> for AxumRequestBody
89where
90    S: Send + Sync,
91{
92    type Rejection = RequestBodyError;
93
94    async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
95        // Extract the body from the request
96        let body = req.into_body();
97
98        // Collect the body with size limit
99        let collected = match body.collect().await {
100            Ok(collected) => collected,
101            Err(e) => return Err(RequestBodyError::ReadError(e.to_string())),
102        };
103
104        let bytes = collected.to_bytes();
105
106        // Check size limit
107        if bytes.len() > REQUEST_BODY_LIMIT {
108            return Err(RequestBodyError::TooLarge);
109        }
110
111        // Convert to UTF-8 string
112        let content = match String::from_utf8(bytes.to_vec()) {
113            Ok(content) => content,
114            Err(_) => return Err(RequestBodyError::InvalidEncoding),
115        };
116
117        Ok(Self(RequestBody { content }))
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[tokio::test]
126    async fn test_empty_body_extraction() {
127        use axum::body::Body;
128        use axum::http::Request;
129
130        let req = Request::builder()
131            .body(Body::empty())
132            .unwrap();
133
134        let body = AxumRequestBody::from_request(req, &()).await.unwrap();
135        assert_eq!("", body.0.content);
136        assert_eq!(None, body.into_optional_string());
137    }
138
139    #[tokio::test]
140    async fn test_body_extraction() {
141        use axum::body::Body;
142        use axum::http::Request;
143
144        let req = Request::builder()
145            .body(Body::from("lorem ipsum dolor"))
146            .unwrap();
147
148        let body = AxumRequestBody::from_request(req, &()).await.unwrap();
149        assert_eq!("lorem ipsum dolor", body.0.content);
150        assert_eq!(Some("lorem ipsum dolor".to_string()), body.into_optional_string());
151    }
152}
limiting_factor_axum/api/guards.rs

limiting_factor_axum/api/
guards.rs
