如何在actix_web中间件中打印响应体？

Question

我想使用actix_web框架编写一个非常简单的中间件，但是这个任务一直在各个方面击败我。

我的代码框架如下：

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;

我可以通过 res.unwrap().response().body() 访问 ResponseBody 类型，但我不知道该怎么使用它。

有什么想法吗？

- Patryk

---

我删除了我的回答，因为我不知道响应的情况，并且我错误地将您的问题误读为“RequestBody”。然而，如果响应是流，则我怀疑您会遇到同样的问题。（访问请求的问题：https://github.com/actix/actix-web/issues/373）。 - Peter Hall

可能的解决方案是将流收集到一些堆分配的数据结构中，然后（以某种方式）重新构建流以生成有效的响应。 - Peter Hall

相关链接：https://dev59.com/p1IG5IYBdhLWcg3wlCO9 - Peter Hall

Answer 1

这是我用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
                        ))
                    }
                }
            }
        })
    }
}

提取请求主体(Request Body)的方法很简单，与Actix示例类似。但是，在更新到Beta 14版本后，由于引入了BoxedBody，直接从AnyBody中提取字节的方法已经发生了变化。幸运的是，我找到了一个实用函数body::to_bytes (use actix_web::body::to_bytes)，它做得很好。它的当前实现如下所示：

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())
}

我认为这种方式提取正文应该是可行的，因为正文是通过to_bytes()从正文流中提取出来的。

如果有更好的方法，请告诉我，但是这有点麻烦，而且当它从Beta 13切换到Beta 14时，我才最近确定如何做到这一点。

这个特定的例子拦截错误并将它们重写为JSON格式（如果它们还不是JSON格式）。例如，如果在处理程序之外发生错误，比如在处理程序函数本身_data：web::Json<Request<'a，LoginRequest>>中解析JSON，而不是在处理程序正文中。 提取请求正文和响应正文并不是实现目标所必需的，只是为了说明问题。

- Chris Andrew