feat(node): add KV session wrapper (#1054)

* feat(node): add kv session wrapper utility

* fix(node): address kv session wrapper comments

Author: wangsijie

packages/node/README.md

@@ -83,38 +83,14 @@ But the cookie size is limited, so you may need to use external storage like Red
 
 ```typescript
 // storage.ts
-import { CookieStorage } from '@logto/node';
-
-class RedisSessionWrapper implements SessionWrapper {
-  private currentSessionId?: string;
-
-  constructor(private readonly redis: Redis) {}
-
-  async wrap(data: unknown, _key: string): Promise<string> {
-    // Reuse existing session ID if available, only generate new one for first-time users.
-    // This is important for environments where cookies cannot be updated (e.g. React Server Components),
-    // as the session ID in the cookie must remain stable while the data in Redis can be updated.
-    const sessionId = this.currentSessionId ?? randomUUID();
-    this.currentSessionId = sessionId;
-    await this.redis.set(`logto_session_${sessionId}`, JSON.stringify(data));
-    return sessionId;
-  }
-
-  async unwrap(value: string, _key: string): Promise<SessionData> {
-    if (!value) {
-      return {};
-    }
-
-    // Store the session ID for potential reuse in wrap()
-    this.currentSessionId = value;
-    const data = await this.redis.get(`logto_session_${value}`);
-    return data ? JSON.parse(data) : {};
-  }
-}
+import { CookieStorage, createKVSessionWrapper } from '@logto/node';
 
-export const storage = new CookieStorage({
+export const createStorage = () => new CookieStorage({
   cookieKey: `<logto_app_xxx>`,
-  sessionWrapper: new RedisSessionWrapper(redis),
+  sessionWrapper: createKVSessionWrapper({
+    get: (key) => redis.get(key),
+    set: (key, value, ttl) => redis.set(key, value, 'EX', ttl),
+  }),
   isSecure: false, // Set to true if you are using HTTPS
   getCookie: (name) => {
     // Example usage, get cookie from the request, depends on your framework

packages/node/src/utils/cookie-storage.test.ts

@@ -114,6 +114,32 @@ describe('CookieStorage', () => {
     const cookie = await storage.config.getCookie('logtoCookies');
     expect(cookie).toBe(JSON.stringify({ [PersistKey.AccessToken]: 'test-token' }));
   });
+
+  it('should pass current session value to custom sessionWrapper when saving', async () => {
+    const mockSessionWrapper = {
+      wrap: vi.fn(
+        async (_data: unknown, _key: string, currentValue?: string) => currentValue ?? ''
+      ),
+      unwrap: vi.fn(async () => ({ [PersistKey.AccessToken]: 'test-token' })),
+    };
+    const storage = new TestCookieStorage({
+      ...createCookieConfig(encryptionKey, {
+        logtoCookies: 'existing-session',
+      }),
+      sessionWrapper: mockSessionWrapper,
+    });
+
+    await storage.init();
+    await storage.setItem(PersistKey.AccessToken, 'updated-token');
+
+    expect(mockSessionWrapper.unwrap).toHaveBeenCalledWith('existing-session', encryptionKey);
+    expect(mockSessionWrapper.wrap).toHaveBeenCalledWith(
+      { [PersistKey.AccessToken]: 'updated-token' },
+      encryptionKey,
+      'existing-session'
+    );
+    await expect(storage.config.getCookie('logtoCookies')).resolves.toBe('existing-session');
+  });
 });
 
 describe('CookieStorage concurrency', () => {

packages/node/src/utils/cookie-storage.ts

@@ -19,7 +19,7 @@ export type CookieConfigBase = {
 };
 
 export type SessionWrapper = {
-  wrap: (data: SessionData, key: string) => Promise<string>;
+  wrap: (data: SessionData, key: string, currentValue?: string) => Promise<string>;
   unwrap: (value: string, key: string) => Promise<SessionData>;
 };
 
@@ -61,6 +61,7 @@ export class CookieStorage implements Storage<PersistKey> {
   }
 
   protected sessionData: SessionData = {};
+  protected sessionValue = '';
   protected saveQueue = new PromiseQueue();
 
   /**
@@ -85,10 +86,9 @@ export class CookieStorage implements Storage<PersistKey> {
 
   async init() {
     const { encryptionKey = '' } = this.config;
-    this.sessionData = await this.sessionWrapper.unwrap(
-      (await this.config.getCookie(this.cookieKey)) ?? '',
-      encryptionKey
-    );
+    const sessionValue = (await this.config.getCookie(this.cookieKey)) ?? '';
+    this.sessionValue = sessionValue;
+    this.sessionData = await this.sessionWrapper.unwrap(sessionValue, encryptionKey);
   }
 
   async getItem(key: PersistKey): Promise<Nullable<string>> {
@@ -117,7 +117,8 @@ export class CookieStorage implements Storage<PersistKey> {
 
   protected async write(data = this.sessionData) {
     const { encryptionKey = '' } = this.config;
-    const value = await this.sessionWrapper.wrap(data, encryptionKey);
+    const value = await this.sessionWrapper.wrap(data, encryptionKey, this.sessionValue);
+    this.sessionValue = value;
     await this.config.setCookie(this.cookieKey, value, this.cookieOptions);
   }
 }

packages/node/src/utils/session.test.ts

@@ -1,6 +1,6 @@
 import { PersistKey } from '@logto/client';
 
-import { unwrapSession, wrapSession } from './session.js';
+import { createKVSessionWrapper, unwrapSession, wrapSession } from './session.js';
 
 const secret = 'secret';
 
@@ -23,4 +23,73 @@ describe('session', () => {
     const session = await unwrapSession(cookie, secret);
     expect(session[PersistKey.IdToken]).toEqual('idToken');
   });
+
+  it('should be able to wrap and unwrap with kv wrapper', async () => {
+    const store = new Map<string, string>();
+    const sessionWrapper = createKVSessionWrapper({
+      get: async (key) => store.get(key),
+      set: async (key, value) => {
+        store.set(key, value);
+      },
+    });
+
+    const cookieValue = await sessionWrapper.wrap({ [PersistKey.IdToken]: 'idToken' }, '');
+    const session = await sessionWrapper.unwrap(cookieValue, '');
+
+    expect(session[PersistKey.IdToken]).toEqual('idToken');
+  });
+
+  it('should reuse the current kv session id and pass options to the adapter', async () => {
+    const store = new Map<string, string>();
+    const set = vi.fn(async (key: string, value: string, ttl?: number) => {
+      store.set(key, value);
+    });
+    const sessionWrapper = createKVSessionWrapper(
+      {
+        get: async (key) => store.get(key) ?? null,
+        set,
+      },
+      {
+        keyPrefix: 'custom_prefix_',
+        ttl: 123,
+      }
+    );
+    const existingSessionId = 'session-id';
+    store.set(
+      `custom_prefix_${existingSessionId}`,
+      JSON.stringify({ [PersistKey.IdToken]: 'idToken' })
+    );
+
+    const session = await sessionWrapper.unwrap(existingSessionId, '');
+    const cookieValue = await sessionWrapper.wrap(
+      { ...session, [PersistKey.AccessToken]: 'accessToken' },
+      '',
+      existingSessionId
+    );
+
+    expect(cookieValue).toBe(existingSessionId);
+    expect(set).toHaveBeenCalledWith(
+      `custom_prefix_${existingSessionId}`,
+      JSON.stringify({
+        [PersistKey.IdToken]: 'idToken',
+        [PersistKey.AccessToken]: 'accessToken',
+      }),
+      123
+    );
+    expect(JSON.parse(store.get(`custom_prefix_${existingSessionId}`) ?? '')).toEqual({
+      [PersistKey.IdToken]: 'idToken',
+      [PersistKey.AccessToken]: 'accessToken',
+    });
+  });
+
+  it.each(['null', '1', '[]'])('should ignore invalid kv session data: %s', async (value) => {
+    const sessionWrapper = createKVSessionWrapper({
+      get: async () => value,
+      set: async () => {
+        await Promise.resolve();
+      },
+    });
+
+    await expect(sessionWrapper.unwrap('session-id', '')).resolves.toEqual({});
+  });
 });

packages/node/src/utils/session.ts

@@ -12,6 +12,19 @@ export type Session = SessionData & {
   getValues?: () => Promise<string>;
 };
 
+// eslint-disable-next-line @typescript-eslint/ban-types
+type NullableKVValue = string | null | undefined;
+
+export type KVAdapter = {
+  get: (key: string) => Promise<NullableKVValue>;
+  set: (key: string, value: string, ttlSeconds?: number) => Promise<void>;
+};
+
+export type KVSessionWrapperOptions = {
+  keyPrefix?: string;
+  ttl?: number;
+};
+
 async function getKeyFromPassword(password: string, crypto: Crypto): Promise<string> {
   const encoder = new TextEncoder();
   const data = encoder.encode(password);
@@ -99,3 +112,45 @@ export const wrapSession = async (session: SessionData, secret: string): Promise
   const { ciphertext, iv } = await encrypt(JSON.stringify(session), secret, crypto);
   return `${ciphertext}.${iv}`;
 };
+
+const isSessionData = (data: unknown): data is SessionData =>
+  typeof data === 'object' && data !== null && !Array.isArray(data);
+
+export const createKVSessionWrapper = (
+  kv: KVAdapter,
+  options: KVSessionWrapperOptions = {}
+): {
+  wrap: (data: SessionData, key: string, currentValue?: string) => Promise<string>;
+  unwrap: (value: string, key: string) => Promise<SessionData>;
+} => {
+  const prefix = options.keyPrefix ?? 'logto_session_';
+  const ttl = options.ttl ?? 14 * 24 * 3600;
+
+  return {
+    async wrap(data: SessionData, _key: string, currentValue?: string): Promise<string> {
+      const sessionId =
+        currentValue === undefined || currentValue === '' ? crypto.randomUUID() : currentValue;
+      await kv.set(`${prefix}${sessionId}`, JSON.stringify(data), ttl);
+      return sessionId;
+    },
+    async unwrap(value: string, _key: string): Promise<SessionData> {
+      if (!value) {
+        return {};
+      }
+
+      const data = await kv.get(`${prefix}${value}`);
+
+      if (!data) {
+        return {};
+      }
+
+      try {
+        // eslint-disable-next-line no-restricted-syntax
+        const sessionData = JSON.parse(data) as unknown;
+        return isSessionData(sessionData) ? sessionData : {};
+      } catch {
+        return {};
+      }
+    },
+  };
+};