axum/middleware/
response_axum_body.rs

1use std::{
2    error::Error,
3    future::Future,
4    pin::Pin,
5    task::{ready, Context, Poll},
6};
7
8use axum_core::{body::Body, response::Response};
9use bytes::Bytes;
10use http_body::Body as HttpBody;
11use pin_project_lite::pin_project;
12use tower::{Layer, Service};
13
14/// Layer that transforms the Response body to [`crate::body::Body`].
15///
16/// This is useful when another layer maps the body to some other type to convert it back.
17#[derive(Debug, Clone)]
18pub struct ResponseAxumBodyLayer;
19
20impl<S> Layer<S> for ResponseAxumBodyLayer {
21    type Service = ResponseAxumBody<S>;
22
23    fn layer(&self, inner: S) -> Self::Service {
24        ResponseAxumBody::<S>(inner)
25    }
26}
27
28/// Service generated by [`ResponseAxumBodyLayer`].
29#[derive(Debug, Clone)]
30pub struct ResponseAxumBody<S>(S);
31
32impl<S, Request, ResBody> Service<Request> for ResponseAxumBody<S>
33where
34    S: Service<Request, Response = Response<ResBody>>,
35    ResBody: HttpBody<Data = Bytes> + Send + 'static,
36    <ResBody as HttpBody>::Error: Error + Send + Sync,
37{
38    type Response = Response;
39
40    type Error = S::Error;
41
42    type Future = ResponseAxumBodyFuture<S::Future>;
43
44    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
45        self.0.poll_ready(cx)
46    }
47
48    fn call(&mut self, req: Request) -> Self::Future {
49        ResponseAxumBodyFuture {
50            inner: self.0.call(req),
51        }
52    }
53}
54
55pin_project! {
56    /// Response future for [`ResponseAxumBody`].
57    pub struct ResponseAxumBodyFuture<Fut> {
58        #[pin]
59        inner: Fut,
60    }
61}
62
63impl<Fut, ResBody, E> Future for ResponseAxumBodyFuture<Fut>
64where
65    Fut: Future<Output = Result<Response<ResBody>, E>>,
66    ResBody: HttpBody<Data = Bytes> + Send + 'static,
67    <ResBody as HttpBody>::Error: Error + Send + Sync,
68{
69    type Output = Result<Response<Body>, E>;
70
71    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
72        let this = self.project();
73        let res = ready!(this.inner.poll(cx)?);
74        Poll::Ready(Ok(res.map(Body::new)))
75    }
76}