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

package/src/auth.spec.ts

@@ -1,57 +1,63 @@
-import { describe, it, expect, beforeEach, jest } from '@jest/globals';
-
 import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization } from './auth.js';
 import * as session from './session.js';
 import * as cache from 'next/cache';
 import * as workosModule from './workos.js';
 
-// These are mocked in jest.setup.ts
+// These are mocked in vitest.setup.ts
 import { cookies, headers } from 'next/headers';
 import { redirect } from 'next/navigation';
 import { generateSession, generateTestToken } from './test-helpers.js';
 import { sealData } from 'iron-session';
 import { getWorkOS } from './workos.js';
+import { getStateFromPKCECookieValue } from './pkce.js';
 
 const workos = getWorkOS();
 
-jest.mock('next/cache', () => {
-  const actual = jest.requireActual<typeof cache>('next/cache');
+vi.mock('next/cache', async () => {
+  const actual = await vi.importActual<typeof cache>('next/cache');
   return {
     ...actual,
-    revalidateTag: jest.fn(),
-    revalidatePath: jest.fn(),
+    revalidateTag: vi.fn(),
+    revalidatePath: vi.fn(),
   };
 });
 
 // Create a fake WorkOS instance that will be used only in the "on error" tests
 const fakeWorkosInstance = {
   userManagement: {
-    authenticateWithRefreshToken: jest.fn(),
-    getAuthorizationUrl: jest.fn(),
-    getJwksUrl: jest.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
-    getLogoutUrl: jest.fn(),
+    authenticateWithRefreshToken: vi.fn(),
+    getAuthorizationUrl: vi.fn(),
+    getJwksUrl: vi.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
+    getLogoutUrl: vi.fn(),
+  },
+  pkce: {
+    generate: vi.fn().mockResolvedValue({
+      codeVerifier: 'test-code-verifier',
+      codeChallenge: 'test-code-challenge',
+      codeChallengeMethod: 'S256' as const,
+    }),
   },
 };
 
-const revalidatePath = jest.mocked(cache.revalidatePath);
-const revalidateTag = jest.mocked(cache.revalidateTag);
+const revalidatePath = vi.mocked(cache.revalidatePath);
+const revalidateTag = vi.mocked(cache.revalidateTag);
 // We'll only use these in the "on error" tests
 const authenticateWithRefreshToken = fakeWorkosInstance.userManagement.authenticateWithRefreshToken;
 const getAuthorizationUrl = fakeWorkosInstance.userManagement.getAuthorizationUrl;
 
-jest.mock('../src/session', () => {
-  const actual = jest.requireActual<typeof session>('../src/session');
+vi.mock('../src/session', async () => {
+  const actual = await vi.importActual<typeof session>('../src/session');
 
   return {
     ...actual,
-    refreshSession: jest.fn(actual.refreshSession),
+    refreshSession: vi.fn(actual.refreshSession),
   };
 });
 
 describe('auth.ts', () => {
   beforeEach(async () => {
     // Clear all mocks between tests
-    jest.clearAllMocks();
+    vi.clearAllMocks();
 
     // Reset the cookie store
     const nextCookies = await cookies();
@@ -76,6 +82,15 @@ describe('auth.ts', () => {
       expect(url).toBeDefined();
       expect(() => new URL(url)).not.toThrow();
     });
+
+    it('should include returnTo as returnPathname in the state parameter', async () => {
+      const url = await getSignInUrl({ returnTo: '/dashboard' });
+      const parsedUrl = new URL(url);
+      const state = parsedUrl.searchParams.get('state');
+      expect(state).toBeDefined();
+      const decoded = await getStateFromPKCECookieValue(state!);
+      expect(decoded.returnPathname).toBe('/dashboard');
+    });
   });
 
   it('should not include prompt when not specified for getSignInUrl', async () => {
@@ -103,6 +118,15 @@ describe('auth.ts', () => {
       const url = await getSignUpUrl({ prompt: 'consent' });
       expect(url).toContain('prompt=consent');
     });
+
+    it('should include returnTo as returnPathname in the state parameter', async () => {
+      const url = await getSignUpUrl({ returnTo: '/welcome' });
+      const parsedUrl = new URL(url);
+      const state = parsedUrl.searchParams.get('state');
+      expect(state).toBeDefined();
+      const decoded = await getStateFromPKCECookieValue(state!);
+      expect(decoded.returnPathname).toBe('/welcome');
+    });
   });
 
   describe('switchToOrganization', () => {
@@ -135,14 +159,21 @@ describe('auth.ts', () => {
         nextHeaders.set('x-url', 'http://localhost/test');
         await generateSession();
 
+        fakeWorkosInstance.pkce.generate.mockResolvedValue({
+          codeVerifier: 'test-code-verifier',
+          codeChallenge: 'test-code-challenge',
+          codeChallengeMethod: 'S256' as const,
+        });
+
         // Create a WorkOS-like object that matches what our tests need
         const mockWorkOS = {
           userManagement: fakeWorkosInstance.userManagement,
+          pkce: fakeWorkosInstance.pkce,
           // Add minimal properties to satisfy TypeScript
-          createHttpClient: jest.fn(),
-          createWebhookClient: jest.fn(),
-          createActionsClient: jest.fn(),
-          createIronSessionProvider: jest.fn(),
+          createHttpClient: vi.fn(),
+          createWebhookClient: vi.fn(),
+          createActionsClient: vi.fn(),
+          createIronSessionProvider: vi.fn(),
           apiKey: 'test',
           clientId: 'test',
           host: 'test',
@@ -154,12 +185,12 @@ describe('auth.ts', () => {
 
         // Apply the mock for these tests only
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        jest.spyOn(workosModule, 'getWorkOS').mockImplementation(() => mockWorkOS as any);
+        vi.spyOn(workosModule, 'getWorkOS').mockImplementation(() => mockWorkOS as any);
       });
 
       afterEach(() => {
         // Restore all mocks after each test
-        jest.restoreAllMocks();
+        vi.restoreAllMocks();
       });
 
       it('should redirect to sign in when error is "sso_required"', async () => {
@@ -243,11 +274,30 @@ describe('auth.ts', () => {
       expect(sessionCookie).toBeUndefined();
     });
 
+    it('should clear lingering PKCE verifier cookies (legacy and per-flow)', async () => {
+      const nextCookies = await cookies();
+      const nextHeaders = await headers();
+
+      nextHeaders.set('x-workos-middleware', 'true');
+      nextCookies.set('wos-session', 'foo');
+      nextCookies.set('wos-auth-verifier', 'legacy-state');
+      nextCookies.set('wos-auth-verifier-a1b2c3d4', 'flow-a-state');
+      nextCookies.set('wos-auth-verifier-deadbeef', 'flow-b-state');
+      nextCookies.set('unrelated-cookie', 'keep-me');
+
+      await signOut();
+
+      expect(nextCookies.get('wos-auth-verifier')).toBeUndefined();
+      expect(nextCookies.get('wos-auth-verifier-a1b2c3d4')).toBeUndefined();
+      expect(nextCookies.get('wos-auth-verifier-deadbeef')).toBeUndefined();
+      expect(nextCookies.get('unrelated-cookie')?.value).toBe('keep-me');
+    });
+
     describe('when given a `returnTo` parameter', () => {
       it('passes the `returnTo` through to the `getLogoutUrl` call', async () => {
-        jest
-          .spyOn(workos.userManagement, 'getLogoutUrl')
-          .mockReturnValue('https://user-management-logout.com/signed-out');
+        vi.spyOn(workos.userManagement, 'getLogoutUrl').mockReturnValue(
+          'https://user-management-logout.com/signed-out',
+        );
         const mockSession = {
           accessToken: await generateTestToken(),
           sessionId: 'session_123',
@@ -306,9 +356,9 @@ describe('auth.ts', () => {
 
         nextCookies.set('wos-session', encryptedSession);
 
-        jest
-          .spyOn(workos.userManagement, 'getLogoutUrl')
-          .mockReturnValue('https://api.workos.com/user_management/sessions/logout?session_id=session_123');
+        vi.spyOn(workos.userManagement, 'getLogoutUrl').mockReturnValue(
+          'https://api.workos.com/user_management/sessions/logout?session_id=session_123',
+        );
 
         await signOut();
 

package/src/auth.ts

@@ -5,41 +5,47 @@ import { revalidatePath, revalidateTag } from 'next/cache';
 import { cookies, headers } from 'next/headers';
 import { redirect } from 'next/navigation';
 import { WORKOS_COOKIE_NAME } from './env-variables.js';
-import { getCookieOptions } from './cookie.js';
+import { getCookieOptions, getPKCECookieOptions } from './cookie.js';
 import { getAuthorizationUrl } from './get-authorization-url.js';
-import type { AccessToken, SwitchToOrganizationOptions, UserInfo } from './interfaces.js';
+import type { AccessToken, GetAuthURLOptions, SwitchToOrganizationOptions, UserInfo } from './interfaces.js';
+import { PKCE_COOKIE_NAME, setPKCECookie } from './pkce.js';
 import { getSessionFromCookie, refreshSession, withAuth } from './session.js';
 import { getWorkOS } from './workos.js';
-export async function getSignInUrl({
-  organizationId,
-  loginHint,
-  redirectUri,
-  prompt,
-  state,
-}: {
-  organizationId?: string;
-  loginHint?: string;
-  redirectUri?: string;
-  prompt?: 'consent';
-  state?: string;
-} = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-in', loginHint, redirectUri, prompt, state });
+
+/**
+ * A wrapper around revalidateTag to provide compatibility with previous versions.
+ * @param tag The tag to revalidate.
+ */
+function revalidateTagCompat(tag: string): void {
+  const fn = revalidateTag as (tag: string, profile: string) => void;
+  return fn(tag, 'max');
+}
+
+async function getAuthURLAndSetPKCECookie(options: GetAuthURLOptions): Promise<string> {
+  const { url, sealedState } = await getAuthorizationUrl(options);
+  await setPKCECookie(sealedState);
+
+  return url;
+}
+
+type GetSignUrlOptions = Omit<GetAuthURLOptions, 'screenHint' | 'returnPathname'> & {
+  returnTo?: string;
+};
+
+export async function getSignInUrl(authUrlOptions: GetSignUrlOptions = {}) {
+  return getAuthURLAndSetPKCECookie({
+    ...authUrlOptions,
+    returnPathname: authUrlOptions.returnTo,
+    screenHint: 'sign-in',
+  });
 }
 
-export async function getSignUpUrl({
-  organizationId,
-  loginHint,
-  redirectUri,
-  prompt,
-  state,
-}: {
-  organizationId?: string;
-  loginHint?: string;
-  redirectUri?: string;
-  prompt?: 'consent';
-  state?: string;
-} = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-up', loginHint, redirectUri, prompt, state });
+export async function getSignUpUrl(authUrlOptions: GetSignUrlOptions = {}) {
+  return getAuthURLAndSetPKCECookie({
+    ...authUrlOptions,
+    returnPathname: authUrlOptions.returnTo,
+    screenHint: 'sign-up',
+  });
 }
 
 /**
@@ -67,7 +73,30 @@ export async function signOut({ returnTo }: { returnTo?: string } = {}) {
     const nextCookies = await cookies();
     const cookieName = WORKOS_COOKIE_NAME || 'wos-session';
     const { domain, path, sameSite, secure } = getCookieOptions();
-    nextCookies.delete({ name: cookieName, domain, path, sameSite, secure });
+    try {
+      nextCookies.delete({ name: cookieName, domain, path, sameSite, secure });
+    } catch {
+      // Some environments (e.g., vinext) only accept a string cookie name
+      nextCookies.delete(cookieName);
+    }
+
+    // Clear any lingering PKCE verifier cookies so orphans from abandoned
+    // flows don't accumulate toward HTTP 431 or confuse future sign-ins.
+    const pkceOptions = getPKCECookieOptions();
+    for (const { name } of nextCookies.getAll()) {
+      if (!name.startsWith(PKCE_COOKIE_NAME)) continue;
+      try {
+        nextCookies.delete({
+          name,
+          domain: pkceOptions.domain,
+          path: pkceOptions.path,
+          sameSite: pkceOptions.sameSite,
+          secure: pkceOptions.secure,
+        });
+      } catch {
+        nextCookies.delete(name);
+      }
+    }
 
     if (sessionId) {
       redirect(getWorkOS().userManagement.getLogoutUrl({ sessionId, returnTo }));
@@ -98,22 +127,25 @@ export async function switchToOrganization(
       redirect(cause.rawData.authkit_redirect_url);
     } else {
       if (cause?.error === 'sso_required' || cause?.error === 'mfa_enrollment') {
-        const url = await getAuthorizationUrl({ organizationId });
-        return redirect(url);
+        return redirect(await getAuthURLAndSetPKCECookie({ organizationId }));
       }
       throw error;
     }
   }
 
-  switch (revalidationStrategy) {
-    case 'path':
-      revalidatePath(pathname);
-      break;
-    case 'tag':
-      for (const tag of revalidationTags) {
-        revalidateTag(tag);
-      }
-      break;
+  try {
+    switch (revalidationStrategy) {
+      case 'path':
+        revalidatePath(pathname);
+        break;
+      case 'tag':
+        for (const tag of revalidationTags) {
+          revalidateTagCompat(tag);
+        }
+        break;
+    }
+  } catch {
+    // revalidatePath/revalidateTag may not be available in non-Next.js environments (e.g., vinext)
   }
   if (revalidationStrategy !== 'none') {
     redirect(pathname);