Middleware Response body manipulation

[0rangeFox]

Hey folks 👋! The title already tells what I am looking for, basically I have this bit of code, but I would say that only the manipulate_body function is what will be used to be analyzed. I've already seen a lot of issues (#2856, #3166, #3227) and examples, but none of them achieve what I want to do. Can anyone help me with this part? Thanks for any suggestions and opinions.

// ...
    forward_ready!(service);

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

        Box::pin(async move {
            let mut res: ServiceResponse<B> = service.call(req).await?;
            res = manipulate_body(&res);
            Ok(res)
        })
    }

// ...

fn manipulate_body<B: MessageBody>(res: &ServiceResponse<B>) -> ServiceResponse<B> {
    // Read body bytes
    let body_bytes = res.response().body();

    let mut new_body: Value = serde_json::from_slice(&body_bytes).unwrap();
    new_body["test"] = Value::String("anything".to_owned());

    // Destructures ServiceResponse into request and response components
    let (req, res) = res.into_parts();

    // Create a new response with the modified body
    let res: HttpResponse<B> = res.set_body(&new_body.to_string()).map_into_boxed_body();

    // Create a new ServiceResponse with the modified response
    ServiceResponse::new(req, res).map_into_right_body()
}

[robjtede]

There's no way to get a body from a shared & reference since bodies could be streams. Middleware isn't typically allowed to make any assumptions like that. You need ownership of the body in the place you process it.

So, more or less...

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

        Box::pin(async move {
            let res = service.call(req).await?;

            let (req, res) = res.into_parts();
            let (res, body) = res.into_parts();

            let body = convert_body(body).await;

            let res = res.set_body(body);
            Ok(ServiceResponse::new(req, res))
        })
    }

// ...

async fn convert_body<B: MessageBody>(body: B) -> String {
    // Read body bytes
    let body_bytes = actix_web::body::to_bytes_limited(
        body,
        2_usize.pow(22), // 4 MB
    )
    .await
    .unwrap()
    .map_err(Into::into)
    .unwrap();

    let mut new_body: Value = serde_json::from_slice(&body_bytes).unwrap();
    new_body["test"] = Value::String("anything".to_owned());

    new_body.to_string()
}

If you are completely sure you'll never return a streaming body, then you could get away with .try_into_bytes() instead of to_bytes_limited.

[robjtede]

If you're comfortable with that last limitation then it could be simplified to:

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

        Box::pin(async move {
            let res = service.call(req).await?;
            Ok(res.map_body(convert_body))
        })
    }

// ...

fn convert_body<B: MessageBody>(
  _head: &mut actix_web::dev::ResponseHead,
  body: B,
) -> String {
    // Read body bytes
    let body_bytes = match body.try_into_bytes() {
        Ok(bytes) => bytes,
        Err(_body) => panic!("streaming bodies are not allowed"),
    };

    let mut new_body: Value = serde_json::from_slice(&body_bytes).unwrap();
    new_body["test"] = Value::String("anything".to_owned());

    new_body.to_string()
}

[0rangeFox]

I'm not getting a result with that, supposedly as B derives from MessageBody, and I've tried doing it this way let res = res.set_body::<B>(BoxBody::new(body.to_string())); and still no results, or I must be asleep and missed something. But isn't this supposed to be what I'm looking for?

error[E0308]: mismatched types
   --> src/middlewares/test_middleware.rs:245:33
    |
224 | async fn manipulate_body<B: MessageBody>(res: &ServiceResponse<B>) -> ServiceResponse<B> {
    |                           - expected this type parameter
...
245 |     let res = res.set_body::<B>(body.to_string());
    |                   ------------- ^^^^^^^^^^^^^^^^ expected type parameter `B`, found `String`
    |                   |
    |                   arguments to this method are incorrect
    |
    = note: expected type parameter `B`
                       found struct `std::string::String`

This is what I have as a middleware declaration:

impl<S, B> Service<ServiceRequest> for InnerTestMiddleware<S>
    where
        S: Service<ServiceRequest, Response=ServiceResponse<B>, Error=Error> + 'static,
        S::Future: 'static,
        B: MessageBody
{

// ...

And the current code at the moment with your suggestions.

async fn manipulate_body<B: MessageBody>(res: &ServiceResponse<B>) -> ServiceResponse<B> {
    let (req, res) = res.into_parts();
    let (res, body) = res.into_parts();

    // Read body bytes
    let body_bytes = actix_web::body::to_bytes_limited(
        body,
        2_usize.pow(22), // 4 MB
    )
        .await
        .unwrap()
        .map_err(Into::into)
        .unwrap();

    let mut body: Value = serde_json::from_slice(&body_bytes).unwrap();
    body["test"] = Value::String("anything".to_owned());

    let res = res.set_body::<B>(body.to_string());
    ServiceResponse::new(req, res)
}

[0rangeFox]

Oh yeah, I kept hitting the same button with the references and I forgot to remove that part since you had mentioned it in the first explanation, but the problem is that it seems to be always in the part of the value type that is returned.

error[E0308]: mismatched types
   --> src/middlewares/test_middleware.rs:242:31
    |
224 | async fn manipulate_body<B1: MessageBody, B2: MessageBody>(res: ServiceResponse<B1>) -> ServiceResponse<B2> {
    |                                           -- expected this type parameter
...
242 |     ServiceResponse::new(req, res)
    |     --------------------      ^^^ expected `HttpResponse<B2>`, found `HttpResponse<String>`
    |     |
    |     arguments to this function are incorrect
    |
    = note: expected struct `HttpResponse<B2>`
               found struct `HttpResponse<std::string::String>`

[robjtede]

Ahh right, in that case:

- async fn manipulate_body<B1: MessageBody, B2: MessageBody>(res: ServiceResponse<B1>) -> ServiceResponse<B2> {
+ async fn manipulate_body<B: MessageBody>(res: ServiceResponse<B>) -> ServiceResponse<String> {

The point is B != String so you can't use one in place for the other.

[0rangeFox]

That I know, but I would be limited only to String, but isn’t possible to make it convertible into B without needing to change the whole code? Or I don’t need to worry just to be string as output? Won’t crash in future?

[robjtede]

You're converting your body type into a string so String seems fine. You don't know what the concrete type of B was so you'll find it hard/impossible to convert back.

Think of it like:

fn stringify<T: ToString>(val: T) -> String {
  val.to_string()
}

It'd be hard/impossible to return a T from this, even if val was a String, because you have no way of asserting that it's the case.

[0rangeFox]

Well, if there was nothing to say that it should be extended or derived from, then it would be impossible to know the type. But in this case, I'm stating that B is derived from MessageBody, so at least we know that the returned type will be composed of MessageBody, that's why in one code example I was trying to convert String to BoxBody.

[0rangeFox]

Ooohhh yeah, I managed to get everything working with the help of @robjtede, if it wasn't for him it might have taken me longer to get there. Basically, most of the problem I had was that I was restricting the generic type B to be a MessageBody, so if we look at the struct ServiceResponse, the generic type is already defined as a BoxBody and that makes it even better than defining our generic types, making it dynamic and being able to manipulate any types we want. So for those interested who want to know what the solution looks like with the full code, here goes! Happy coding! 💪

use std::future::{Future, ready, Ready};
use std::pin::Pin;
use std::rc::Rc;
use actix_web::{dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform}, Error};
use actix_web::body::{BoxBody, MessageBody, to_bytes};
use serde_json::Value;

#[derive(Clone, Default)]
pub struct BodyManipulationMiddleware;

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

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(InnerBodyManipulationMiddleware {
            service: Rc::new(service)
        }))
    }
}

pub struct InnerBodyManipulationMiddleware<S> {
    service: Rc<S>
}

impl<S> Service<ServiceRequest> for InnerBodyManipulationMiddleware<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>>>>;

    forward_ready!(service);

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

        Box::pin(async move {
            let res = service.call(req).await?;
            Ok(manipulate_body(res).await)
        })
    }
}

async fn manipulate_body(res: ServiceResponse) -> ServiceResponse {
    let (req, res) = res.into_parts();
    let (res, body) = res.into_parts();

    let body_bytes = to_bytes(body).await.unwrap();
    let mut body: Value = serde_json::from_slice(&body_bytes).unwrap();
    body["test"] = Value::String("anything".to_owned());

    let res = res.set_body(BoxBody::new(body.to_string()));
    ServiceResponse::new(req, res)
}