Search code examples
rustcookiesmiddlewarerust-axum

Best way to create a session cookie if it's not present and persist to a database using axum and tower_cookies?


I'm relatively new to axum and I'm trying to use it as the webserver for a service I'm creating.

I have some state which contains a database I would like to write session / session metadata to.

I want all users which access any route to set a session cookie if it's not already present in the header. I then want to write this session cookie to the database (which is stored in some state).

I am using the tower_cookies crate in an attempt to make cookie reading and writing more ergonomic, so I'd prefer if the middleware were able to access the Cookies struct rather than reading and writing the raw Request<...> and Response<...>.

All my attempts to do this have been insufficient in some fashion -- either being a standalone handler (instead of a middleware automatically applied to all routes), or not using the Cookie class (and dealing with the raw Cookie / Set-Cookie headers directly), or something similarly not ergonomic. I feel that this is a common thing to want to do so I must be missing something. For a distilled example of one of these not-ergonomic but compiling approaches:

// Problem: not using Cookies from the CookieManager, instead dealing with the headers directly.
pub async fn set_and_store_cookie_if_absent<B>(
    State(state): State<MyState>,
    request: Request<B>,
    next: Next<B>,
) -> Response {
    // Read request cookies, create / write new session to state.db if needed
    let response = next.run(request).await;
    // Write to response cookies.
    response
}

let state = MyState::new();
Route(...)
    .layer(
        ServiceBuilder::new()
            .layer(CookieManagerLayer::new())
            .layer(middleware::from_fn_with_state(state.clone(), set_and_store_cookie_if_absent)));

But this isn't using the CookieManagerLayer that I want to use.


Solution

  • Ah, after doing some research, some of my attempted approaches were correct, except my State was wrapping a type that doesn't play well with async: https://github.com/tokio-rs/axum/discussions/964#discussioncomment-2629585

    So the right thing to do would be something like this:

    pub async fn set_and_store_cookie_if_absent<B>(
        State(state): State<MyState>,
        cookies: Cookies,
        request: Request<B>,
        next: Next<B>,
    ) -> Response {
        if cookies.get(SESSION_COOKIE).is_none() {
            // handle, write to State, etc.
        }
        
        next.run(request).await
    }
    

    But this was giving me a TON of errors that were hard to wade through before -- because I was using rusqlite instead of something like deadpool-sqlite.