Reapply feat: add optional update() to SessionDataStore (#2590) (#2616)

Author: Piyush-85

EXAMPLES.md

@@ -2651,6 +2651,28 @@ export const auth0 = new Auth0Client({
     async delete(id) {
       // delete the session using its ID
     },
+
+    async update(id, sessionData) {
+      // Optional: update the session by its ID only if it already exists; return true if updated, false if not found.
+      // Prevents a session deleted by logout from being re-created by a concurrent in-flight rolling write.
+      //
+      // MUST be a single atomic operation — do not use get() + set() here,
+      // as that would recreate the TOCTOU race condition this method is designed to prevent.
+      //
+      // Example (PostgreSQL):
+      //   const r = await db.query(
+      //     "UPDATE sessions SET data=$2, expires_at=NOW()+interval'1 day' WHERE id=$1",
+      //     [id, sessionData]
+      //   );
+      //   return r.rowCount > 0;
+      //
+      // Example (Redis — Lua script for atomicity):
+      //   const lua = `if redis.call('exists',KEYS[1])==1 then
+      //     redis.call('set',KEYS[1],ARGV[1],'EX',ARGV[2]) return 1
+      //     else return 0 end`;
+      //   return await redis.eval(lua, 1, id, JSON.stringify(sessionData), ttl) === 1;
+    },
+
     async deleteByLogoutToken({ sid, sub }: { sid?: string; sub?: string }) {
       // optional method to be implemented when using Back-Channel Logout
     }

src/server/session/stateful-session-store.test.ts

@@ -463,6 +463,101 @@ describe("Stateful Session Store", async () => {
     });
 
     describe("rolling session race condition", async () => {
+      it("should use store.update() atomically when implemented and bail if it returns false", async () => {
+        // store.update() returns false → session deleted by concurrent logout.
+        // The SDK must not fall through to store.set() — zero TOCTOU gap.
+        const sessionId = "ses_atomic";
+        const secret = await generateSecret(32);
+        const session: SessionData = {
+          user: { sub: "user_123" },
+          tokenSet: {
+            accessToken: "at_123",
+            refreshToken: "rt_123",
+            expiresAt: 123456
+          },
+          internal: {
+            sid: "auth0-sid",
+            createdAt: Math.floor(Date.now() / 1000)
+          }
+        };
+        const store = {
+          get: vi.fn(),
+          set: vi.fn(),
+          delete: vi.fn(),
+          update: vi.fn().mockResolvedValue(false) // session already gone
+        };
+        const maxAge = 60 * 60;
+        const expiration = Math.floor(Date.now() / 1000 + maxAge);
+        const encryptedCookieValue = await encrypt(
+          { id: sessionId },
+          secret,
+          expiration
+        );
+        const headers = new Headers();
+        headers.append("cookie", `__session=${encryptedCookieValue}`);
+        const requestCookies = new RequestCookies(headers);
+        const responseCookies = new ResponseCookies(new Headers());
+
+        const sessionStore = new StatefulSessionStore({
+          secret,
+          store,
+          rolling: true
+        });
+        await sessionStore.set(requestCookies, responseCookies, session);
+
+        expect(store.update).toHaveBeenCalledWith(sessionId, session);
+        expect(store.set).not.toHaveBeenCalled();
+        expect(store.get).not.toHaveBeenCalled(); // no fallback read
+        expect(responseCookies.get("__session")).toBeUndefined();
+      });
+
+      it("should use store.update() atomically and refresh the cookie when it returns true", async () => {
+        // store.update() returns true → session found and updated in one DB op.
+        const sessionId = "ses_atomic_alive";
+        const secret = await generateSecret(32);
+        const session: SessionData = {
+          user: { sub: "user_123" },
+          tokenSet: {
+            accessToken: "at_123",
+            refreshToken: "rt_123",
+            expiresAt: 123456
+          },
+          internal: {
+            sid: "auth0-sid",
+            createdAt: Math.floor(Date.now() / 1000)
+          }
+        };
+        const store = {
+          get: vi.fn(),
+          set: vi.fn(),
+          delete: vi.fn(),
+          update: vi.fn().mockResolvedValue(true)
+        };
+        const maxAge = 60 * 60;
+        const expiration = Math.floor(Date.now() / 1000 + maxAge);
+        const encryptedCookieValue = await encrypt(
+          { id: sessionId },
+          secret,
+          expiration
+        );
+        const headers = new Headers();
+        headers.append("cookie", `__session=${encryptedCookieValue}`);
+        const requestCookies = new RequestCookies(headers);
+        const responseCookies = new ResponseCookies(new Headers());
+
+        const sessionStore = new StatefulSessionStore({
+          secret,
+          store,
+          rolling: true
+        });
+        await sessionStore.set(requestCookies, responseCookies, session);
+
+        expect(store.update).toHaveBeenCalledWith(sessionId, session);
+        expect(store.set).not.toHaveBeenCalled();
+        expect(store.get).not.toHaveBeenCalled();
+        expect(responseCookies.get("__session")).toBeDefined();
+      });
+
       it("should not re-create a session that was deleted by a concurrent logout", async () => {
         const sessionId = "ses_123";
         const secret = await generateSecret(32);
@@ -609,6 +704,45 @@ describe("Stateful Session Store", async () => {
         expect(store.set).toHaveBeenCalledOnce();
         expect(responseCookies.get("__session")).toBeDefined();
       });
+
+      it("should call store.set() and not store.update() for a brand-new session even when update() is implemented", async () => {
+        // When there is no existing session cookie, existingSessionId is null and the guard
+        // is bypassed entirely. store.set() must be called to create the new session;
+        // store.update() must NOT be called — it would immediately return false (nothing
+        // exists yet) and silently swallow the new session creation.
+        const secret = await generateSecret(32);
+        const session: SessionData = {
+          user: { sub: "user_123" },
+          tokenSet: {
+            accessToken: "at",
+            refreshToken: "rt",
+            expiresAt: 123456
+          },
+          internal: {
+            sid: "sid",
+            createdAt: Math.floor(Date.now() / 1000)
+          }
+        };
+        const store = {
+          get: vi.fn(),
+          set: vi.fn(),
+          delete: vi.fn(),
+          update: vi.fn() // must NOT be called for new sessions
+        };
+        const requestCookies = new RequestCookies(new Headers()); // no existing cookie
+        const responseCookies = new ResponseCookies(new Headers());
+
+        const sessionStore = new StatefulSessionStore({
+          secret,
+          store,
+          rolling: true
+        });
+        await sessionStore.set(requestCookies, responseCookies, session);
+
+        expect(store.update).not.toHaveBeenCalled();
+        expect(store.set).toHaveBeenCalledOnce();
+        expect(responseCookies.get("__session")).toBeDefined();
+      });
     });
 
     describe("session fixation", async () => {

src/server/session/stateful-session-store.ts

@@ -157,11 +157,25 @@ export class StatefulSessionStore extends AbstractSessionStore {
     // The guard only applies when we found an existing session ID in the request cookie
     // (existingSessionId !== null). Brand-new sessions (no cookie, or isNew login) bypass
     // the check because there is nothing in the store to verify against.
+    //
+    // If the store implements the optional update() method we use it as an atomic
+    // check-and-write (UPDATE WHERE id = $1). Otherwise we fall back to a non-atomic
+    // get() + set() pair.
     if (existingSessionId !== null) {
-      const existingSession = await this.store.get(existingSessionId);
-      if (!existingSession) {
-        return;
+      if (typeof this.store.update === "function") {
+        const updated = await this.store.update(existingSessionId, session);
+        if (!updated) {
+          return;
+        }
+      } else {
+        const existingSession = await this.store.get(existingSessionId);
+        if (!existingSession) {
+          return;
+        }
+        await this.store.set(existingSessionId, session);
       }
+    } else {
+      await this.store.set(sessionId, session);
     }
 
     const maxAge = this.calculateMaxAge(session.internal.createdAt);
@@ -180,7 +194,6 @@ export class StatefulSessionStore extends AbstractSessionStore {
       ...this.cookieConfig,
       maxAge
     });
-    await this.store.set(sessionId, session);
 
     // to enable read-after-write in the same request for middleware
     reqCookies.set(this.sessionCookieName, jwe.toString());

src/types/index.ts

@@ -37,6 +37,12 @@ export interface SessionData {
   [key: string]: unknown;
 }
 
+/**
+ * Interface for a custom session data store.
+ *
+ * **TTL contract:** every successful write method (`set`, `update`) must reset the session
+ * TTL/expiry so that active sessions are not silently expired between requests.
+ */
 export interface SessionDataStore {
   /**
    * Gets the session from the store given a session ID.
@@ -48,6 +54,13 @@ export interface SessionDataStore {
    */
   set(id: string, session: SessionData): Promise<void>;
 
+  /**
+   * Optional: update the session by its ID only if it already exists.
+   * Return `true` if updated, `false` if not found.
+   *
+   */
+  update?(id: string, session: SessionData): Promise<boolean>;
+
   /**
    * Destroys the session with the given session ID.
    */