From 5c5bb8d2f3a9926581c2b768383b2a4fd7f2576c Mon Sep 17 00:00:00 2001
From: Adam Mcgrath <adam.mcgrath@auth0.com>
Date: Fri, 19 Feb 2021 11:34:46 +0000
Subject: [PATCH] [SDK-2330] New tokens should be applied to existing session
 (#307)

---
 src/session/get-access-token.ts | 17 ++++++-----------
 tests/fixtures/setup.ts         |  6 ++++--
 tests/handlers/profile.test.ts  | 26 ++++++++++++++++++++++++--
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/src/session/get-access-token.ts b/src/session/get-access-token.ts
index cf7df0abb..f7cf85463 100644
--- a/src/session/get-access-token.ts
+++ b/src/session/get-access-token.ts
@@ -3,7 +3,7 @@ import { NextApiRequest, NextApiResponse } from 'next';
 import { ClientFactory } from '../auth0-session';
 import { AccessTokenError } from '../utils/errors';
 import { intersect, match } from '../utils/array';
-import { SessionCache, fromTokenSet, fromJson } from '../session';
+import { SessionCache, fromTokenSet } from '../session';
 import { NextConfig } from '../config';
 
 /**
@@ -114,16 +114,11 @@ export default function accessTokenFactory(
 
       // Update the session.
       const newSession = fromTokenSet(tokenSet, config);
-      sessionCache.set(
-        req,
-        res,
-        fromJson({
-          ...session,
-          ...newSession,
-          refreshToken: newSession.refreshToken || session.refreshToken,
-          user: { ...session.user, ...newSession.user }
-        })
-      );
+      Object.assign(session, {
+        ...newSession,
+        refreshToken: newSession.refreshToken || session.refreshToken,
+        user: { ...session.user, ...newSession.user }
+      });
 
       // Return the new access token.
       return {
diff --git a/tests/fixtures/setup.ts b/tests/fixtures/setup.ts
index 0fe7882f9..ad3e215f7 100644
--- a/tests/fixtures/setup.ts
+++ b/tests/fixtures/setup.ts
@@ -29,6 +29,7 @@ export type SetupOptions = {
   getAccessTokenOptions?: AccessTokenRequest;
   discoveryOptions?: object;
   userInfoPayload?: object;
+  userInfoToken?: string;
 };
 
 export const setup = async (
@@ -42,13 +43,14 @@ export const setup = async (
     withPageAuthRequiredOptions,
     getAccessTokenOptions,
     discoveryOptions,
-    userInfoPayload = {}
+    userInfoPayload = {},
+    userInfoToken = 'eyJz93a...k4laUWw'
   }: SetupOptions = {}
 ): Promise<string> => {
   discovery(config, discoveryOptions);
   jwksEndpoint(config, jwks);
   codeExchange(config, makeIdToken({ iss: 'https://acme.auth0.local/', ...idTokenClaims }));
-  userInfo(config, 'eyJz93a...k4laUWw', userInfoPayload);
+  userInfo(config, userInfoToken, userInfoPayload);
   const {
     handleAuth,
     handleCallback,
diff --git a/tests/handlers/profile.test.ts b/tests/handlers/profile.test.ts
index 39bb13532..8cd97c192 100644
--- a/tests/handlers/profile.test.ts
+++ b/tests/handlers/profile.test.ts
@@ -1,6 +1,6 @@
 import nock from 'nock';
-import { withoutApi } from '../fixtures/default-settings';
-import { userInfo } from '../fixtures/oidc-nocks';
+import { withApi, withoutApi } from '../fixtures/default-settings';
+import { refreshTokenRotationExchange, userInfo } from '../fixtures/oidc-nocks';
 import { get } from '../auth0-session/fixtures/helpers';
 import { setup, teardown, login } from '../fixtures/setup';
 import { Session, AfterCallback } from '../../src';
@@ -91,6 +91,28 @@ describe('profile handler', () => {
     );
   });
 
+  test('should refetch the user and preserve new tokens', async () => {
+    const afterCallback: AfterCallback = (_req, _res, session: Session): Session => {
+      session.accessTokenExpiresAt = -60;
+      return session;
+    };
+    const baseUrl = await setup(withApi, {
+      profileOptions: { refetch: true },
+      userInfoPayload: { foo: 'bar' },
+      callbackOptions: {
+        afterCallback
+      },
+      userInfoToken: 'new-access-token'
+    });
+    refreshTokenRotationExchange(withApi, 'GEbRxBN...edjnXbL', {}, 'new-access-token', 'new-refresh-token');
+    const cookieJar = await login(baseUrl);
+    const profile = await get(baseUrl, '/api/auth/me', { cookieJar });
+    expect(profile).toMatchObject({ foo: 'bar' });
+    const session = await get(baseUrl, '/api/session', { cookieJar });
+    expect(session.accessToken).toEqual('new-access-token');
+    expect(session.refreshToken).toEqual('new-refresh-token');
+  });
+
   test('should update the session in the afterRefetch hook', async () => {
     const baseUrl = await setup(withoutApi, {
       profileOptions: {