fix: do not return token to driver which is being rejected MONGOSH-2147

Ensure that we do not return tokens from a token set that the driver
is currently rejecting (in the sense of calling the OIDC callback
again and referring to it via the driver `refreshToken` property).

Author: addaleax

src/log-hook.ts

@@ -241,7 +241,15 @@ export function hookLoggerToMongoLogWriter(
 
   emitter.on(
     'mongodb-oidc-plugin:auth-succeeded',
-    ({ tokenType, refreshToken, expiresAt, passIdTokenAsAccessToken }) => {
+    ({
+      tokenType,
+      refreshToken,
+      expiresAt,
+      passIdTokenAsAccessToken,
+      forceRefreshOrReauth,
+      willRetryWithForceRefreshOrReauth,
+      tokenSetId,
+    }) => {
       log.info(
         'OIDC-PLUGIN',
         mongoLogId(1_002_000_017),
@@ -252,6 +260,9 @@ export function hookLoggerToMongoLogWriter(
           refreshToken,
           expiresAt,
           passIdTokenAsAccessToken,
+          forceRefreshOrReauth,
+          willRetryWithForceRefreshOrReauth,
+          tokenSetId,
         }
       );
     }

src/plugin.spec.ts

@@ -56,13 +56,15 @@ function requestToken(
   plugin: MongoDBOIDCPlugin,
   oidcParams: IdPServerInfo,
   abortSignal?: OIDCAbortSignal,
-  username?: string
+  username?: string,
+  refreshToken?: string
 ): ReturnType<OIDCCallbackFunction> {
   return plugin.mongoClientOptions.authMechanismProperties.OIDC_HUMAN_CALLBACK({
     timeoutContext: abortSignal,
     version: 1,
     idpInfo: oidcParams,
     username,
+    refreshToken,
   });
 }
 
@@ -390,6 +392,7 @@ describe('OIDC plugin (local OIDC provider)', function () {
         // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
         expect(Object.keys(serializedData.state[0][1]).sort()).to.deep.equal([
           'currentTokenSet',
+          'discardingTokenSets',
           'lastIdTokenClaims',
           'serverOIDCMetadata',
         ]);
@@ -1277,6 +1280,93 @@ describe('OIDC plugin (mock OIDC provider)', function () {
     });
   });
 
+  context('when drivers re-request tokens early', function () {
+    let plugin: MongoDBOIDCPlugin;
+
+    beforeEach(function () {
+      plugin = createMongoDBOIDCPlugin({
+        openBrowserTimeout: 60_000,
+        openBrowser: fetchBrowser,
+        allowedFlows: ['auth-code'],
+        redirectURI: 'http://localhost:0/callback',
+      });
+    });
+
+    it('will return a different token even if the existing one is not yet expired', async function () {
+      const result1 = await requestToken(plugin, {
+        issuer: provider.issuer,
+        clientId: 'mockclientid',
+        requestScopes: [],
+      });
+      const result2 = await requestToken(plugin, {
+        issuer: provider.issuer,
+        clientId: 'mockclientid',
+        requestScopes: [],
+      });
+      const result3 = await requestToken(
+        plugin,
+        {
+          issuer: provider.issuer,
+          clientId: 'mockclientid',
+          requestScopes: [],
+        },
+        undefined,
+        undefined,
+        result2.refreshToken
+      );
+      const result4 = await requestToken(
+        plugin,
+        {
+          issuer: provider.issuer,
+          clientId: 'mockclientid',
+          requestScopes: [],
+        },
+        undefined,
+        undefined,
+        result2.refreshToken
+      );
+      expect(result1).to.deep.equal(result2);
+      expect(result2.accessToken).not.to.equal(result3.accessToken);
+      expect(result2.refreshToken).not.to.equal(result3.refreshToken);
+      expect(result3).to.deep.equal(result4);
+    });
+
+    it('will return only one new token per expired token even when called in parallel', async function () {
+      const result1 = await requestToken(plugin, {
+        issuer: provider.issuer,
+        clientId: 'mockclientid',
+        requestScopes: [],
+      });
+      const [result2, result3] = await Promise.all([
+        requestToken(
+          plugin,
+          {
+            issuer: provider.issuer,
+            clientId: 'mockclientid',
+            requestScopes: [],
+          },
+          undefined,
+          undefined,
+          result1.refreshToken
+        ),
+        requestToken(
+          plugin,
+          {
+            issuer: provider.issuer,
+            clientId: 'mockclientid',
+            requestScopes: [],
+          },
+          undefined,
+          undefined,
+          result1.refreshToken
+        ),
+      ]);
+      expect(result1.accessToken).not.to.equal(result2.accessToken);
+      expect(result1.refreshToken).not.to.equal(result2.refreshToken);
+      expect(result2).to.deep.equal(result3);
+    });
+  });
+
   context('HTTP request tracking', function () {
     it('will log all outgoing HTTP requests', async function () {
       const pluginHttpRequests: string[] = [];

src/plugin.ts

@@ -10,6 +10,7 @@ import { MongoDBOIDCError } from './types';
 import {
   errorString,
   getRefreshTokenId,
+  getStableTokenSetId,
   messageFromError,
   normalizeObject,
   throwIfAborted,
@@ -69,6 +70,10 @@ interface UserOIDCAuthState {
   // A cached Client instance that uses the issuer metadata as discovered
   // through serverOIDCMetadata.
   client?: Client;
+  // A set of refresh token IDs which are currently being rejected, i.e.
+  // where the driver has called our callback indicating that the corresponding
+  // access token has become invalid.
+  discardingTokenSets?: string[];
 }
 
 // eslint-disable-next-line @typescript-eslint/consistent-type-imports
@@ -239,6 +244,7 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
           lastIdTokenClaims: serializedState.lastIdTokenClaims
             ? { ...serializedState.lastIdTokenClaims }
             : undefined,
+          discardingTokenSets: serializedState.discardingTokenSets,
         };
         this.updateStateWithTokenSet(
           state,
@@ -274,6 +280,7 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
               lastIdTokenClaims: state.lastIdTokenClaims
                 ? { ...state.lastIdTokenClaims }
                 : undefined,
+              discardingTokenSets: state.discardingTokenSets ?? [],
             },
           ] as const;
         }),
@@ -652,6 +659,7 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
     this.logger.emit('mongodb-oidc-plugin:state-updated', {
       updateId,
       timerDuration,
+      tokenSetId: getStableTokenSetId(tokenSet),
     });
   }
 
@@ -821,7 +829,8 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
 
   private async initiateAuthAttempt(
     state: UserOIDCAuthState,
-    driverAbortSignal?: OIDCAbortSignal
+    driverAbortSignal?: OIDCAbortSignal,
+    { forceRefreshOrReauth = false } = {}
   ): Promise<IdPServerResponse> {
     throwIfAborted(this.options.signal);
     throwIfAborted(driverAbortSignal);
@@ -843,11 +852,12 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
     try {
       get_tokens: {
         if (
+          !forceRefreshOrReauth &&
           tokenExpiryInSeconds(
             state.currentTokenSet?.set,
             passIdTokenAsAccessToken
           ) >
-          5 * 60
+            5 * 60
         ) {
           this.logger.emit('mongodb-oidc-plugin:skip-auth-attempt', {
             reason: 'not-expired',
@@ -915,6 +925,13 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
 
     const { token_type, expires_at, access_token, id_token, refresh_token } =
       state.currentTokenSet.set;
+    const tokenSetId = getStableTokenSetId(state.currentTokenSet.set);
+
+    // We would not want to return the access token or ID token of a token set whose
+    // accompanying refresh token was passed to us by
+    const willRetryWithForceRefreshOrReauth =
+      !forceRefreshOrReauth &&
+      !!state.discardingTokenSets?.includes(tokenSetId);
 
     this.logger.emit('mongodb-oidc-plugin:auth-succeeded', {
       tokenType: token_type ?? null, // DPoP or Bearer
@@ -926,11 +943,20 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
         idToken: id_token,
         refreshToken: refresh_token,
       },
+      tokenSetId,
+      forceRefreshOrReauth,
+      willRetryWithForceRefreshOrReauth,
     });
 
+    if (willRetryWithForceRefreshOrReauth) {
+      return await this.initiateAuthAttempt(state, driverAbortSignal, {
+        forceRefreshOrReauth: true,
+      });
+    }
+
     return {
       accessToken: passIdTokenAsAccessToken ? id_token || '' : access_token,
-      refreshToken: refresh_token,
+      refreshToken: tokenSetId,
       // Passing `expiresInSeconds: 0` results in the driver not caching the token.
       // We perform our own caching here inside the plugin, so interactions with the
       // cache of the driver are not really required or necessarily helpful.
@@ -969,20 +995,39 @@ export class MongoDBOIDCPluginImpl implements MongoDBOIDCPlugin {
       username: params.username,
     });
 
-    if (state.currentAuthAttempt) {
-      return await state.currentAuthAttempt;
+    // If the driver called us with a refresh token, that means that its corresponding
+    // access token has become invalid and we should always return a new one.
+    if (params.refreshToken) {
+      (state.discardingTokenSets ??= []).push(params.refreshToken);
+      this.logger.emit('mongodb-oidc-plugin:discarding-token-set', {
+        tokenSetId: params.refreshToken,
+      });
     }
 
-    const newAuthAttempt = this.initiateAuthAttempt(
-      state,
-      params.timeoutContext
-    );
     try {
-      state.currentAuthAttempt = newAuthAttempt;
-      return await newAuthAttempt;
+      if (state.currentAuthAttempt) {
+        return await state.currentAuthAttempt;
+      }
+
+      const newAuthAttempt = this.initiateAuthAttempt(
+        state,
+        params.timeoutContext
+      );
+      try {
+        state.currentAuthAttempt = newAuthAttempt;
+        return await newAuthAttempt;
+      } finally {
+        if (state.currentAuthAttempt === newAuthAttempt)
+          state.currentAuthAttempt = null;
+      }
     } finally {
-      if (state.currentAuthAttempt === newAuthAttempt)
-        state.currentAuthAttempt = null;
+      if (params.refreshToken) {
+        const index =
+          state.discardingTokenSets?.indexOf(params.refreshToken) ?? -1;
+        if (index > 0) {
+          state.discardingTokenSets?.splice(index, 1);
+        }
+      }
     }
   }
 

src/types.ts

@@ -5,6 +5,7 @@ export interface MongoDBOIDCLogEventsMap {
   }) => void;
   'mongodb-oidc-plugin:state-updated': (event: {
     updateId: number;
+    tokenSetId: string;
     timerDuration: number | undefined;
   }) => void;
   'mongodb-oidc-plugin:local-redirect-accessed': (event: {
@@ -82,6 +83,12 @@ export interface MongoDBOIDCLogEventsMap {
       idToken: string | undefined;
       refreshToken: string | undefined;
     };
+    forceRefreshOrReauth: boolean;
+    willRetryWithForceRefreshOrReauth: boolean;
+    tokenSetId: string;
+  }) => void;
+  'mongodb-oidc-plugin:discarding-token-set': (event: {
+    tokenSetId: string;
   }) => void;
   'mongodb-oidc-plugin:destroyed': () => void;
   'mongodb-oidc-plugin:missing-id-token': () => void;

src/util.ts

@@ -1,3 +1,4 @@
+import type { TokenSet } from 'openid-client';
 import type { OIDCAbortSignal } from './types';
 import { createHash, randomBytes } from 'crypto';
 
@@ -135,3 +136,19 @@ export function getRefreshTokenId(
     'debugid:' + createHash('sha256').update(salt).update(token).digest('hex')
   );
 }
+
+export function getStableTokenSetId(tokenSet: TokenSet): string {
+  const { access_token, id_token, refresh_token, token_type, expires_at } =
+    tokenSet;
+  return createHash('sha256')
+    .update(
+      JSON.stringify({
+        access_token,
+        id_token,
+        refresh_token,
+        token_type,
+        expires_at,
+      })
+    )
+    .digest('hex');
+}