feat(NODE-6391): Add timeoutMS support to explicit encryption

src/client-side-encryption/client_encryption.ts

@@ -24,7 +24,7 @@ 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 { TimeoutContext } from '../timeout';
+import { type CSOTTimeoutContext, TimeoutContext } from '../timeout';
 import { MongoDBCollectionNamespace, resolveTimeoutOptions } from '../utils';
 import * as cryptoCallbacks from './crypto_callbacks';
 import {
@@ -220,7 +220,13 @@ export class ClientEncryption {
       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
@@ -229,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;
   }
@@ -511,6 +522,7 @@ export class ClientEncryption {
         }
       }
     ];
+
     const value = await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
@@ -555,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) =>
@@ -582,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) {
@@ -667,7 +691,12 @@ export class ClientEncryption {
       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;
   }
@@ -747,7 +776,11 @@ export class ClientEncryption {
     });
     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;
   }
 }
@@ -833,7 +866,8 @@ export interface ClientEncryptionOptions {
    */
   tlsOptions?: CSFLEKMSTlsOptions;
 
-  /**
+  /** @internal TODO(NODE-5688): make this public
+   *
    * The timeout setting to be used for all the operations on ClientEncryption.
    */
   timeoutMS?: number;
@@ -965,6 +999,9 @@ export interface ClientEncryptionCreateDataKeyProviderOptions {
 
   /** @experimental */
   keyMaterial?: Buffer | Binary;
+
+  /** @internal */
+  timeoutContext?: CSOTTimeoutContext;
 }
 
 /**

src/timeout.ts

@@ -323,6 +323,20 @@ export class CSOTTimeoutContext extends TimeoutContext {
     return remainingTimeMS;
   }
 
+  /**
+   * @internal
+   * This method is intended to be used in situations where concurrent operation are on the same deadline, but cannot share a single `TimeoutContext` instance.
+   * Returns a new instance of `CSOTTimeoutContext` constructed with identical options, but setting the `start` property to `this.start`.
+   */
+  clone(): CSOTTimeoutContext {
+    const timeoutContext = new CSOTTimeoutContext({
+      timeoutMS: this.timeoutMS,
+      serverSelectionTimeoutMS: this.serverSelectionTimeoutMS
+    });
+    timeoutContext.start = this.start;
+    return timeoutContext;
+  }
+
   override refreshed(): CSOTTimeoutContext {
     return new CSOTTimeoutContext(this);
   }

test/integration/client-side-encryption/driver.test.ts

@@ -12,6 +12,8 @@ import {
   Connection,
   CSOTTimeoutContext,
   type MongoClient,
+  MongoCryptCreateDataKeyError,
+  MongoCryptCreateEncryptedCollectionError,
   MongoOperationTimeoutError,
   StateMachine
 } from '../../mongodb';
@@ -1050,4 +1052,165 @@ describe('CSOT', function () {
       );
     });
   });
+
+  describe('Explicit Encryption', function () {
+    describe('#createEncryptedCollection', function () {
+      let client: MongoClient;
+      let clientEncryption: ClientEncryption;
+      let local_key;
+      const timeoutMS = 1000;
+
+      const encryptedCollectionMetadata: MongoDBMetadataUI = {
+        requires: {
+          clientSideEncryption: true,
+          mongodb: '>=7.0.0',
+          topology: '!single'
+        }
+      };
+
+      beforeEach(async function () {
+        local_key = { local: EJSON.parse(process.env.CSFLE_KMS_PROVIDERS).local };
+        client = this.configuration.newClient({ timeoutMS });
+        await client.connect();
+        await client.db('keyvault').createCollection('datakeys');
+        clientEncryption = new ClientEncryption(client, {
+          keyVaultNamespace: 'keyvault.datakeys',
+          keyVaultClient: client,
+          kmsProviders: local_key
+        });
+      });
+
+      afterEach(async function () {
+        await client
+          .db()
+          .admin()
+          .command({
+            configureFailPoint: 'failCommand',
+            mode: 'off'
+          } as FailPoint);
+        await client
+          .db('db')
+          .collection('newnew')
+          .drop()
+          .catch(() => null);
+        await client
+          .db('keyvault')
+          .collection('datakeys')
+          .drop()
+          .catch(() => null);
+        await client.close();
+      });
+
+      async function runCreateEncryptedCollection() {
+        const createCollectionOptions = {
+          encryptedFields: { fields: [{ path: 'ssn', bsonType: 'string', keyId: null }] }
+        };
+
+        const db = client.db('db');
+
+        return await measureDuration(() =>
+          clientEncryption
+            .createEncryptedCollection(db, 'newnew', {
+              provider: 'local',
+              createCollectionOptions,
+              masterKey: null
+            })
+            .catch(err => err)
+        );
+      }
+
+      context(
+        'when `createDataKey` hangs longer than timeoutMS and `createCollection` does not hang',
+        () => {
+          it(
+            '`createEncryptedCollection throws `MongoCryptCreateDataKeyError` due to a timeout error',
+            encryptedCollectionMetadata,
+            async function () {
+              await client
+                .db()
+                .admin()
+                .command({
+                  configureFailPoint: 'failCommand',
+                  mode: {
+                    times: 1
+                  },
+                  data: {
+                    failCommands: ['insert'],
+                    blockConnection: true,
+                    blockTimeMS: timeoutMS * 1.2
+                  }
+                } as FailPoint);
+
+              const { duration, result: err } = await runCreateEncryptedCollection();
+              expect(err).to.be.instanceOf(MongoCryptCreateDataKeyError);
+              expect(err.cause).to.be.instanceOf(MongoOperationTimeoutError);
+              expect(duration).to.be.within(timeoutMS - 100, timeoutMS + 100);
+            }
+          );
+        }
+      );
+
+      context(
+        'when `createDataKey` does not hang and `createCollection` hangs longer than timeoutMS',
+        () => {
+          it(
+            '`createEncryptedCollection throws `MongoCryptCreateEncryptedCollectionError` due to a timeout error',
+            encryptedCollectionMetadata,
+            async function () {
+              await client
+                .db()
+                .admin()
+                .command({
+                  configureFailPoint: 'failCommand',
+                  mode: {
+                    times: 1
+                  },
+                  data: {
+                    failCommands: ['create'],
+                    blockConnection: true,
+                    blockTimeMS: timeoutMS * 1.2
+                  }
+                } as FailPoint);
+
+              const { duration, result: err } = await runCreateEncryptedCollection();
+              expect(err).to.be.instanceOf(MongoCryptCreateEncryptedCollectionError);
+              expect(err.cause).to.be.instanceOf(MongoOperationTimeoutError);
+              expect(duration).to.be.within(timeoutMS - 100, timeoutMS + 100);
+            }
+          );
+        }
+      );
+
+      context(
+        'when `createDataKey` and `createCollection` cumulatively hang longer than timeoutMS',
+        () => {
+          it(
+            '`createEncryptedCollection throws `MongoCryptCreateEncryptedCollectionError` due to a timeout error',
+            encryptedCollectionMetadata,
+            async function () {
+              await client
+                .db()
+                .admin()
+                .command({
+                  configureFailPoint: 'failCommand',
+                  mode: {
+                    times: 2
+                  },
+                  data: {
+                    failCommands: ['insert', 'create'],
+                    blockConnection: true,
+                    blockTimeMS: timeoutMS * 0.6
+                  }
+                } as FailPoint);
+
+              const { duration, result: err } = await runCreateEncryptedCollection();
+              expect(err).to.be.instanceOf(MongoCryptCreateEncryptedCollectionError);
+              expect(err.cause).to.be.instanceOf(MongoOperationTimeoutError);
+              expect(duration).to.be.within(timeoutMS - 100, timeoutMS + 100);
+            }
+          );
+        }
+      );
+    });
+  });
 });