itsscb/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2025-09-27 04:50:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

axum: add ResponseAxumBodyLayer for mapping response body to axum::body::Body

Browse Source
This commit is contained in:
David Mládek 2025-09-12 14:33:10 +02:00
parent 4ab8df5e42
commit 853d7e707a
2 changed files with 80 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
axum/src/middleware/mod.rs
View File

@ -6,6 +6,7 @@ mod from_extractor;
mod from_fn;
mod map_request;
mod map_response;
mod response_axum_body;
pub use self::from_extractor::{
from_extractor, from_extractor_with_state, FromExtractor, FromExtractorLayer,
@ -17,6 +18,9 @@ pub use self::map_request::{
pub use self::map_response::{
map_response, map_response_with_state, MapResponse, MapResponseLayer,
};
pub use self::response_axum_body::{
ResponseAxumBody, ResponseAxumBodyFuture, ResponseAxumBodyLayer,
};
pub use crate::extension::AddExtension;
pub mod future {

76
axum/src/middleware/response_axum_body.rs Normal file
View File

@ -0,0 +1,76 @@
use std::{
error::Error,
future::Future,
pin::Pin,
task::{ready, Context, Poll},
};
use axum_core::{body::Body, response::Response};
use bytes::Bytes;
use http_body::Body as HttpBody;
use pin_project_lite::pin_project;
use tower::{Layer, Service};
/// Layer that transforms the Response body to [`crate::body::Body`].
///
/// This is useful when another layer maps the body to some other type to convert it back.
#[derive(Debug, Clone)]
pub struct ResponseAxumBodyLayer;
impl<S> Layer<S> for ResponseAxumBodyLayer {
type Service = ResponseAxumBody<S>;
fn layer(&self, inner: S) -> Self::Service {
ResponseAxumBody::<S>(inner)
}
}
/// Service generated by [`ResponseAxumBodyLayer`].
#[derive(Debug, Clone)]
pub struct ResponseAxumBody<S>(S);
impl<S, Request, ResBody> Service<Request> for ResponseAxumBody<S>
where
S: Service<Request, Response = Response<ResBody>>,
ResBody: HttpBody<Data = Bytes> + Send + 'static,
<ResBody as HttpBody>::Error: Error + Send + Sync,
{
type Response = Response;
type Error = S::Error;
type Future = ResponseAxumBodyFuture<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.0.poll_ready(cx)
}
fn call(&mut self, req: Request) -> Self::Future {
ResponseAxumBodyFuture {
inner: self.0.call(req),
}
}
}
pin_project! {
/// Response future for [`ResponseAxumBody`].
pub struct ResponseAxumBodyFuture<Fut> {
#[pin]
inner: Fut,
}
}
impl<Fut, ResBody, E> Future for ResponseAxumBodyFuture<Fut>
where
Fut: Future<Output = Result<Response<ResBody>, E>>,
ResBody: HttpBody<Data = Bytes> + Send + 'static,
<ResBody as HttpBody>::Error: Error + Send + Sync,
{
type Output = Result<Response<Body>, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let res = ready!(this.inner.poll(cx)?);
Poll::Ready(Ok(res.map(Body::new)))
}
}