refactor: setting auto select family options

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

Hide file tree

Showing 6 changed files with 66 additions and 58 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,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);
  }
  }
@@ -388,7 +392,7 @@ export class AutoEncrypter {
  promoteLongs: false,
  proxyOptions: this._proxyOptions,
  tlsOptions: this._tlsOptions,
- socketOptions: this._socketOptions
+ socketOptions: autoSelectSocketOptions(this._client.options)
  });
 
  return deserialize(await stateMachine.execute(this, context), {
@@ -409,7 +413,7 @@ export class AutoEncrypter {
  ...options,
  proxyOptions: this._proxyOptions,
  tlsOptions: this._tlsOptions,
- socketOptions: this._socketOptions
+ socketOptions: autoSelectSocketOptions(this._client.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.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.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.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.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 {
+ options: any;
+
+ constructor(options?: any) {
+ this.options = { 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);
  });
  });
  });

diff --git a/test/unit/client-side-encryption/client_encryption.test.ts b/test/unit/client-side-encryption/client_encryption.test.ts
@@ -19,10 +19,10 @@ import { Binary, BSON, deserialize } from '../../mongodb';
 const { EJSON } = BSON;
 
 class MockClient {
- s: any;
+ options: any;
 
  constructor(options?: any) {
- this.s = { options: options || {} };
+ this.options = { options: options || {} };
  }
  db(dbName) {
  return {