Merge pull request #2481 from contentstack/fix/dx-3878-oauth-session-expire

Refactor OAuth token refresh logic in AuthHandler to prevent concurre…

Author: harshitha-cstk

packages/contentstack-utilities/src/auth-handler.ts

@@ -35,7 +35,10 @@ class AuthHandler {
   private allAuthConfigItems: any;
   private oauthHandler: any;
   private managementAPIClient: ContentstackClient;
+  /** True while an OAuth access-token refresh is running (for logging/diagnostics; correctness uses `oauthRefreshInFlight`). */
   private isRefreshingToken: boolean = false; // Flag to track if a refresh operation is in progress
+  /** Serialize OAuth refresh so concurrent API calls await the same refresh instead of proceeding with a stale token. */
+  private oauthRefreshInFlight: Promise<void> | null = null;
   private cmaHost: string;
 
   set host(contentStackHost) {
@@ -376,42 +379,42 @@ class AuthHandler {
   checkExpiryAndRefresh = (force: boolean = false) => this.compareOAuthExpiry(force);
 
   async compareOAuthExpiry(force: boolean = false) {
-    // Avoid recursive refresh operations
-    if (this.isRefreshingToken) {
-      cliux.print('Refresh operation already in progress');
-      return Promise.resolve();
-    }
     const oauthDateTime = configHandler.get(this.oauthDateTimeKeyName);
     const authorisationType = configHandler.get(this.authorisationTypeKeyName);
     if (oauthDateTime && authorisationType === this.authorisationTypeOAUTHValue) {
       const now = new Date();
       const oauthDate = new Date(oauthDateTime);
-      const oauthValidUpto = new Date();
-      oauthValidUpto.setTime(oauthDate.getTime() + 59 * 60 * 1000);
-      if (force) {
-        cliux.print('Forcing token refresh...');
-        return this.refreshToken();
-      } else {
-        if (oauthValidUpto > now) {
-          return Promise.resolve();
-        } else {
-          cliux.print('Token expired, refreshing the token');
-          // Set the flag before refreshing the token
-          this.isRefreshingToken = true;
-
-          try {
-            await this.refreshToken();
-          } catch (error) {
-            cliux.error('Error refreshing token');
-            throw error;
-          } finally {
-            // Reset the flag after refresh operation is completed
-            this.isRefreshingToken = false;
-          }
+      const oauthValidUpto = new Date(oauthDate.getTime() + 59 * 60 * 1000);
+      const tokenExpired = oauthValidUpto <= now;
+      const shouldRefresh = force || tokenExpired;
 
-          return Promise.resolve();
-        }
+      if (!shouldRefresh) {
+        return Promise.resolve();
+      }
+
+      if (this.oauthRefreshInFlight) {
+        return this.oauthRefreshInFlight;
       }
+
+      this.isRefreshingToken = true;
+      this.oauthRefreshInFlight = (async () => {
+        try {
+          if (force) {
+            cliux.print('Forcing token refresh...');
+          } else {
+            cliux.print('Token expired, refreshing the token');
+          }
+          await this.refreshToken();
+        } catch (error) {
+          cliux.error('Error refreshing token');
+          throw error;
+        } finally {
+          this.isRefreshingToken = false;
+          this.oauthRefreshInFlight = null;
+        }
+      })();
+
+      return this.oauthRefreshInFlight;
     } else {
       cliux.print('No OAuth configuration set.');
       this.unsetConfigData();

packages/contentstack-utilities/test/unit/auth-handler.test.ts

@@ -456,6 +456,8 @@ describe('Auth Handler', () => {
 
     beforeEach(() => {
       sandbox = createSandbox();
+      authHandler.oauthRefreshInFlight = null;
+      authHandler.isRefreshingToken = false;
       configHandlerGetStub = sandbox.stub(configHandler, 'get');
       cliuxPrintStub = sandbox.stub(cliux, 'print');
       refreshTokenStub = sandbox.stub(authHandler, 'refreshToken').resolves();
@@ -467,40 +469,64 @@ describe('Auth Handler', () => {
     });
 
     it('should resolve if the OAuth token is valid and not expired', async () => {
-      const expectedOAuthDateTime = '2023-05-30T12:00:00Z';
-      const expectedAuthorisationType = 'oauth';
-      const now = new Date('2023-05-30T12:30:00Z');
+      const expectedOAuthDateTime = new Date(Date.now() - 30 * 60 * 1000).toISOString();
+      const expectedAuthorisationType = 'OAUTH';
 
       configHandlerGetStub.withArgs(authHandler.oauthDateTimeKeyName).returns(expectedOAuthDateTime);
       configHandlerGetStub.withArgs(authHandler.authorisationTypeKeyName).returns(expectedAuthorisationType);
 
-      sandbox.stub(Date, 'now').returns(now.getTime());
+      await authHandler.compareOAuthExpiry();
+      expect(cliuxPrintStub.called).to.be.false;
+      expect(refreshTokenStub.called).to.be.false;
+      expect(unsetConfigDataStub.called).to.be.false;
+    });
 
-      try {
-        await authHandler.compareOAuthExpiry();
-      } catch (error) {
-        expect(error).to.be.undefined;
-        expect(cliuxPrintStub.called).to.be.false;
-        expect(refreshTokenStub.called).to.be.false;
-        expect(unsetConfigDataStub.called).to.be.false;
-      }
+    it('should refresh when force is true even if token is not expired', async () => {
+      const expectedOAuthDateTime = new Date(Date.now() - 30 * 60 * 1000).toISOString();
+      const expectedAuthorisationType = 'OAUTH';
+
+      configHandlerGetStub.withArgs(authHandler.oauthDateTimeKeyName).returns(expectedOAuthDateTime);
+      configHandlerGetStub.withArgs(authHandler.authorisationTypeKeyName).returns(expectedAuthorisationType);
+
+      await authHandler.compareOAuthExpiry(true);
+      expect(cliuxPrintStub.calledOnceWithExactly('Forcing token refresh...')).to.be.true;
+      expect(refreshTokenStub.calledOnce).to.be.true;
+      expect(unsetConfigDataStub.called).to.be.false;
     });
 
-    it('should resolve if force is true and refreshToken is called', async () => {
-      const expectedOAuthDateTime = '2023-05-30T12:00:00Z';
-      const expectedAuthorisationType = 'oauth';
+    it('should refresh when token is expired', async () => {
+      const expectedOAuthDateTime = new Date(Date.now() - 2 * 60 * 60 * 1000).toISOString();
+      const expectedAuthorisationType = 'OAUTH';
 
       configHandlerGetStub.withArgs(authHandler.oauthDateTimeKeyName).returns(expectedOAuthDateTime);
       configHandlerGetStub.withArgs(authHandler.authorisationTypeKeyName).returns(expectedAuthorisationType);
 
-      try {
-        await authHandler.compareOAuthExpiry();
-      } catch (error) {
-        expect(error).to.be.undefined;
-        expect(cliuxPrintStub.calledOnceWithExactly('Forcing token refresh...')).to.be.true;
-        expect(refreshTokenStub.calledOnce).to.be.true;
-        expect(unsetConfigDataStub.called).to.be.false;
-      }
+      await authHandler.compareOAuthExpiry(false);
+      expect(cliuxPrintStub.calledOnceWithExactly('Token expired, refreshing the token')).to.be.true;
+      expect(refreshTokenStub.calledOnce).to.be.true;
+    });
+
+    it('should run a single refresh when compareOAuthExpiry is called concurrently', async () => {
+      const expectedOAuthDateTime = new Date(Date.now() - 2 * 60 * 60 * 1000).toISOString();
+      const expectedAuthorisationType = 'OAUTH';
+      let resolveRefresh;
+      const refreshDone = new Promise((r) => {
+        resolveRefresh = r;
+      });
+
+      configHandlerGetStub.withArgs(authHandler.oauthDateTimeKeyName).returns(expectedOAuthDateTime);
+      configHandlerGetStub.withArgs(authHandler.authorisationTypeKeyName).returns(expectedAuthorisationType);
+
+      refreshTokenStub.callsFake(async () => {
+        await refreshDone;
+      });
+
+      const p1 = authHandler.compareOAuthExpiry(false);
+      const p2 = authHandler.compareOAuthExpiry(false);
+      resolveRefresh();
+      await Promise.all([p1, p2]);
+
+      expect(refreshTokenStub.callCount).to.equal(1);
     });
   });
 });