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
82            // Only 404/405/422 are rewritten. Every other status (notably the 200s that make up
83            // the bulk of the traffic) is returned untouched, without buffering: this preserves
84            // streaming and avoids a full-body copy + UTF-8 scan, plus the spurious 413 a legitimate
85            // large body used to trigger by overflowing `to_bytes`.
86            match parts.status {
87                // The 405 body is discarded, so there is no need to read it.
88                StatusCode::METHOD_NOT_ALLOWED => Ok(ApiError::MethodNotAllowed.into_response()),
89                // 422/404 are the only statuses whose body we still need to inspect.
90                StatusCode::UNPROCESSABLE_ENTITY | StatusCode::NOT_FOUND => {
91                    match axum::body::to_bytes(body, config.body_max_size).await {
92                        Ok(bytes) => match String::from_utf8(bytes.to_vec()) {
93                            Ok(body) => match parts.status {
94                                StatusCode::UNPROCESSABLE_ENTITY => {
95                                    Ok(ApiError::UnprocessableEntity(body).into_response())
96                                }
97                                StatusCode::NOT_FOUND if body.is_empty() => {
98                                    Ok(ApiError::NotFound("Resource Not Found".to_owned()).into_response())
99                                }
100                                _ => Ok(Response::from_parts(parts, Body::from(body))),
101                            },
102                            Err(err) => Ok(ApiError::InternalServerError(err.to_string()).into_response()),
103                        },
104                        Err(_) => Ok(ApiError::PayloadTooLarge.into_response()),
105                    }
106                }
107                _ => Ok(Response::from_parts(parts, body)),
108            }
109        })
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use axum::http::header;
117    use std::convert::Infallible;
118    use tower::{ServiceBuilder, ServiceExt};
119
120    fn layer() -> HttpErrorsLayer {
121        HttpErrorsLayer::new(&HttpErrorsConfig { body_max_size: 1024 })
122    }
123
124    async fn read_body(response: Response) -> String {
125        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
126        String::from_utf8(body.to_vec()).unwrap()
127    }
128
129    #[tokio::test]
130    async fn ok_response_passes_through_unchanged() {
131        let svc = ServiceBuilder::new()
132            .layer(layer())
133            .service(tower::service_fn(|_req: Request<Body>| async {
134                Ok::<_, Infallible>(
135                    Response::builder()
136                        .status(StatusCode::OK)
137                        .body(Body::from("hello"))
138                        .unwrap(),
139                )
140            }));
141
142        let response = svc
143            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
144            .await
145            .unwrap();
146
147        assert_eq!(response.status(), StatusCode::OK);
148        assert_eq!(read_body(response).await, "hello");
149    }
150
151    #[tokio::test]
152    async fn empty_404_is_rewritten_as_json_api_error() {
153        let svc = ServiceBuilder::new()
154            .layer(layer())
155            .service(tower::service_fn(|_req: Request<Body>| async {
156                Ok::<_, Infallible>(
157                    Response::builder()
158                        .status(StatusCode::NOT_FOUND)
159                        .body(Body::empty())
160                        .unwrap(),
161                )
162            }));
163
164        let response = svc
165            .oneshot(Request::builder().uri("/missing").body(Body::empty()).unwrap())
166            .await
167            .unwrap();
168
169        assert_eq!(response.status(), StatusCode::NOT_FOUND);
170        assert_eq!(
171            response.headers().get(header::CONTENT_TYPE).unwrap(),
172            "application/json",
173        );
174
175        let body = read_body(response).await;
176        assert!(body.contains("\"code\":404"), "body was: {body}");
177        assert!(body.contains("Resource Not Found"), "body was: {body}");
178    }
179
180    /// 404 with a non-empty body must be left untouched — the inner service
181    /// already provided a response body.
182    #[tokio::test]
183    async fn non_empty_404_is_passed_through() {
184        let svc = ServiceBuilder::new()
185            .layer(layer())
186            .service(tower::service_fn(|_req: Request<Body>| async {
187                Ok::<_, Infallible>(
188                    Response::builder()
189                        .status(StatusCode::NOT_FOUND)
190                        .body(Body::from("custom 404"))
191                        .unwrap(),
192                )
193            }));
194
195        let response = svc
196            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
197            .await
198            .unwrap();
199
200        assert_eq!(response.status(), StatusCode::NOT_FOUND);
201        assert_eq!(read_body(response).await, "custom 404");
202    }
203
204    #[tokio::test]
205    async fn method_not_allowed_is_rewritten_as_json_api_error() {
206        let svc = ServiceBuilder::new()
207            .layer(layer())
208            .service(tower::service_fn(|_req: Request<Body>| async {
209                Ok::<_, Infallible>(
210                    Response::builder()
211                        .status(StatusCode::METHOD_NOT_ALLOWED)
212                        .body(Body::empty())
213                        .unwrap(),
214                )
215            }));
216
217        let response = svc
218            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
219            .await
220            .unwrap();
221
222        assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
223        assert_eq!(
224            response.headers().get(header::CONTENT_TYPE).unwrap(),
225            "application/json",
226        );
227
228        let body = read_body(response).await;
229        assert!(body.contains("\"code\":405"), "body was: {body}");
230        assert!(body.contains("Method not allowed"), "body was: {body}");
231    }
232
233    #[tokio::test]
234    async fn unprocessable_entity_body_is_wrapped_into_json_message() {
235        let svc = ServiceBuilder::new()
236            .layer(layer())
237            .service(tower::service_fn(|_req: Request<Body>| async {
238                Ok::<_, Infallible>(
239                    Response::builder()
240                        .status(StatusCode::UNPROCESSABLE_ENTITY)
241                        .body(Body::from("validation failed"))
242                        .unwrap(),
243                )
244            }));
245
246        let response = svc
247            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
248            .await
249            .unwrap();
250
251        assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
252
253        let body = read_body(response).await;
254        assert!(body.contains("\"code\":422"), "body was: {body}");
255        assert!(body.contains("validation failed"), "body was: {body}");
256    }
257
258    /// Binary content-types (image/audio/video) must short-circuit before any
259    /// body reading: rewriting them as JSON would corrupt the payload.
260    #[tokio::test]
261    async fn image_content_type_short_circuits_without_touching_body() {
262        let payload = vec![0u8, 1, 2, 3, 4];
263        let payload_clone = payload.clone();
264
265        let svc = ServiceBuilder::new()
266            .layer(layer())
267            .service(tower::service_fn(move |_req: Request<Body>| {
268                let payload = payload_clone.clone();
269                async move {
270                    Ok::<_, Infallible>(
271                        Response::builder()
272                            .status(StatusCode::NOT_FOUND)
273                            .header(header::CONTENT_TYPE, "image/png")
274                            .body(Body::from(payload))
275                            .unwrap(),
276                    )
277                }
278            }));
279
280        let response = svc
281            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
282            .await
283            .unwrap();
284
285        assert_eq!(response.status(), StatusCode::NOT_FOUND);
286        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
287        assert_eq!(body.to_vec(), payload);
288    }
289
290    /// The size limit is still enforced on the statuses whose body is actually
291    /// read (`422`): an oversized body there is turned into a `413`.
292    #[tokio::test]
293    async fn unprocessable_entity_over_limit_still_413() {
294        let small_layer = HttpErrorsLayer::new(&HttpErrorsConfig { body_max_size: 4 });
295        let svc = ServiceBuilder::new()
296            .layer(small_layer)
297            .service(tower::service_fn(|_req: Request<Body>| async {
298                Ok::<_, Infallible>(
299                    Response::builder()
300                        .status(StatusCode::UNPROCESSABLE_ENTITY)
301                        .body(Body::from("validation failed for many fields"))
302                        .unwrap(),
303                )
304            }));
305
306        let response = svc
307            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
308            .await
309            .unwrap();
310
311        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
312        let body = read_body(response).await;
313        assert!(body.contains("\"code\":413"), "body was: {body}");
314    }
315
316    /// A `200 OK` whose body exceeds `body_max_size` must pass through untouched:
317    /// the body is never buffered for passthrough statuses, so it can no longer
318    /// be turned into a spurious `413`.
319    #[tokio::test]
320    async fn large_ok_body_is_not_truncated_to_413() {
321        let small_layer = HttpErrorsLayer::new(&HttpErrorsConfig { body_max_size: 4 });
322        let svc = ServiceBuilder::new()
323            .layer(small_layer)
324            .service(tower::service_fn(|_req: Request<Body>| async {
325                Ok::<_, Infallible>(
326                    Response::builder()
327                        .status(StatusCode::OK)
328                        .body(Body::from("hello world"))
329                        .unwrap(),
330                )
331            }));
332
333        let response = svc
334            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
335            .await
336            .unwrap();
337
338        assert_eq!(response.status(), StatusCode::OK);
339        assert_eq!(read_body(response).await, "hello world");
340    }
341
342    /// A non-rewritten status (`500`) carrying a body must be returned
343    /// unchanged, without being buffered.
344    #[tokio::test]
345    async fn passthrough_status_body_is_not_buffered() {
346        let svc = ServiceBuilder::new()
347            .layer(layer())
348            .service(tower::service_fn(|_req: Request<Body>| async {
349                Ok::<_, Infallible>(
350                    Response::builder()
351                        .status(StatusCode::INTERNAL_SERVER_ERROR)
352                        .body(Body::from("upstream failure"))
353                        .unwrap(),
354                )
355            }));
356
357        let response = svc
358            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
359            .await
360            .unwrap();
361
362        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
363        assert_eq!(read_body(response).await, "upstream failure");
364    }
365}
api_tools/server/axum/layers/http_errors.rs

api_tools/server/axum/layers/
http_errors.rs
