Search code examples
rustmiddlewarerust-tokiorust-tonic

Retrieving the request body in a tower layer to sign GRPC requests


I'm attempting to implement an authentication layer on top of grpc (through tonic) with the tower middleware layer functionality. In order to do that I need to get the body of the request, consisting of the protobuf payload that is being sent to the server, authenticate it with an HMAC, and then set the authenticating HMAC in the metadata / headers.

However I am having some issues retrieving the request body via the API, which seems to not wait for the entire request to be generalized for small as well as large streaming requests. The latter would require a large body to be buffered in memory before actually issuing the request to the server. Since I am in control of the requests that I want to intercept, and I know my requests are small enough to be bufferable in memory without much overhead, this generalization adds some complexity that I do not know how to handle.

My Layer currently looks like this:

use http::Request;
use std::task::{Context, Poll};
use tonic::body::BoxBody;
use tonic::codegen::Body;
use tower::{Layer, Service};

pub struct AuthLayer {
    hmac_key: Vec<u8>,
}

impl AuthLayer {
    pub fn new(hmac_key: Vec<u8>) -> Self {
        AuthLayer { hmac_key: hmac_key }
    }
}

impl<S> Layer<S> for AuthLayer {
    type Service = AuthService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        AuthService {
            hmac_key: self.hmac_key.clone(),
            inner,
        }
    }
}

// This service implements the Log behavior
pub struct AuthService<S> {
    hmac_key: Vec<u8>,
    inner: S,
}

impl<S> Service<Request<BoxBody>> for AuthService<S>
where
    S: Service<Request<BoxBody>>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

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

    fn call(&mut self, request: Request<BoxBody>) -> Self::Future {
        let _body = dbg!(request.body().data());
        // TODO Compute an HMAC over _body
        // TODO Set HMAC in metadata / headers
        self.inner.call(request)
    }
}

It fails because therequest is not mutable, hence we cannot call data() on the body. Since the requests are small, I'd be ok with cloning the entire request, buffer the body and then compute the HMAC over it, before forwarding the request to the inner service, but everything I tried fails. Ideally however it'd be possible to buffer in place and forward the original.

How can I get the body from the http::Request?


Solution

  • The method signature accepts the request by value:

        fn call(&mut self, request: Request<BoxBody>) -> Self::Future 
    

    Thus you can re-bind it in order to make it mutable. By rebinding it you move the data to a new variable (with the same name), that is mutable:

    let mut request = request; 
    let request_body = request.body_mut(); 
    let body = request_body.data();
    

    Alternatively you can add mut directly on the method:

        fn call(&mut self, mut request: Request<BoxBody>) -> Self::Future 
    

    This will not change the signature, because the caller does not care if you will modify it or not, because it's passed by value.

    how is upgrading the non-mut borrow to a mut borrow safe?

    It is safe (and possible) because it is passed by value, not by reference. Being the sole owner of the value grants you the ability to make it mutable or immutable