test(NODE-4363): Add support for ClientEncryption in the unified test runner (#3314)

Author: baileympearson

src/encrypter.ts

@@ -4,7 +4,7 @@ import { MONGO_CLIENT_EVENTS } from './constants';
 import type { AutoEncrypter, AutoEncryptionOptions } from './deps';
 import { MongoInvalidArgumentError, MongoMissingDependencyError } from './error';
 import { MongoClient, MongoClientOptions } from './mongo_client';
-import type { Callback } from './utils';
+import { Callback, getMongoDBClientEncryption } from './utils';
 
 let AutoEncrypterClass: AutoEncrypter;
 
@@ -123,23 +123,15 @@ export class Encrypter {
   }
 
   static checkForMongoCrypt(): void {
-    let mongodbClientEncryption = undefined;
-    // Ensure you always wrap an optional require in the try block NODE-3199
     try {
-      // Note (NODE-4254): This is to get around the circular dependency between
-      // mongodb-client-encryption and the driver in the test scenarios.
-      if (process.env.MONGODB_CLIENT_ENCRYPTION_OVERRIDE) {
-        mongodbClientEncryption = require(process.env.MONGODB_CLIENT_ENCRYPTION_OVERRIDE);
-      } else {
-        mongodbClientEncryption = require('mongodb-client-encryption');
-      }
-    } catch (err) {
+      // NOTE(NODE-3199): Ensure you always wrap an optional require in the try block
+      const mongodbClientEncryption = getMongoDBClientEncryption();
+      AutoEncrypterClass = mongodbClientEncryption.extension(require('../lib/index')).AutoEncrypter;
+    } catch {
       throw new MongoMissingDependencyError(
         'Auto-encryption requested, but the module is not installed. ' +
           'Please add `mongodb-client-encryption` as a dependency of your project'
       );
     }
-
-    AutoEncrypterClass = mongodbClientEncryption.extension(require('../lib/index')).AutoEncrypter;
   }
 }

src/utils.ts

@@ -1407,3 +1407,26 @@ export function commandSupportsReadConcern(command: Document, options?: Document
 
   return false;
 }
+
+/**
+ * A utility function to get the instance of mongodb-client-encryption, if it exists.
+ *
+ * @throws MongoMissingDependencyError if mongodb-client-encryption isn't installed.
+ * @returns
+ */
+export function getMongoDBClientEncryption() {
+  let mongodbClientEncryption;
+
+  // NOTE(NODE-4254): This is to get around the circular dependency between
+  // mongodb-client-encryption and the driver in the test scenarios.
+  if (
+    typeof process.env.MONGODB_CLIENT_ENCRYPTION_OVERRIDE === 'string' &&
+    process.env.MONGODB_CLIENT_ENCRYPTION_OVERRIDE.length > 0
+  ) {
+    mongodbClientEncryption = require(process.env.MONGODB_CLIENT_ENCRYPTION_OVERRIDE);
+  } else {
+    mongodbClientEncryption = require('mongodb-client-encryption');
+  }
+
+  return mongodbClientEncryption;
+}

test/integration/client-side-encryption/client_side_encryption.spec.test.ts

@@ -7,6 +7,7 @@ import {
   TestRunnerContext
 } from '../../tools/spec-runner';
 import { runUnifiedSuite } from '../../tools/unified-spec-runner/runner';
+import { TestFilter } from '../../tools/unified-spec-runner/schema';
 
 const isAuthEnabled = process.env.AUTH === 'auth';
 
@@ -80,8 +81,19 @@ describe('Client Side Encryption (Legacy)', function () {
 });
 
 describe('Client Side Encryption (Unified)', function () {
-  runUnifiedSuite(
-    loadSpecTests(path.join('client-side-encryption', 'tests', 'unified')),
-    () => 'NODE-4330 - implement the key management API'
-  );
+  const filter: TestFilter = ({ description }) => {
+    if (
+      description.includes('create datakey with') ||
+      description.includes('create data key with') ||
+      description.includes('rewrap')
+    ) {
+      if (description === 'no keys to rewrap due to no filter matches') {
+        return 'TODO(NODE-4330): implement the key management API';
+      }
+      return false;
+    }
+
+    return 'TODO(NODE-4330): implement the key management API';
+  };
+  runUnifiedSuite(loadSpecTests(path.join('client-side-encryption', 'tests', 'unified')), filter);
 });

test/integration/unified-test-format/unified_test_format.spec.test.ts

@@ -37,14 +37,6 @@ const filter: TestFilter = ({ description }) => {
     return 'TODO(NODE-3891): fix tests broken when AUTH enabled';
   }
 
-  if (description.length === 0) {
-    // This may seem weird, but the kmsProvider valid pass tests really test that the new
-    // client encryption entity is constructed correctly so the "test" section of each
-    // unified test is empty (save the required properties) and the test description
-    // is just the empty string
-    return 'TODO(NODE-4363): add support for client encryption entity to unified runner';
-  }
-
   return false;
 };
 

test/readme.md

@@ -402,6 +402,12 @@ The following steps will walk you through how to run the tests for CSFLE.
   set -e CREDS
   ```
 
+1. If you are running the unified tests, you must set the following environment variable as well:
+
+```shell
+export TEST_CSFLE=true
+```
+
 1. Run the functional tests:
 
    `npm run check:test`

test/tools/unified-spec-runner/entities.ts

@@ -35,8 +35,13 @@ import { WriteConcern } from '../../../src/write_concern';
 import { ejson, getEnvironmentalOptions } from '../../tools/utils';
 import type { TestConfiguration } from '../runner/config';
 import { trace } from './runner';
-import type { ClientEntity, EntityDescription } from './schema';
-import { makeConnectionString, patchCollectionOptions, patchDbOptions } from './unified-utils';
+import type { ClientEncryption, ClientEntity, EntityDescription } from './schema';
+import {
+  createClientEncryption,
+  makeConnectionString,
+  patchCollectionOptions,
+  patchDbOptions
+} from './unified-utils';
 
 export interface UnifiedChangeStream extends ChangeStream {
   eventCollector: InstanceType<typeof import('../../tools/utils')['EventCollector']>;
@@ -231,6 +236,7 @@ export type Entity =
   | AbstractCursor
   | UnifiedChangeStream
   | GridFSBucket
+  | ClientEncryption
   | Document; // Results from operations
 
 export type EntityCtor =
@@ -240,7 +246,8 @@ export type EntityCtor =
   | typeof ClientSession
   | typeof ChangeStream
   | typeof AbstractCursor
-  | typeof GridFSBucket;
+  | typeof GridFSBucket
+  | ClientEncryption;
 
 export type EntityTypeId =
   | 'client'
@@ -249,7 +256,8 @@ export type EntityTypeId =
   | 'session'
   | 'bucket'
   | 'cursor'
-  | 'stream';
+  | 'stream'
+  | 'clientEncryption';
 
 const ENTITY_CTORS = new Map<EntityTypeId, EntityCtor>();
 ENTITY_CTORS.set('client', UnifiedMongoClient);
@@ -275,6 +283,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
   mapOf(type: 'bucket'): EntitiesMap<GridFSBucket>;
   mapOf(type: 'cursor'): EntitiesMap<AbstractCursor>;
   mapOf(type: 'stream'): EntitiesMap<UnifiedChangeStream>;
+  mapOf(type: 'clientEncryption'): EntitiesMap<ClientEncryption>;
   mapOf(type: EntityTypeId): EntitiesMap<Entity> {
     const ctor = ENTITY_CTORS.get(type);
     if (!ctor) {
@@ -290,12 +299,17 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
   getEntity(type: 'bucket', key: string, assertExists?: boolean): GridFSBucket;
   getEntity(type: 'cursor', key: string, assertExists?: boolean): AbstractCursor;
   getEntity(type: 'stream', key: string, assertExists?: boolean): UnifiedChangeStream;
+  getEntity(type: 'clientEncryption', key: string, assertExists?: boolean): ClientEncryption;
   getEntity(type: EntityTypeId, key: string, assertExists = true): Entity {
     const entity = this.get(key);
     if (!entity) {
       if (assertExists) throw new Error(`Entity '${key}' does not exist`);
       return;
     }
+    if (type === 'clientEncryption') {
+      // we do not have instanceof checking here since csfle might not be installed
+      return entity;
+    }
     const ctor = ENTITY_CTORS.get(type);
     if (!ctor) {
       throw new Error(`Unknown type ${type}`);
@@ -423,6 +437,10 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
         map.set(entity.bucket.id, new GridFSBucket(db, options));
       } else if ('stream' in entity) {
         throw new Error(`Unsupported Entity ${JSON.stringify(entity)}`);
+      } else if ('clientEncryption' in entity) {
+        const clientEncryption = createClientEncryption(map, entity.clientEncryption);
+
+        map.set(entity.clientEncryption.id, clientEncryption);
       } else {
         throw new Error(`Unsupported Entity ${JSON.stringify(entity)}`);
       }

test/tools/unified-spec-runner/match.ts

@@ -204,6 +204,7 @@ export function resultCheck(
 
       if (checkExtraKeys) {
         expect(actual, `Expected actual to exist at ${path.join('')}`).to.exist;
+        // by using `Object.keys`, we ignore non-enumerable properties. This is intentional.
         const actualKeys = Object.keys(actual);
         const expectedKeys = Object.keys(expected);
         // Don't check for full key set equality because some of the actual keys

test/tools/unified-spec-runner/operations.ts

@@ -464,6 +464,76 @@ operations.set('rename', async ({ entities, operation }) => {
   return collection.rename(to, options);
 });
 
+operations.set('createDataKey', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { kmsProvider, opts } = operation.arguments ?? {};
+
+  return clientEncryption.createDataKey(kmsProvider, opts);
+});
+
+operations.set('rewrapManyDataKey', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { filter, opts } = operation.arguments ?? {};
+
+  const rewrapManyDataKeyResult = await clientEncryption.rewrapManyDataKey(filter, opts);
+
+  if (rewrapManyDataKeyResult.bulkWriteResult != null) {
+    // TODO(NODE-4393): refactor BulkWriteResult to not have a 'result' property
+    //
+    // The unified spec runner match function will assert that documents have no extra
+    // keys.  For `rewrapManyDataKey` operations, our unifed tests will fail because
+    // our BulkWriteResult class has an extra property - "result".  We explicitly make it
+    // non-enumerable for the purposes of testing so that the tests can pass.
+    const { bulkWriteResult } = rewrapManyDataKeyResult;
+    Object.defineProperty(bulkWriteResult, 'result', {
+      value: bulkWriteResult.result,
+      enumerable: false
+    });
+  }
+  return rewrapManyDataKeyResult;
+});
+
+operations.set('deleteKey', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { id } = operation.arguments ?? {};
+
+  return clientEncryption.deleteKey(id);
+});
+
+operations.set('getKey', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { id } = operation.arguments ?? {};
+
+  return clientEncryption.getKey(id);
+});
+
+operations.set('getKeys', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+
+  return clientEncryption.getKeys();
+});
+
+operations.set('addKeyAltName', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { id, keyAltName } = operation.arguments ?? {};
+
+  return clientEncryption.addKeyAltName(id, keyAltName);
+});
+
+operations.set('removeKeyAltName', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { id, keyAltName } = operation.arguments ?? {};
+
+  return clientEncryption.removeKeyAltName(id, keyAltName);
+});
+
+operations.set('getKeyByAltName', async ({ entities, operation }) => {
+  const clientEncryption = entities.getEntity('clientEncryption', operation.object);
+  const { keyAltName } = operation.arguments ?? {};
+
+  return clientEncryption.getKeyByAltName(keyAltName);
+});
+
 export async function executeOperationAndCheck(
   operation: OperationDescription,
   entities: EntitiesMap,

test/tools/unified-spec-runner/schema.ts

@@ -1,4 +1,4 @@
-import { ServerApiVersion } from '../../../src';
+import { MongoClient, ServerApiVersion } from '../../../src';
 import type { Document, ObjectId } from '../../../src/bson';
 import type { ReadConcernLevel } from '../../../src/read_concern';
 import type { ReadPreferenceMode } from '../../../src/read_preference';
@@ -57,7 +57,15 @@ export const OperationNames = [
   'runCommand',
   'updateMany',
   'updateOne',
-  'rename'
+  'rename',
+  'createDataKey',
+  'rewrapManyDataKey',
+  'deleteKey',
+  'getKey',
+  'getKeys',
+  'addKeyAltName',
+  'removeKeyAltName',
+  'getKeyByAltName'
 ] as const;
 export type OperationName = typeof OperationNames[number];
 
@@ -95,6 +103,7 @@ export interface RunOnRequirement {
   minServerVersion?: string;
   topologies?: TopologyId[];
   serverParameters?: Document;
+  csfle?: boolean;
 }
 export type ObservableCommandEventId =
   | 'commandStartedEvent'
@@ -146,13 +155,53 @@ export interface StreamEntity {
   id: string;
   hexBytes: string;
 }
+
+export type StringOrPlaceholder = string | { $$placeholder: number };
+
+export interface ClientEncryptionEntity {
+  id: string;
+  clientEncryptionOpts: {
+    /** this is the id of the client entity to use as the keyvault client */
+    keyVaultClient: string;
+    keyVaultNamespace: string;
+    kmsProviders: {
+      aws?: {
+        accessKeyId: StringOrPlaceholder;
+        secretAccessKey: StringOrPlaceholder;
+        sessionToken: StringOrPlaceholder;
+      };
+      azure?: {
+        tenantId: StringOrPlaceholder;
+        clientId: StringOrPlaceholder;
+        clientSecret: StringOrPlaceholder;
+        identityPlatformEndpoint: StringOrPlaceholder;
+      };
+      gcp?: {
+        email: StringOrPlaceholder;
+        privateKey: StringOrPlaceholder;
+        endPoint: StringOrPlaceholder;
+      };
+      kmip?: {
+        endpoint: StringOrPlaceholder;
+      };
+      local?: {
+        key: StringOrPlaceholder;
+      };
+    };
+  };
+}
+
+export type KMSProvidersEntity = ClientEncryptionEntity['clientEncryptionOpts']['kmsProviders'];
+
 export type EntityDescription =
   | { client: ClientEntity }
   | { database: DatabaseEntity }
   | { collection: CollectionEntity }
   | { bucket: BucketEntity }
   | { stream: StreamEntity }
-  | { session: SessionEntity };
+  | { session: SessionEntity }
+  | { clientEncryption: ClientEncryptionEntity };
+
 export interface ServerApi {
   version: ServerApiVersion;
   strict?: boolean;
@@ -241,3 +290,20 @@ export interface ExpectedError {
  * A type that represents the test filter provided to the unifed runner.
  */
 export type TestFilter = (test: Test) => string | false;
+
+/**
+ * This interface represents the bare minimum of type information needed to get *some* type
+ * safety on the client encryption object in unified tests.
+ */
+export interface ClientEncryption {
+  // eslint-disable-next-line @typescript-eslint/no-misused-new
+  new (client: MongoClient, options: any): ClientEncryption;
+  createDataKey(provider, options): Promise<any>;
+  rewrapManyDataKey(filter, options): Promise<any>;
+  deleteKey(id): Promise<any>;
+  getKey(id): Promise<any>;
+  getKeys(): Promise<any>;
+  addKeyAltName(id, keyAltName): Promise<any>;
+  removeKeyAltName(id, keyAltName): Promise<any>;
+  getKeyByAltName(keyAltName): Promise<any>;
+}