Regenerating session from request handler doesn't work

[Ptrskay3]

Calling regenerate from a request handler does not refresh the cookie for the client. Having a reliable way to refresh cookie values is essential in most applications: secure session based authentication systems rotate the cookie value on any privilege level change to prevent session fixation attacks.

This is happening because we insert the cloned session to the request handler as an extension

axum-sessions/src/session.rs

Lines 279 to 280 in f4b2301

	request.extensions_mut().insert(session.clone());
	let mut response = inner.call(request).await?;

and the original session in the middleware is still holding onto the old cookie_value

axum-sessions/src/session.rs

Lines 294 to 295 in f4b2301

	} else if session_layer.save_unchanged || session.data_changed() {
	match session_layer.store.store_session(session).await {

so it won't see the regenerate call from the request handler.

I have a workaround for this, but it's by no means pretty, and requires an extra round-trip to the session store.

Basically I defined an extension trait for async_session::Session to add two methods:

pub trait SessionExt {
    const REGENERATION_MARK_KEY: &'static str;

    fn mark_for_regenerate(&mut self);
    fn should_regenerate(&mut self) -> bool;
}

impl SessionExt for async_session::Session {
    const REGENERATION_MARK_KEY: &'static str = "__regenerate_key";

    fn mark_for_regenerate(&mut self) {
        self.insert(Self::REGENERATION_MARK_KEY, true)
            .expect("bool is serializable");
    }

    fn should_regenerate(&mut self) -> bool {
        let previously_changed = self.data_changed();
        let regenerate = self.get(Self::REGENERATION_MARK_KEY).unwrap_or_default();

        self.remove(Self::REGENERATION_MARK_KEY);
        if !previously_changed {
            self.reset_data_changed();
        }

        regenerate
    }
}

In the middleware we could insert

     if session.should_regenerate() {
            session.regenerate();
         }

and in request handlers consumers should call mark_for_regenerate instead of regenerate.
However, this solution is definitely a little ambiguous/deceptive for users..

[maxcountryman]

I wonder if we should try to avoid cloning the session in some way? (I haven't really thought about the mechanics of if that's even feasible.)

[Ptrskay3]

I don't think we can avoid cloning there. I guess someone should look at it who's smarter than me 🤷‍♂️

[maxcountryman]

That likely isn't me! 😁

In all seriousness, I'm happy to explore whatever direction you think makes the most sense for addressing this issue.

[maxcountryman]

Thinking about this a little more, I took a look at wrapping async_session::Session with Arc<RwLock<_>>. Here's what that could look like:

diff --git a/src/session.rs b/src/session.rs
index 15a460b..54c53a2 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -2,6 +2,7 @@
 // `tide::sessions::middleware::SessionMiddleware`. See: https://github.com/http-rs/tide/blob/20fe435a9544c10f64245e883847fc3cd1d50538/src/sessions/middleware.rs
 
 use std::{
+    sync::Arc,
     task::{Context, Poll},
     time::Duration,
 };
@@ -21,10 +22,13 @@ use axum::{
 };
 use axum_extra::extract::cookie::{Cookie, Key, SameSite};
 use futures::future::BoxFuture;
+use tokio::sync::RwLock;
 use tower::{Layer, Service};
 
 const BASE64_DIGEST_LEN: usize = 44;
 
+pub type SessionHandle = Arc<RwLock<async_session::Session>>;
+
 #[derive(Clone)]
 pub struct SessionLayer<Store> {
     store: Store,
@@ -134,15 +138,17 @@ impl<Store: SessionStore> SessionLayer<Store> {
         self
     }
 
-    async fn load_or_create(&self, cookie_value: Option<String>) -> async_session::Session {
+    async fn load_or_create(&self, cookie_value: Option<String>) -> SessionHandle {
         let session = match cookie_value {
             Some(cookie_value) => self.store.load_session(cookie_value).await.ok().flatten(),
             None => None,
         };
 
-        session
-            .and_then(|session| session.validate())
-            .unwrap_or_default()
+        Arc::new(RwLock::new(
+            session
+                .and_then(async_session::Session::validate)
+                .unwrap_or_default(),
+        ))
     }
 
     fn build_cookie(&self, secure: bool, cookie_value: String) -> Cookie<'static> {
@@ -270,16 +276,24 @@ where
 
         let mut inner = self.inner.clone();
         Box::pin(async move {
-            let mut session = session_layer.load_or_create(cookie_value).await;
+            let session_handle = session_layer.load_or_create(cookie_value).await;
 
+            let mut session = session_handle.write().await;
             if let Some(ttl) = session_layer.session_ttl {
-                session.expire_in(ttl);
+                (*session).expire_in(ttl);
             }
+            drop(session);
 
-            request.extensions_mut().insert(session.clone());
+            request.extensions_mut().insert(session_handle.clone());
             let mut response = inner.call(request).await?;
 
+            let session = session_handle.read().await;
             if session.is_destroyed() {
+                drop(session);
+                let session = RwLock::into_inner(
+                    Arc::try_unwrap(session_handle).expect("Session handle still has owners."),
+                );
+
                 if let Err(e) = session_layer.store.destroy_session(session).await {
                     tracing::error!("Failed to destroy session: {:?}", e);
                     *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
@@ -292,6 +306,11 @@ where
                     HeaderValue::from_str(&removal_cookie.to_string()).unwrap(),
                 );
             } else if session_layer.save_unchanged || session.data_changed() {
+                drop(session);
+                let session = RwLock::into_inner(
+                    Arc::try_unwrap(session_handle).expect("Session handle still has owners."),
+                );
+
                 match session_layer.store.store_session(session).await {
                     Ok(Some(cookie_value)) => {
                         let cookie = session_layer.build_cookie(session_layer.secure, cookie_value);
@@ -300,7 +319,11 @@ where
                             HeaderValue::from_str(&cookie.to_string()).unwrap(),
                         );
                     }
-                    Ok(None) => {}
+
+                    Ok(None) => {
+                        tracing::warn!("The cookie value is missing.");
+                    }
+
                     Err(e) => {
                         tracing::error!("Failed to reach session storage: {:?}", e);
                         *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;

Because async_session::Session's Clone implementation blows away the cookie value, it's necessary to do some extra shenanigans with drop, as seen above: specifically we can't pass the interior session to store_session and get back the expected cookie value otherwise.

Then in the application code, you might write a handler like so:

use axum::{routing::get, Extension, Router};
use axum_sessions::{async_session::MemoryStore, SessionHandle, SessionLayer};

#[tokio::main]
async fn main() {
    let store = MemoryStore::new();
    let secret = b"please do not use this secret in a real application, you must change it";
    let session_layer = SessionLayer::new(store, secret);

    async fn regenerate_handler(Extension(session_handle): Extension<SessionHandle>) {
        let mut session = session_handle.write().await;

        // NB: This DOES NOT update the store, meaning that both sessions will still be
        // found.
        session.regenerate();
    }

    async fn handler() -> &'static str {
        "Hello, world!"
    }

    let app = Router::new()
        .route("/", get(handler))
        .route("/regenerate", get(regenerate_handler))
        .layer(session_layer);

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

Note that callers will still need to purge the session from the underlying store when calling regenerate if they want to protect against the old session being found again later.

@Ptrskay3 I'd love your thoughts on this approach.

[Ptrskay3]

I've tested your solution, and seems to be working correctly. Also ran a quick loadtest with a Redis store backend in the application I'm using this library, and the performance seems to be the same as before, so the implementation is surely not the bottleneck.
One thing I'm not particularly keen on is the way you have to write application code: users now need to care about the RwLock in the handler. Of course you're allowed to do that in a FromRequest implementation, but then you would need to split the API for readable and writable session handles.

Meanwhile, I've been playing with another approach. I made a wrapper around async_session::Session like so:

pub struct Session {
    inner: async_session::Session,
    regenerate: Arc<AtomicBool>,
}

This way, we don't need to care about cloning the session, as the regenerate field is shared. The regenerate functionality becomes:

impl Session {
     // ...

    /// Marks the session for id and cookie value regeneration. Unlike `async_session::Session`,
    /// the actual work is done in the middleware __after__ your handler is called.
    pub fn regenerate(&mut self) {
        self.regenerate.store(true, Ordering::Relaxed);
    }

    pub(crate) fn inner_regenerate(&mut self) {
        self.inner.regenerate();
    }

    pub(crate) fn should_regenerate(&self) -> bool {
        self.regenerate.load(Ordering::Relaxed)
    }
}

and in the middleware, we can just write

if session.should_regenerate() {
    session.inner_regenerate();
}

Application code stays the same as before. I believe this way it should be easy for us to purge the old session from the storage layer after regeneration directly in the middleware, so users would not need to worry about that. This is similar to what actix_session does.

On the other hand, it seems to be a little awkward to wrap the whole async_session::Session's public API. Also, it might be confusing to use our axum_sessions::Session instead of async_session::Session with nearly identical functionality, but not quite (see regenerate documentation I wrote above). The implementation is just a draft, I haven't reviewed very thoroughly, and haven't changed the documentation to match the new API.

Sadly all workarounds have their own drawbacks 😞.

I'm curious what do you think about this approach.

[maxcountryman]

Of course you're allowed to do that in a FromRequest implementation, but then you would need to split the API for readable and writable session handles.

To explore this idea a little more I went ahead and added an extractors module:

use std::ops::{Deref, DerefMut};

use axum::{
    async_trait,
    extract::{FromRequest, RequestParts},
    http, Extension,
};
use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard};

use crate::SessionHandle;

/// An extractor which provides a readable session. Sessions may have many
/// readers.
pub struct ReadableSession {
    session: OwnedRwLockReadGuard<async_session::Session>,
}

impl Deref for ReadableSession {
    type Target = OwnedRwLockReadGuard<async_session::Session>;

    fn deref(&self) -> &Self::Target {
        &self.session
    }
}

#[async_trait]
impl<B> FromRequest<B> for ReadableSession
where
    B: Send,
{
    type Rejection = http::StatusCode;

    async fn from_request(request: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        let Extension(session_handle): Extension<SessionHandle> = Extension::from_request(request)
            .await
            .expect("Session extension missing. Is the session layer installed?");
        let session = session_handle.read_owned().await;

        Ok(Self { session })
    }
}

/// An extractor which provides a writable session. Sessions may have only one
/// writer.
pub struct WritableSession {
    session: OwnedRwLockWriteGuard<async_session::Session>,
}

impl Deref for WritableSession {
    type Target = OwnedRwLockWriteGuard<async_session::Session>;

    fn deref(&self) -> &Self::Target {
        &self.session
    }
}

impl DerefMut for WritableSession {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.session
    }
}

#[async_trait]
impl<B> FromRequest<B> for WritableSession
where
    B: Send,
{
    type Rejection = http::StatusCode;

    async fn from_request(request: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        let Extension(session_handle): Extension<SessionHandle> = Extension::from_request(request)
            .await
            .expect("Session extension missing. Is the session layer installed?");
        let session = session_handle.write_owned().await;

        Ok(Self { session })
    }
}

One of the side effects of doing this is application code no longer needs to wrap the session in Extension, but as you point out, the downside is now we have to pick explicitly between two extractors. (Depending on your perspective, it could be a good thing that this choice becomes an explicit feature of our public type interface.)

Here's what that could look like in an admittedly contrived application example:

use axum::{routing::get, Router};
use axum_sessions::{
    async_session::MemoryStore,
    extractors::{ReadableSession, WritableSession},
    SessionLayer,
};
use rand::Rng;

#[tokio::main]
async fn main() {
    let store = MemoryStore::new();
    let secret = rand::thread_rng().gen::<[u8; 128]>();
    let session_layer = SessionLayer::new(store, &secret);

    async fn signin_handler(mut session: WritableSession) {
        session
            .insert("signed_in", true)
            .expect("Could not sign in.");
    }

    async fn signout_handler(mut session: WritableSession) {
        session.destroy();
    }

    async fn protected_handler(session: ReadableSession) -> &'static str {
        if session
            .get::<bool>("signed_in")
            .map_or(false, |signed_in| signed_in)
        {
            "Shh, it's secret!"
        } else {
            "Nothing to see here."
        }
    }

    let app = Router::new()
        .route("/signin", get(signin_handler))
        .route("/signout", get(signout_handler))
        .route("/protected", get(protected_handler))
        .layer(session_layer);

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

This feels somewhat closer to the original implementation, where we could manipulate the session directly. (Note I haven't tested this extensively, so it probably needs a little more vetting before we can consider it viable.)

I believe this way it should be easy for us to purge the old session from the storage layer after regeneration directly in the middleware, so users would not need to worry about that.

Curious what this looks like: would we update the store in the middleware as we check for should_generate? Perhaps something like:

if session.should_regenerate() {
    session.inner_regenerate();
    // Remove the old session here?
}

On the other hand, it seems to be a little awkward to wrap the whole async_session::Session's public API. Also, it might be confusing to use our axum_sessions::Session instead of async_session::Session with nearly identical functionality

It seems like a good bet that async_session::Session's API is pretty stable. So while it's a shame to have to wrap each method, we'd probably be pretty safe in that regard. That said, I agree that the slight semantic difference with regenerate is potentially confusing and surprising for our users. In practice, it's not clear to me if that will really matter tho.

Overall, I don't have a very strong preference, but am starting to lean a bit towards the interior mutability pattern with readable and writable extractors if we agree that's both viable and not the worst API. I'll look to your input to guide how we should move forward. :)

[Ptrskay3]

I don't have a strong preference either, I'm fine with the interior mutability approach. Actually, having you spelled out the extractor module might convinced me fully to move forward with that.

[maxcountryman]

I'll put together a PR and we can review it from there.