Repository: auth0/nextjs-auth0
Pull request #838 (2022-10-04 10:21:41)

[SDK-3567] Allow response customization in afterCallback
### 📋 Changes

Allow customisation of the response in `afterCallback` to enabled a couple of use cases

1. Terminating the request if a condition in the claim is not met.

```js
const afterCallback = (req, res, session, state) => {
  if (session.user.isAdmin) {
    return session;
  } else {
    res.status(401).end('User is not admin');
  }
}; // Terminates the request with 401 if user is not admin
```

2. Customising the return to url based on the claims

```js
const afterCallback = (req, res, session, state) => {
  if (!session.user.isAdmin) {
    res.setHeader('Location', '/admin');
  }
  return session;
}; // Redirects to `/admin` if user is admin
```



Patch:
diff --git a/V2_MIGRATION_GUIDE.md b/V2_MIGRATION_GUIDE.md
index f2fffd1a..44e359dc 100644
--- a/V2_MIGRATION_GUIDE.md
+++ b/V2_MIGRATION_GUIDE.md
@@ -8,6 +8,7 @@ Guide to migrating from `1.x` to `2.x`
 - [Profile API route no longer returns a 401](#profile-api-route-no-longer-returns-a-401)
 - [The ID token is no longer stored by default](#the-id-token-is-no-longer-stored-by-default)
 - [Override default error handler](#override-default-error-handler)
+- [afterCallback can write to the response](#afterCallback-can-write-to-the-response)
 
 ## `getSession` now returns a `Promise`
 
@@ -156,3 +157,45 @@ export default handleAuth({
   }
 });
 ```
+
+## `afterCallback` can write to the response
+
+You can now write your own redirect header or terminate the request in `afterCallback`
+
+### Before
+
+```js
+const afterCallback = (req, res, session, state) => {
+  if (session.user.isAdmin) {
+    return session;
+  } else {
+    res.status(401).end('User is not admin');
+  }
+}; // 💥Fail with ERR_HTTP_HEADERS_SENT
+
+const afterCallback = (req, res, session, state) => {
+  if (!session.user.isAdmin) {
+    res.setHeader('Location', '/admin');
+  }
+  return session;
+}; // 💥Fail with ERR_HTTP_HEADERS_SENT
+```
+
+### After
+
+```js
+const afterCallback = (req, res, session, state) => {
+  if (session.user.isAdmin) {
+    return session;
+  } else {
+    res.status(401).end('User is not admin');
+  }
+}; // Terminates the request with 401 if user is not admin
+
+const afterCallback = (req, res, session, state) => {
+  if (!session.user.isAdmin) {
+    res.setHeader('Location', '/admin');
+  }
+  return session;
+}; // Redirects to `/admin` if user is admin
+```
diff --git a/src/auth0-session/handlers/callback.ts b/src/auth0-session/handlers/callback.ts
index deeff61d..d0447be1 100644
--- a/src/auth0-session/handlers/callback.ts
+++ b/src/auth0-session/handlers/callback.ts
@@ -19,7 +19,12 @@ function getRedirectUri(config: Config): string {
   return urlJoin(config.baseURL, config.routes.callback);
 }
 
-export type AfterCallback = (req: any, res: any, session: any, state?: Record<string, any>) => Promise<any> | any;
+export type AfterCallback = (
+  req: any,
+  res: any,
+  session: any,
+  state?: Record<string, any>
+) => Promise<any> | any | undefined;
 
 export type CallbackOptions = {
   afterCallback?: AfterCallback;
@@ -87,14 +92,18 @@ export default function callbackHandlerFactory(
     let session = await sessionCache.fromTokenSet(tokenSet);
 
     if (options?.afterCallback) {
-      session = await options.afterCallback(req as any, res as any, session, openidState);
+      session = await options.afterCallback(req, res, session, openidState);
     }
 
-    await sessionCache.create(req, res, session);
+    if (session) {
+      await sessionCache.create(req, res, session);
+    }
 
-    res.writeHead(302, {
-      Location: openidState.returnTo || config.baseURL
-    });
-    res.end(htmlSafe(openidState.returnTo || config.baseURL));
+    if (!res.writableEnded) {
+      res.writeHead(302, {
+        Location: res.getHeader('Location') || openidState.returnTo || config.baseURL
+      });
+      res.end(htmlSafe(openidState.returnTo || config.baseURL));
+    }
   };
 }
diff --git a/src/handlers/auth.ts b/src/handlers/auth.ts
index 4b22ee74..c6b001d8 100644
--- a/src/handlers/auth.ts
+++ b/src/handlers/auth.ts
@@ -119,7 +119,7 @@ export default function handlerFactory({
         }
       } catch (error) {
         await (onError || defaultOnError)(req, res, error);
-        if (!res.finished) {
+        if (!res.writableEnded) {
           // 200 is the default, so we assume it has not been set in the custom error handler if it equals 200
           res.status(res.statusCode === 200 ? 500 : res.statusCode).end();
         }
diff --git a/src/handlers/callback.ts b/src/handlers/callback.ts
index 570a4cfa..577c16de 100644
--- a/src/handlers/callback.ts
+++ b/src/handlers/callback.ts
@@ -17,10 +17,11 @@ import { CallbackHandlerError, HandlerErrorCause } from '../utils/errors';
  * import { handleAuth, handleCallback } from '@auth0/nextjs-auth0';
  *
  * const afterCallback = (req, res, session, state) => {
- *   if (!session.user.isAdmin) {
- *     throw new UnauthorizedError('User is not admin');
+ *   if (session.user.isAdmin) {
+ *     return session;
+ *   } else {
+ *     res.status(401).end('User is not admin');
  *   }
- *   return session;
  * };
  *
  * export default handleAuth({
@@ -57,6 +58,30 @@ import { CallbackHandlerError, HandlerErrorCause } from '../utils/errors';
  * });
  * ```
  *
+ * @example Redirect successful login based on claim
+ *
+ * ```js
+ * // pages/api/auth/[...auth0].js
+ * import { handleAuth, handleCallback } from '@auth0/nextjs-auth0';
+ *
+ * const afterCallback = (req, res, session, state) => {
+ *   if (!session.user.isAdmin) {
+ *     res.setHeader('Location', '/admin');
+ *   }
+ *   return session;
+ * };
+ *
+ * export default handleAuth({
+ *   async callback(req, res) {
+ *     try {
+ *       await handleCallback(req, res, { afterCallback });
+ *     } catch (error) {
+ *       res.status(error.status || 500).end(error.message);
+ *     }
+ *   }
+ * });
+ * ```
+ *
  * @throws {@link HandlerError}
  *
  * @category Server
@@ -66,7 +91,7 @@ export type AfterCallback = (
   res: NextApiResponse,
   session: Session,
   state?: { [key: string]: any }
-) => Promise<Session> | Session;
+) => Promise<Session> | Session | undefined;
 
 /**
  * Options to customize the callback handler.


Test patch:
diff --git a/tests/auth0-session/fixtures/helpers.ts b/tests/auth0-session/fixtures/helpers.ts
index 7f744316..41af9e3a 100644
--- a/tests/auth0-session/fixtures/helpers.ts
+++ b/tests/auth0-session/fixtures/helpers.ts
@@ -59,6 +59,9 @@ const request = (
         rejectUnauthorized: false
       },
       (res) => {
+        if (cookieJar) {
+          (res.headers['set-cookie'] || []).forEach((cookie: string) => cookieJar.setCookieSync(cookie, url));
+        }
         if (res.statusCode && (res.statusCode < 200 || res.statusCode >= 400)) {
           return reject(new Error(res.statusMessage));
         }
@@ -80,9 +83,6 @@ const request = (
             resolve(data);
           }
         });
-        if (cookieJar) {
-          (res.headers['set-cookie'] || []).forEach((cookie: string) => cookieJar.setCookieSync(cookie, url));
-        }
       }
     );
     req.setHeader('content-type', 'application/json');
diff --git a/tests/auth0-session/handlers/callback.test.ts b/tests/auth0-session/handlers/callback.test.ts
index 949ed931..4bedf882 100644
--- a/tests/auth0-session/handlers/callback.test.ts
+++ b/tests/auth0-session/handlers/callback.test.ts
@@ -6,6 +6,7 @@ import { encodeState } from '../../../src/auth0-session/hooks/get-login-state';
 import { SessionResponse, setup, teardown } from '../fixtures/server';
 import { makeIdToken } from '../fixtures/cert';
 import { toSignedCookieJar, get, post, defaultConfig } from '../fixtures/helpers';
+import { ServerResponse } from 'http';
 
 const expectedDefaultState = encodeState({ returnTo: 'https://example.org' });
 
@@ -395,4 +396,68 @@ describe('callback', () => {
     expect(res.statusCode).toEqual(302);
     expect(res.headers.location).toEqual(baseURL);
   });
+
+  it('should not overwrite location header if set in after callback', async () => {
+    const baseURL = await setup(defaultConfig, {
+      callbackOptions: {
+        afterCallback(_req, res: ServerResponse, session) {
+          res.setHeader('Location', '/foo');
+          return session;
+        }
+      }
+    });
+
+    const state = encodeState({ foo: 'bar' });
+    const cookieJar = await toSignedCookieJar(
+      {
+        state: state,
+        nonce: '__test_nonce__'
+      },
+      baseURL
+    );
+
+    const { res } = await post(baseURL, '/callback', {
+      body: {
+        state: state,
+        id_token: await makeIdToken()
+      },
+      cookieJar,
+      fullResponse: true
+    });
+
+    expect(res.statusCode).toEqual(302);
+    expect(res.headers.location).toEqual('/foo');
+    expect(cookieJar.getCookieStringSync(baseURL)).toMatch(/^appSession=.*/);
+  });
+
+  it('should terminate the request in after callback and not set session if none returned', async () => {
+    const baseURL = await setup(defaultConfig, {
+      callbackOptions: {
+        afterCallback(_req, res: ServerResponse) {
+          res.writeHead(401).end();
+        }
+      }
+    });
+
+    const state = encodeState({ foo: 'bar' });
+    const cookieJar = await toSignedCookieJar(
+      {
+        state: state,
+        nonce: '__test_nonce__'
+      },
+      baseURL
+    );
+
+    await expect(
+      post(baseURL, '/callback', {
+        body: {
+          state: state,
+          id_token: await makeIdToken()
+        },
+        cookieJar,
+        fullResponse: true
+      })
+    ).rejects.toThrow('Unauthorized');
+    expect(cookieJar.getCookieStringSync(baseURL)).toBeFalsy();
+  });
 });