gatel_core/hoops/
buffer.rs

1use salvo::http::StatusCode;
2use salvo::{Depot, FlowCtrl, Request, Response, async_trait};
3use tracing::debug;
4
5/// Middleware that enforces request and response body size limits.
6///
7/// - When `max_request_body` is set: if the `Content-Length` header of the incoming request exceeds
8///   the limit, a 413 Payload Too Large response is returned immediately without forwarding to the
9///   upstream.
10/// - When `max_response_body` is set: if the `Content-Length` header of the upstream response
11///   exceeds the limit, a 502 Bad Gateway response is returned instead.
12pub struct BufferLimitHoop {
13    max_request_body: Option<usize>,
14    max_response_body: Option<usize>,
15}
16
17impl BufferLimitHoop {
18    pub fn new(max_request_body: Option<usize>, max_response_body: Option<usize>) -> Self {
19        Self {
20            max_request_body,
21            max_response_body,
22        }
23    }
24}
25
26#[async_trait]
27impl salvo::Handler for BufferLimitHoop {
28    async fn handle(
29        &self,
30        req: &mut Request,
31        depot: &mut Depot,
32        res: &mut Response,
33        ctrl: &mut FlowCtrl,
34    ) {
35        // Check the request body size via Content-Length header.
36        if let Some(max_req) = self.max_request_body
37            && let Some(content_length) = req
38                .headers()
39                .get(http::header::CONTENT_LENGTH)
40                .and_then(|v| v.to_str().ok())
41                .and_then(|s| s.parse::<usize>().ok())
42            && content_length > max_req
43        {
44            debug!(
45                content_length,
46                limit = max_req,
47                "request body exceeds limit, returning 413"
48            );
49            res.status_code(StatusCode::PAYLOAD_TOO_LARGE);
50            res.body("Payload Too Large");
51            ctrl.skip_rest();
52            return;
53        }
54
55        // Forward the request and check the response body size.
56        ctrl.call_next(req, depot, res).await;
57
58        if let Some(max_resp) = self.max_response_body
59            && let Some(content_length) = res
60                .headers()
61                .get(http::header::CONTENT_LENGTH)
62                .and_then(|v| v.to_str().ok())
63                .and_then(|s| s.parse::<usize>().ok())
64            && content_length > max_resp
65        {
66            debug!(
67                content_length,
68                limit = max_resp,
69                "response body exceeds limit, returning 502"
70            );
71            res.status_code(StatusCode::BAD_GATEWAY);
72            res.body("Bad Gateway: response body too large");
73        }
74    }
75}
gatel_core/hoops/buffer.rs

gatel_core/hoops/
buffer.rs
