api_tools/server/axum/layers/
http_errors.rs

1//! Override some HTTP errors
2
3use crate::server::axum::response::ApiError;
4use axum::body::Body;
5use axum::http::{Request, StatusCode};
6use axum::response::{IntoResponse, Response};
7use futures::future::BoxFuture;
8use std::task::{Context, Poll};
9use tower::{Layer, Service};
10
11/// Configuration for the `HttpErrorsLayer`
12#[derive(Clone, Debug)]
13pub struct HttpErrorsConfig {
14    /// Maximum size of the body in bytes
15    pub body_max_size: usize,
16}
17
18#[derive(Clone)]
19pub struct HttpErrorsLayer {
20    pub config: HttpErrorsConfig,
21}
22
23impl HttpErrorsLayer {
24    /// Create a new `HttpErrorsLayer`
25    pub fn new(config: &HttpErrorsConfig) -> Self {
26        Self { config: config.clone() }
27    }
28}
29
30impl<S> Layer<S> for HttpErrorsLayer {
31    type Service = HttpErrorsMiddleware<S>;
32
33    fn layer(&self, inner: S) -> Self::Service {
34        HttpErrorsMiddleware {
35            inner,
36            config: self.config.clone(),
37        }
38    }
39}
40
41#[derive(Clone)]
42pub struct HttpErrorsMiddleware<S> {
43    inner: S,
44    config: HttpErrorsConfig,
45}
46
47impl<S> Service<Request<Body>> for HttpErrorsMiddleware<S>
48where
49    S: Service<Request<Body>, Response = Response> + Send + Clone + 'static,
50    S::Future: Send + 'static,
51{
52    type Response = S::Response;
53    type Error = S::Error;
54    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
55    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
56
57    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
58        self.inner.poll_ready(cx)
59    }
60
61    fn call(&mut self, request: Request<Body>) -> Self::Future {
62        let mut inner = self.inner.clone();
63        let config = self.config.clone();
64
65        Box::pin(async move {
66            let response: Response = inner.call(request).await?;
67
68            // Check the content-type
69            let headers = response.headers();
70            if let Some(content_type) = headers.get("content-type") {
71                let content_type = content_type.to_str().unwrap_or_default();
72                if content_type.starts_with("image/")
73                    || content_type.starts_with("audio/")
74                    || content_type.starts_with("video/")
75                {
76                    return Ok(response);
77                }
78            }
79
80            let (parts, body) = response.into_parts();
81            match axum::body::to_bytes(body, config.body_max_size).await {
82                Ok(body) => match String::from_utf8(body.to_vec()) {
83                    Ok(body) => match parts.status {
84                        StatusCode::METHOD_NOT_ALLOWED => Ok(ApiError::MethodNotAllowed.into_response()),
85                        StatusCode::UNPROCESSABLE_ENTITY => Ok(ApiError::UnprocessableEntity(body).into_response()),
86                        StatusCode::NOT_FOUND if body.is_empty() => {
87                            Ok(ApiError::NotFound("Resource Not Found".to_owned()).into_response())
88                        }
89                        _ => Ok(Response::from_parts(parts, Body::from(body))),
90                    },
91                    Err(err) => Ok(ApiError::InternalServerError(err.to_string()).into_response()),
92                },
93                Err(_) => Ok(ApiError::PayloadTooLarge.into_response()),
94            }
95        })
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use axum::http::header;
103    use std::convert::Infallible;
104    use tower::{ServiceBuilder, ServiceExt};
105
106    fn layer() -> HttpErrorsLayer {
107        HttpErrorsLayer::new(&HttpErrorsConfig { body_max_size: 1024 })
108    }
109
110    async fn read_body(response: Response) -> String {
111        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
112        String::from_utf8(body.to_vec()).unwrap()
113    }
114
115    #[tokio::test]
116    async fn ok_response_passes_through_unchanged() {
117        let svc = ServiceBuilder::new()
118            .layer(layer())
119            .service(tower::service_fn(|_req: Request<Body>| async {
120                Ok::<_, Infallible>(
121                    Response::builder()
122                        .status(StatusCode::OK)
123                        .body(Body::from("hello"))
124                        .unwrap(),
125                )
126            }));
127
128        let response = svc
129            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
130            .await
131            .unwrap();
132
133        assert_eq!(response.status(), StatusCode::OK);
134        assert_eq!(read_body(response).await, "hello");
135    }
136
137    #[tokio::test]
138    async fn empty_404_is_rewritten_as_json_api_error() {
139        let svc = ServiceBuilder::new()
140            .layer(layer())
141            .service(tower::service_fn(|_req: Request<Body>| async {
142                Ok::<_, Infallible>(
143                    Response::builder()
144                        .status(StatusCode::NOT_FOUND)
145                        .body(Body::empty())
146                        .unwrap(),
147                )
148            }));
149
150        let response = svc
151            .oneshot(Request::builder().uri("/missing").body(Body::empty()).unwrap())
152            .await
153            .unwrap();
154
155        assert_eq!(response.status(), StatusCode::NOT_FOUND);
156        assert_eq!(
157            response.headers().get(header::CONTENT_TYPE).unwrap(),
158            "application/json",
159        );
160
161        let body = read_body(response).await;
162        assert!(body.contains("\"code\":404"), "body was: {body}");
163        assert!(body.contains("Resource Not Found"), "body was: {body}");
164    }
165
166    /// 404 with a non-empty body must be left untouched — the inner service
167    /// already provided a response body.
168    #[tokio::test]
169    async fn non_empty_404_is_passed_through() {
170        let svc = ServiceBuilder::new()
171            .layer(layer())
172            .service(tower::service_fn(|_req: Request<Body>| async {
173                Ok::<_, Infallible>(
174                    Response::builder()
175                        .status(StatusCode::NOT_FOUND)
176                        .body(Body::from("custom 404"))
177                        .unwrap(),
178                )
179            }));
180
181        let response = svc
182            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
183            .await
184            .unwrap();
185
186        assert_eq!(response.status(), StatusCode::NOT_FOUND);
187        assert_eq!(read_body(response).await, "custom 404");
188    }
189
190    #[tokio::test]
191    async fn method_not_allowed_is_rewritten_as_json_api_error() {
192        let svc = ServiceBuilder::new()
193            .layer(layer())
194            .service(tower::service_fn(|_req: Request<Body>| async {
195                Ok::<_, Infallible>(
196                    Response::builder()
197                        .status(StatusCode::METHOD_NOT_ALLOWED)
198                        .body(Body::empty())
199                        .unwrap(),
200                )
201            }));
202
203        let response = svc
204            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
205            .await
206            .unwrap();
207
208        assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
209        assert_eq!(
210            response.headers().get(header::CONTENT_TYPE).unwrap(),
211            "application/json",
212        );
213
214        let body = read_body(response).await;
215        assert!(body.contains("\"code\":405"), "body was: {body}");
216        assert!(body.contains("Method not allowed"), "body was: {body}");
217    }
218
219    #[tokio::test]
220    async fn unprocessable_entity_body_is_wrapped_into_json_message() {
221        let svc = ServiceBuilder::new()
222            .layer(layer())
223            .service(tower::service_fn(|_req: Request<Body>| async {
224                Ok::<_, Infallible>(
225                    Response::builder()
226                        .status(StatusCode::UNPROCESSABLE_ENTITY)
227                        .body(Body::from("validation failed"))
228                        .unwrap(),
229                )
230            }));
231
232        let response = svc
233            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
234            .await
235            .unwrap();
236
237        assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
238
239        let body = read_body(response).await;
240        assert!(body.contains("\"code\":422"), "body was: {body}");
241        assert!(body.contains("validation failed"), "body was: {body}");
242    }
243
244    /// Binary content-types (image/audio/video) must short-circuit before any
245    /// body reading: rewriting them as JSON would corrupt the payload.
246    #[tokio::test]
247    async fn image_content_type_short_circuits_without_touching_body() {
248        let payload = vec![0u8, 1, 2, 3, 4];
249        let payload_clone = payload.clone();
250
251        let svc = ServiceBuilder::new()
252            .layer(layer())
253            .service(tower::service_fn(move |_req: Request<Body>| {
254                let payload = payload_clone.clone();
255                async move {
256                    Ok::<_, Infallible>(
257                        Response::builder()
258                            .status(StatusCode::NOT_FOUND)
259                            .header(header::CONTENT_TYPE, "image/png")
260                            .body(Body::from(payload))
261                            .unwrap(),
262                    )
263                }
264            }));
265
266        let response = svc
267            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
268            .await
269            .unwrap();
270
271        assert_eq!(response.status(), StatusCode::NOT_FOUND);
272        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
273        assert_eq!(body.to_vec(), payload);
274    }
275
276    #[tokio::test]
277    async fn body_exceeding_max_size_returns_payload_too_large() {
278        let small_layer = HttpErrorsLayer::new(&HttpErrorsConfig { body_max_size: 4 });
279        let svc = ServiceBuilder::new()
280            .layer(small_layer)
281            .service(tower::service_fn(|_req: Request<Body>| async {
282                Ok::<_, Infallible>(
283                    Response::builder()
284                        .status(StatusCode::OK)
285                        .body(Body::from("hello world"))
286                        .unwrap(),
287                )
288            }));
289
290        let response = svc
291            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
292            .await
293            .unwrap();
294
295        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
296        let body = read_body(response).await;
297        assert!(body.contains("\"code\":413"), "body was: {body}");
298    }
299}
api_tools/server/axum/layers/http_errors.rs

api_tools/server/axum/layers/
http_errors.rs
