@workos-inc/authkit-nextjs 3.0.0-beta.1 → 3.0.1

package/src/session.spec.ts

@@ -1,3 +1,4 @@
+import type { Mock, MockInstance } from 'vitest';
 import { NextRequest, NextResponse } from 'next/server';
 import { cookies, headers } from 'next/headers';
 import { redirect } from 'next/navigation';
@@ -7,15 +8,24 @@ import { getWorkOS } from './workos.js';
 import * as envVariables from './env-variables.js';
 
 import { jwtVerify } from 'jose';
+
+// Helper to override env variable exports without triggering no-import-assign on the import binding
+function setEnvVar(mod: Record<string, unknown>, key: string, value: unknown) {
+  Object.defineProperty(mod, key, { value, configurable: true });
+}
 import { sealData } from 'iron-session';
 import { User } from '@workos-inc/node';
-
-jest.mock('jose', () => ({
-  jwtVerify: jest.fn(),
-  createRemoteJWKSet: jest.fn(),
-  SignJWT: jest.requireActual('jose').SignJWT,
-  decodeJwt: jest.requireActual('jose').decodeJwt,
-}));
+import { getStateFromPKCECookieValue } from './pkce.js';
+
+vi.mock('jose', async () => {
+  const actual = await vi.importActual<typeof import('jose')>('jose');
+  return {
+    jwtVerify: vi.fn(),
+    createRemoteJWKSet: vi.fn(),
+    SignJWT: actual.SignJWT,
+    decodeJwt: actual.decodeJwt,
+  };
+});
 
 // logging is disabled by default, flip this to true to still have logs in the console
 const DEBUG = false;
@@ -42,20 +52,16 @@ describe('session.ts', () => {
       profilePictureUrl: null,
       firstName: null,
       lastName: null,
-      lastSignInAt: null,
-      locale: null,
-      externalId: null,
-      metadata: {},
       createdAt: '2024-01-01',
       updatedAt: '2024-01-01',
     } as User,
   };
 
-  let consoleLogSpy: jest.SpyInstance;
+  let consoleLogSpy: MockInstance;
 
   beforeEach(async () => {
     // Clear all mocks between tests
-    jest.clearAllMocks();
+    vi.clearAllMocks();
 
     // Reset the cookie store
     const nextCookies = await cookies();
@@ -67,9 +73,9 @@ describe('session.ts', () => {
     nextHeaders._reset();
     nextHeaders.set('x-workos-middleware', 'true');
 
-    (jwtVerify as jest.Mock).mockReset();
+    (jwtVerify as Mock).mockReset();
 
-    consoleLogSpy = jest.spyOn(console, 'log').mockImplementation((...args) => {
+    consoleLogSpy = vi.spyOn(console, 'log').mockImplementation((...args) => {
       if (DEBUG) {
         console.info(...args);
       }
@@ -78,7 +84,7 @@ describe('session.ts', () => {
 
   afterEach(() => {
     consoleLogSpy.mockRestore();
-    jest.resetModules();
+    vi.resetModules();
   });
 
   describe('withAuth', () => {
@@ -120,7 +126,7 @@ describe('session.ts', () => {
       await expect(async () => {
         await withAuth();
       }).rejects.toThrow(
-        "You are calling 'withAuth' on https://example.com/ that isn’t covered by the AuthKit middleware. Make sure it is running on all paths you are calling 'withAuth' from by updating your middleware config in 'middleware.(js|ts)'.",
+        /You are calling 'withAuth' on https:\/\/example\.com\/ that isn't covered by the AuthKit middleware/,
       );
     });
 
@@ -130,9 +136,7 @@ describe('session.ts', () => {
 
       await expect(async () => {
         await withAuth({ ensureSignedIn: true });
-      }).rejects.toThrow(
-        "You are calling 'withAuth' on a route that isn’t covered by the AuthKit middleware. Make sure it is running on all paths you are calling 'withAuth' from by updating your middleware config in 'middleware.(js|ts)'.",
-      );
+      }).rejects.toThrow(/You are calling 'withAuth' on a route that isn't covered by the AuthKit middleware/);
     });
 
     it('should throw an error if the URL is not found in the headers', async () => {
@@ -150,14 +154,12 @@ describe('session.ts', () => {
 
       await withAuth({ ensureSignedIn: true });
 
-      // URL-safe base64 encoding
-      const pathname = encodeURIComponent(
-        btoa(JSON.stringify({ returnPathname: '/protected?test=123' }))
-          .replace(/\+/g, '-')
-          .replace(/\//g, '_'),
-      );
+      // The state is now sealed, se we need to unseal it
+      const redirectUrl = new URL((redirect as unknown as Mock).mock.calls[0][0]);
+      const sealedState = redirectUrl.searchParams.get('state')!;
+      const { returnPathname } = await getStateFromPKCECookieValue(sealedState);
 
-      expect(redirect).toHaveBeenCalledWith(expect.stringContaining(pathname));
+      expect(returnPathname).toBe('/protected?test=123');
     });
   });
 
@@ -165,7 +167,7 @@ describe('session.ts', () => {
     it('should throw an error if the redirect URI is not set', async () => {
       const originalWorkosRedirectUri = envVariables.WORKOS_REDIRECT_URI;
 
-      jest.replaceProperty(envVariables, 'WORKOS_REDIRECT_URI', '');
+      setEnvVar(envVariables, 'WORKOS_REDIRECT_URI', '');
 
       await expect(async () => {
         await updateSessionMiddleware(
@@ -180,13 +182,13 @@ describe('session.ts', () => {
         );
       }).rejects.toThrow('You must provide a redirect URI in the AuthKit middleware or in the environment variables.');
 
-      jest.replaceProperty(envVariables, 'WORKOS_REDIRECT_URI', originalWorkosRedirectUri);
+      setEnvVar(envVariables, 'WORKOS_REDIRECT_URI', originalWorkosRedirectUri);
     });
 
     it('should throw an error if the cookie password is not set', async () => {
       const originalWorkosCookiePassword = envVariables.WORKOS_COOKIE_PASSWORD;
 
-      jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', '');
+      setEnvVar(envVariables, 'WORKOS_COOKIE_PASSWORD', '');
 
       await expect(async () => {
         await updateSessionMiddleware(
@@ -203,13 +205,13 @@ describe('session.ts', () => {
         'You must provide a valid cookie password that is at least 32 characters in the environment variables.',
       );
 
-      jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', originalWorkosCookiePassword);
+      setEnvVar(envVariables, 'WORKOS_COOKIE_PASSWORD', originalWorkosCookiePassword);
     });
 
     it('should throw an error if the cookie password is less than 32 characters', async () => {
       const originalWorkosCookiePassword = envVariables.WORKOS_COOKIE_PASSWORD;
 
-      jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', 'short');
+      setEnvVar(envVariables, 'WORKOS_COOKIE_PASSWORD', 'short');
 
       await expect(async () => {
         await updateSessionMiddleware(
@@ -226,7 +228,7 @@ describe('session.ts', () => {
         'You must provide a valid cookie password that is at least 32 characters in the environment variables.',
       );
 
-      jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', originalWorkosCookiePassword);
+      setEnvVar(envVariables, 'WORKOS_COOKIE_PASSWORD', originalWorkosCookiePassword);
     });
 
     it('should return early if there is no session', async () => {
@@ -247,7 +249,7 @@ describe('session.ts', () => {
     });
 
     it('should return 200 if the session is valid', async () => {
-      jest.spyOn(console, 'log').mockImplementation(() => {});
+      vi.spyOn(console, 'log').mockImplementation(() => {});
 
       const nextCookies = await cookies();
       nextCookies.set(
@@ -255,7 +257,7 @@ describe('session.ts', () => {
         await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
       );
 
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         return true;
       });
 
@@ -278,11 +280,11 @@ describe('session.ts', () => {
     it('should attempt to refresh the session when the access token is invalid', async () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
         accessToken: await generateTestToken(),
         refreshToken: 'new-refresh-token',
         user: mockSession.user,
@@ -314,17 +316,15 @@ describe('session.ts', () => {
     });
 
     it('should delete the cookie when refreshing fails', async () => {
-      jest.spyOn(console, 'log').mockImplementation(() => {});
+      vi.spyOn(console, 'log').mockImplementation(() => {});
 
       mockSession.accessToken = await generateTestToken({}, true);
 
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
-      jest
-        .spyOn(workos.userManagement, 'authenticateWithRefreshToken')
-        .mockRejectedValue(new Error('Failed to refresh'));
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('Failed to refresh'));
 
       const request = new NextRequest(new URL('http://example.com'));
 
@@ -360,7 +360,7 @@ describe('session.ts', () => {
 
     describe('middleware auth', () => {
       it('should redirect unauthenticated users on protected routes', async () => {
-        jest.spyOn(console, 'log').mockImplementation(() => {});
+        vi.spyOn(console, 'log').mockImplementation(() => {});
 
         const request = new NextRequest(new URL('http://example.com/protected'));
         const result = await updateSessionMiddleware(
@@ -380,10 +380,7 @@ describe('session.ts', () => {
         );
       });
 
-      it('should use Response if NextResponse.redirect is not available', async () => {
-        const originalRedirect = NextResponse.redirect;
-        (NextResponse as Partial<typeof NextResponse>).redirect = undefined;
-
+      it('should return a redirect response when middlewareAuth is enabled and user is not authenticated', async () => {
         const request = new NextRequest(new URL('http://example.com/protected'));
         const result = await updateSessionMiddleware(
           request,
@@ -396,10 +393,9 @@ describe('session.ts', () => {
           [],
         );
 
-        expect(result).toBeInstanceOf(Response);
-
-        // Restore the original redirect method
-        (NextResponse as Partial<typeof NextResponse>).redirect = originalRedirect;
+        expect(result).toBeInstanceOf(NextResponse);
+        expect(result.status).toBe(307);
+        expect(result.headers.get('Location')).toContain('workos.com');
       });
 
       it('should automatically add the redirect URI to unauthenticatedPaths when middleware is enabled', async () => {
@@ -435,7 +431,7 @@ describe('session.ts', () => {
         expect(result.headers.get('Location')).toContain('screen_hint=sign-up');
       });
 
-      it('should set the sign up paths in the headers', async () => {
+      it('should not leak sign-up paths header to the browser', async () => {
         const request = new NextRequest(new URL('http://example.com/protected-signup'));
         const result = await updateSessionMiddleware(
           request,
@@ -448,7 +444,8 @@ describe('session.ts', () => {
           ['/protected-signup'],
         );
 
-        expect(result.headers.get('x-sign-up-paths')).toBe('/protected-signup');
+        // x-sign-up-paths is an internal header that should not leak to the browser
+        expect(result.headers.get('x-sign-up-paths')).toBeNull();
       });
 
       it('should allow logged out users on unauthenticated paths', async () => {
@@ -485,11 +482,11 @@ describe('session.ts', () => {
 
       it('should throw an error if the provided regex is invalid and a non-Error object is thrown', async () => {
         // Reset modules to ensure clean import state
-        jest.resetModules();
+        vi.resetModules();
 
         // Import first, then spy
         const pathToRegexp = await import('path-to-regexp');
-        const parseSpy = jest.spyOn(pathToRegexp, 'parse').mockImplementation(() => {
+        const parseSpy = vi.spyOn(pathToRegexp, 'parse').mockImplementation(() => {
           throw 'invalid regex';
         });
 
@@ -513,6 +510,9 @@ describe('session.ts', () => {
 
         // Verify the mock was called
         expect(parseSpy).toHaveBeenCalled();
+
+        // Restore the spy to prevent leaking to subsequent tests
+        parseSpy.mockRestore();
       });
 
       it('should default to the WORKOS_REDIRECT_URI environment variable if no redirect URI is provided', async () => {
@@ -532,17 +532,17 @@ describe('session.ts', () => {
       });
 
       it('should delete the cookie and redirect when refreshing fails', async () => {
-        jest.spyOn(console, 'log').mockImplementation(() => {});
+        vi.spyOn(console, 'log').mockImplementation(() => {});
 
         mockSession.accessToken = await generateTestToken({}, true);
 
-        (jwtVerify as jest.Mock).mockImplementation(() => {
+        (jwtVerify as Mock).mockImplementation(() => {
           throw new Error('Invalid token');
         });
 
-        jest
-          .spyOn(workos.userManagement, 'authenticateWithRefreshToken')
-          .mockRejectedValue(new Error('Failed to refresh'));
+        vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(
+          new Error('Failed to refresh'),
+        );
 
         const request = new NextRequest(new URL('http://example.com'));
 
@@ -643,12 +643,12 @@ describe('session.ts', () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
       // Mock token verification to fail
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
       // Mock successful refresh
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
         accessToken: await generateTestToken(),
         refreshToken: 'new-refresh-token',
         user: mockSession.user,
@@ -676,12 +676,12 @@ describe('session.ts', () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
       // Mock token verification to fail
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
       // Mock refresh failure
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('Refresh failed'));
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('Refresh failed'));
 
       const request = new NextRequest(new URL('http://example.com/protected'));
       request.cookies.set(
@@ -703,15 +703,15 @@ describe('session.ts', () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
       // Mock token verification to fail
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
       const newAccessToken = await generateTestToken();
-      const mockSuccessCallback = jest.fn();
+      const mockSuccessCallback = vi.fn();
 
       // Mock successful refresh
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
         accessToken: newAccessToken,
         refreshToken: 'new-refresh-token',
         user: mockSession.user,
@@ -742,15 +742,15 @@ describe('session.ts', () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
       // Mock token verification to fail
-      (jwtVerify as jest.Mock).mockImplementation(() => {
+      (jwtVerify as Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
 
       const mockError = new Error('Refresh failed');
-      const mockErrorCallback = jest.fn();
+      const mockErrorCallback = vi.fn();
 
       // Mock refresh failure
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(mockError);
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(mockError);
 
       const request = new NextRequest(new URL('http://example.com/protected'));
       request.cookies.set(
@@ -773,15 +773,15 @@ describe('session.ts', () => {
 
   describe('refreshSession', () => {
     it('should refresh session successfully', async () => {
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
         accessToken: await generateTestToken(),
         refreshToken: 'new-refresh-token',
         user: mockSession.user,
       });
 
-      jest
-        .spyOn(workos.userManagement, 'getJwksUrl')
-        .mockReturnValue('https://api.workos.com/sso/jwks/client_1234567890');
+      vi.spyOn(workos.userManagement, 'getJwksUrl').mockReturnValue(
+        'https://api.workos.com/sso/jwks/client_1234567890',
+      );
 
       const nextCookies = await cookies();
       nextCookies.set(
@@ -817,15 +817,15 @@ describe('session.ts', () => {
         await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
       );
 
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
         accessToken: await generateTestToken({ org_id: 'org_456' }),
         refreshToken: 'new-refresh-token',
         user: mockSession.user,
       });
 
-      jest
-        .spyOn(workos.userManagement, 'getJwksUrl')
-        .mockReturnValue('https://api.workos.com/sso/jwks/client_1234567890');
+      vi.spyOn(workos.userManagement, 'getJwksUrl').mockReturnValue(
+        'https://api.workos.com/sso/jwks/client_1234567890',
+      );
 
       const result = await refreshSession({ organizationId: 'org_456' });
 
@@ -844,8 +844,8 @@ describe('session.ts', () => {
         'wos-session',
         await sealData(mockSessionWithValidJWT, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
       );
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue('fail');
-      expect(refreshSession({ ensureSignedIn: false })).rejects.toThrow('Failed to refresh session: fail');
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue('fail');
+      await expect(refreshSession({ ensureSignedIn: false })).rejects.toThrow('Failed to refresh session: fail');
     });
 
     it('throws if authenticateWithRefreshToken fails with error', async () => {
@@ -859,7 +859,7 @@ describe('session.ts', () => {
         'wos-session',
         await sealData(mockSessionWithValidJWT, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
       );
-      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('error'));
+      vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('error'));
       await expect(refreshSession()).rejects.toThrow('Failed to refresh session: error');
     });
   });
@@ -869,7 +869,7 @@ describe('session.ts', () => {
       const nextCookies = await cookies();
       // @ts-expect-error - _reset is part of the mock
       nextCookies._reset();
-      jest.clearAllMocks();
+      vi.clearAllMocks();
     });
 
     it('should return all token claims when accessToken is provided', async () => {
@@ -963,7 +963,7 @@ describe('session.ts', () => {
 
   describe('eager auth functionality', () => {
     beforeEach(() => {
-      jest.clearAllMocks();
+      vi.clearAllMocks();
     });
 
     describe('isInitialDocumentRequest', () => {
@@ -1057,12 +1057,12 @@ describe('session.ts', () => {
         // Setup invalid session that needs refresh
         mockSession.accessToken = await generateTestToken({}, true);
 
-        (jwtVerify as jest.Mock).mockImplementation(() => {
+        (jwtVerify as Mock).mockImplementation(() => {
           throw new Error('Invalid token');
         });
 
         const newAccessToken = await generateTestToken();
-        jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+        vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
           accessToken: newAccessToken,
           refreshToken: 'new-refresh-token',
           user: mockSession.user,
@@ -1088,13 +1088,11 @@ describe('session.ts', () => {
         // Setup invalid session
         mockSession.accessToken = await generateTestToken({}, true);
 
-        (jwtVerify as jest.Mock).mockImplementation(() => {
+        (jwtVerify as Mock).mockImplementation(() => {
           throw new Error('Invalid token');
         });
 
-        jest
-          .spyOn(workos.userManagement, 'authenticateWithRefreshToken')
-          .mockRejectedValue(new Error('Refresh failed'));
+        vi.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('Refresh failed'));
 
         const request = new NextRequest(new URL('http://example.com/page'));
         request.headers.set('accept', 'text/html');