how to share data between the request and response?

[kyle-mccarthy]

I am trying to use layers to create a "middleware" that allows me to load a session for a handler and then after the handler, check if the session was modified and if so update it in session store.

What I've tried so far is the method shown in the print-request-response example, using the AsyncFilterLayer and AndThenLayer. I load the session in AsyncFilterLayer and then store it in the response's extensions and can access it in my handler using axum's Extension. However, when I try to access it in my AndThenLayer, the session doesn't exist in the response's extensions.

1. How can I share data between the request and response?
2. Is there a better or "right way" (I'm not familiar with tower) to do what I am trying?

Thanks in advance for any help or guidance!

[davidpdrsn]

Awesome sequence diagram! You don't see a lot of those these days 😅

This is possible but arguably not as easy as it could be. The problem is that request extensions are not automatically propagated to response extensions. You could write a middleware that does that and then combine it with AsyncFilterLayer and AndThenLayer but I think the best solution is to write a single middleware that does what you need. Something like this:

use axum::{
    extract::Extension,
    handler::get,
    http::{Request, Response},
    Router,
};
use std::{
    future::Future,
    net::SocketAddr,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    task::{Context, Poll},
};
use tower::Service;

#[tokio::main]
async fn main() {
    let app = Router::new()
        // endpoint that access the session and changes it
        .route(
            "/change",
            get(|Extension(session): Extension<Session>| async move {
                println!("handler called, changing session");
                session.changed.store(true, Ordering::Relaxed);
            }),
        )
        // endpoint that does not change the session
        .route("/change-not", get(|| async {}))
        // add our `HandleSession` middleware. `layer_fn` is a convenience for
        // creating a `tower::Layer` from a `fn(Service) -> OtherService`
        // function
        .layer(tower::layer::layer_fn(HandleSession::new));

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

// example session type
#[derive(Clone, Debug)]
struct Session {
    changed: Arc<AtomicBool>,
}

// middleware that
// 1. inserts a session into request extensions so handlers can access it
// 2. calls the inner service and awaits the response
// 3. prints the session so we can see if its changed
#[derive(Clone)]
struct HandleSession<S> {
    inner: S,
}

impl<S> HandleSession<S> {
    fn new(inner: S) -> Self {
        Self { inner }
    }
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for HandleSession<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
    S::Future: Send + 'static,
    ReqBody: Send + 'static,
    ResBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;

    #[allow(clippy::type_complexity)]
    type Future =
        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

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

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        // best practice is to clone the inner service like this
        // see https://github.com/tower-rs/tower/issues/547 for details
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);

        Box::pin(async move {
            let session = Session {
                changed: Arc::new(AtomicBool::new(false)),
            };
            // we have to clone the session here because we need to access it
            // once we have the response
            req.extensions_mut().insert(session.clone());

            let res = inner.call(req).await?;
            println!("inner service return. session={:?}", session);
            Ok(res)
        })
    }
}

If you run this and call each endpoint you'll see this:

# GET /change
handler called, changing session
inner service return. session=Session { changed: true }

# GET /change-not
inner service return. session=Session { changed: false }

[kyle-mccarthy]

Thanks for the detailed answer; this is extremely helpful and is exactly what I was trying to accomplish! And yeah I find that sequence diagrams help me clearly explain topics or things that I am trying to do without overloading people with unnecessary or confusing details :)