How to fetch response body from tower::util::MapResponseLayer ?

[xiuno]

I try to return a friendly page when StatusCode == StatusCode::BAD_REQUEST, found this function is not a async, but it must be async If want to fetch body content from Response<BoxBody>.

file examples/404.rs:

//! Run with
//!
//! ```not_rust
//! cargo run --example 404
//! ```

use axum::{
    body::{box_body, Body, BoxBody},
    prelude::*,
};
use http::{Response, StatusCode};
use std::net::SocketAddr;
use tower::util::MapResponseLayer;

#[tokio::main]
async fn main() {
    // Set the RUST_LOG, if it hasn't been explicitly defined
    if std::env::var("RUST_LOG").is_err() {
        std::env::set_var("RUST_LOG", "404=debug")
    }
    //tracing_subscriber::fmt::init();

    // build our application with a route
    let app = route("/", get(handler))
        // make sure this is added as the very last thing
        .layer(MapResponseLayer::new(map_500_400));

    // run it
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    //tracing::debug!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn handler() -> response::Html<&'static str> {
    response::Html("<h1>Hello, World!</h1>")
}

fn map_500_400(response: Response<BoxBody>) -> Response<BoxBody> {
    if response.status() == StatusCode::INTERNAL_SERVER_ERROR ||  response.status() == StatusCode::BAD_REQUEST {
        let body = response.body();

        // todo: how to fetch body string from Response<BoxBody> in this function ?
        let body = hyper::body::to_bytes(body).await.unwrap();
        let body = String::from_utf8(body.to_vec()).unwrap();

        return Response::builder()
            .status(StatusCode::NOT_FOUND)
            .body(box_body(Body::from(format!("Internal Error! {:?}", body()))))
            .unwrap();
    }
    response
}

[davidpdrsn]

You should be able to use AndThenLayer from tower. Its basically an async map_response that can return a result as well.

Something like this:

use axum::{
    body::{box_body, Body, BoxBody},
    handler::get,
    http::{Response, StatusCode},
    route,
    routing::RoutingDsl,
};
use std::net::SocketAddr;

#[tokio::main]
async fn main() {
    let app = route("/", get(|| async {}))
        // make sure this is added as the very last thing
        .layer(tower::util::AndThenLayer::new(map_404));

    // run it
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn map_404(
    response: Response<BoxBody>,
) -> Result<Response<BoxBody>, std::convert::Infallible> {
    if response.status() == StatusCode::NOT_FOUND
        || response.status() == StatusCode::METHOD_NOT_ALLOWED
    {
        return Ok(Response::builder()
            .status(StatusCode::NOT_FOUND)
            .body(box_body(Body::from("nothing to see here")))
            .unwrap());
    }

    Ok(response)
}