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

package/src/authkit-callback-route.spec.ts

@@ -1,21 +1,29 @@
+import type { Mock } from 'vitest';
 import { getWorkOS } from './workos.js';
 import { handleAuth } from './authkit-callback-route.js';
 import { getSessionFromCookie, saveSession } from './session.js';
 import { NextRequest, NextResponse } from 'next/server';
+import { sealData } from 'iron-session';
 
-// Mocked in jest.setup.ts
+// Mocked in vitest.setup.ts
 import { cookies, headers } from 'next/headers';
+import { State } from './interfaces.js';
 
 // Mock dependencies
-const fakeWorkosInstance = {
-  userManagement: {
-    authenticateWithCode: jest.fn(),
-    getJwksUrl: jest.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
+const { fakeWorkosInstance } = vi.hoisted(() => ({
+  fakeWorkosInstance: {
+    userManagement: {
+      authenticateWithCode: vi.fn(),
+      getJwksUrl: vi.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
+    },
+    pkce: {
+      generate: vi.fn(),
+    },
   },
-};
+}));
 
-jest.mock('../src/workos', () => ({
-  getWorkOS: jest.fn(() => fakeWorkosInstance),
+vi.mock('../src/workos', () => ({
+  getWorkOS: vi.fn(() => fakeWorkosInstance),
 }));
 
 describe('authkit-callback-route', () => {
@@ -34,9 +42,9 @@ describe('authkit-callback-route', () => {
       createdAt: '2024-01-01T00:00:00Z',
       updatedAt: '2024-01-01T00:00:00Z',
       lastSignInAt: '2024-01-01T00:00:00Z',
-      locale: null,
       externalId: null,
       metadata: {},
+      locale: null,
     },
     oauthTokens: {
       accessToken: 'access123',
@@ -46,17 +54,23 @@ describe('authkit-callback-route', () => {
     },
   };
 
+  async function setAuthCookie(req: NextRequest, state: State): Promise<string> {
+    const sealedState = await sealData(state, { password: process.env.WORKOS_COOKIE_PASSWORD! });
+    req.cookies.set('wos-auth-verifier', sealedState);
+    return sealedState;
+  }
+
   describe('handleAuth', () => {
     let request: NextRequest;
 
     beforeAll(() => {
       // Silence console.error during tests
-      jest.spyOn(console, 'error').mockImplementation(() => {});
+      vi.spyOn(console, 'error').mockImplementation(() => {});
     });
 
     beforeEach(async () => {
       // Reset all mocks
-      jest.clearAllMocks();
+      vi.clearAllMocks();
 
       // Create a new request with searchParams
       request = new NextRequest(new URL('http://example.com/callback'));
@@ -72,10 +86,12 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle successful authentication', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Set up request with code
+      // Set up request with code & state
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -83,15 +99,17 @@ describe('authkit-callback-route', () => {
       expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith({
         clientId: process.env.WORKOS_CLIENT_ID,
         code: 'test-code',
+        codeVerifier: 'test-verifier',
       });
       expect(response).toBeInstanceOf(NextResponse);
     });
 
     it('should handle authentication failure', async () => {
-      // Mock authentication failure
-      (workos.userManagement.authenticateWithCode as jest.Mock).mockRejectedValue(new Error('Auth failed'));
+      (workos.userManagement.authenticateWithCode as Mock).mockRejectedValue(new Error('Auth failed'));
 
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'invalid-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -102,10 +120,11 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle authentication failure if a non-Error object is thrown', async () => {
-      // Mock authentication failure
-      jest.mocked(workos.userManagement.authenticateWithCode).mockRejectedValue('Auth failed');
+      vi.mocked(workos.userManagement.authenticateWithCode).mockRejectedValue('Auth failed');
 
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'invalid-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -116,9 +135,11 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle authentication failure with custom onError handler', async () => {
-      // Mock authentication failure
-      jest.mocked(workos.userManagement.authenticateWithCode).mockRejectedValue('Auth failed');
+      vi.mocked(workos.userManagement.authenticateWithCode).mockRejectedValue('Auth failed');
+
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'invalid-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth({
         onError: () => {
@@ -145,9 +166,11 @@ describe('authkit-callback-route', () => {
     });
 
     it('should respect custom returnPathname', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth({ returnPathname: '/dashboard' });
       const response = await handler(request);
@@ -156,11 +179,15 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle state parameter with returnPathname', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      const state = btoa(JSON.stringify({ returnPathname: '/custom-path' }));
+      const sealedState = await setAuthCookie(request, {
+        nonce: 'foo',
+        codeVerifier: 'test-verifier',
+        returnPathname: '/custom-path',
+      });
       request.nextUrl.searchParams.set('code', 'test-code');
-      request.nextUrl.searchParams.set('state', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -169,11 +196,15 @@ describe('authkit-callback-route', () => {
     });
 
     it('should extract custom search params from returnPathname', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      const state = btoa(JSON.stringify({ returnPathname: '/custom-path?foo=bar&baz=qux' }));
+      const sealedState = await setAuthCookie(request, {
+        nonce: 'foo',
+        codeVerifier: 'test-verifier',
+        returnPathname: '/custom-path?foo=bar&baz=qux',
+      });
       request.nextUrl.searchParams.set('code', 'test-code');
-      request.nextUrl.searchParams.set('state', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -181,14 +212,35 @@ describe('authkit-callback-route', () => {
       expect(response.headers.get('Location')).toContain('/custom-path?foo=bar&baz=qux');
     });
 
+    it('should handle full URL in returnPathname by extracting only the pathname', async () => {
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      const sealedState = await setAuthCookie(request, {
+        nonce: 'foo',
+        codeVerifier: 'test-verifier',
+        returnPathname: 'https://example.com/invite/k0123456789',
+      });
+
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
+
+      const handler = handleAuth();
+      const response = await handler(request);
+
+      const location = response.headers.get('Location');
+      expect(location).toContain('/invite/k0123456789');
+      expect(location).not.toContain('https://example.com/invite');
+    });
+
     it('should use Response if NextResponse.redirect is not available', async () => {
       const originalRedirect = NextResponse.redirect;
       (NextResponse as Partial<typeof NextResponse>).redirect = undefined;
 
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Set up request with code
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -203,6 +255,10 @@ describe('authkit-callback-route', () => {
       const originalJson = NextResponse.json;
       (NextResponse as Partial<typeof NextResponse>).json = undefined;
 
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
+
       const handler = handleAuth();
       const response = await handler(request);
 
@@ -217,10 +273,12 @@ describe('authkit-callback-route', () => {
     });
 
     it('should use baseURL if provided', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Set up request with code
+      // Set up request with code & state
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth({ baseURL: 'https://base.com' });
       const response = await handler(request);
@@ -229,14 +287,15 @@ describe('authkit-callback-route', () => {
     });
 
     it('should throw an error if response is missing tokens', async () => {
-      const mockAuthResponse = {
+      const incompleteAuthResponse = {
         user: { id: 'user_123' },
       };
 
-      (workos.userManagement.authenticateWithCode as jest.Mock).mockResolvedValue(mockAuthResponse);
+      (workos.userManagement.authenticateWithCode as Mock).mockResolvedValue(incompleteAuthResponse);
 
-      // Set up request with code
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -245,12 +304,14 @@ describe('authkit-callback-route', () => {
     });
 
     it('should call onSuccess if provided', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Set up request with code
+      // Set up request with code & state
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
-      const onSuccess = jest.fn();
+      const onSuccess = vi.fn();
       const handler = handleAuth({ onSuccess: onSuccess });
       await handler(request);
 
@@ -261,10 +322,11 @@ describe('authkit-callback-route', () => {
 
     it('should allow onSuccess to update session', async () => {
       const newAccessToken = 'new-access-token';
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Set up request with code
+      const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
       request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth({
         onSuccess: async (data) => {
@@ -278,19 +340,19 @@ describe('authkit-callback-route', () => {
     });
 
     it('should pass custom state data to onSuccess callback', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      // Create state with new format: internal.user
-      const internalState = btoa(JSON.stringify({ returnPathname: '/dashboard' }))
-        .replace(/\+/g, '-')
-        .replace(/\//g, '_');
-      const userState = 'custom-user-state-string';
-      const state = `${internalState}.${userState}`;
+      const sealedState = await setAuthCookie(request, {
+        nonce: 'foo',
+        codeVerifier: 'test-verifier',
+        returnPathname: '/dashboard',
+        customState: 'custom-user-state-string',
+      });
 
       request.nextUrl.searchParams.set('code', 'test-code');
-      request.nextUrl.searchParams.set('state', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
-      const onSuccess = jest.fn();
+      const onSuccess = vi.fn();
       const handler = handleAuth({ onSuccess });
       await handler(request);
 
@@ -308,15 +370,19 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle state without custom data', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
       // State with only returnPathname
-      const state = btoa(JSON.stringify({ returnPathname: '/profile' }));
+      const sealedState = await setAuthCookie(request, {
+        nonce: 'foo',
+        codeVerifier: 'test-verifier',
+        returnPathname: '/profile',
+      });
 
       request.nextUrl.searchParams.set('code', 'test-code');
-      request.nextUrl.searchParams.set('state', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
-      const onSuccess = jest.fn();
+      const onSuccess = vi.fn();
       const handler = handleAuth({ onSuccess });
       await handler(request);
 
@@ -329,20 +395,206 @@ describe('authkit-callback-route', () => {
       );
     });
 
-    it('should handle backward compatibility with old state format', async () => {
-      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+    it('should NOT handle backward compatibility with old state format', async () => {
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
       // Old format: just returnPathname
-      const state = btoa(JSON.stringify({ returnPathname: '/old-path' }));
-
+      // @ts-expect-error we're purposely testing backward compatibility with an old format that doesn't match the current State interface
+      const sealedState = await setAuthCookie(request, { returnPathname: '/old-path' });
       request.nextUrl.searchParams.set('code', 'test-code');
-      request.nextUrl.searchParams.set('state', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
 
-      // Should still redirect correctly
-      expect(response.headers.get('Location')).toContain('/old-path');
+      // Should error
+      expect(response.status).toBe(500);
+      expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+    });
+
+    it('should not leak nonce-only state as custom state in onSuccess', async () => {
+      vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      // Simulate a nonce-only state (no returnPathname, no custom state)
+      const nonceState = await setAuthCookie(request, { nonce: 'test-nonce', codeVerifier: 'test-verifier' });
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', nonceState);
+
+      const onSuccess = vi.fn();
+      const handler = handleAuth({ onSuccess });
+      await handler(request);
+
+      expect(onSuccess).toHaveBeenCalledWith(expect.objectContaining({ state: undefined }));
+    });
+
+    describe('state verification', () => {
+      it('should reject callback when state does not match stored state', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        const state = 'attacker-state';
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', state);
+        await setAuthCookie(request, { nonce: 'legitimate-state', codeVerifier: 'test-verifier' });
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(response.status).toBe(500);
+        expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+      });
+
+      it('should reject when state is present but no cookie exists', async () => {
+        const sealedState = await sealData(
+          { nonce: 'foo', codeVerifier: 'test-verifier' },
+          { password: process.env.WORKOS_COOKIE_PASSWORD! },
+        );
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+        expect(response.status).toBe(500);
+      });
+
+      it('should pass when state matches stored state', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier' });
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalled();
+        expect(response.status).not.toBe(500);
+      });
+
+      it('should return 500 when neither state nor cookie exist', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        request.nextUrl.searchParams.set('code', 'test-code');
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+        expect(response.status).toBe(500);
+      });
+    });
+
+    describe('PKCE', () => {
+      it('should pass codeVerifier and verify state when both are in the cookie', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        const sealedState = await setAuthCookie(request, {
+          codeVerifier: 'test-verifier-456',
+          returnPathname: '/dashboard',
+          nonce: 'foo',
+        });
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
+          expect.objectContaining({
+            code: 'test-code',
+            codeVerifier: 'test-verifier-456',
+          }),
+        );
+        expect(response.headers.get('Location')).toContain('/dashboard');
+      });
+
+      it('should pass codeVerifier from cookie to authenticateWithCode', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+        const sealedState = await setAuthCookie(request, {
+          nonce: 'foo',
+          codeVerifier: 'test-verifier-123',
+        });
+
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
+          expect.objectContaining({
+            code: 'test-code',
+            codeVerifier: 'test-verifier-123',
+          }),
+        );
+      });
+
+      it('should reject when cookie is missing even if state contains valid sealed data', async () => {
+        const sealedState = await sealData(
+          { nonce: 'foo', codeVerifier: 'test-verifier-123' },
+          { password: process.env.WORKOS_COOKIE_PASSWORD! },
+        );
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+        expect(response.status).toBe(500);
+      });
+
+      it('should return an error response when PKCE cookie is corrupted', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        // Set a corrupted cookie
+        request.cookies.set('wos-auth-verifier', 'not-a-valid-sealed-value');
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', 'not-a-valid-sealed-value');
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(response.status).toBe(500);
+        expect(workos.userManagement.authenticateWithCode).not.toHaveBeenCalled();
+      });
+
+      it('should delete PKCE cookie after successful authentication', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier-123' });
+        request.nextUrl.searchParams.set('code', 'test-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        // The response should be a redirect (success) and have a Set-Cookie header to delete the PKCE cookie
+        expect(response.status).toBe(307);
+
+        const setCookieHeaders = response.headers.getSetCookie();
+        const pkceDeletionCookie = setCookieHeaders.find((c: string) => c.startsWith('wos-auth-verifier='));
+        expect(pkceDeletionCookie).toBeDefined();
+        expect(pkceDeletionCookie).toContain('Max-Age=0');
+      });
+
+      it('should delete PKCE cookie after failed authentication', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockRejectedValue(new Error('Auth failed'));
+
+        const sealedState = await setAuthCookie(request, { nonce: 'foo', codeVerifier: 'test-verifier-123' });
+        request.nextUrl.searchParams.set('code', 'bad-code');
+        request.nextUrl.searchParams.set('state', sealedState);
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        expect(response.status).toBe(500);
+        const setCookieHeaders = response.headers.getSetCookie();
+        const pkceDeletionCookie = setCookieHeaders.find((c: string) => c.startsWith('wos-auth-verifier='));
+        expect(pkceDeletionCookie).toBeDefined();
+        expect(pkceDeletionCookie).toContain('Max-Age=0');
+      });
     });
   });
 });