Add composable middleware

README.md

@@ -233,27 +233,36 @@ In the above example the `/admin` page will require a user to be signed in, wher
 
 `unauthenticatedPaths` uses the same glob logic as the [Next.js matcher](https://nextjs.org/docs/pages/building-your-application/routing/middleware#matcher).
 
-### Retrieve session in middleware
+### Composing middleware
 
-Sometimes it's useful to check the user session if you want to compose custom middleware. The `getSession` helper method will retrieve the session from the cookie and verify the access token.
+If you don't want to use `authkitMiddleware` and instead want to compose your own middleware, you can use the `authkit` method. In this mode you are responsible to handling what to do when there's no session on a protected route.
 
 ```ts
-import { authkitMiddleware, getSession } from '@workos-inc/authkit-nextjs';
-import { NextRequest, NextFetchEvent } from 'next/server';
+export default async function middleware(request: NextRequest) {
+  // Perform logic before or after AuthKit
 
-export default async function middleware(request: NextRequest, event: NextFetchEvent) {
-  // authkitMiddleware will handle refreshing the session if the access token has expired
-  const response = await authkitMiddleware()(request, event);
+  // Auth object contains the session, response headers and an auhorization URL in the case that the session isn't valid
+  // This method will automatically handle setting the cookie and refreshing the session
+  const { session, headers, authorizationUrl } = await authkit(request, {
+    debug: true,
+  });
 
-  // If session is undefined, the user is not authenticated
-  const session = await getSession(response);
+  // Control of what to do when there's no session on a protected route is left to the developer
+  if (request.url.includes('/account') && !session.user) {
+    console.log('No session on protected path');
+    return NextResponse.redirect(authorizationUrl);
 
-  // ...add additional middleware logic here
+    // Alternatively you could redirect to your own login page, for example if you want to use your own UI instead of hosted AuthKit
+    return NextResponse.redirect('/login');
+  }
 
-  return response;
+  // Headers from the authkit response need to be included in every non-redirect response to ensure that `withAuth` works as expected
+  return NextResponse.next({
+    headers: headers,
+  });
 }
 
-// Match against pages that require auth
+// Match against the pages
 export const config = { matcher: ['/', '/account/:path*'] };
 ```
 

__tests__/session.spec.ts

@@ -2,7 +2,7 @@ import { NextRequest, NextResponse } from 'next/server';
 import { cookies, headers } from 'next/headers';
 import { redirect } from 'next/navigation';
 import { generateTestToken } from './test-helpers.js';
-import { withAuth, updateSession, refreshSession, getSession, terminateSession } from '../src/session.js';
+import { withAuth, updateSession, refreshSession, terminateSession, updateSessionMiddleware } from '../src/session.js';
 import { workos } from '../src/workos.js';
 import * as envVariables from '../src/env-variables.js';
 
@@ -137,14 +137,14 @@ describe('session.ts', () => {
     });
   });
 
-  describe('updateSession', () => {
+  describe('updateSessionMiddleware', () => {
     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', '');
 
       await expect(async () => {
-        await updateSession(
+        await updateSessionMiddleware(
           new NextRequest(new URL('http://example.com')),
           false,
           {
@@ -165,7 +165,7 @@ describe('session.ts', () => {
       jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', '');
 
       await expect(async () => {
-        await updateSession(
+        await updateSessionMiddleware(
           new NextRequest(new URL('http://example.com')),
           false,
           {
@@ -188,7 +188,7 @@ describe('session.ts', () => {
       jest.replaceProperty(envVariables, 'WORKOS_COOKIE_PASSWORD', 'short');
 
       await expect(async () => {
-        await updateSession(
+        await updateSessionMiddleware(
           new NextRequest(new URL('http://example.com')),
           false,
           {
@@ -207,7 +207,7 @@ describe('session.ts', () => {
 
     it('should return early if there is no session', async () => {
       const request = new NextRequest(new URL('http://example.com'));
-      const result = await updateSession(
+      const result = await updateSessionMiddleware(
         request,
         false,
         {
@@ -236,7 +236,7 @@ describe('session.ts', () => {
       });
 
       const request = new NextRequest(new URL('http://example.com'));
-      const result = await updateSession(
+      const result = await updateSessionMiddleware(
         request,
         true,
         {
@@ -273,7 +273,7 @@ describe('session.ts', () => {
 
       const request = new NextRequest(new URL('http://example.com'));
 
-      const result = await updateSession(
+      const result = await updateSessionMiddleware(
         request,
         true,
         {
@@ -312,7 +312,7 @@ describe('session.ts', () => {
 
       const request = new NextRequest(new URL('http://example.com'));
 
-      const result = await updateSession(
+      const result = await updateSessionMiddleware(
         request,
         true,
         {
@@ -342,7 +342,7 @@ describe('session.ts', () => {
         jest.spyOn(console, 'log').mockImplementation(() => {});
 
         const request = new NextRequest(new URL('http://example.com/protected'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           true,
           {
@@ -364,7 +364,7 @@ describe('session.ts', () => {
         (NextResponse as Partial<typeof NextResponse>).redirect = undefined;
 
         const request = new NextRequest(new URL('http://example.com/protected'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           false,
           {
@@ -383,7 +383,7 @@ describe('session.ts', () => {
 
       it('should automatically add the redirect URI to unauthenticatedPaths when middleware is enabled', async () => {
         const request = new NextRequest(new URL('http://example.com/protected'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           false,
           {
@@ -399,7 +399,7 @@ describe('session.ts', () => {
 
       it('should redirect unauthenticated users to sign up page on protected routes included in signUpPaths', async () => {
         const request = new NextRequest(new URL('http://example.com/protected-signup'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           false,
           {
@@ -416,7 +416,7 @@ describe('session.ts', () => {
 
       it('should allow logged out users on unauthenticated paths', async () => {
         const request = new NextRequest(new URL('http://example.com/unauthenticated'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           false,
           {
@@ -433,7 +433,7 @@ describe('session.ts', () => {
       it('should throw an error if the provided regex is invalid', async () => {
         const request = new NextRequest(new URL('http://example.com/invalid-regex'));
         await expect(async () => {
-          await updateSession(
+          await updateSessionMiddleware(
             request,
             false,
             {
@@ -457,12 +457,12 @@ describe('session.ts', () => {
         });
 
         // Import session after setting up the spy
-        const { updateSession } = await import('../src/session.js');
+        const { updateSessionMiddleware } = await import('../src/session.js');
 
         const request = new NextRequest(new URL('http://example.com/invalid-regex'));
 
         await expect(async () => {
-          await updateSession(
+          await updateSessionMiddleware(
             request,
             false,
             {
@@ -480,7 +480,7 @@ describe('session.ts', () => {
 
       it('should default to the WORKOS_REDIRECT_URI environment variable if no redirect URI is provided', async () => {
         const request = new NextRequest(new URL('http://example.com/protected'));
-        const result = await updateSession(
+        const result = await updateSessionMiddleware(
           request,
           false,
           {
@@ -498,7 +498,7 @@ describe('session.ts', () => {
         it('should redirect to sign up when unauthenticated user is on a sign up path', async () => {
           const request = new NextRequest(new URL('http://example.com/signup'));
 
-          const result = await updateSession(
+          const result = await updateSessionMiddleware(
             request,
             false,
             {
@@ -526,6 +526,92 @@ describe('session.ts', () => {
     });
   });
 
+  describe('updateSession', () => {
+    it('should return an authorization url if the session is invalid', async () => {
+      const result = await updateSession(new NextRequest(new URL('http://example.com/protected')), {
+        debug: true,
+      });
+
+      expect(result.authorizationUrl).toBeDefined();
+      expect(result.session.user).toBeNull();
+    });
+
+    it('should return a session if the session is valid', async () => {
+      const nextCookies = await cookies();
+      nextCookies.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')), {
+        debug: true,
+      });
+      console.log('result', result);
+
+      expect(result.session).toBeDefined();
+    });
+
+    it('should attempt to refresh an invalid session', async () => {
+      // 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');
+      });
+
+      // Mock successful refresh
+      jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
+        accessToken: await generateTestToken(),
+        refreshToken: 'new-refresh-token',
+        user: mockSession.user,
+      });
+
+      const result = await updateSession(new NextRequest(new URL('http://example.com/protected')), {
+        debug: true,
+      });
+
+      expect(result.session).toBeDefined();
+      expect(result.session.user).toBeDefined();
+      expect(console.log).toHaveBeenCalledWith(
+        expect.stringContaining('Session invalid. Refreshing access token that ends in'),
+      );
+    });
+
+    it('should handle refresh failure by returning auth URL', async () => {
+      // 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');
+      });
+
+      // 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')), {
+        debug: true,
+      });
+
+      expect(result.session.user).toBeNull();
+      expect(result.authorizationUrl).toBeDefined();
+      expect(console.log).toHaveBeenCalledWith('Failed to refresh. Deleting cookie.', expect.any(Error));
+    });
+  });
+
   describe('refreshSession', () => {
     it('should refresh session successfully', async () => {
       jest.spyOn(workos.userManagement, 'authenticateWithRefreshToken').mockResolvedValue({
@@ -589,54 +675,6 @@ describe('session.ts', () => {
     });
   });
 
-  describe('getSession', () => {
-    it('should return session info when valid', async () => {
-      const nextCookies = await cookies();
-      nextCookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
-      const result = await getSession();
-      expect(result).toHaveProperty('user');
-    });
-
-    it('should return null user when no session exists', async () => {
-      const result = await getSession();
-      expect(result).toEqual({ user: null });
-    });
-
-    it('should return undefined if the access token is invalid', async () => {
-      mockSession.accessToken = 'invalid-token';
-
-      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');
-      });
-
-      const result = await getSession();
-      expect(result).toEqual(undefined);
-    });
-
-    it('should return cookie from a response object if provided', async () => {
-      mockSession.accessToken = await generateTestToken();
-
-      const response = new NextResponse();
-      response.cookies.set(
-        'wos-session',
-        await sealData(mockSession, { password: process.env.WORKOS_COOKIE_PASSWORD as string }),
-      );
-
-      const result = await getSession(response);
-      expect(result).toEqual(mockSession);
-    });
-  });
-
   describe('terminateSession', () => {
     it('should redirect to logout url when there is a session', async () => {
       const nextHeaders = await headers();

jest.setup.ts

@@ -26,7 +26,10 @@ jest.mock('next/headers', () => {
       get: jest.fn((name: string) => cookieStore.get(name)),
       getAll: jest.fn(() => Array.from(cookieStore.entries())),
       set: jest.fn((name: string, value: string | { [key: string]: string | number | boolean }) =>
-        cookieStore.set(name, value),
+        cookieStore.set(name, {
+          name,
+          value,
+        }),
       ),
       _reset: () => {
         cookieStore.clear();

src/index.ts

@@ -1,6 +1,6 @@
 import { handleAuth } from './authkit-callback-route.js';
-import { authkitMiddleware } from './middleware.js';
-import { withAuth, refreshSession, getSession } from './session.js';
+import { authkit, authkitMiddleware } from './middleware.js';
+import { withAuth, refreshSession } from './session.js';
 import { getSignInUrl, getSignUpUrl, signOut } from './auth.js';
 import { Impersonation } from './impersonation.js';
 import { AuthKitProvider } from './authkit-provider.js';
@@ -9,7 +9,7 @@ export {
   handleAuth,
   //
   authkitMiddleware,
-  getSession,
+  authkit,
   //
   getSignInUrl,
   getSignUpUrl,