How to print a response body in actix_web middleware?

Question

I'd like to write a very simple middleware using actix_web framework but it's so far beating me on every front.

I have a skeleton like this:

let result = actix_web::HttpServer::new(move || {
    actix_web::App::new()
        .wrap_fn(move |req, srv| {
            srv.call(req).map(move|res| {
                println!("Got response");
                // let s = res.unwrap().response().body();
                // ???
                res
            })
        })
})
.bind("0.0.0.0:8080")?
.run()
.await;

and I can access ResponseBody type via res.unwrap().response().body() but I don't know what can I do with this.

Any ideas?

Answer 1

This is an example of how I was able to accomplish this with 4.0.0-beta.14:

use std::cell::RefCell;
use std::pin::Pin;
use std::rc::Rc;
use std::collections::HashMap;
use std::str;

use erp_contrib::{actix_http, actix_web, futures, serde_json};

use actix_web::dev::{Service,  ServiceRequest, ServiceResponse, Transform};
use actix_web::{HttpMessage, body, http::StatusCode, error::Error ,HttpResponseBuilder};
use actix_http::{h1::Payload, header};
use actix_web::web::{BytesMut};

use futures::future::{ok, Future, Ready};
use futures::task::{Context, Poll};
use futures::StreamExt;

use crate::response::ErrorResponse;

pub struct UnhandledErrorResponse;

impl<S: 'static> Transform<S, ServiceRequest> for UnhandledErrorResponse
    where
        S: Service<ServiceRequest, Response = ServiceResponse, Error = Error>,
        S::Future: 'static,
{
    type Response = ServiceResponse;
    type Error = Error;
    type Transform = UnhandledErrorResponseMiddleware<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ok(UnhandledErrorResponseMiddleware { service: Rc::new(RefCell::new(service)), })
    }
}

pub struct UnhandledErrorResponseMiddleware<S> {
    service: Rc<RefCell<S>>,
}

impl<S> Service<ServiceRequest> for UnhandledErrorResponseMiddleware<S>
    where
        S: Service<ServiceRequest, Response = ServiceResponse, Error = Error> + 'static,
        S::Future: 'static,
{
    type Response = ServiceResponse;
    type Error = Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;

    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&self, mut req: ServiceRequest) -> Self::Future {
        let svc = self.service.clone();

        Box::pin(async move {

            /* EXTRACT THE BODY OF REQUEST */
            let mut request_body = BytesMut::new();
            while let Some(chunk) = req.take_payload().next().await {
                request_body.extend_from_slice(&chunk?);
            }

            let mut orig_payload = Payload::empty();
            orig_payload.unread_data(request_body.freeze());
            req.set_payload(actix_http::Payload::from(orig_payload));

            /* now process the response */
            let res: ServiceResponse = svc.call(req).await?;

            let content_type = match res.headers().get("content-type") {
                None => { "unknown"}
                Some(header) => {
                    match header.to_str() {
                        Ok(value) => {value}
                        Err(_) => { "unknown"}
                    }
                }
            };

            return match res.response().error() {
                None => {
                    Ok(res)
                }
                Some(error) => {
                    if content_type.to_uppercase().contains("APPLICATION/JSON") {
                        Ok(res)
                    } else {

                        let error = error.to_string();
                        let new_request = res.request().clone();

                        /* EXTRACT THE BODY OF RESPONSE */
                        let _body_data =
                            match str::from_utf8(&body::to_bytes(res.into_body()).await?){
                            Ok(str) => {
                                str
                            }
                            Err(_) => {
                                "Unknown"
                            }
                        };

                        let mut errors = HashMap::new();
                        errors.insert("general".to_string(), vec![error]);

                        let new_response = match ErrorResponse::new(&false, errors) {
                            Ok(response) => {
                                HttpResponseBuilder::new(StatusCode::BAD_REQUEST)
                                    .insert_header((header::CONTENT_TYPE, "application/json"))
                                    .body(serde_json::to_string(&response).unwrap())
                            }
                            Err(_error) => {
                                HttpResponseBuilder::new(StatusCode::BAD_REQUEST)
                                    .insert_header((header::CONTENT_TYPE, "application/json"))
                                    .body("An unknown error occurred.")
                            }
                        };

                        Ok(ServiceResponse::new(
                            new_request,
                            new_response
                        ))
                    }
                }
            }
        })
    }
}

The extraction of the Request Body is straightforward and similar to how Actix example's illustrate. However, with the update to version Beta 14, pulling the bytes directly from AnyBody has changed with the introduction of BoxedBody. Fortunately, I found a utility function body::to_bytes (use actix_web::body::to_bytes) which does a good job. It's current implementation looks like this:

pub async fn to_bytes<B: MessageBody>(body: B) -> Result<Bytes, B::Error> {
    let cap = match body.size() {
        BodySize::None | BodySize::Sized(0) => return Ok(Bytes::new()),
        BodySize::Sized(size) => size as usize,
        // good enough first guess for chunk size
        BodySize::Stream => 32_768,
    };

    let mut buf = BytesMut::with_capacity(cap);

    pin!(body);

    poll_fn(|cx| loop {
        let body = body.as_mut();

        match ready!(body.poll_next(cx)) {
            Some(Ok(bytes)) => buf.extend_from_slice(&*bytes),
            None => return Poll::Ready(Ok(())),
            Some(Err(err)) => return Poll::Ready(Err(err)),
        }
    })
    .await?;

    Ok(buf.freeze())
}

which I believe should be fine to extract the body in this way, as the body is extracted from the body stream by to_bytes().

If someone has a better way, let me know, but it was a little bit of pain, and I only had recently determined how to do it in Beta 13 when it switched to Beta 14.

This particular example intercepts errors and rewrites them to JSON format if they're not already json format. This would be the case, as an example, if an error occurs outside of a handler, such as parsing JSON in the handler function itself _data: web::Json<Request<'a, LoginRequest>> and not in the handler body. Extracting the Request Body and Response Body is not necessary to accomplish the goal, and is just here for illustration.