Skip to content

Commit

Permalink
refactor: setting auto select family options
Browse files Browse the repository at this point in the history
  • Loading branch information
durran committed Aug 1, 2024
1 parent 7c92254 commit fd5039b
Show file tree
Hide file tree
Showing 5 changed files with 60 additions and 56 deletions.
24 changes: 12 additions & 12 deletions src/client-side-encryption/auto_encrypter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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 {
Expand Down Expand Up @@ -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;
}

/**
Expand Down Expand Up @@ -156,7 +152,6 @@ export class AutoEncrypter {
_kmsProviders: KMSProviders;
_bypassMongocryptdAndCryptShared: boolean;
_contextCounter: number;
_socketOptions: ClientEncryptionSocketOptions;

_mongocryptdManager?: MongocryptdManager;
_mongocryptdClient?: MongoClient;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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), {
Expand All @@ -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);
Expand Down
39 changes: 23 additions & 16 deletions src/client-side-encryption/client_encryption.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -66,8 +66,6 @@ export class ClientEncryption {
_tlsOptions: CSFLEKMSTlsOptions;
/** @internal */
_kmsProviders: KMSProviders;
/** @internal */
_socketOptions: ClientEncryptionSocketOptions;

/** @internal */
_mongoCrypt: MongoCrypt;
Expand Down Expand Up @@ -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`');
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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;
}
22 changes: 9 additions & 13 deletions src/client-side-encryption/state_machine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -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 {
Expand Down
10 changes: 0 additions & 10 deletions src/encrypter.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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);
}

Expand Down
21 changes: 16 additions & 5 deletions test/unit/client-side-encryption/auto_encrypter.test.ts
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
});
});
});
Expand Down

0 comments on commit fd5039b

Please sign in to comment.