feat(NODE-5036): reauthenticate OIDC and retry

Author: durran

src/cmap/auth/auth_provider.ts

@@ -5,14 +5,20 @@ import type { HandshakeDocument } from '../connect';
 import type { Connection, ConnectionOptions } from '../connection';
 import type { MongoCredentials } from './mongo_credentials';
 
+/** @internal */
 export type AuthContextOptions = ConnectionOptions & ClientMetadataOptions;
 
-/** Context used during authentication */
+/**
+ * Context used during authentication
+ * @internal
+ */
 export class AuthContext {
   /** The connection to authenticate */
   connection: Connection;
   /** The credentials to use for authentication */
   credentials?: MongoCredentials;
+  /** If the context if for reauthentication. */
+  reauthenticating = false;
   /** The options passed to the `connect` method */
   options: AuthContextOptions;
 

src/cmap/auth/mongodb_oidc.ts

@@ -65,7 +65,7 @@ export class MongoDBOIDC extends AuthProvider {
    * Authenticate using OIDC
    */
   override auth(authContext: AuthContext, callback: Callback): void {
-    const { connection, credentials, response } = authContext;
+    const { connection, credentials, response, reauthenticating } = authContext;
 
     if (response?.speculativeAuthenticate) {
       return callback();
@@ -86,7 +86,7 @@ export class MongoDBOIDC extends AuthProvider {
           )
         );
       }
-      workflow.execute(connection, credentials).then(
+      workflow.execute(connection, credentials, reauthenticating).then(
         result => {
           return callback(undefined, result);
         },

src/cmap/auth/mongodb_oidc/callback_workflow.ts

@@ -58,7 +58,11 @@ export class CallbackWorkflow implements Workflow {
    *   - put the new entry in the cache.
    *   - execute step two.
    */
-  async execute(connection: Connection, credentials: MongoCredentials): Promise<Document> {
+  async execute(
+    connection: Connection,
+    credentials: MongoCredentials,
+    reauthenticate = false
+  ): Promise<Document> {
     const request = credentials.mechanismProperties.REQUEST_TOKEN_CALLBACK;
     const refresh = credentials.mechanismProperties.REFRESH_TOKEN_CALLBACK;
 
@@ -69,8 +73,8 @@ export class CallbackWorkflow implements Workflow {
       refresh || null
     );
     if (entry) {
-      // Check if the entry is not expired.
-      if (entry.isValid()) {
+      // Check if the entry is not expired and if we are reauthenticating.
+      if (!reauthenticate && entry.isValid()) {
         // Skip step one and execute the step two saslContinue.
         try {
           const result = await finishAuth(entry.tokenResult, undefined, connection, credentials);

src/cmap/auth/mongodb_oidc/workflow.ts

@@ -8,7 +8,11 @@ export interface Workflow {
    * All device workflows must implement this method in order to get the access
    * token and then call authenticate with it.
    */
-  execute(connection: Connection, credentials: MongoCredentials): Promise<Document>;
+  execute(
+    connection: Connection,
+    credentials: MongoCredentials,
+    reauthenticate?: boolean
+  ): Promise<Document>;
 
   /**
    * Get the document to add for speculative authentication.

src/cmap/auth/scram.ts

@@ -53,8 +53,8 @@ class ScramSHA extends AuthProvider {
   }
 
   override auth(authContext: AuthContext, callback: Callback) {
-    const response = authContext.response;
-    if (response && response.speculativeAuthenticate) {
+    const { reauthenticating, response } = authContext;
+    if (response?.speculativeAuthenticate && !reauthenticating) {
       continueScramConversation(
         this.cryptoMethod,
         response.speculativeAuthenticate,

src/cmap/connect.ts

@@ -36,7 +36,8 @@ import {
   MIN_SUPPORTED_WIRE_VERSION
 } from './wire_protocol/constants';
 
-const AUTH_PROVIDERS = new Map<AuthMechanism | string, AuthProvider>([
+/** @internal */
+export const AUTH_PROVIDERS = new Map<AuthMechanism | string, AuthProvider>([
   [AuthMechanism.MONGODB_AWS, new MongoDBAWS()],
   [AuthMechanism.MONGODB_CR, new MongoCR()],
   [AuthMechanism.MONGODB_GSSAPI, new GSSAPI()],
@@ -117,6 +118,7 @@ function performInitialHandshake(
   }
 
   const authContext = new AuthContext(conn, credentials, options);
+  conn.authContext = authContext;
   prepareHandshakeDocument(authContext, (err, handshakeDoc) => {
     if (err || !handshakeDoc) {
       return callback(err);

src/cmap/connection.ts

@@ -37,6 +37,7 @@ import {
   uuidV4
 } from '../utils';
 import type { WriteConcern } from '../write_concern';
+import type { AuthContext } from './auth/auth_provider';
 import type { MongoCredentials } from './auth/mongo_credentials';
 import {
   CommandFailedEvent,
@@ -127,7 +128,6 @@ export interface ConnectionOptions
   noDelay?: boolean;
   socketTimeoutMS?: number;
   cancellationToken?: CancellationToken;
-
   metadata: ClientMetadata;
 }
 
@@ -165,6 +165,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     cmd: Document,
     options: CommandOptions | undefined
   ) => Promise<Document>;
+  /** @internal */
+  authContext?: AuthContext;
 
   /**@internal */
   [kDelayedTimeoutId]: NodeJS.Timeout | null;

src/cmap/connection_pool.ts

@@ -16,16 +16,18 @@ import {
   CONNECTION_READY
 } from '../constants';
 import {
+  MONGODB_ERROR_CODES,
   MongoError,
   MongoInvalidArgumentError,
+  MongoMissingCredentialsError,
   MongoNetworkError,
   MongoRuntimeError,
   MongoServerError
 } from '../error';
 import { CancellationToken, TypedEventEmitter } from '../mongo_types';
 import type { Server } from '../sdam/server';
 import { Callback, eachAsync, List, makeCounter } from '../utils';
-import { connect } from './connect';
+import { AUTH_PROVIDERS, connect } from './connect';
 import { Connection, ConnectionEvents, ConnectionOptions } from './connection';
 import {
   ConnectionCheckedInEvent,
@@ -544,7 +546,17 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       fn(undefined, conn, (fnErr, result) => {
         if (typeof callback === 'function') {
           if (fnErr) {
-            callback(fnErr);
+            if ((fnErr as MongoError).code === MONGODB_ERROR_CODES.Reauthenticate) {
+              this.reauthenticate(conn, fn, (error, res) => {
+                if (error) {
+                  callback(error);
+                } else {
+                  callback(undefined, res);
+                }
+              });
+            } else {
+              callback(fnErr);
+            }
           } else {
             callback(undefined, result);
           }
@@ -559,7 +571,17 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       fn(err as MongoError, conn, (fnErr, result) => {
         if (typeof callback === 'function') {
           if (fnErr) {
-            callback(fnErr);
+            if (conn && (fnErr as MongoError).code === MONGODB_ERROR_CODES.Reauthenticate) {
+              this.reauthenticate(conn, fn, (error, res) => {
+                if (error) {
+                  callback(error);
+                } else {
+                  callback(undefined, res);
+                }
+              });
+            } else {
+              callback(fnErr);
+            }
           } else {
             callback(undefined, result);
           }
@@ -572,6 +594,50 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     });
   }
 
+  /**
+   * Reauthenticate on the same connection and then retry the operation.
+   */
+  private reauthenticate(
+    connection: Connection,
+    fn: WithConnectionCallback,
+    callback: Callback
+  ): void {
+    const authContext = connection.authContext;
+    if (!authContext) {
+      return callback(new MongoRuntimeError('No auth context found on connection.'));
+    }
+    authContext.reauthenticating = true;
+    const credentials = authContext.credentials;
+    if (!credentials) {
+      return callback(
+        new MongoMissingCredentialsError(
+          'Connection is missing credentials when asked to reauthenticate'
+        )
+      );
+    }
+    const resolvedCredentials = credentials.resolveAuthMechanism(connection.hello || undefined);
+    const provider = AUTH_PROVIDERS.get(resolvedCredentials.mechanism);
+    if (!provider) {
+      return callback(
+        new MongoMissingCredentialsError(
+          `Reauthenticate failed due to no auth provider for ${credentials.mechanism}`
+        )
+      );
+    }
+    provider.auth(authContext, error => {
+      authContext.reauthenticating = false;
+      if (error) {
+        return callback(error);
+      }
+      return fn(undefined, connection, (fnErr, fnResult) => {
+        if (fnErr) {
+          return callback(fnErr);
+        }
+        callback(undefined, fnResult);
+      });
+    });
+  }
+
   /** Clear the min pool size timer */
   private clearMinPoolSizeTimer(): void {
     const minPoolSizeTimer = this[kMinPoolSizeTimer];

src/error.ts

@@ -58,7 +58,8 @@ export const MONGODB_ERROR_CODES = Object.freeze({
   IllegalOperation: 20,
   MaxTimeMSExpired: 50,
   UnknownReplWriteConcern: 79,
-  UnsatisfiableWriteConcern: 100
+  UnsatisfiableWriteConcern: 100,
+  Reauthenticate: 391
 } as const);
 
 // From spec@https://github.com/mongodb/specifications/blob/f93d78191f3db2898a59013a7ed5650352ef6da8/source/change-streams/change-streams.rst#resumable-error

src/index.ts

@@ -197,6 +197,7 @@ export type {
   ResumeToken,
   UpdateDescription
 } from './change_stream';
+export { AuthContext, AuthContextOptions } from './cmap/auth/auth_provider';
 export type {
   AuthMechanismProperties,
   MongoCredentials,