Skip to main content

siderite_core/middleware/
limits.rs

1//! Timeouts, concurrency, body-size and rate limits.
2
3use super::adapt::{BoxService, Next, from_fn, impl_layer, wrap_engine};
4use crate::error::ApiError;
5use crate::extract::Request;
6use crate::response::IntoResponse;
7use http::{StatusCode, header};
8use std::time::Duration;
9use tower::Layer;
10use tower_http::limit::RequestBodyLimitLayer;
11
12/// Fails requests that take longer than a deadline (default status 504).
13///
14/// The deadline covers producing the response, not streaming its body. The
15/// status is 504 Gateway Timeout by default (the server, not the client, was
16/// too slow); [`Timeout::status`] changes it, e.g. to 408.
17#[derive(Debug, Clone)]
18pub struct Timeout {
19    limit: Duration,
20    status: StatusCode,
21}
22
23impl Timeout {
24    /// Time out after `limit`.
25    pub fn new(limit: Duration) -> Self {
26        Self {
27            limit,
28            status: StatusCode::GATEWAY_TIMEOUT,
29        }
30    }
31
32    /// Status code of the problem returned on timeout.
33    #[must_use]
34    pub fn status(mut self, status: StatusCode) -> Self {
35        self.status = status;
36        self
37    }
38
39    fn wrap(&self, inner: BoxService) -> BoxService {
40        let (limit, status) = (self.limit, self.status);
41        from_fn(move |req: Request, next: Next| async move {
42            match tokio::time::timeout(limit, next.run(req)).await {
43                Ok(response) => response,
44                Err(_) => ApiError::new(status, "The request timed out.").into_response(),
45            }
46        })
47        .layer(inner)
48    }
49}
50
51/// Rejects request bodies larger than a limit with a 413 problem.
52///
53/// Checked eagerly via `Content-Length`; bodies without it (chunked) are
54/// capped while being read, which surfaces as a body error in the extractor.
55#[derive(Debug, Clone)]
56pub struct BodyLimit {
57    max_bytes: usize,
58}
59
60impl BodyLimit {
61    /// Allow bodies up to `max_bytes`.
62    pub fn new(max_bytes: usize) -> Self {
63        Self { max_bytes }
64    }
65
66    fn wrap(&self, inner: BoxService) -> BoxService {
67        let max = self.max_bytes;
68        let capped = wrap_engine(&RequestBodyLimitLayer::new(max), inner);
69        from_fn(move |req: Request, next: Next| async move {
70            let declared = req
71                .headers()
72                .get(header::CONTENT_LENGTH)
73                .and_then(|v| v.to_str().ok())
74                .and_then(|v| v.parse::<usize>().ok());
75            if declared.is_some_and(|len| len > max) {
76                return ApiError::new(
77                    StatusCode::PAYLOAD_TOO_LARGE,
78                    format!("Request body exceeds the limit of {max} bytes."),
79                )
80                .into_response();
81            }
82            next.run(req).await
83        })
84        .layer(capped)
85    }
86}
87
88impl_layer!(Timeout, BodyLimit);