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

src/client-side-encryption/client_encryption.ts

@@ -24,6 +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 { CSOTTimeoutContext } from '../timeout';
 import { MongoDBCollectionNamespace } from '../utils';
 import * as cryptoCallbacks from './crypto_callbacks';
 import {
@@ -74,6 +75,8 @@ export class ClientEncryption {
   _tlsOptions: CSFLEKMSTlsOptions;
   /** @internal */
   _kmsProviders: KMSProviders;
+  /** @internal */
+  _timeoutMS?: number;
 
   /** @internal */
   _mongoCrypt: MongoCrypt;
@@ -120,6 +123,7 @@ export class ClientEncryption {
     this._proxyOptions = options.proxyOptions ?? {};
     this._tlsOptions = options.tlsOptions ?? {};
     this._kmsProviders = options.kmsProviders || {};
+    this._timeoutMS = options.timeoutMS ?? client.options.timeoutMS;
 
     if (options.keyVaultNamespace == null) {
       throw new MongoCryptInvalidArgumentError('Missing required option `keyVaultNamespace`');
@@ -215,7 +219,17 @@ export class ClientEncryption {
       socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
-    const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
+    const timeoutContext = options?.timeoutContext
+      ? options?.timeoutContext
+      : this._timeoutMS
+        ? new CSOTTimeoutContext({
+            timeoutMS: this._timeoutMS,
+            serverSelectionTimeoutMS: this._client.options.serverSelectionTimeoutMS
+          })
+        : undefined;
+    const dataKey = deserialize(
+      await stateMachine.execute(this, context, timeoutContext)
+    ) as DataKey;
 
     const { db: dbName, collection: collectionName } = MongoDBCollectionNamespace.fromString(
       this._keyVaultNamespace
@@ -224,7 +238,10 @@ 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?.remainingTimeMS : undefined
+      });
 
     return insertedId;
   }
@@ -498,6 +515,7 @@ export class ClientEncryption {
         }
       }
     ];
+
     const value = await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
@@ -541,13 +559,20 @@ export class ClientEncryption {
       }
     } = options;
 
+    const timeoutContext = this._timeoutMS
+      ? new CSOTTimeoutContext({
+          timeoutMS: this._timeoutMS,
+          serverSelectionTimeoutMS: this._client.options.serverSelectionTimeoutMS
+        })
+      : 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, timeoutContext })
             }
       );
 
@@ -568,7 +593,8 @@ export class ClientEncryption {
     try {
       const collection = await db.createCollection<TSchema>(name, {
         ...createCollectionOptions,
-        encryptedFields
+        encryptedFields,
+        timeoutMS: timeoutContext?.remainingTimeMS
       });
       return { collection, encryptedFields };
     } catch (cause) {
@@ -653,7 +679,13 @@ export class ClientEncryption {
       socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
-    const { v } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext = this._timeoutMS
+      ? new CSOTTimeoutContext({
+          timeoutMS: this._timeoutMS,
+          serverSelectionTimeoutMS: this._client.options.serverSelectionTimeoutMS
+        })
+      : undefined;
+    const { v } = deserialize(await stateMachine.execute(this, context, timeoutContext));
 
     return v;
   }
@@ -733,7 +765,13 @@ export class ClientEncryption {
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
-    const { v } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext = this._timeoutMS
+      ? new CSOTTimeoutContext({
+          timeoutMS: this._timeoutMS,
+          serverSelectionTimeoutMS: this._client.options.serverSelectionTimeoutMS
+        })
+      : undefined;
+    const { v } = deserialize(await stateMachine.execute(this, context, timeoutContext));
     return v;
   }
 }
@@ -818,6 +856,9 @@ export interface ClientEncryptionOptions {
    * TLS options for kms providers to use.
    */
   tlsOptions?: CSFLEKMSTlsOptions;
+
+  /** @internal TODO(NODE-5688): make this public */
+  timeoutMS?: number;
 }
 
 /**
@@ -946,6 +987,9 @@ export interface ClientEncryptionCreateDataKeyProviderOptions {
 
   /** @experimental */
   keyMaterial?: Buffer | Binary;
+
+  /** @internal */
+  timeoutContext?: CSOTTimeoutContext;
 }
 
 /**

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

@@ -4,9 +4,14 @@ import * as crypto from 'crypto';
 
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
 import { ClientEncryption } from '../../../src/client-side-encryption/client_encryption';
-import { type Collection, type CommandStartedEvent, type MongoClient } from '../../mongodb';
+import {
+  type Collection,
+  type CommandStartedEvent,
+  type MongoClient,
+  MongoCryptCreateEncryptedCollectionError
+} from '../../mongodb';
 import * as BSON from '../../mongodb';
-import { getEncryptExtraOptions } from '../../tools/utils';
+import { type FailPoint, getEncryptExtraOptions } from '../../tools/utils';
 
 const metadata = {
   requires: {
@@ -471,3 +476,85 @@ describe('Range Explicit Encryption with JS native types', function () {
     });
   });
 });
+
+describe('CSOT', function () {
+  describe('Explicit Encryption', function () {
+    describe('#createEncryptedCollection', function () {
+      let keyVaultClient: MongoClient;
+      let internalClient: MongoClient;
+      let clientEncryption: ClientEncryption;
+      const LOCAL_KEY = Buffer.from(
+        'Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN3YkQ5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk',
+        'base64'
+      );
+
+      beforeEach(async function () {
+        internalClient = this.configuration.newClient();
+        await internalClient.connect();
+        await internalClient.db('keyvault').createCollection('datakeys');
+        keyVaultClient = this.configuration.newClient({ timeoutMS: 100 });
+        clientEncryption = new ClientEncryption(keyVaultClient, {
+          keyVaultNamespace: 'keyvault.datakeys',
+          kmsProviders: { local: { key: LOCAL_KEY } }
+        });
+        await internalClient
+          .db()
+          .admin()
+          .command({
+            configureFailPoint: 'failCommand',
+            mode: {
+              times: 1
+            },
+            data: {
+              failCommands: ['create'],
+              blockConnection: true,
+              blockTimeMS: 1000
+            }
+          } as FailPoint);
+      });
+
+      afterEach(async function () {
+        await internalClient
+          .db()
+          .admin()
+          .command({
+            configureFailPoint: 'failCommand',
+            mode: 'off'
+          } as FailPoint);
+        internalClient.db('db').collection('newnew').drop();
+        await internalClient.close();
+        await keyVaultClient.close();
+      });
+
+      it(
+        'times out due to timeoutMS',
+        {
+          requires: {
+            clientSideEncryption: true,
+            mongodb: '>=7.0.0',
+            topology: '!single'
+          }
+        },
+        async function () {
+          const createCollectionOptions = {
+            encryptedFields: { fields: [{ path: 'ssn', bsonType: 'string', keyId: null }] }
+          };
+
+          const db = internalClient.db('db');
+          const err = await clientEncryption
+            .createEncryptedCollection(db, 'newnew', {
+              provider: 'local',
+              createCollectionOptions,
+              masterKey: null
+            })
+            .catch(err => err);
+
+          expect(err).to.be.instanceOf(MongoCryptCreateEncryptedCollectionError);
+          expect(err.message).to.contain(
+            'Unable to create collection: Timed out during socket read'
+          );
+        }
+      );
+    });
+  });
+});