feat(NODE-5754): allow auto select family options (#4185)

Author: durran

Diff for: .evergreen/run-typescript.sh

@@ -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,6 +12,7 @@ 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';
@@ -297,10 +299,20 @@ 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;
       }
 
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore: TS complains as this always returns true on versions where it is present.
+      if (net.getDefaultAutoSelectFamily) {
+        Object.assign(clientOptions, autoSelectSocketOptions(this._client.options));
+      }
+
       this._mongocryptdClient = new MongoClient(this._mongocryptdManager.uri, clientOptions);
     }
   }
@@ -379,7 +391,8 @@ export class AutoEncrypter {
       promoteValues: false,
       promoteLongs: false,
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
     return deserialize(await stateMachine.execute(this, context), {
@@ -399,7 +412,8 @@ export class AutoEncrypter {
     const stateMachine = new StateMachine({
       ...options,
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
     return await stateMachine.execute(this, context);

Diff for: src/client-side-encryption/auto_encrypter.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';
@@ -28,7 +28,11 @@ import {
   type KMSProviders,
   refreshKMSCredentials
 } from './providers/index';
-import { type CSFLEKMSTlsOptions, StateMachine } from './state_machine';
+import {
+  type ClientEncryptionSocketOptions,
+  type CSFLEKMSTlsOptions,
+  StateMachine
+} from './state_machine';
 
 /**
  * @public
@@ -199,7 +203,8 @@ export class ClientEncryption {
 
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
     const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
@@ -256,7 +261,8 @@ export class ClientEncryption {
     const context = this._mongoCrypt.makeRewrapManyDataKeyContext(filterBson, keyEncryptionKeyBson);
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
     const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
@@ -637,7 +643,8 @@ export class ClientEncryption {
 
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
 
     const { v } = deserialize(await stateMachine.execute(this, context));
@@ -715,7 +722,8 @@ export class ClientEncryption {
     const valueBuffer = serialize({ v: value });
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions
+      tlsOptions: this._tlsOptions,
+      socketOptions: autoSelectSocketOptions(this._client.options)
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
@@ -957,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 for: src/client-side-encryption/client_encryption.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';
@@ -114,6 +114,16 @@ export type CSFLEKMSTlsOptions = {
   [key: string]: ClientEncryptionTlsOptions | undefined;
 };
 
+/**
+ * @public
+ *
+ * Socket options to use for KMS requests.
+ */
+export type ClientEncryptionSocketOptions = Pick<
+  MongoClientOptions,
+  'autoSelectFamily' | 'autoSelectFamilyAttemptTimeout'
+>;
+
 /**
  * This is kind of a hack.  For `rewrapManyDataKey`, we have tests that
  * guarantee that when there are no matching keys, `rewrapManyDataKey` returns
@@ -153,6 +163,9 @@ export type StateMachineOptions = {
 
   /** TLS options for KMS requests, if set. */
   tlsOptions: CSFLEKMSTlsOptions;
+
+  /** Socket specific options we support. */
+  socketOptions: ClientEncryptionSocketOptions;
 } & Pick<BSONSerializeOptions, 'promoteLongs' | 'promoteValues'>;
 
 /**
@@ -289,10 +302,17 @@ 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 options: tls.ConnectionOptions & { host: string; port: number } = {
+    const socketOptions = autoSelectSocketOptions(this.options.socketOptions || {});
+    const options: tls.ConnectionOptions & {
+      host: string;
+      port: number;
+      autoSelectFamily?: boolean;
+      autoSelectFamilyAttemptTimeout?: number;
+    } = {
       host: parsedUrl[0],
       servername: parsedUrl[0],
-      port
+      port,
+      ...socketOptions
     };
     const message = request.message;
     const buffer = new BufferPool();
@@ -351,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 for: src/client-side-encryption/state_machine.ts

@@ -269,6 +269,8 @@ export const LEGAL_TLS_SOCKET_OPTIONS = [
 
 /** @public */
 export const LEGAL_TCP_SOCKET_OPTIONS = [
+  'autoSelectFamily',
+  'autoSelectFamilyAttemptTimeout',
   'family',
   'hints',
   'localAddress',

Diff for: src/cmap/connect.ts

@@ -740,6 +740,13 @@ export const OPTIONS = {
   autoEncryption: {
     type: 'record'
   },
+  autoSelectFamily: {
+    type: 'boolean',
+    default: true
+  },
+  autoSelectFamilyAttemptTimeout: {
+    type: 'uint'
+  },
   bsonRegExp: {
     type: 'boolean'
   },

Diff for: src/connection_string.ts

@@ -248,6 +248,7 @@ export type {
   LocalKMSProviderConfiguration
 } from './client-side-encryption/providers/index';
 export type {
+  ClientEncryptionSocketOptions,
   ClientEncryptionTlsOptions,
   CSFLEKMSTlsOptions,
   StateMachineExecutable

Diff for: src/index.ts

@@ -104,7 +104,7 @@ export type SupportedTLSSocketOptions = Pick<
 
 /** @public */
 export type SupportedSocketOptions = Pick<
-  TcpNetConnectOpts,
+  TcpNetConnectOpts & { autoSelectFamily?: boolean; autoSelectFamilyAttemptTimeout?: number },
   (typeof LEGAL_TCP_SOCKET_OPTIONS)[number]
 >;
 

Diff for: src/mongo_client.ts

@@ -1,5 +1,6 @@
 import { expect } from 'chai';
 import { once } from 'events';
+import * as net from 'net';
 import * as sinon from 'sinon';
 
 import {
@@ -721,4 +722,56 @@ describe('class MongoClient', function () {
       });
     });
   });
+
+  context('when connecting', function () {
+    let netSpy;
+
+    beforeEach(function () {
+      netSpy = sinon.spy(net, 'createConnection');
+    });
+
+    afterEach(function () {
+      sinon.restore();
+    });
+
+    context('when auto select options are provided', function () {
+      beforeEach(function () {
+        client = this.configuration.newClient({
+          autoSelectFamily: false,
+          autoSelectFamilyAttemptTimeout: 100
+        });
+      });
+
+      it('sets the provided options', {
+        metadata: { requires: { topology: ['single'] } },
+        test: async function () {
+          await client.connect();
+          expect(netSpy).to.have.been.calledWith({
+            autoSelectFamily: false,
+            autoSelectFamilyAttemptTimeout: 100,
+            host: 'localhost',
+            port: 27017
+          });
+        }
+      });
+    });
+
+    context('when auto select options are not provided', function () {
+      beforeEach(function () {
+        client = this.configuration.newClient();
+      });
+
+      it('sets the default options', {
+        metadata: { requires: { topology: ['single'] } },
+        test: async function () {
+          await client.connect();
+          expect(netSpy).to.have.been.calledWith({
+            autoSelectFamily: true,
+            host: 'localhost',
+            port: 27017
+          });
+        }
+      });
+    });
+  });
 });

Diff for: test/integration/node-specific/mongo_client.test.ts

@@ -1,5 +1,8 @@
 {
-  "require": "ts-node/register",
+  "require": [
+    "ts-node/register",
+    "test/tools/runner/chai_addons.ts"
+  ],
   "reporter": "test/tools/reporter/mongodb_reporter.js",
   "failZero": true,
   "color": true,