Fix cookie bugs (#171)

* Fix cookie bugs * Address feedback
workos · Jan 14, 2025 · ce5f0a3 · ce5f0a3
1 parent 5e9ac80
commit ce5f0a3
Show file tree

Hide file tree

Showing 8 changed files with 228 additions and 129 deletions.
diff --git a/__tests__/cookie.spec.ts b/__tests__/cookie.spec.ts
@@ -46,6 +46,40 @@ describe('cookie.ts', () => {
           domain: 'foobar.com',
         }),
       );
+
+      Object.defineProperty(envVars, 'WORKOS_COOKIE_DOMAIN', { value: '' });
+
+      const options2 = getCookieOptions('http://example.com');
+      expect(options2).toEqual(
+        expect.objectContaining({
+          secure: false,
+          maxAge: 1000,
+          domain: '',
+        }),
+      );
+
+      const options3 = getCookieOptions('https://example.com', true);
+
+      expect(options3).toEqual(expect.stringContaining('Domain='));
+    });
+
+    it('should return the cookie options with expired set to true', async () => {
+      const { getCookieOptions } = await import('../src/cookie');
+      const options = getCookieOptions('http://example.com', false, true);
+      expect(options).toEqual(expect.objectContaining({ maxAge: 0 }));
+    });
+
+    it('should return the cookie options as a string', async () => {
+      const { getCookieOptions } = await import('../src/cookie');
+      const options = getCookieOptions('http://example.com', true, false);
+      expect(options).toEqual(
+        expect.stringContaining('Path=/; HttpOnly; Secure=false; SameSite="Lax"; Max-Age=34560000; Domain=example.com'),
+      );
+
+      const options2 = getCookieOptions('https://example.com', true, true);
+      expect(options2).toEqual(
+        expect.stringContaining('Path=/; HttpOnly; Secure=true; SameSite="Lax"; Max-Age=0; Domain=example.com'),
+      );
     });
   });
 });
diff --git a/__tests__/session.spec.ts b/__tests__/session.spec.ts
@@ -265,12 +265,6 @@ describe('session.ts', () => {
     it('should attempt to refresh the session when the access token is invalid', async () => {
       mockSession.accessToken = await generateTestToken({}, true);
 
-      const nextCookies = await cookies();
-      nextCookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
       (jwtVerify as jest.Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
@@ -283,6 +277,11 @@ describe('session.ts', () => {
 
       const request = new NextRequest(new URL('http://example.com'));
 
+      request.cookies.set(
+        'wos-session',
+        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
+      );
+
       const result = await updateSessionMiddleware(
         request,
         true,
@@ -306,12 +305,6 @@ describe('session.ts', () => {
 
       mockSession.accessToken = await generateTestToken({}, true);
 
-      const nextCookies = await cookies();
-      nextCookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
       (jwtVerify as jest.Mock).mockImplementation(() => {
         throw new Error('Invalid token');
       });
@@ -322,7 +315,12 @@ describe('session.ts', () => {
 
       const request = new NextRequest(new URL('http://example.com'));
 
-      const result = await updateSessionMiddleware(
+      request.cookies.set(
+        'wos-session',
+        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
+      );
+
+      const response = await updateSessionMiddleware(
         request,
         true,
         {
@@ -333,8 +331,8 @@ describe('session.ts', () => {
         [],
       );
 
-      expect(result.status).toBe(200);
-      expect(nextCookies.get('wos-session')).toBeUndefined();
+      expect(response.status).toBe(200);
+      expect(response.headers.get('Set-Cookie')).toContain('wos-session=;');
       expect(console.log).toHaveBeenCalledTimes(2);
       expect(console.log).toHaveBeenNthCalledWith(
         1,
@@ -343,7 +341,7 @@ describe('session.ts', () => {
       expect(console.log).toHaveBeenNthCalledWith(
         2,
         'Failed to refresh. Deleting cookie.',
-        new Error('Failed to refresh session: Failed to refresh'),
+        new Error('Failed to refresh'),
       );
     });
 
@@ -437,9 +435,7 @@ describe('session.ts', () => {
           ['/protected-signup'],
         );
 
-        console.log('result headers:', result.headers);
-
-        expect(result.headers.get('x-middleware-request-x-sign-up-paths')).toBe('/protected-signup');
+        expect(result.headers.get('x-sign-up-paths')).toBe('/protected-signup');
       });
 
       it('should allow logged out users on unauthenticated paths', async () => {
@@ -527,12 +523,6 @@ describe('session.ts', () => {
 
         mockSession.accessToken = await generateTestToken({}, true);
 
-        const nextCookies = await cookies();
-        nextCookies.set(
-          'wos-session',
-          await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-        );
-
         (jwtVerify as jest.Mock).mockImplementation(() => {
           throw new Error('Invalid token');
         });
@@ -543,7 +533,12 @@ describe('session.ts', () => {
 
         const request = new NextRequest(new URL('http://example.com'));
 
-        const result = await updateSessionMiddleware(
+        request.cookies.set(
+          'wos-session',
+          await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
+        );
+
+        const response = await updateSessionMiddleware(
           request,
           true,
           {
@@ -554,8 +549,8 @@ describe('session.ts', () => {
           [],
         );
 
-        expect(result.status).toBe(307);
-        expect(nextCookies.get('wos-session')).toBeUndefined();
+        expect(response.status).toBe(307);
+        expect(response.headers.get('Set-Cookie')).toContain('wos-session=;');
         expect(console.log).toHaveBeenCalledTimes(3);
         expect(console.log).toHaveBeenNthCalledWith(
           1,
@@ -564,7 +559,7 @@ describe('session.ts', () => {
         expect(console.log).toHaveBeenNthCalledWith(
           2,
           'Failed to refresh. Deleting cookie.',
-          new Error('Failed to refresh session: Failed to refresh'),
+          new Error('Failed to refresh'),
         );
 
         expect(console.log).toHaveBeenNthCalledWith(
@@ -619,13 +614,13 @@ describe('session.ts', () => {
     });
 
     it('should return a session if the session is valid', async () => {
-      const nextCookies = await cookies();
-      nextCookies.set(
+      const request = new NextRequest(new URL('http://example.com/protected'));
+      request.cookies.set(
         'wos-session',
         await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
       );
 
-      const result = await updateSession(new NextRequest(new URL('http://example.com/protected')));
+      const result = await updateSession(request);
 
       expect(result.session).toBeDefined();
     });
@@ -634,12 +629,6 @@ describe('session.ts', () => {
       // Setup invalid session
       mockSession.accessToken = await generateTestToken({}, true);
 
-      const nextCookies = await cookies();
-      nextCookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
       // Mock token verification to fail
       (jwtVerify as jest.Mock).mockImplementation(() => {
         throw new Error('Invalid token');
@@ -652,12 +641,18 @@ describe('session.ts', () => {
         user: mockSession.user,
       });
 
-      const result = await updateSession(new NextRequest(new URL('http://example.com/protected')), {
+      const request = new NextRequest(new URL('http://example.com/protected'));
+      request.cookies.set(
+        'wos-session',
+        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
+      );
+
+      const response = await updateSession(request, {
         debug: true,
       });
 
-      expect(result.session).toBeDefined();
-      expect(result.session.user).toBeDefined();
+      expect(response.session).toBeDefined();
+      expect(response.session.user).toBeDefined();
       expect(console.log).toHaveBeenCalledWith(
         expect.stringContaining('Session invalid. Refreshing access token that ends in'),
       );
@@ -667,12 +662,6 @@ describe('session.ts', () => {
       // Setup invalid session
       mockSession.accessToken = await generateTestToken({}, true);
 
-      const nextCookies = await cookies();
-      nextCookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
       // Mock token verification to fail
       (jwtVerify as jest.Mock).mockImplementation(() => {
         throw new Error('Invalid token');
@@ -681,12 +670,18 @@ describe('session.ts', () => {
       // Mock refresh failure
       jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockRejectedValue(new Error('Refresh failed'));
 
-      const result = await updateSession(new NextRequest(new URL('http://example.com/protected')), {
+      const request = new NextRequest(new URL('http://example.com/protected'));
+      request.cookies.set(
+        'wos-session',
+        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
+      );
+
+      const response = await updateSession(request, {
         debug: true,
       });
 
-      expect(result.session.user).toBeNull();
-      expect(result.authorizationUrl).toBeDefined();
+      expect(response.session.user).toBeNull();
+      expect(response.authorizationUrl).toBeDefined();
       expect(console.log).toHaveBeenCalledWith('Failed to refresh. Deleting cookie.', expect.any(Error));
     });
   });

diff --git a/__tests__/utils.spec.ts b/__tests__/utils.spec.ts
@@ -16,12 +16,22 @@ describe('utils', () => {
 
       const result = redirectWithFallback(redirectUrl);
 
-      expect(mockRedirect).toHaveBeenCalledWith(redirectUrl);
+      expect(mockRedirect).toHaveBeenCalledWith(redirectUrl, { headers: undefined });
       expect(result).toBe('redirected');
 
       NextResponse.redirect = originalRedirect;
     });
 
+    it('uses headers when provided', () => {
+      const redirectUrl = 'https://example.com';
+      const headers = new Headers();
+      headers.set('Set-Cookie', 'test=1');
+
+      const result = redirectWithFallback(redirectUrl, headers);
+
+      expect(result.headers.get('Set-Cookie')).toBe('test=1');
+    });
+
     it('falls back to standard Response when NextResponse exists but redirect is undefined', async () => {
       const redirectUrl = 'https://example.com';
 

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@workos-inc/authkit-nextjs",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "Authentication and session helpers for using WorkOS & AuthKit with Next.js",
   "sideEffects": false,
   "type": "module",

diff --git a/src/cookie.ts b/src/cookie.ts
@@ -1,18 +1,39 @@
 import { WORKOS_REDIRECT_URI, WORKOS_COOKIE_MAX_AGE, WORKOS_COOKIE_DOMAIN } from './env-variables.js';
 import { CookieOptions } from './interfaces.js';
 
-export function getCookieOptions(redirectUri?: string | null): CookieOptions {
+export function getCookieOptions(): CookieOptions;
+export function getCookieOptions(redirectUri?: string | null): CookieOptions;
+export function getCookieOptions(redirectUri: string | null | undefined, asString: true, expired?: boolean): string;
+export function getCookieOptions(
+  redirectUri: string | null | undefined,
+  asString: false,
+  expired?: boolean,
+): CookieOptions;
+export function getCookieOptions(
+  redirectUri?: string | null,
+  asString?: boolean,
+  expired?: boolean,
+): CookieOptions | string;
+export function getCookieOptions(
+  redirectUri?: string | null,
+  asString: boolean = false,
+  expired: boolean = false,
+): CookieOptions | string {
   const url = new URL(redirectUri || WORKOS_REDIRECT_URI);
 
-  return {
-    path: '/',
-    httpOnly: true,
-    secure: url.protocol === 'https:',
-    sameSite: 'lax' as const,
-    // Defaults to 400 days, the maximum allowed by Chrome
-    // It's fine to have a long cookie expiry date as the access/refresh tokens
-    // act as the actual time-limited aspects of the session.
-    maxAge: WORKOS_COOKIE_MAX_AGE ? parseInt(WORKOS_COOKIE_MAX_AGE, 10) : 60 * 60 * 24 * 400,
-    domain: WORKOS_COOKIE_DOMAIN,
-  };
+  const maxAge = expired ? 0 : WORKOS_COOKIE_MAX_AGE ? parseInt(WORKOS_COOKIE_MAX_AGE, 10) : 60 * 60 * 24 * 400;
+
+  return asString
+    ? `Path=/; HttpOnly; Secure=${url.protocol === 'https:'}; SameSite="Lax"; Max-Age=${maxAge}; Domain=${WORKOS_COOKIE_DOMAIN || ''}`
+    : {
+        path: '/',
+        httpOnly: true,
+        secure: url.protocol === 'https:',
+        sameSite: 'lax' as const,
+        // Defaults to 400 days, the maximum allowed by Chrome
+        // It's fine to have a long cookie expiry date as the access/refresh tokens
+        // act as the actual time-limited aspects of the session.
+        maxAge,
+        domain: WORKOS_COOKIE_DOMAIN || '',
+      };
 }