Allow for configurable Session Storage (#40)

* add configuration option for session storage

Use a promise to ensure the value is set

* fix tests

* rename cookie to sessionStorage

* adjust build to minimum supported version

* fixes and adjustments

* Update src/sessionStorage.spec.ts

Co-authored-by: Dan Dorman <[email protected]>

* fix typo

* allow for configurable cookie name

* go deeper on testing session storage configuration

* remove unused import

* fix issue where cookieName isn't always used, fix types

---------

Co-authored-by: Dan Dorman <[email protected]>

Author: nicknisi

src/authkit-callback-route.spec.ts

@@ -6,6 +6,7 @@ import {
   createAuthWithCodeResponse,
   assertIsResponse,
 } from './test-utils/test-helpers.js';
+import { configureSessionStorage } from './sessionStorage.js';
 
 // Mock dependencies
 jest.mock('../src/workos.js', () => ({
@@ -25,6 +26,7 @@ describe('authLoader', () => {
   beforeAll(() => {
     // Silence console.error during tests
     jest.spyOn(console, 'error').mockImplementation(() => {});
+    configureSessionStorage();
   });
 
   beforeEach(async () => {

src/authkit-callback-route.ts

@@ -2,11 +2,12 @@ import { HandleAuthOptions } from './interfaces.js';
 import { WORKOS_CLIENT_ID } from './env-variables.js';
 import { workos } from './workos.js';
 import { encryptSession } from './session.js';
-import { getSession, commitSession, cookieName } from './cookie.js';
+import { getSessionStorage } from './sessionStorage.js';
 import { redirect, json, LoaderFunctionArgs } from '@remix-run/node';
 
 export function authLoader(options: HandleAuthOptions = {}) {
   return async function loader({ request }: LoaderFunctionArgs) {
+    const { getSession, commitSession, cookieName } = await getSessionStorage();
     const { returnPathname: returnPathnameOption = '/', onSuccess } = options;
 
     const url = new URL(request.url);

src/cookie.spec.ts

@@ -20,6 +20,7 @@ const WORKOS_API_HOSTNAME = getOptionalEnvVariable('WORKOS_API_HOSTNAME');
 const WORKOS_API_HTTPS = getOptionalEnvVariable('WORKOS_API_HTTPS');
 const WORKOS_API_PORT = getOptionalEnvVariable('WORKOS_API_PORT');
 const WORKOS_COOKIE_MAX_AGE = getOptionalEnvVariable('WORKOS_COOKIE_MAX_AGE');
+const WORKOS_COOKIE_NAME = getOptionalEnvVariable('WORKOS_COOKIE_NAME');
 
 if (WORKOS_COOKIE_PASSWORD.length < 32) {
   throw new Error('WORKOS_COOKIE_PASSWORD must be at least 32 characters long');
@@ -34,4 +35,5 @@ export {
   WORKOS_API_HTTPS,
   WORKOS_API_PORT,
   WORKOS_COOKIE_MAX_AGE,
+  WORKOS_COOKIE_NAME,
 };

src/cookie.ts

@@ -1,13 +1,13 @@
+import { getSignInUrl, getSignUpUrl, signOut } from './auth.js';
 import { authLoader } from './authkit-callback-route.js';
 import { authkitLoader } from './session.js';
-import { getSignInUrl, getSignUpUrl, signOut } from './auth.js';
 
 export {
   authLoader,
   //
+  authkitLoader,
+  //
   getSignInUrl,
   getSignUpUrl,
   signOut,
-  //
-  authkitLoader,
 };

src/env-variables.ts

@@ -1,4 +1,5 @@
-import { OauthTokens, User } from '@workos-inc/node';
+import type { SessionStorage, SessionIdStorageStrategy } from '@remix-run/node';
+import type { OauthTokens, User } from '@workos-inc/node';
 
 export interface HandleAuthOptions {
   returnPathname?: string;
@@ -39,10 +40,19 @@ export interface GetAuthURLOptions {
   returnPathname?: string;
 }
 
-export interface AuthKitLoaderOptions {
+export type AuthKitLoaderOptions = {
   ensureSignedIn?: boolean;
   debug?: boolean;
-}
+} & (
+  | {
+      storage?: never;
+      cookie?: SessionIdStorageStrategy['cookie'];
+    }
+  | {
+      storage: SessionStorage;
+      cookie: SessionIdStorageStrategy['cookie'];
+    }
+);
 
 export interface AuthorizedData {
   user: User;

src/index.ts

@@ -1,26 +1,28 @@
 import { LoaderFunctionArgs, Session as RemixSession, redirect } from '@remix-run/node';
+import { AuthenticationResponse } from '@workos-inc/node';
 import * as ironSession from 'iron-session';
-import * as cookie from './cookie.js';
+import * as jose from 'jose';
+import {
+  configureSessionStorage as configureSessionStorageMock,
+  getSessionStorage as getSessionStorageMock,
+} from './sessionStorage.js';
 import { WORKOS_COOKIE_PASSWORD } from './env-variables.js';
 import { Session } from './interfaces.js';
-import { encryptSession, terminateSession, authkitLoader } from './session.js';
-import { workos } from './workos.js';
-import * as jose from 'jose';
-import { AuthenticationResponse } from '@workos-inc/node';
+import { authkitLoader, encryptSession, terminateSession } from './session.js';
 import { assertIsResponse } from './test-utils/test-helpers.js';
+import { workos } from './workos.js';
 
-const getSession = jest.mocked(cookie.getSession);
-const destroySession = jest.mocked(cookie.destroySession);
 const unsealData = jest.mocked(ironSession.unsealData);
 const sealData = jest.mocked(ironSession.sealData);
 const getLogoutUrl = jest.mocked(workos.userManagement.getLogoutUrl);
 const authenticateWithRefreshToken = jest.mocked(workos.userManagement.authenticateWithRefreshToken);
+const getSessionStorage = jest.mocked(getSessionStorageMock);
+const configureSessionStorage = jest.mocked(configureSessionStorageMock);
 const jwtVerify = jest.mocked(jose.jwtVerify);
 
-jest.mock('./cookie', () => ({
-  getSession: jest.fn(),
-  destroySession: jest.fn().mockResolvedValue('destroyed-session-cookie'),
-  commitSession: jest.fn(),
+jest.mock('./sessionStorage.js', () => ({
+  configureSessionStorage: jest.fn(),
+  getSessionStorage: jest.fn(),
 }));
 
 jest.mock('./workos.js', () => ({
@@ -68,6 +70,30 @@ describe('session', () => {
       }),
     });
 
+  let getSession: jest.Mock;
+  let destroySession: jest.Mock;
+  let commitSession: jest.Mock;
+
+  beforeEach(async () => {
+    getSession = jest.fn();
+    destroySession = jest.fn().mockResolvedValue('destroyed-session-cookie');
+    commitSession = jest.fn();
+
+    getSessionStorage.mockResolvedValue({
+      cookieName: 'wos-cookie',
+      getSession,
+      destroySession,
+      commitSession,
+    });
+
+    configureSessionStorage.mockResolvedValue({
+      cookieName: 'wos-cookie',
+      getSession,
+      destroySession,
+      commitSession,
+    });
+  });
+
   describe('encryptSession', () => {
     it('should encrypt session data with correct parameters', async () => {
       const mockSession = {
@@ -371,7 +397,7 @@ describe('session', () => {
         };
         unsealData.mockResolvedValue(expiredSessionData);
         sealData.mockResolvedValue('new-encrypted-jwt');
-        (cookie.commitSession as jest.Mock).mockResolvedValue('new-session-cookie');
+        commitSession.mockResolvedValue('new-session-cookie');
 
         // Token verification fails
         jwtVerify.mockRejectedValue(new Error('Token expired'));

src/interfaces.ts

@@ -1,18 +1,19 @@
-import { json, redirect } from '@remix-run/node';
 import type { LoaderFunctionArgs, SessionData, TypedResponse } from '@remix-run/node';
-import { WORKOS_CLIENT_ID, WORKOS_COOKIE_PASSWORD } from './env-variables.js';
-import type { AccessToken, AuthorizedData, UnauthorizedData, AuthKitLoaderOptions, Session } from './interfaces.js';
-import { getSession, destroySession, commitSession } from './cookie.js';
+import { json, redirect } from '@remix-run/node';
+import { WORKOS_CLIENT_ID, WORKOS_COOKIE_NAME, WORKOS_COOKIE_PASSWORD } from './env-variables.js';
 import { getAuthorizationUrl } from './get-authorization-url.js';
+import type { AccessToken, AuthKitLoaderOptions, AuthorizedData, Session, UnauthorizedData } from './interfaces.js';
 import { workos } from './workos.js';
 
 import { sealData, unsealData } from 'iron-session';
-import { jwtVerify, createRemoteJWKSet, decodeJwt } from 'jose';
+import { createRemoteJWKSet, decodeJwt, jwtVerify } from 'jose';
+import { configureSessionStorage, getSessionStorage } from './sessionStorage.js';
 
 const JWKS = createRemoteJWKSet(new URL(workos.userManagement.getJwksUrl(WORKOS_CLIENT_ID)));
 
 async function updateSession(request: Request, debug: boolean) {
   const session = await getSessionFromCookie(request.headers.get('Cookie') as string);
+  const { commitSession, getSession, destroySession } = await getSessionStorage();
 
   // If no session, just continue
   if (!session) {
@@ -115,7 +116,15 @@ async function authkitLoader<Data = unknown>(
   options: AuthKitLoaderOptions = {},
 ) {
   const loader = typeof loaderOrOptions === 'function' ? loaderOrOptions : undefined;
-  const { ensureSignedIn = false, debug = false } = typeof loaderOrOptions === 'object' ? loaderOrOptions : options;
+  const {
+    ensureSignedIn = false,
+    debug = false,
+    storage,
+    cookie,
+  } = typeof loaderOrOptions === 'object' ? loaderOrOptions : options;
+
+  const cookieName = cookie?.name ?? WORKOS_COOKIE_NAME ?? 'wos-session';
+  const { getSession, destroySession } = await configureSessionStorage({ storage, cookieName });
 
   const { request } = loaderArgs;
   const session = await updateSession(request, debug);
@@ -215,6 +224,7 @@ async function handleAuthLoader(
 }
 
 async function terminateSession(request: Request) {
+  const { getSession, destroySession } = await getSessionStorage();
   const encryptedSession = await getSession(request.headers.get('Cookie'));
   const { accessToken } = (await getSessionFromCookie(
     request.headers.get('Cookie') as string,
@@ -257,6 +267,7 @@ function getClaimsFromAccessToken(accessToken: string) {
 }
 
 async function getSessionFromCookie(cookie: string, session?: SessionData) {
+  const { getSession } = await getSessionStorage();
   if (!session) {
     session = await getSession(cookie);
   }
@@ -286,4 +297,4 @@ function getReturnPathname(url: string): string {
   return `${newUrl.pathname}${newUrl.searchParams.size > 0 ? '?' + newUrl.searchParams.toString() : ''}`;
 }
 
-export { encryptSession, terminateSession, authkitLoader };
+export { authkitLoader, encryptSession, terminateSession };