Adding afterRefetch hook (#276)

src/handlers/callback.ts

@@ -4,7 +4,7 @@ import { Session } from '../session';
 import { assertReqRes } from '../utils/assert';
 
 /**
- * Use this function for validating additional claims on the user's access token or adding removing items from
+ * Use this function for validating additional claims on the user's ID Token or adding removing items from
  * the session after login, eg
  *
  * ### Validate additional claims

src/handlers/index.ts

@@ -1,5 +1,5 @@
 export { default as callbackHandler, HandleCallback, CallbackOptions, AfterCallback } from './callback';
 export { default as loginHandler, HandleLogin, LoginOptions, GetLoginState } from './login';
 export { default as logoutHandler, HandleLogout, LogoutOptions } from './logout';
-export { default as profileHandler, HandleProfile, ProfileOptions } from './profile';
+export { default as profileHandler, HandleProfile, ProfileOptions, AfterRefetch } from './profile';
 export { default as handlerFactory, Handlers, HandleAuth } from './auth';

src/handlers/profile.ts

@@ -3,6 +3,8 @@ import { ClientFactory } from '../auth0-session';
 import { SessionCache, Session, fromJson, GetAccessToken } from '../session';
 import { assertReqRes } from '../utils/assert';
 
+export type AfterRefetch = (req: NextApiRequest, res: NextApiResponse, session: Session) => Promise<Session> | Session;
+
 /**
  * Custom options for {@link HandleProfile}
  *
@@ -13,6 +15,12 @@ export type ProfileOptions = {
    * If set to `true` this will refetch the user profile information from `/userinfo` and save it to the session.
    */
   refetch?: boolean;
+
+  /**
+   * Like {@AfterCallback} when a session is created, you can use this function to validate or add/remove claims
+   * after the session is updated. Will only run if {@link ProfileOptions.refetch} is `true`
+   */
+  afterRefetch?: AfterRefetch;
 };
 
 /**
@@ -42,8 +50,9 @@ export default function profileHandler(
     }
 
     const session = sessionCache.get(req, res) as Session;
+    res.setHeader('Cache-Control', 'no-store');
 
-    if (options && options.refetch) {
+    if (options?.refetch) {
       const { accessToken } = await getAccessToken(req, res);
       if (!accessToken) {
         throw new Error('No access token available to refetch the profile');
@@ -52,25 +61,24 @@ export default function profileHandler(
       const client = await getClient();
       const userInfo = await client.userinfo(accessToken);
 
-      const updatedUser = {
-        ...session.user,
-        ...userInfo
-      };
+      let newSession = fromJson({
+        ...session,
+        user: {
+          ...session.user,
+          ...userInfo
+        }
+      }) as Session;
+
+      if (options.afterRefetch) {
+        newSession = await options.afterRefetch(req, res, newSession);
+      }
 
-      sessionCache.set(
-        req,
-        res,
-        fromJson({
-          ...session,
-          user: updatedUser
-        })
-      );
+      sessionCache.set(req, res, newSession);
 
-      res.json(updatedUser);
+      res.json(newSession.user);
       return;
     }
 
-    res.setHeader('Cache-Control', 'no-store');
     res.json(session.user);
   };
 }

src/index.ts

@@ -23,7 +23,8 @@ import {
   GetLoginState,
   ProfileOptions,
   CallbackOptions,
-  AfterCallback
+  AfterCallback,
+  AfterRefetch
 } from './handlers';
 import {
   sessionFactory,
@@ -131,6 +132,7 @@ export {
   GetAccessTokenResult,
   CallbackOptions,
   AfterCallback,
+  AfterRefetch,
   LoginOptions,
   LogoutOptions,
   GetLoginState

tests/handlers/profile.test.ts

@@ -31,6 +31,14 @@ describe('profile handler', () => {
     expect(res.headers['cache-control']).toEqual('no-store');
   });
 
+  test('should not allow caching the profile response when refetch is true', async () => {
+    const baseUrl = await setup(withoutApi, { profileOptions: { refetch: true } });
+    const cookieJar = await login(baseUrl);
+
+    const { res } = await get(baseUrl, '/api/auth/me', { cookieJar, fullResponse: true });
+    expect(res.headers['cache-control']).toEqual('no-store');
+  });
+
   test('should throw if re-fetching with no Access Token', async () => {
     const afterCallback: AfterCallback = (_req, _res, session: Session): Session => {
       delete session.accessToken;
@@ -82,4 +90,34 @@ describe('profile handler', () => {
       'No access token available to refetch the profile'
     );
   });
+
+  test('should update the session in the afterRefetch hook', async () => {
+    const baseUrl = await setup(withoutApi, {
+      profileOptions: {
+        refetch: true,
+        afterRefetch(_req, _res, session) {
+          session.user.foo = 'bar';
+          return session;
+        }
+      }
+    });
+    const cookieJar = await login(baseUrl);
+
+    const user = await get(baseUrl, '/api/auth/me', { cookieJar });
+    expect(user.foo).toEqual('bar');
+  });
+
+  test('should throw from the afterRefetch hook', async () => {
+    const baseUrl = await setup(withoutApi, {
+      profileOptions: {
+        refetch: true,
+        afterRefetch() {
+          throw new Error('some validation error');
+        }
+      }
+    });
+    const cookieJar = await login(baseUrl);
+
+    await expect(get(baseUrl, '/api/auth/me', { cookieJar })).rejects.toThrowError('some validation error');
+  });
 });