Skip to content

Commit fd5039b

Browse files
committed
refactor: setting auto select family options
1 parent 7c92254 commit fd5039b

File tree

5 files changed

+60
-56
lines changed

5 files changed

+60
-56
lines changed

src/client-side-encryption/auto_encrypter.ts

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import {
33
type MongoCryptConstructor,
44
type MongoCryptOptions
55
} from 'mongodb-client-encryption';
6+
import * as net from 'net';
67

78
import { deserialize, type Document, serialize } from '../bson';
89
import { type CommandOptions, type ProxyOptions } from '../cmap/connection';
@@ -11,15 +12,12 @@ import { getMongoDBClientEncryption } from '../deps';
1112
import { MongoRuntimeError } from '../error';
1213
import { MongoClient, type MongoClientOptions } from '../mongo_client';
1314
import { MongoDBCollectionNamespace } from '../utils';
15+
import { autoSelectSocketOptions } from './client_encryption';
1416
import * as cryptoCallbacks from './crypto_callbacks';
1517
import { MongoCryptInvalidArgumentError } from './errors';
1618
import { MongocryptdManager } from './mongocryptd_manager';
1719
import { type KMSProviders, refreshKMSCredentials } from './providers';
18-
import {
19-
type ClientEncryptionSocketOptions,
20-
type CSFLEKMSTlsOptions,
21-
StateMachine
22-
} from './state_machine';
20+
import { type CSFLEKMSTlsOptions, StateMachine } from './state_machine';
2321

2422
/** @public */
2523
export interface AutoEncryptionOptions {
@@ -105,8 +103,6 @@ export interface AutoEncryptionOptions {
105103
proxyOptions?: ProxyOptions;
106104
/** The TLS options to use connecting to the KMS provider */
107105
tlsOptions?: CSFLEKMSTlsOptions;
108-
/** Options for KMS socket requests. */
109-
socketOptions?: ClientEncryptionSocketOptions;
110106
}
111107

112108
/**
@@ -156,7 +152,6 @@ export class AutoEncrypter {
156152
_kmsProviders: KMSProviders;
157153
_bypassMongocryptdAndCryptShared: boolean;
158154
_contextCounter: number;
159-
_socketOptions: ClientEncryptionSocketOptions;
160155

161156
_mongocryptdManager?: MongocryptdManager;
162157
_mongocryptdClient?: MongoClient;
@@ -241,7 +236,6 @@ export class AutoEncrypter {
241236
this._proxyOptions = options.proxyOptions || {};
242237
this._tlsOptions = options.tlsOptions || {};
243238
this._kmsProviders = options.kmsProviders || {};
244-
this._socketOptions = options.socketOptions || {};
245239

246240
const mongoCryptOptions: MongoCryptOptions = {
247241
cryptoCallbacks
@@ -305,8 +299,14 @@ export class AutoEncrypter {
305299
serverSelectionTimeoutMS: 10000
306300
};
307301

308-
if (options.extraOptions == null || typeof options.extraOptions.mongocryptdURI !== 'string') {
302+
if (
303+
(options.extraOptions == null || typeof options.extraOptions.mongocryptdURI !== 'string') &&
304+
!net.getDefaultAutoSelectFamily
305+
) {
306+
// Only set family if autoSelectFamily options are not supported.
309307
clientOptions.family = 4;
308+
} else {
309+
Object.assign(clientOptions, autoSelectSocketOptions(this._client.s.options));
310310
}
311311

312312
this._mongocryptdClient = new MongoClient(this._mongocryptdManager.uri, clientOptions);
@@ -388,7 +388,7 @@ export class AutoEncrypter {
388388
promoteLongs: false,
389389
proxyOptions: this._proxyOptions,
390390
tlsOptions: this._tlsOptions,
391-
socketOptions: this._socketOptions
391+
socketOptions: autoSelectSocketOptions(this._client.s.options)
392392
});
393393

394394
return deserialize(await stateMachine.execute(this, context), {
@@ -409,7 +409,7 @@ export class AutoEncrypter {
409409
...options,
410410
proxyOptions: this._proxyOptions,
411411
tlsOptions: this._tlsOptions,
412-
socketOptions: this._socketOptions
412+
socketOptions: autoSelectSocketOptions(this._client.s.options)
413413
});
414414

415415
return await stateMachine.execute(this, context);

src/client-side-encryption/client_encryption.ts

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { type Collection } from '../collection';
1212
import { type FindCursor } from '../cursor/find_cursor';
1313
import { type Db } from '../db';
1414
import { getMongoDBClientEncryption } from '../deps';
15-
import { type MongoClient } from '../mongo_client';
15+
import { type MongoClient, type MongoClientOptions } from '../mongo_client';
1616
import { type Filter, type WithId } from '../mongo_types';
1717
import { type CreateCollectionOptions } from '../operations/create_collection';
1818
import { type DeleteResult } from '../operations/delete';
@@ -66,8 +66,6 @@ export class ClientEncryption {
6666
_tlsOptions: CSFLEKMSTlsOptions;
6767
/** @internal */
6868
_kmsProviders: KMSProviders;
69-
/** @internal */
70-
_socketOptions: ClientEncryptionSocketOptions;
7169

7270
/** @internal */
7371
_mongoCrypt: MongoCrypt;
@@ -114,15 +112,6 @@ export class ClientEncryption {
114112
this._proxyOptions = options.proxyOptions ?? {};
115113
this._tlsOptions = options.tlsOptions ?? {};
116114
this._kmsProviders = options.kmsProviders || {};
117-
this._socketOptions = {};
118-
119-
if ('autoSelectFamily' in client.s.options) {
120-
this._socketOptions.autoSelectFamily = client.s.options.autoSelectFamily;
121-
}
122-
if ('autoSelectFamilyAttemptTimeout' in client.s.options) {
123-
this._socketOptions.autoSelectFamilyAttemptTimeout =
124-
client.s.options.autoSelectFamilyAttemptTimeout;
125-
}
126115

127116
if (options.keyVaultNamespace == null) {
128117
throw new MongoCryptInvalidArgumentError('Missing required option `keyVaultNamespace`');
@@ -215,7 +204,7 @@ export class ClientEncryption {
215204
const stateMachine = new StateMachine({
216205
proxyOptions: this._proxyOptions,
217206
tlsOptions: this._tlsOptions,
218-
socketOptions: this._socketOptions
207+
socketOptions: autoSelectSocketOptions(this._client.s.options)
219208
});
220209

221210
const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
@@ -273,7 +262,7 @@ export class ClientEncryption {
273262
const stateMachine = new StateMachine({
274263
proxyOptions: this._proxyOptions,
275264
tlsOptions: this._tlsOptions,
276-
socketOptions: this._socketOptions
265+
socketOptions: autoSelectSocketOptions(this._client.s.options)
277266
});
278267

279268
const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
@@ -655,7 +644,7 @@ export class ClientEncryption {
655644
const stateMachine = new StateMachine({
656645
proxyOptions: this._proxyOptions,
657646
tlsOptions: this._tlsOptions,
658-
socketOptions: this._socketOptions
647+
socketOptions: autoSelectSocketOptions(this._client.s.options)
659648
});
660649

661650
const { v } = deserialize(await stateMachine.execute(this, context));
@@ -734,7 +723,7 @@ export class ClientEncryption {
734723
const stateMachine = new StateMachine({
735724
proxyOptions: this._proxyOptions,
736725
tlsOptions: this._tlsOptions,
737-
socketOptions: this._socketOptions
726+
socketOptions: autoSelectSocketOptions(this._client.s.options)
738727
});
739728
const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
740729

@@ -976,3 +965,21 @@ export interface RangeOptions {
976965
sparsity: Long;
977966
precision?: number;
978967
}
968+
969+
/**
970+
* Get the socket options from the client.
971+
* @param baseOptions - The mongo client options.
972+
* @returns ClientEncryptionSocketOptions
973+
*/
974+
export function autoSelectSocketOptions(
975+
baseOptions: MongoClientOptions
976+
): ClientEncryptionSocketOptions {
977+
const options: ClientEncryptionSocketOptions = { autoSelectFamily: true };
978+
if ('autoSelectFamily' in baseOptions) {
979+
options.autoSelectFamily = baseOptions.autoSelectFamily;
980+
}
981+
if ('autoSelectFamilyAttemptTimeout' in baseOptions) {
982+
options.autoSelectFamilyAttemptTimeout = baseOptions.autoSelectFamilyAttemptTimeout;
983+
}
984+
return options;
985+
}

src/client-side-encryption/state_machine.ts

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import { type ProxyOptions } from '../cmap/connection';
1414
import { getSocks, type SocksLib } from '../deps';
1515
import { type MongoClient, type MongoClientOptions } from '../mongo_client';
1616
import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
17-
import { type DataKey } from './client_encryption';
17+
import { autoSelectSocketOptions, type DataKey } from './client_encryption';
1818
import { MongoCryptError } from './errors';
1919
import { type MongocryptdManager } from './mongocryptd_manager';
2020
import { type KMSProviders } from './providers';
@@ -302,6 +302,7 @@ export class StateMachine {
302302
async kmsRequest(request: MongoCryptKMSRequest): Promise<void> {
303303
const parsedUrl = request.endpoint.split(':');
304304
const port = parsedUrl[1] != null ? Number.parseInt(parsedUrl[1], 10) : HTTPS_PORT;
305+
const socketOptions = autoSelectSocketOptions(this.options.socketOptions || {});
305306
const options: tls.ConnectionOptions & {
306307
host: string;
307308
port: number;
@@ -310,19 +311,12 @@ export class StateMachine {
310311
} = {
311312
host: parsedUrl[0],
312313
servername: parsedUrl[0],
313-
port
314+
port,
315+
...socketOptions
314316
};
315317
const message = request.message;
316318
const buffer = new BufferPool();
317319

318-
const socketOptions = this.options.socketOptions || {};
319-
if ('autoSelectFamily' in socketOptions) {
320-
options.autoSelectFamily = socketOptions.autoSelectFamily;
321-
}
322-
if ('autoSelectFamilyAttemptTimeout' in socketOptions) {
323-
options.autoSelectFamilyAttemptTimeout = socketOptions.autoSelectFamilyAttemptTimeout;
324-
}
325-
326320
const netSocket: net.Socket = new net.Socket();
327321
let socket: tls.TLSSocket;
328322

@@ -377,10 +371,12 @@ export class StateMachine {
377371

378372
try {
379373
if (this.options.proxyOptions && this.options.proxyOptions.proxyHost) {
380-
netSocket.connect({
374+
const netSocketOptions = {
381375
host: this.options.proxyOptions.proxyHost,
382-
port: this.options.proxyOptions.proxyPort || 1080
383-
});
376+
port: this.options.proxyOptions.proxyPort || 1080,
377+
...socketOptions
378+
};
379+
netSocket.connect(netSocketOptions);
384380
await willConnect;
385381

386382
try {

src/encrypter.ts

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { callbackify } from 'util';
22

33
import { AutoEncrypter, type AutoEncryptionOptions } from './client-side-encryption/auto_encrypter';
4-
import { type ClientEncryptionSocketOptions } from './client-side-encryption/state_machine';
54
import { MONGO_CLIENT_EVENTS } from './constants';
65
import { getMongoDBClientEncryption } from './deps';
76
import { MongoInvalidArgumentError, MongoMissingDependencyError } from './error';
@@ -57,15 +56,6 @@ export class Encrypter {
5756
};
5857
}
5958

60-
const socketOptions: ClientEncryptionSocketOptions = {};
61-
if ('autoSelectFamily' in options) {
62-
socketOptions.autoSelectFamily = options.autoSelectFamily;
63-
}
64-
if ('autoSelectFamilyAttemptTimeout' in options) {
65-
socketOptions.autoSelectFamilyAttemptTimeout = options.autoSelectFamilyAttemptTimeout;
66-
}
67-
options.autoEncryption.socketOptions = socketOptions;
68-
6959
this.autoEncrypter = new AutoEncrypter(client, options.autoEncryption);
7060
}
7161

test/unit/client-side-encryption/auto_encrypter.test.ts

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { expect } from 'chai';
22
import * as fs from 'fs';
3+
import * as net from 'net';
34
import * as sinon from 'sinon';
45

56
// eslint-disable-next-line @typescript-eslint/no-restricted-imports
@@ -37,7 +38,13 @@ const MOCK_MONGOCRYPTD_RESPONSE = readExtendedJsonToBuffer(
3738
const MOCK_KEYDOCUMENT_RESPONSE = readExtendedJsonToBuffer(`${__dirname}/data/key-document.json`);
3839
const MOCK_KMS_DECRYPT_REPLY = readHttpResponse(`${__dirname}/data/kms-decrypt-reply.txt`);
3940

40-
class MockClient {}
41+
class MockClient {
42+
s: any;
43+
44+
constructor(options?: any) {
45+
this.s = { options: options || {} };
46+
}
47+
}
4148

4249
const originalAccessKeyId = process.env.AWS_ACCESS_KEY_ID;
4350
const originalSecretAccessKey = process.env.AWS_SECRET_ACCESS_KEY;
@@ -105,23 +112,27 @@ describe('AutoEncrypter', function () {
105112
});
106113

107114
context('when mongocryptdURI is not specified', () => {
108-
it('sets the ip address family to ipv4', function () {
115+
it('sets family options', function () {
109116
expect(autoEncrypter).to.have.nested.property('_mongocryptdClient.s.options');
110117
const options = autoEncrypter._mongocryptdClient?.s.options;
111-
expect(options).to.have.property('family', 4);
118+
if (net.getDefaultAutoSelectFamily) {
119+
expect(options).to.have.property('autoSelectFamily', true);
120+
} else {
121+
expect(options).to.have.property('family', 4);
122+
}
112123
});
113124
});
114125

115126
context('when mongocryptdURI is specified', () => {
116-
it('does not set the ip address family to ipv4', function () {
127+
it('sets autoSelectFamily options', function () {
117128
const autoEncrypter = new AutoEncrypter(client, {
118129
...autoEncrypterOptions,
119130
extraOptions: { mongocryptdURI: MongocryptdManager.DEFAULT_MONGOCRYPTD_URI }
120131
});
121132

122133
expect(autoEncrypter).to.have.nested.property('_mongocryptdClient.s.options');
123134
const options = autoEncrypter._mongocryptdClient?.s.options;
124-
expect(options).not.to.have.property('family', 4);
135+
expect(options).to.have.property('autoSelectFamily', true);
125136
});
126137
});
127138
});

0 commit comments

Comments
 (0)