refactor: setting auto select family options

mongodb · Aug 1, 2024 · fd5039b · fd5039b
1 parent 7c92254
commit fd5039b
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 56 deletions.
diff --git a/src/client-side-encryption/auto_encrypter.ts b/src/client-side-encryption/auto_encrypter.ts
@@ -3,6 +3,7 @@ import {
   type MongoCryptConstructor,
   type MongoCryptOptions
 } from 'mongodb-client-encryption';
+import * as net from 'net';
 
 import { deserialize, type Document, serialize } from '../bson';
 import { type CommandOptions, type ProxyOptions } from '../cmap/connection';
@@ -11,15 +12,12 @@ import { getMongoDBClientEncryption } from '../deps';
 import { MongoRuntimeError } from '../error';
 import { MongoClient, type MongoClientOptions } from '../mongo_client';
 import { MongoDBCollectionNamespace } from '../utils';
+import { autoSelectSocketOptions } from './client_encryption';
 import * as cryptoCallbacks from './crypto_callbacks';
 import { MongoCryptInvalidArgumentError } from './errors';
 import { MongocryptdManager } from './mongocryptd_manager';
 import { type KMSProviders, refreshKMSCredentials } from './providers';
-import {
-  type ClientEncryptionSocketOptions,
-  type CSFLEKMSTlsOptions,
-  StateMachine
-} from './state_machine';
+import { type CSFLEKMSTlsOptions, StateMachine } from './state_machine';
 
 /** @public */
 export interface AutoEncryptionOptions {
@@ -105,8 +103,6 @@ export interface AutoEncryptionOptions {
   proxyOptions?: ProxyOptions;
   /** The TLS options to use connecting to the KMS provider */
   tlsOptions?: CSFLEKMSTlsOptions;
-  /** Options for KMS socket requests. */
-  socketOptions?: ClientEncryptionSocketOptions;
 }
 
 /**
@@ -156,7 +152,6 @@ export class AutoEncrypter {
   _kmsProviders: KMSProviders;
   _bypassMongocryptdAndCryptShared: boolean;
   _contextCounter: number;
-  _socketOptions: ClientEncryptionSocketOptions;
 
   _mongocryptdManager?: MongocryptdManager;
   _mongocryptdClient?: MongoClient;
@@ -241,7 +236,6 @@ export class AutoEncrypter {
     this._proxyOptions = options.proxyOptions || {};
     this._tlsOptions = options.tlsOptions || {};
     this._kmsProviders = options.kmsProviders || {};
-    this._socketOptions = options.socketOptions || {};
 
     const mongoCryptOptions: MongoCryptOptions = {
       cryptoCallbacks
@@ -305,8 +299,14 @@ export class AutoEncrypter {
         serverSelectionTimeoutMS: 10000
       };
 
-      if (options.extraOptions == null || typeof options.extraOptions.mongocryptdURI !== 'string') {
+      if (
+        (options.extraOptions == null || typeof options.extraOptions.mongocryptdURI !== 'string') &&
+        !net.getDefaultAutoSelectFamily
+      ) {
+        // Only set family if autoSelectFamily options are not supported.
         clientOptions.family = 4;
+      } else {
+        Object.assign(clientOptions, autoSelectSocketOptions(this._client.s.options));
       }
 
       this._mongocryptdClient = new MongoClient(this._mongocryptdManager.uri, clientOptions);
@@ -388,7 +388,7 @@ export class AutoEncrypter {
       promoteLongs: false,
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     return deserialize(await stateMachine.execute(this, context), {
@@ -409,7 +409,7 @@ export class AutoEncrypter {
       ...options,
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     return await stateMachine.execute(this, context);

diff --git a/src/client-side-encryption/client_encryption.ts b/src/client-side-encryption/client_encryption.ts
@@ -12,7 +12,7 @@ import { type Collection } from '../collection';
 import { type FindCursor } from '../cursor/find_cursor';
 import { type Db } from '../db';
 import { getMongoDBClientEncryption } from '../deps';
-import { type MongoClient } from '../mongo_client';
+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';
@@ -66,8 +66,6 @@ export class ClientEncryption {
   _tlsOptions: CSFLEKMSTlsOptions;
   /** @internal */
   _kmsProviders: KMSProviders;
-  /** @internal */
-  _socketOptions: ClientEncryptionSocketOptions;
 
   /** @internal */
   _mongoCrypt: MongoCrypt;
@@ -114,15 +112,6 @@ export class ClientEncryption {
     this._proxyOptions = options.proxyOptions ?? {};
     this._tlsOptions = options.tlsOptions ?? {};
     this._kmsProviders = options.kmsProviders || {};
-    this._socketOptions = {};
-
-    if ('autoSelectFamily' in client.s.options) {
-      this._socketOptions.autoSelectFamily = client.s.options.autoSelectFamily;
-    }
-    if ('autoSelectFamilyAttemptTimeout' in client.s.options) {
-      this._socketOptions.autoSelectFamilyAttemptTimeout =
-        client.s.options.autoSelectFamilyAttemptTimeout;
-    }
 
     if (options.keyVaultNamespace == null) {
       throw new MongoCryptInvalidArgumentError('Missing required option `keyVaultNamespace`');
@@ -215,7 +204,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
@@ -273,7 +262,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
@@ -655,7 +644,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     const { v } = deserialize(await stateMachine.execute(this, context));
@@ -734,7 +723,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: this._socketOptions
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
@@ -976,3 +965,21 @@ export interface RangeOptions {
   sparsity: Long;
   precision?: number;
 }
+
+/**
+ * Get the socket options from the client.
+ * @param baseOptions - The mongo client options.
+ * @returns ClientEncryptionSocketOptions
+ */
+export function autoSelectSocketOptions(
+  baseOptions: MongoClientOptions
+): ClientEncryptionSocketOptions {
+  const options: ClientEncryptionSocketOptions = { autoSelectFamily: true };
+  if ('autoSelectFamily' in baseOptions) {
+    options.autoSelectFamily = baseOptions.autoSelectFamily;
+  }
+  if ('autoSelectFamilyAttemptTimeout' in baseOptions) {
+    options.autoSelectFamilyAttemptTimeout = baseOptions.autoSelectFamilyAttemptTimeout;
+  }
+  return options;
+}
diff --git a/src/client-side-encryption/state_machine.ts b/src/client-side-encryption/state_machine.ts
@@ -14,7 +14,7 @@ import { type ProxyOptions } from '../cmap/connection';
 import { getSocks, type SocksLib } from '../deps';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
 import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
-import { type DataKey } from './client_encryption';
+import { autoSelectSocketOptions, type DataKey } from './client_encryption';
 import { MongoCryptError } from './errors';
 import { type MongocryptdManager } from './mongocryptd_manager';
 import { type KMSProviders } from './providers';
@@ -302,6 +302,7 @@ export class StateMachine {
   async kmsRequest(request: MongoCryptKMSRequest): 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 || {});
     const options: tls.ConnectionOptions & {
       host: string;
       port: number;
@@ -310,19 +311,12 @@ export class StateMachine {
     } = {
       host: parsedUrl[0],
       servername: parsedUrl[0],
-      port
+      port,
+      ...socketOptions
     };
     const message = request.message;
     const buffer = new BufferPool();
 
-    const socketOptions = this.options.socketOptions || {};
-    if ('autoSelectFamily' in socketOptions) {
-      options.autoSelectFamily = socketOptions.autoSelectFamily;
-    }
-    if ('autoSelectFamilyAttemptTimeout' in socketOptions) {
-      options.autoSelectFamilyAttemptTimeout = socketOptions.autoSelectFamilyAttemptTimeout;
-    }
-
     const netSocket: net.Socket = new net.Socket();
     let socket: tls.TLSSocket;
 
@@ -377,10 +371,12 @@ export class StateMachine {
 
     try {
       if (this.options.proxyOptions && this.options.proxyOptions.proxyHost) {
-        netSocket.connect({
+        const netSocketOptions = {
           host: this.options.proxyOptions.proxyHost,
-          port: this.options.proxyOptions.proxyPort || 1080
-        });
+          port: this.options.proxyOptions.proxyPort || 1080,
+          ...socketOptions
+        };
+        netSocket.connect(netSocketOptions);
         await willConnect;
 
         try {

diff --git a/src/encrypter.ts b/src/encrypter.ts
@@ -1,7 +1,6 @@
 import { callbackify } from 'util';
 
 import { AutoEncrypter, type AutoEncryptionOptions } from './client-side-encryption/auto_encrypter';
-import { type ClientEncryptionSocketOptions } from './client-side-encryption/state_machine';
 import { MONGO_CLIENT_EVENTS } from './constants';
 import { getMongoDBClientEncryption } from './deps';
 import { MongoInvalidArgumentError, MongoMissingDependencyError } from './error';
@@ -57,15 +56,6 @@ export class Encrypter {
       };
     }
 
-    const socketOptions: ClientEncryptionSocketOptions = {};
-    if ('autoSelectFamily' in options) {
-      socketOptions.autoSelectFamily = options.autoSelectFamily;
-    }
-    if ('autoSelectFamilyAttemptTimeout' in options) {
-      socketOptions.autoSelectFamilyAttemptTimeout = options.autoSelectFamilyAttemptTimeout;
-    }
-    options.autoEncryption.socketOptions = socketOptions;
-
     this.autoEncrypter = new AutoEncrypter(client, options.autoEncryption);
   }
 

diff --git a/test/unit/client-side-encryption/auto_encrypter.test.ts b/test/unit/client-side-encryption/auto_encrypter.test.ts
@@ -1,5 +1,6 @@
 import { expect } from 'chai';
 import * as fs from 'fs';
+import * as net from 'net';
 import * as sinon from 'sinon';
 
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
@@ -37,7 +38,13 @@ const MOCK_MONGOCRYPTD_RESPONSE = readExtendedJsonToBuffer(
 const MOCK_KEYDOCUMENT_RESPONSE = readExtendedJsonToBuffer(`${__dirname}/data/key-document.json`);
 const MOCK_KMS_DECRYPT_REPLY = readHttpResponse(`${__dirname}/data/kms-decrypt-reply.txt`);
 
-class MockClient {}
+class MockClient {
+  s: any;
+
+  constructor(options?: any) {
+    this.s = { options: options || {} };
+  }
+}
 
 const originalAccessKeyId = process.env.AWS_ACCESS_KEY_ID;
 const originalSecretAccessKey = process.env.AWS_SECRET_ACCESS_KEY;
@@ -105,23 +112,27 @@ describe('AutoEncrypter', function () {
       });
 
       context('when mongocryptdURI is not specified', () => {
-        it('sets the ip address family to ipv4', function () {
+        it('sets family options', function () {
           expect(autoEncrypter).to.have.nested.property('_mongocryptdClient.s.options');
           const options = autoEncrypter._mongocryptdClient?.s.options;
-          expect(options).to.have.property('family', 4);
+          if (net.getDefaultAutoSelectFamily) {
+            expect(options).to.have.property('autoSelectFamily', true);
+          } else {
+            expect(options).to.have.property('family', 4);
+          }
         });
       });
 
       context('when mongocryptdURI is specified', () => {
-        it('does not set the ip address family to ipv4', function () {
+        it('sets autoSelectFamily options', function () {
           const autoEncrypter = new AutoEncrypter(client, {
             ...autoEncrypterOptions,
             extraOptions: { mongocryptdURI: MongocryptdManager.DEFAULT_MONGOCRYPTD_URI }
           });
 
           expect(autoEncrypter).to.have.nested.property('_mongocryptdClient.s.options');
           const options = autoEncrypter._mongocryptdClient?.s.options;
-          expect(options).not.to.have.property('family', 4);
+          expect(options).to.have.property('autoSelectFamily', true);
         });
       });
     });