From 3032958c4a736b45b536107b9207c0d544952c9d Mon Sep 17 00:00:00 2001
From: Paulo Margarido <paulo.margarido@shopify.com>
Date: Thu, 14 Jan 2021 11:07:31 -0500
Subject: [PATCH] Clone the JWT session before overriding the expiration

---
 src/auth/oauth/oauth.ts           |  9 +++++----
 src/auth/oauth/test/oauth.test.ts | 10 +++++++---
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/auth/oauth/oauth.ts b/src/auth/oauth/oauth.ts
index b41dfff4e..c14a3afe4 100644
--- a/src/auth/oauth/oauth.ts
+++ b/src/auth/oauth/oauth.ts
@@ -143,14 +143,15 @@ const ShopifyOAuth = {
       oauthSessionExpiration = new Date();
     }
     else if (Context.IS_EMBEDDED_APP) {
-      oauthSessionExpiration = new Date(Date.now() + 30000);
-      currentSession.expires = oauthSessionExpiration;
-
-      // If this is an online session for an embedded app, prepare a JWT session to be used from here on out
+      // If this is an online session for an embedded app, prepare a JWT session to be used going forward
       const onlineInfo = currentSession.onlineAccesInfo as OnlineAccessInfo;
       const jwtSessionId = this.getJwtSessionId(currentSession.shop, '' + onlineInfo.associated_user.id);
       const jwtSession = Session.cloneSession(currentSession, jwtSessionId);
       await Context.storeSession(jwtSession);
+
+      // Make sure the current OAuth session expires along with the cookie
+      oauthSessionExpiration = new Date(Date.now() + 30000);
+      currentSession.expires = oauthSessionExpiration;
     }
 
     cookies.set(ShopifyOAuth.SESSION_COOKIE_NAME, currentSession.id, {
diff --git a/src/auth/oauth/test/oauth.test.ts b/src/auth/oauth/test/oauth.test.ts
index e1cd59179..49ee2f7f6 100644
--- a/src/auth/oauth/test/oauth.test.ts
+++ b/src/auth/oauth/test/oauth.test.ts
@@ -259,7 +259,7 @@ describe('validateAuthCallback', () => {
     const successResponse = {
       access_token: 'some access token',
       scope: 'pet_kitties, walk_dogs',
-      expires_in: '525600',
+      expires_in: 525600,
       associated_user_scope: 'pet_kitties',
       associated_user: {
         id: '1',
@@ -299,7 +299,7 @@ describe('validateAuthCallback', () => {
       dest: `https://${shop}`,
       aud: Context.API_KEY,
       sub: '1',
-      exp: Date.now() / 1000 + 3600,
+      exp: new Date(Date.now() + successResponse.expires_in * 1000).getTime() / 1000,
       nbf: 1234,
       iat: 1234,
       jti: '4321',
@@ -307,7 +307,11 @@ describe('validateAuthCallback', () => {
     };
 
     const jwtSessionId = `${shop}_${jwtPayload.sub}`;
-    await expect(Context.loadSession(jwtSessionId)).resolves.not.toBeUndefined();
+    const actualJwtSession = await Context.loadSession(jwtSessionId);
+    expect(actualJwtSession).not.toBeUndefined();
+
+    const actualJwtExpiration = actualJwtSession?.expires ? actualJwtSession.expires.getTime() / 1000 : 0;
+    expect(Math.abs(actualJwtExpiration - jwtPayload.exp)).toBeLessThan(1); // 1-second grace period
 
     // Simulate a subsequent JWT request to see if the session is loaded as the current one