mongodb 6.10.0 → 6.11.0-dev.20241123.sha.32f7ac63

package/src/client-side-encryption/client_encryption.ts

@@ -24,7 +24,8 @@ import { type MongoClient, type MongoClientOptions } from '../mongo_client';
 import { type Filter, type WithId } from '../mongo_types';
 import { type CreateCollectionOptions } from '../operations/create_collection';
 import { type DeleteResult } from '../operations/delete';
-import { MongoDBCollectionNamespace } from '../utils';
+import { type CSOTTimeoutContext, TimeoutContext } from '../timeout';
+import { MongoDBCollectionNamespace, resolveTimeoutOptions } from '../utils';
 import * as cryptoCallbacks from './crypto_callbacks';
 import {
   MongoCryptCreateDataKeyError,
@@ -74,6 +75,8 @@ export class ClientEncryption {
   _tlsOptions: CSFLEKMSTlsOptions;
   /** @internal */
   _kmsProviders: KMSProviders;
+  /** @internal */
+  _timeoutMS?: number;
 
   /** @internal */
   _mongoCrypt: MongoCrypt;
@@ -120,6 +123,8 @@ export class ClientEncryption {
     this._proxyOptions = options.proxyOptions ?? {};
     this._tlsOptions = options.tlsOptions ?? {};
     this._kmsProviders = options.kmsProviders || {};
+    const { timeoutMS } = resolveTimeoutOptions(client, options);
+    this._timeoutMS = timeoutMS;
 
     if (options.keyVaultNamespace == null) {
       throw new MongoCryptInvalidArgumentError('Missing required option `keyVaultNamespace`');
@@ -212,10 +217,16 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
-    const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
+    const timeoutContext =
+      options?.timeoutContext ??
+      TimeoutContext.create(resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS }));
+
+    const dataKey = deserialize(
+      await stateMachine.execute(this, context, timeoutContext)
+    ) as DataKey;
 
     const { db: dbName, collection: collectionName } = MongoDBCollectionNamespace.fromString(
       this._keyVaultNamespace
@@ -224,7 +235,12 @@ export class ClientEncryption {
     const { insertedId } = await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .insertOne(dataKey, { writeConcern: { w: 'majority' } });
+      .insertOne(dataKey, {
+        writeConcern: { w: 'majority' },
+        timeoutMS: timeoutContext?.csotEnabled()
+          ? timeoutContext?.getRemainingTimeMSOrThrow()
+          : undefined
+      });
 
     return insertedId;
   }
@@ -270,10 +286,14 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
-    const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext = TimeoutContext.create(
+      resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS })
+    );
+
+    const { v: dataKeys } = deserialize(await stateMachine.execute(this, context, timeoutContext));
     if (dataKeys.length === 0) {
       return {};
     }
@@ -303,7 +323,8 @@ export class ClientEncryption {
       .db(dbName)
       .collection<DataKey>(collectionName)
       .bulkWrite(replacements, {
-        writeConcern: { w: 'majority' }
+        writeConcern: { w: 'majority' },
+        timeoutMS: timeoutContext.csotEnabled() ? timeoutContext?.remainingTimeMS : undefined
       });
 
     return { bulkWriteResult: result };
@@ -332,7 +353,7 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .deleteOne({ _id }, { writeConcern: { w: 'majority' } });
+      .deleteOne({ _id }, { writeConcern: { w: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -355,7 +376,7 @@ export class ClientEncryption {
     return this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .find({}, { readConcern: { level: 'majority' } });
+      .find({}, { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -381,7 +402,7 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .findOne({ _id }, { readConcern: { level: 'majority' } });
+      .findOne({ _id }, { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -408,7 +429,10 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .findOne({ keyAltNames: keyAltName }, { readConcern: { level: 'majority' } });
+      .findOne(
+        { keyAltNames: keyAltName },
+        { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS }
+      );
   }
 
   /**
@@ -442,7 +466,7 @@ export class ClientEncryption {
       .findOneAndUpdate(
         { _id },
         { $addToSet: { keyAltNames: keyAltName } },
-        { writeConcern: { w: 'majority' }, returnDocument: 'before' }
+        { writeConcern: { w: 'majority' }, returnDocument: 'before', timeoutMS: this._timeoutMS }
       );
 
     return value;
@@ -498,12 +522,14 @@ export class ClientEncryption {
         }
       }
     ];
+
     const value = await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
       .findOneAndUpdate({ _id }, pipeline, {
         writeConcern: { w: 'majority' },
-        returnDocument: 'before'
+        returnDocument: 'before',
+        timeoutMS: this._timeoutMS
       });
 
     return value;
@@ -541,16 +567,25 @@ export class ClientEncryption {
       }
     } = options;
 
+    const timeoutContext =
+      this._timeoutMS != null
+        ? TimeoutContext.create(resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS }))
+        : undefined;
+
     if (Array.isArray(encryptedFields.fields)) {
       const createDataKeyPromises = encryptedFields.fields.map(async field =>
         field == null || typeof field !== 'object' || field.keyId != null
           ? field
           : {
               ...field,
-              keyId: await this.createDataKey(provider, { masterKey })
+              keyId: await this.createDataKey(provider, {
+                masterKey,
+                // clone the timeoutContext
+                // in order to avoid sharing the same timeout for server selection and connection checkout across different concurrent operations
+                timeoutContext: timeoutContext?.csotEnabled() ? timeoutContext?.clone() : undefined
+              })
             }
       );
-
       const createDataKeyResolutions = await Promise.allSettled(createDataKeyPromises);
 
       encryptedFields.fields = createDataKeyResolutions.map((resolution, index) =>
@@ -568,7 +603,10 @@ export class ClientEncryption {
     try {
       const collection = await db.createCollection<TSchema>(name, {
         ...createCollectionOptions,
-        encryptedFields
+        encryptedFields,
+        timeoutMS: timeoutContext?.csotEnabled()
+          ? timeoutContext?.getRemainingTimeMSOrThrow()
+          : undefined
       });
       return { collection, encryptedFields };
     } catch (cause) {
@@ -650,10 +688,15 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
-    const { v } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext =
+      this._timeoutMS != null
+        ? TimeoutContext.create(resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS }))
+        : undefined;
+
+    const { v } = deserialize(await stateMachine.execute(this, context, timeoutContext));
 
     return v;
   }
@@ -729,11 +772,15 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
-    const { v } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext =
+      this._timeoutMS != null
+        ? TimeoutContext.create(resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS }))
+        : undefined;
+    const { v } = deserialize(await stateMachine.execute(this, context, timeoutContext));
     return v;
   }
 }
@@ -818,6 +865,39 @@ export interface ClientEncryptionOptions {
    * TLS options for kms providers to use.
    */
   tlsOptions?: CSFLEKMSTlsOptions;
+
+  /**
+   * @experimental
+   *
+   * The timeout setting to be used for all the operations on ClientEncryption.
+   *
+   * When provided, `timeoutMS` is used as the timeout for each operation executed on
+   * the ClientEncryption object.  For example:
+   *
+   * ```typescript
+   * const clientEncryption = new ClientEncryption(client, {
+   *  timeoutMS: 1_000
+   *  kmsProviders: { local: { key: '<KEY>' } }
+   * });
+   *
+   * // `1_000` is used as the timeout for createDataKey call
+   * await clientEncryption.createDataKey('local');
+   * ```
+   *
+   * If `timeoutMS` is configured on the provided client, the client's `timeoutMS` value
+   * will be used unless `timeoutMS` is also provided as a client encryption option.
+   *
+   * ```typescript
+   * const client = new MongoClient('<uri>', { timeoutMS: 2_000 });
+   *
+   * // timeoutMS is set to 1_000 on clientEncryption
+   * const clientEncryption = new ClientEncryption(client, {
+   *  timeoutMS: 1_000
+   *  kmsProviders: { local: { key: '<KEY>' } }
+   * });
+   * ```
+   */
+  timeoutMS?: number;
 }
 
 /**
@@ -946,6 +1026,9 @@ export interface ClientEncryptionCreateDataKeyProviderOptions {
 
   /** @experimental */
   keyMaterial?: Buffer | Binary;
+
+  /** @internal */
+  timeoutContext?: CSOTTimeoutContext;
 }
 
 /**

package/src/client-side-encryption/state_machine.ts

@@ -11,8 +11,11 @@ import {
   serialize
 } from '../bson';
 import { type ProxyOptions } from '../cmap/connection';
+import { CursorTimeoutContext } from '../cursor/abstract_cursor';
 import { getSocks, type SocksLib } from '../deps';
+import { MongoOperationTimeoutError } from '../error';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
+import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';
 import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
 import { autoSelectSocketOptions, type DataKey } from './client_encryption';
 import { MongoCryptError } from './errors';
@@ -173,6 +176,7 @@ export type StateMachineOptions = {
  * An internal class that executes across a MongoCryptContext until either
  * a finishing state or an error is reached. Do not instantiate directly.
  */
+// TODO(DRIVERS-2671): clarify CSOT behavior for FLE APIs
 export class StateMachine {
   constructor(
     private options: StateMachineOptions,
@@ -182,7 +186,11 @@ export class StateMachine {
   /**
    * Executes the state machine according to the specification
    */
-  async execute(executor: StateMachineExecutable, context: MongoCryptContext): Promise<Uint8Array> {
+  async execute(
+    executor: StateMachineExecutable,
+    context: MongoCryptContext,
+    timeoutContext?: TimeoutContext
+  ): Promise<Uint8Array> {
     const keyVaultNamespace = executor._keyVaultNamespace;
     const keyVaultClient = executor._keyVaultClient;
     const metaDataClient = executor._metaDataClient;
@@ -201,8 +209,13 @@ export class StateMachine {
               'unreachable state machine state: entered MONGOCRYPT_CTX_NEED_MONGO_COLLINFO but metadata client is undefined'
             );
           }
-          const collInfo = await this.fetchCollectionInfo(metaDataClient, context.ns, filter);
 
+          const collInfo = await this.fetchCollectionInfo(
+            metaDataClient,
+            context.ns,
+            filter,
+            timeoutContext
+          );
           if (collInfo) {
             context.addMongoOperationResponse(collInfo);
           }
@@ -222,9 +235,9 @@ export class StateMachine {
           // When we are using the shared library, we don't have a mongocryptd manager.
           const markedCommand: Uint8Array = mongocryptdManager
             ? await mongocryptdManager.withRespawn(
-                this.markCommand.bind(this, mongocryptdClient, context.ns, command)
+                this.markCommand.bind(this, mongocryptdClient, context.ns, command, timeoutContext)
               )
-            : await this.markCommand(mongocryptdClient, context.ns, command);
+            : await this.markCommand(mongocryptdClient, context.ns, command, timeoutContext);
 
           context.addMongoOperationResponse(markedCommand);
           context.finishMongoOperation();
@@ -233,7 +246,12 @@ export class StateMachine {
 
         case MONGOCRYPT_CTX_NEED_MONGO_KEYS: {
           const filter = context.nextMongoOperation();
-          const keys = await this.fetchKeys(keyVaultClient, keyVaultNamespace, filter);
+          const keys = await this.fetchKeys(
+            keyVaultClient,
+            keyVaultNamespace,
+            filter,
+            timeoutContext
+          );
 
           if (keys.length === 0) {
             // See docs on EMPTY_V
@@ -255,9 +273,7 @@ export class StateMachine {
         }
 
         case MONGOCRYPT_CTX_NEED_KMS: {
-          const requests = Array.from(this.requests(context));
-          await Promise.all(requests);
-
+          await Promise.all(this.requests(context, timeoutContext));
           context.finishKMSRequests();
           break;
         }
@@ -299,7 +315,7 @@ export class StateMachine {
    * @param kmsContext - A C++ KMS context returned from the bindings
    * @returns A promise that resolves when the KMS reply has be fully parsed
    */
-  async kmsRequest(request: MongoCryptKMSRequest): Promise<void> {
+  async kmsRequest(request: MongoCryptKMSRequest, timeoutContext?: TimeoutContext): Promise<void> {
     const parsedUrl = request.endpoint.split(':');
     const port = parsedUrl[1] != null ? Number.parseInt(parsedUrl[1], 10) : HTTPS_PORT;
     const socketOptions = autoSelectSocketOptions(this.options.socketOptions || {});
@@ -329,10 +345,6 @@ export class StateMachine {
       }
     }
 
-    function ontimeout() {
-      return new MongoCryptError('KMS request timed out');
-    }
-
     function onerror(cause: Error) {
       return new MongoCryptError('KMS request failed', { cause });
     }
@@ -364,7 +376,6 @@ export class StateMachine {
       resolve: resolveOnNetSocketConnect
     } = promiseWithResolvers<void>();
     netSocket
-      .once('timeout', () => rejectOnNetSocketError(ontimeout()))
       .once('error', err => rejectOnNetSocketError(onerror(err)))
       .once('close', () => rejectOnNetSocketError(onclose()))
       .once('connect', () => resolveOnNetSocketConnect());
@@ -410,8 +421,8 @@ export class StateMachine {
         reject: rejectOnTlsSocketError,
         resolve
       } = promiseWithResolvers<void>();
+
       socket
-        .once('timeout', () => rejectOnTlsSocketError(ontimeout()))
         .once('error', err => rejectOnTlsSocketError(onerror(err)))
         .once('close', () => rejectOnTlsSocketError(onclose()))
         .on('data', data => {
@@ -425,20 +436,26 @@ export class StateMachine {
             resolve();
           }
         });
-      await willResolveKmsRequest;
+      await (timeoutContext?.csotEnabled()
+        ? Promise.all([willResolveKmsRequest, Timeout.expires(timeoutContext?.remainingTimeMS)])
+        : willResolveKmsRequest);
+    } catch (error) {
+      if (error instanceof TimeoutError)
+        throw new MongoOperationTimeoutError('KMS request timed out');
+      throw error;
     } finally {
       // There's no need for any more activity on this socket at this point.
       destroySockets();
     }
   }
 
-  *requests(context: MongoCryptContext) {
+  *requests(context: MongoCryptContext, timeoutContext?: TimeoutContext) {
     for (
       let request = context.nextKMSRequest();
       request != null;
       request = context.nextKMSRequest()
     ) {
-      yield this.kmsRequest(request);
+      yield this.kmsRequest(request, timeoutContext);
     }
   }
 
@@ -498,17 +515,21 @@ export class StateMachine {
   async fetchCollectionInfo(
     client: MongoClient,
     ns: string,
-    filter: Document
+    filter: Document,
+    timeoutContext?: TimeoutContext
   ): Promise<Uint8Array | null> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
 
-    const collections = await client
-      .db(db)
-      .listCollections(filter, {
-        promoteLongs: false,
-        promoteValues: false
-      })
-      .toArray();
+    const cursor = client.db(db).listCollections(filter, {
+      promoteLongs: false,
+      promoteValues: false,
+      timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
+    });
+
+    // There is always exactly zero or one matching documents, so this should always exhaust the cursor
+    // in a single batch.  We call `toArray()` just to be safe and ensure that the cursor is always
+    // exhausted and closed.
+    const collections = await cursor.toArray();
 
     const info = collections.length > 0 ? serialize(collections[0]) : null;
     return info;
@@ -522,12 +543,22 @@ export class StateMachine {
    * @param command - The command to execute.
    * @param callback - Invoked with the serialized and marked bson command, or with an error
    */
-  async markCommand(client: MongoClient, ns: string, command: Uint8Array): Promise<Uint8Array> {
-    const options = { promoteLongs: false, promoteValues: false };
+  async markCommand(
+    client: MongoClient,
+    ns: string,
+    command: Uint8Array,
+    timeoutContext?: TimeoutContext
+  ): Promise<Uint8Array> {
     const { db } = MongoDBCollectionNamespace.fromString(ns);
-    const rawCommand = deserialize(command, options);
+    const bsonOptions = { promoteLongs: false, promoteValues: false };
+    const rawCommand = deserialize(command, bsonOptions);
 
-    const response = await client.db(db).command(rawCommand, options);
+    const response = await client.db(db).command(rawCommand, {
+      ...bsonOptions,
+      ...(timeoutContext?.csotEnabled()
+        ? { timeoutMS: timeoutContext?.remainingTimeMS }
+        : undefined)
+    });
 
     return serialize(response, this.bsonOptions);
   }
@@ -543,7 +574,8 @@ export class StateMachine {
   fetchKeys(
     client: MongoClient,
     keyVaultNamespace: string,
-    filter: Uint8Array
+    filter: Uint8Array,
+    timeoutContext?: TimeoutContext
   ): Promise<Array<DataKey>> {
     const { db: dbName, collection: collectionName } =
       MongoDBCollectionNamespace.fromString(keyVaultNamespace);
@@ -551,7 +583,9 @@ export class StateMachine {
     return client
       .db(dbName)
       .collection<DataKey>(collectionName, { readConcern: { level: 'majority' } })
-      .find(deserialize(filter))
+      .find(deserialize(filter), {
+        timeoutContext: timeoutContext && new CursorTimeoutContext(timeoutContext, Symbol())
+      })
       .toArray();
   }
 }

package/src/cmap/auth/mongo_credentials.ts

@@ -10,7 +10,9 @@ import { GSSAPICanonicalizationValue } from './gssapi';
 import type { OIDCCallbackFunction } from './mongodb_oidc';
 import { AUTH_MECHS_AUTH_SRC_EXTERNAL, AuthMechanism } from './providers';
 
-// https://github.com/mongodb/specifications/blob/master/source/auth/auth.rst
+/**
+ * @see https://github.com/mongodb/specifications/blob/master/source/auth/auth.md
+ */
 function getDefaultAuthMechanism(hello: Document | null): AuthMechanism {
   if (hello) {
     // If hello contains saslSupportedMechs, use scram-sha-256
@@ -29,7 +31,8 @@ function getDefaultAuthMechanism(hello: Document | null): AuthMechanism {
 const ALLOWED_ENVIRONMENT_NAMES: AuthMechanismProperties['ENVIRONMENT'][] = [
   'test',
   'azure',
-  'gcp'
+  'gcp',
+  'k8s'
 ];
 const ALLOWED_HOSTS_ERROR = 'Auth mechanism property ALLOWED_HOSTS must be an array of strings.';
 
@@ -60,7 +63,7 @@ export interface AuthMechanismProperties extends Document {
   /** A user provided OIDC human interacted callback function. */
   OIDC_HUMAN_CALLBACK?: OIDCCallbackFunction;
   /** The OIDC environment. Note that 'test' is for internal use only. */
-  ENVIRONMENT?: 'test' | 'azure' | 'gcp';
+  ENVIRONMENT?: 'test' | 'azure' | 'gcp' | 'k8s';
   /** Allowed hosts that OIDC auth can connect to. */
   ALLOWED_HOSTS?: string[];
   /** The resource token for OIDC auth in Azure and GCP. */

package/src/cmap/auth/mongodb_aws.ts

@@ -107,7 +107,7 @@ export class MongoDBAWS extends AuthProvider {
 
     if (!ByteUtils.equals(serverNonce.subarray(0, nonce.byteLength), nonce)) {
       // throw because the serverNonce's leading 32 bytes must equal the client nonce's 32 bytes
-      // https://github.com/mongodb/specifications/blob/875446db44aade414011731840831f38a6c668df/source/auth/auth.rst#id11
+      // https://github.com/mongodb/specifications/blob/master/source/auth/auth.md#conversation-5
 
       // TODO(NODE-3483)
       throw new MongoRuntimeError('Server nonce does not begin with client nonce');

package/src/cmap/auth/mongodb_oidc/k8s_machine_workflow.ts

@@ -0,0 +1,38 @@
+import { readFile } from 'fs/promises';
+
+import { type AccessToken, MachineWorkflow } from './machine_workflow';
+import { type TokenCache } from './token_cache';
+
+/** The fallback file name */
+const FALLBACK_FILENAME = '/var/run/secrets/kubernetes.io/serviceaccount/token';
+
+/** The azure environment variable for the file name. */
+const AZURE_FILENAME = 'AZURE_FEDERATED_TOKEN_FILE';
+
+/** The AWS environment variable for the file name. */
+const AWS_FILENAME = 'AWS_WEB_IDENTITY_TOKEN_FILE';
+
+export class K8SMachineWorkflow extends MachineWorkflow {
+  /**
+   * Instantiate the machine workflow.
+   */
+  constructor(cache: TokenCache) {
+    super(cache);
+  }
+
+  /**
+   * Get the token from the environment.
+   */
+  async getToken(): Promise<AccessToken> {
+    let filename: string;
+    if (process.env[AZURE_FILENAME]) {
+      filename = process.env[AZURE_FILENAME];
+    } else if (process.env[AWS_FILENAME]) {
+      filename = process.env[AWS_FILENAME];
+    } else {
+      filename = FALLBACK_FILENAME;
+    }
+    const token = await readFile(filename, 'utf8');
+    return { access_token: token };
+  }
+}

package/src/cmap/auth/mongodb_oidc.ts

@@ -6,6 +6,7 @@ import { type AuthContext, AuthProvider } from './auth_provider';
 import type { MongoCredentials } from './mongo_credentials';
 import { AzureMachineWorkflow } from './mongodb_oidc/azure_machine_workflow';
 import { GCPMachineWorkflow } from './mongodb_oidc/gcp_machine_workflow';
+import { K8SMachineWorkflow } from './mongodb_oidc/k8s_machine_workflow';
 import { TokenCache } from './mongodb_oidc/token_cache';
 import { TokenMachineWorkflow } from './mongodb_oidc/token_machine_workflow';
 
@@ -88,7 +89,7 @@ export type OIDCCallbackFunction = (params: OIDCCallbackParams) => Promise<OIDCR
 /** The current version of OIDC implementation. */
 export const OIDC_VERSION = 1;
 
-type EnvironmentName = 'test' | 'azure' | 'gcp' | undefined;
+type EnvironmentName = 'test' | 'azure' | 'gcp' | 'k8s' | undefined;
 
 /** @internal */
 export interface Workflow {
@@ -118,6 +119,7 @@ export const OIDC_WORKFLOWS: Map<EnvironmentName, () => Workflow> = new Map();
 OIDC_WORKFLOWS.set('test', () => new TokenMachineWorkflow(new TokenCache()));
 OIDC_WORKFLOWS.set('azure', () => new AzureMachineWorkflow(new TokenCache()));
 OIDC_WORKFLOWS.set('gcp', () => new GCPMachineWorkflow(new TokenCache()));
+OIDC_WORKFLOWS.set('k8s', () => new K8SMachineWorkflow(new TokenCache()));
 
 /**
  * OIDC auth provider.

package/src/cmap/connect.ts

@@ -113,7 +113,8 @@ export async function performInitialHandshake(
   }
 
   const start = new Date().getTime();
-  const response = await conn.command(ns('admin.$cmd'), handshakeDoc, handshakeOptions);
+
+  const response = await executeHandshake(handshakeDoc, handshakeOptions);
 
   if (!('isWritablePrimary' in response)) {
     // Provide hello-style response document.
@@ -175,6 +176,22 @@ export async function performInitialHandshake(
   // Connection establishment is socket creation (tcp handshake, tls handshake, MongoDB handshake (saslStart, saslContinue))
   // Once connection is established, command logging can log events (if enabled)
   conn.established = true;
+
+  async function executeHandshake(handshakeDoc: Document, handshakeOptions: CommandOptions) {
+    try {
+      const handshakeResponse = await conn.command(
+        ns('admin.$cmd'),
+        handshakeDoc,
+        handshakeOptions
+      );
+      return handshakeResponse;
+    } catch (error) {
+      if (error instanceof MongoError) {
+        error.addErrorLabel(MongoErrorLabel.HandshakeError);
+      }
+      throw error;
+    }
+  }
 }
 
 /**