diff --git a/src/index.ts b/src/index.ts index 7d038c1..2a086c7 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,7 +1,7 @@ import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization, withAuth } from './auth.js'; import { authLoader } from './authkit-callback-route.js'; import { configure, getConfig } from './config.js'; -import { authkitLoader, refreshSession } from './session.js'; +import { authkitLoader, refreshSession, saveSession } from './session.js'; import { getWorkOS } from './workos.js'; export { @@ -14,6 +14,7 @@ export { getSignUpUrl, getWorkOS, refreshSession, + saveSession, signOut, switchToOrganization, }; diff --git a/src/session.spec.ts b/src/session.spec.ts index 67df068..1670c61 100644 --- a/src/session.spec.ts +++ b/src/session.spec.ts @@ -7,7 +7,7 @@ import { getSessionStorage as getSessionStorageMock, } from './sessionStorage.js'; import { Session } from './interfaces.js'; -import { authkitLoader, encryptSession, terminateSession, refreshSession } from './session.js'; +import { authkitLoader, encryptSession, terminateSession, refreshSession, saveSession } from './session.js'; import { assertIsResponse } from './test-utils/test-helpers.js'; import { getWorkOS } from './workos.js'; import { getConfig } from './config.js'; @@ -818,4 +818,96 @@ describe('session', () => { ); }); }); + + describe('saveSession', () => { + const sessionData = { + accessToken: 'new.valid.token', + refreshToken: 'new.refresh.token', + user: { + object: 'user', + id: 'user-1', + email: 'test@example.com', + emailVerified: true, + profilePictureUrl: null, + firstName: 'Test', + lastName: 'User', + lastSignInAt: '2021-01-01T00:00:00Z', + createdAt: '2021-01-01T00:00:00Z', + updatedAt: '2021-01-01T00:00:00Z', + externalId: null, + locale: null, + metadata: {}, + } satisfies User, + impersonator: undefined, + headers: {}, + } satisfies Session; + + const createMockRequest = (cookie = 'test-cookie', url = 'http://example.com./some-path') => + new Request(url, { + headers: new Headers({ + Cookie: cookie, + }), + }); + + let getSession: jest.Mock; + let destroySession: jest.Mock; + let commitSession: jest.Mock; + let mockSession: ReactRouterSession; + + beforeEach(() => { + getSession = jest.fn(); + destroySession = jest.fn().mockResolvedValue('destroyed-session-cookie'); + commitSession = jest.fn().mockResolvedValue('new-session-cookie'); + + mockSession = createMockSession({ + has: jest.fn().mockReturnValue(true), + get: jest.fn().mockReturnValue('encrypted-jwt'), + set: jest.fn(), + }); + + getSessionStorage.mockResolvedValue({ + cookieName: 'wos-cookie', + getSession, + destroySession, + commitSession, + }); + + getSession.mockResolvedValue(mockSession); + + const validSessionData = { + accessToken: 'valid.token', + refreshToken: 'refresh.token', + user: { + id: 'user-1', + email: 'test@example.com', + firstName: 'Test', + lastName: 'User', + object: 'user', + }, + impersonator: null, + }; + unsealData.mockResolvedValue(validSessionData); + sealData.mockResolvedValue('new-encrypted-jwt'); + + authenticateWithRefreshToken.mockResolvedValue(sessionData); + + // Mock JWT decoding + (jose.decodeJwt as jest.Mock).mockReturnValue({ + sid: 'new-session-id', + org_id: 'org-123', + role: 'user', + roles: ['user'], + permissions: ['read'], + entitlements: ['basic'], + feature_flags: ['flag-1'], + }); + }); + + it('should save the session to the cookie', async () => { + await saveSession(sessionData, createMockRequest()); + expect(getSessionStorage).toHaveBeenCalled(); + expect(commitSession).toHaveBeenCalledWith(mockSession); + expect(mockSession.set).toHaveBeenCalledWith('jwt', 'new-encrypted-jwt'); + }); + }); }); diff --git a/src/session.ts b/src/session.ts index 4f3e73a..ae34023 100644 --- a/src/session.ts +++ b/src/session.ts @@ -16,6 +16,7 @@ import { createRemoteJWKSet, decodeJwt, jwtVerify } from 'jose'; import { getConfig } from './config.js'; import { configureSessionStorage, getSessionStorage } from './sessionStorage.js'; import { isDataWithResponseInit, isJsonResponse, isRedirect, isResponse } from './utils.js'; +import type { AuthenticationResponse } from '@workos-inc/node'; // must be a type since this is a subtype of response // interfaces must conform to the types they extend @@ -37,37 +38,24 @@ export class SessionRefreshError extends Error { * @param options - Optional configuration options * @returns A promise that resolves to the new session object */ -export async function refreshSession(request: Request, { organizationId }: { organizationId?: string } = {}) { - const { getSession, commitSession } = await getSessionStorage(); - const session = await getSessionFromCookie(request.headers.get('Cookie') as string); - +export async function refreshSession(request: Request, options: { organizationId?: string } = {}) { + const { organizationId } = options; + const { getSession } = await getSessionStorage(); + const cookie = request.headers.get('Cookie'); + const session = cookie ? await getSessionFromCookie(cookie) : null; if (!session) { throw redirect(await getAuthorizationUrl()); } try { - const { accessToken, refreshToken, user, impersonator } = - await getWorkOS().userManagement.authenticateWithRefreshToken({ - clientId: getConfig('clientId'), - refreshToken: session.refreshToken, - organizationId, - }); - - const newSession = { - accessToken, - refreshToken, - user, - impersonator, - headers: {} as Record, - }; - - const cookieSession = await getSession(request.headers.get('Cookie')); - cookieSession.set('jwt', await encryptSession(newSession)); - const cookie = await commitSession(cookieSession); - - newSession.headers = { - 'Set-Cookie': cookie, - }; + const refreshResult = await getWorkOS().userManagement.authenticateWithRefreshToken({ + clientId: getConfig('clientId'), + refreshToken: session.refreshToken, + organizationId, + }); + const { headers } = await saveSession(refreshResult, request); + const cookieSession = await getSession(cookie); + const { accessToken, user, impersonator } = refreshResult; const { sessionId, @@ -91,7 +79,7 @@ export async function refreshSession(request: Request, { organizationId }: { org featureFlags, impersonator: impersonator ?? null, sealedSession: cookieSession.get('jwt'), - headers: newSession.headers, + headers, }; } catch (error) { throw new Error(`Failed to refresh session: ${error instanceof Error ? error.message : String(error)}`, { @@ -100,7 +88,54 @@ export async function refreshSession(request: Request, { organizationId }: { org } } -async function updateSession(request: Request, debug: boolean) { +/** + * Saves a WorkOS session to a cookie for use with AuthKit. + * + * This function is intended for advanced use cases where you need to manually + * manage sessions, such as custom authentication flows (email verification, + * etc.) that don't use the standard AuthKit authentication flow. + * + * @param sessionOrResponse The WorkOS session or AuthenticationResponse + * containing access token, refresh token, and user information. + * @param request A Request object, used to determine cookie settings. + * + * @example + * import { saveSession } from '@workos-inc/authkit-react-router'; + * + * async function handleEmailVerification(req: Request) { + * const { code } = await req.json(); + * const authResponse = await workos.userManagement.authenticateWithEmailVerification({ + * clientId: process.env.WORKOS_CLIENT_ID, + * code, + * }); + * + * await saveSession(authResponse, req); + * } + */ +export async function saveSession( + sessionOrResponse: Session | AuthenticationResponse, + request: Request, +): Promise { + const { getSession, commitSession } = await getSessionStorage(); + const { accessToken, refreshToken, user, impersonator } = sessionOrResponse; + const newSession: Session = { + accessToken, + refreshToken, + user, + impersonator, + headers: {}, + }; + const cookieSession = await getSession(request.headers.get('Cookie')); + cookieSession.set('jwt', await encryptSession(newSession)); + const cookie = await commitSession(cookieSession); + newSession.headers = { + 'Set-Cookie': cookie, + }; + + return newSession; +} + +async function updateSession(request: Request, debug: boolean): Promise { const session = await getSessionFromCookie(request.headers.get('Cookie') as string); const { commitSession, getSession } = await getSessionStorage(); @@ -158,7 +193,7 @@ async function updateSession(request: Request, debug: boolean) { } } -export async function encryptSession(session: Session) { +export async function encryptSession(session: Session | AuthenticationResponse) { return sealData(session, { password: getConfig('cookiePassword'), ttl: 0,