siderite_core/middleware/
limits.rs1use 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#[derive(Debug, Clone)]
18pub struct Timeout {
19 limit: Duration,
20 status: StatusCode,
21}
22
23impl Timeout {
24 pub fn new(limit: Duration) -> Self {
26 Self {
27 limit,
28 status: StatusCode::GATEWAY_TIMEOUT,
29 }
30 }
31
32 #[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#[derive(Debug, Clone)]
56pub struct BodyLimit {
57 max_bytes: usize,
58}
59
60impl BodyLimit {
61 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);