feat(NODE-5036): reauthenticate OIDC and retry (#3589)

durran · web-flow · commit a41846db3d76 · 2023-03-17T15:34:21.000-04:00
diff --git a/.evergreen/setup-oidc-roles.sh b/.evergreen/setup-oidc-roles.sh
@@ -5,4 +5,4 @@ set -o xtrace   # Write all commands first to stderr
 cd ${DRIVERS_TOOLS}/.evergreen/auth_oidc
 . ./activate-authoidcvenv.sh
 
-${DRIVERS_TOOLS}/mongodb/bin/mongosh setup_oidc.js
+${DRIVERS_TOOLS}/mongodb/bin/mongosh "mongodb://localhost:27017,localhost:27018/?replicaSet=oidc-repl0&readPreference=primary" setup_oidc.js
diff --git a/src/cmap/auth/auth_provider.ts b/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 is for reauthentication. */
+  reauthenticating = false;
   /** The options passed to the `connect` method */
   options: AuthContextOptions;
 
@@ -57,4 +63,22 @@ export class AuthProvider {
     // TODO(NODE-3483): Replace this with MongoMethodOverrideError
     callback(new MongoRuntimeError('`auth` method must be overridden by subclass'));
   }
+
+  /**
+   * Reauthenticate.
+   * @param context - The shared auth context.
+   * @param callback - The callback.
+   */
+  reauth(context: AuthContext, callback: Callback): void {
+    // If we are already reauthenticating this is a no-op.
+    if (context.reauthenticating) {
+      return callback(new MongoRuntimeError('Reauthentication already in progress.'));
+    }
+    context.reauthenticating = true;
+    const cb: Callback = (error, result) => {
+      context.reauthenticating = false;
+      callback(error, result);
+    };
+    this.auth(context, cb);
+  }
 }
diff --git a/src/cmap/auth/mongo_credentials.ts b/src/cmap/auth/mongo_credentials.ts
@@ -37,8 +37,11 @@ export interface AuthMechanismProperties extends Document {
   SERVICE_REALM?: string;
   CANONICALIZE_HOST_NAME?: GSSAPICanonicalizationValue;
   AWS_SESSION_TOKEN?: string;
+  /** @experimental */
   REQUEST_TOKEN_CALLBACK?: OIDCRequestFunction;
+  /** @experimental */
   REFRESH_TOKEN_CALLBACK?: OIDCRefreshFunction;
+  /** @experimental */
   PROVIDER_NAME?: 'aws';
 }
 
diff --git a/src/cmap/auth/mongodb_oidc.ts b/src/cmap/auth/mongodb_oidc.ts
@@ -11,7 +11,10 @@ import { AwsServiceWorkflow } from './mongodb_oidc/aws_service_workflow';
 import { CallbackWorkflow } from './mongodb_oidc/callback_workflow';
 import type { Workflow } from './mongodb_oidc/workflow';
 
-/** @public */
+/**
+ * @public
+ * @experimental
+ */
 export interface OIDCMechanismServerStep1 {
   authorizationEndpoint?: string;
   tokenEndpoint?: string;
@@ -21,21 +24,30 @@ export interface OIDCMechanismServerStep1 {
   requestScopes?: string[];
 }
 
-/** @public */
+/**
+ * @public
+ * @experimental
+ */
 export interface OIDCRequestTokenResult {
   accessToken: string;
   expiresInSeconds?: number;
   refreshToken?: string;
 }
 
-/** @public */
+/**
+ * @public
+ * @experimental
+ */
 export type OIDCRequestFunction = (
   principalName: string,
   serverResult: OIDCMechanismServerStep1,
   timeout: AbortSignal | number
 ) => Promise<OIDCRequestTokenResult>;
 
-/** @public */
+/**
+ * @public
+ * @experimental
+ */
 export type OIDCRefreshFunction = (
   principalName: string,
   serverResult: OIDCMechanismServerStep1,
@@ -52,6 +64,7 @@ OIDC_WORKFLOWS.set('aws', new AwsServiceWorkflow());
 
 /**
  * OIDC auth provider.
+ * @experimental
  */
 export class MongoDBOIDC extends AuthProvider {
   /**
@@ -65,7 +78,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 +99,7 @@ export class MongoDBOIDC extends AuthProvider {
           )
         );
       }
-      workflow.execute(connection, credentials).then(
+      workflow.execute(connection, credentials, reauthenticating).then(
         result => {
           return callback(undefined, result);
         },
diff --git a/src/cmap/auth/mongodb_oidc/callback_workflow.ts b/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);
diff --git a/src/cmap/auth/mongodb_oidc/workflow.ts b/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.
diff --git a/src/cmap/auth/providers.ts b/src/cmap/auth/providers.ts
@@ -8,6 +8,7 @@ export const AuthMechanism = Object.freeze({
   MONGODB_SCRAM_SHA1: 'SCRAM-SHA-1',
   MONGODB_SCRAM_SHA256: 'SCRAM-SHA-256',
   MONGODB_X509: 'MONGODB-X509',
+  /** @experimental */
   MONGODB_OIDC: 'MONGODB-OIDC'
 } as const);
 
diff --git a/src/cmap/auth/scram.ts b/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,
diff --git a/src/cmap/connect.ts b/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);
diff --git a/src/cmap/connection.ts b/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,
@@ -126,7 +127,6 @@ export interface ConnectionOptions
   noDelay?: boolean;
   socketTimeoutMS?: number;
   cancellationToken?: CancellationToken;
-
   metadata: ClientMetadata;
 }
 
@@ -164,6 +164,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     cmd: Document,
     options: CommandOptions | undefined
   ) => Promise<Document>;
+  /** @internal */
+  authContext?: AuthContext;
 
   /**@internal */
   [kDelayedTimeoutId]: NodeJS.Timeout | null;
diff --git a/src/cmap/connection_pool.ts b/src/cmap/connection_pool.ts
@@ -16,16 +16,19 @@ import {
   CONNECTION_READY
 } from '../constants';
 import {
+  AnyError,
+  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,
@@ -537,32 +540,30 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
   withConnection(
     conn: Connection | undefined,
     fn: WithConnectionCallback,
-    callback?: Callback<Connection>
+    callback: Callback<Connection>
   ): void {
     if (conn) {
       // use the provided connection, and do _not_ check it in after execution
       fn(undefined, conn, (fnErr, result) => {
-        if (typeof callback === 'function') {
-          if (fnErr) {
-            callback(fnErr);
-          } else {
-            callback(undefined, result);
-          }
+        if (fnErr) {
+          return this.withReauthentication(fnErr, conn, fn, callback);
         }
+        callback(undefined, result);
       });
-
       return;
     }
 
     this.checkOut((err, conn) => {
       // don't callback with `err` here, we might want to act upon it inside `fn`
       fn(err as MongoError, conn, (fnErr, result) => {
-        if (typeof callback === 'function') {
-          if (fnErr) {
-            callback(fnErr);
+        if (fnErr) {
+          if (conn) {
+            this.withReauthentication(fnErr, conn, fn, callback);
           } else {
-            callback(undefined, result);
+            callback(fnErr);
           }
+        } else {
+          callback(undefined, result);
         }
 
         if (conn) {
@@ -572,6 +573,66 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     });
   }
 
+  private withReauthentication(
+    fnErr: AnyError,
+    conn: Connection,
+    fn: WithConnectionCallback,
+    callback: Callback<Connection>
+  ) {
+    if (fnErr instanceof MongoError && fnErr.code === MONGODB_ERROR_CODES.Reauthenticate) {
+      this.reauthenticate(conn, fn, (error, res) => {
+        if (error) {
+          return callback(error);
+        }
+        callback(undefined, res);
+      });
+    } else {
+      callback(fnErr);
+    }
+  }
+
+  /**
+   * 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.'));
+    }
+    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.reauth(authContext, error => {
+      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];
diff --git a/src/error.ts b/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
diff --git a/src/index.ts b/src/index.ts
@@ -197,6 +197,7 @@ export type {
   ResumeToken,
   UpdateDescription
 } from './change_stream';
+export type { AuthContext, AuthContextOptions } from './cmap/auth/auth_provider';
 export type {
   AuthMechanismProperties,
   MongoCredentials,
diff --git a/test/integration/auth/auth.spec.test.ts b/test/integration/auth/auth.spec.test.ts
@@ -0,0 +1,8 @@
+import * as path from 'path';
+
+import { loadSpecTests } from '../../spec';
+import { runUnifiedSuite } from '../../tools/unified-spec-runner/runner';
+
+describe('Auth (unified)', function () {
+  runUnifiedSuite(loadSpecTests(path.join('auth', 'unified')));
+});
diff --git a/test/manual/mongodb_oidc.prose.test.ts b/test/manual/mongodb_oidc.prose.test.ts
diff --git a/test/spec/auth/unified/reauthenticate_with_retry.json b/test/spec/auth/unified/reauthenticate_with_retry.json
diff --git a/test/spec/auth/unified/reauthenticate_with_retry.yml b/test/spec/auth/unified/reauthenticate_with_retry.yml
diff --git a/test/spec/auth/unified/reauthenticate_without_retry.json b/test/spec/auth/unified/reauthenticate_without_retry.json
diff --git a/test/spec/auth/unified/reauthenticate_without_retry.yml b/test/spec/auth/unified/reauthenticate_without_retry.yml
diff --git a/test/unit/cmap/auth/auth_provider.test.ts b/test/unit/cmap/auth/auth_provider.test.ts
diff --git a/test/unit/cmap/connection_pool.test.js b/test/unit/cmap/connection_pool.test.js
