@workos-inc/authkit-nextjs 2.17.0 → 3.0.0

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

@@ -1,3 +1,4 @@
+import type { Mock } from 'vitest';
 import { getWorkOS } from './workos.js';
 import { handleAuth } from './authkit-callback-route.js';
 import { getSessionFromCookie, saveSession } from './session.js';
@@ -6,6 +7,7 @@ import { sealData } from 'iron-session';
 
 // Mocked in vitest.setup.ts
 import { cookies, headers } from 'next/headers';
+import { State } from './interfaces.js';
 
 // Mock dependencies
 const { fakeWorkosInstance } = vi.hoisted(() => ({
@@ -52,6 +54,12 @@ 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;
 
@@ -80,8 +88,10 @@ describe('authkit-callback-route', () => {
     it('should handle successful authentication', async () => {
       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);
@@ -89,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 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);
@@ -108,10 +120,11 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle authentication failure if a non-Error object is thrown', async () => {
-      // Mock authentication failure
       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);
@@ -122,9 +135,11 @@ describe('authkit-callback-route', () => {
     });
 
     it('should handle authentication failure with custom onError handler', async () => {
-      // Mock authentication failure
       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: () => {
@@ -153,7 +168,9 @@ describe('authkit-callback-route', () => {
     it('should respect custom returnPathname', 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({ returnPathname: '/dashboard' });
       const response = await handler(request);
@@ -164,9 +181,13 @@ describe('authkit-callback-route', () => {
     it('should handle state parameter with returnPathname', async () => {
       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);
@@ -177,9 +198,13 @@ describe('authkit-callback-route', () => {
     it('should extract custom search params from returnPathname', async () => {
       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);
@@ -190,9 +215,14 @@ describe('authkit-callback-route', () => {
     it('should handle full URL in returnPathname by extracting only the pathname', async () => {
       vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-      const state = btoa(JSON.stringify({ returnPathname: 'https://example.com/invite/k0123456789' }));
+      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', state);
+      request.nextUrl.searchParams.set('state', sealedState);
 
       const handler = handleAuth();
       const response = await handler(request);
@@ -208,8 +238,9 @@ describe('authkit-callback-route', () => {
 
       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);
@@ -224,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);
 
@@ -240,8 +275,10 @@ describe('authkit-callback-route', () => {
     it('should use baseURL if provided', async () => {
       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);
@@ -250,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 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);
@@ -268,8 +306,10 @@ describe('authkit-callback-route', () => {
     it('should call onSuccess if provided', async () => {
       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 = vi.fn();
       const handler = handleAuth({ onSuccess: onSuccess });
@@ -284,8 +324,9 @@ describe('authkit-callback-route', () => {
       const newAccessToken = 'new-access-token';
       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) => {
@@ -301,15 +342,15 @@ describe('authkit-callback-route', () => {
     it('should pass custom state data to onSuccess callback', async () => {
       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 = vi.fn();
       const handler = handleAuth({ onSuccess });
@@ -332,10 +373,14 @@ describe('authkit-callback-route', () => {
       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 = vi.fn();
       const handler = handleAuth({ onSuccess });
@@ -350,51 +395,129 @@ describe('authkit-callback-route', () => {
       );
     });
 
-    it('should handle backward compatibility with old state format', async () => {
+    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('PKCE', () => {
-      it('should pass codeVerifier from cookie to authenticateWithCode', async () => {
+    describe('state verification', () => {
+      it('should reject callback when state does not match stored state', async () => {
         vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
 
-        // Seal a verifier into a cookie value
-        const sealedVerifier = await sealData(
-          { codeVerifier: 'test-verifier-123' },
+        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);
 
-        // Set the PKCE cookie on the request
-        request.cookies.set('wos-pkce-verifier', sealedVerifier);
+        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();
-        await handler(request);
+        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-123',
+            codeVerifier: 'test-verifier-456',
           }),
         );
+        expect(response.headers.get('Location')).toContain('/dashboard');
       });
 
-      it('should proceed without codeVerifier when PKCE cookie is missing', async () => {
+      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);
@@ -402,46 +525,73 @@ describe('authkit-callback-route', () => {
         expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
           expect.objectContaining({
             code: 'test-code',
-            codeVerifier: undefined,
+            codeVerifier: 'test-verifier-123',
           }),
         );
       });
 
-      it('should proceed without codeVerifier when PKCE cookie is corrupted', async () => {
+      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-pkce-verifier', 'not-a-valid-sealed-value');
+        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();
-        await handler(request);
+        const response = await handler(request);
 
-        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
-          expect.objectContaining({
-            code: 'test-code',
-            codeVerifier: undefined,
-          }),
-        );
+        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 sealedVerifier = await sealData(
-          { codeVerifier: 'test-verifier-123' },
-          { password: process.env.WORKOS_COOKIE_PASSWORD! },
-        );
-
-        request.cookies.set('wos-pkce-verifier', sealedVerifier);
+        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);
 
-        // The response should have a Set-Cookie header to delete the PKCE cookie
+        expect(response.status).toBe(500);
         const setCookieHeaders = response.headers.getSetCookie();
-        const pkceDeletionCookie = setCookieHeaders.find((c: string) => c.startsWith('wos-pkce-verifier='));
+        const pkceDeletionCookie = setCookieHeaders.find((c: string) => c.startsWith('wos-auth-verifier='));
         expect(pkceDeletionCookie).toBeDefined();
         expect(pkceDeletionCookie).toContain('Max-Age=0');
       });