axum/middleware/
response_axum_body.rs1use 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#[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#[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 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}