Module tower_http::map_response_body

source ·
Available on crate feature map-response-body only.
Expand description

Apply a transformation to the response body.

Example

use bytes::Bytes;
use http::{Request, Response};
use hyper::Body;
use std::convert::Infallible;
use std::{pin::Pin, task::{Context, Poll}};
use tower::{ServiceBuilder, service_fn, ServiceExt, Service};
use tower_http::map_response_body::MapResponseBodyLayer;
use futures::ready;

// A wrapper for a `hyper::Body` that prints the size of data chunks
struct PrintChunkSizesBody {
    inner: Body,
}

impl PrintChunkSizesBody {
    fn new(inner: Body) -> Self {
        Self { inner }
    }
}

impl http_body::Body for PrintChunkSizesBody {
    type Data = Bytes;
    type Error = hyper::Error;

    fn poll_data(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
        if let Some(chunk) = ready!(Pin::new(&mut self.inner).poll_data(cx)?) {
            println!("chunk size = {}", chunk.len());
            Poll::Ready(Some(Ok(chunk)))
        } else {
            Poll::Ready(None)
        }
    }

    fn poll_trailers(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<Option<hyper::HeaderMap>, Self::Error>> {
        Pin::new(&mut self.inner).poll_trailers(cx)
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> http_body::SizeHint {
        self.inner.size_hint()
    }
}

async fn handle<B>(_: Request<B>) -> Result<Response<Body>, Infallible> {
    // ...
}

let mut svc = ServiceBuilder::new()
    // Wrap response bodies in `PrintChunkSizesBody`
    .layer(MapResponseBodyLayer::new(PrintChunkSizesBody::new))
    .service_fn(handle);

// Call the service
let request = Request::new(Body::from("foobar"));

svc.ready().await?.call(request).await?;

Structs