fix: FLE

Author: nbbeeken

src/client-side-encryption/auto_encrypter.ts

@@ -6,6 +6,7 @@ import {
 
 import { deserialize, type Document, serialize } from '../bson';
 import { type CommandOptions, type ProxyOptions } from '../cmap/connection';
+import { kDecorateResult } from '../constants';
 import { getMongoDBClientEncryption } from '../deps';
 import { MongoRuntimeError } from '../error';
 import { MongoClient, type MongoClientOptions } from '../mongo_client';
@@ -212,15 +213,6 @@ export const AutoEncryptionLoggerLevel = Object.freeze({
 export type AutoEncryptionLoggerLevel =
   (typeof AutoEncryptionLoggerLevel)[keyof typeof AutoEncryptionLoggerLevel];
 
-// Typescript errors if we index objects with `Symbol.for(...)`, so
-// to avoid TS errors we pull them out into variables.  Then we can type
-// the objects (and class) that we expect to see them on and prevent TS
-// errors.
-/** @internal */
-const kDecorateResult = Symbol.for('@@mdb.decorateDecryptionResult');
-/** @internal */
-const kDecoratedKeys = Symbol.for('@@mdb.decryptedKeys');
-
 /**
  * @internal An internal class to be used by the driver for auto encryption
  * **NOTE**: Not meant to be instantiated directly, this is for internal use only.
@@ -467,16 +459,18 @@ export class AutoEncrypter {
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions
     });
-    return await stateMachine.execute<Document>(this, context);
+
+    return deserialize(await stateMachine.execute(this, context), {
+      promoteValues: false,
+      promoteLongs: false
+    });
   }
 
   /**
    * Decrypt a command response
    */
-  async decrypt(response: Uint8Array | Document, options: CommandOptions = {}): Promise<Document> {
-    const buffer = Buffer.isBuffer(response) ? response : serialize(response, options);
-
-    const context = this._mongocrypt.makeDecryptionContext(buffer);
+  async decrypt(response: Uint8Array, options: CommandOptions = {}): Promise<Uint8Array> {
+    const context = this._mongocrypt.makeDecryptionContext(response);
 
     context.id = this._contextCounter++;
 
@@ -486,12 +480,7 @@ export class AutoEncrypter {
       tlsOptions: this._tlsOptions
     });
 
-    const decorateResult = this[kDecorateResult];
-    const result = await stateMachine.execute<Document>(this, context);
-    if (decorateResult) {
-      decorateDecryptionResult(result, response);
-    }
-    return result;
+    return await stateMachine.execute(this, context);
   }
 
   /**
@@ -518,53 +507,3 @@ export class AutoEncrypter {
     return AutoEncrypter.getMongoCrypt().libmongocryptVersion;
   }
 }
-
-/**
- * Recurse through the (identically-shaped) `decrypted` and `original`
- * objects and attach a `decryptedKeys` property on each sub-object that
- * contained encrypted fields. Because we only call this on BSON responses,
- * we do not need to worry about circular references.
- *
- * @internal
- */
-function decorateDecryptionResult(
-  decrypted: Document & { [kDecoratedKeys]?: Array<string> },
-  original: Document,
-  isTopLevelDecorateCall = true
-): void {
-  if (isTopLevelDecorateCall) {
-    // The original value could have been either a JS object or a BSON buffer
-    if (Buffer.isBuffer(original)) {
-      original = deserialize(original);
-    }
-    if (Buffer.isBuffer(decrypted)) {
-      throw new MongoRuntimeError('Expected result of decryption to be deserialized BSON object');
-    }
-  }
-
-  if (!decrypted || typeof decrypted !== 'object') return;
-  for (const k of Object.keys(decrypted)) {
-    const originalValue = original[k];
-
-    // An object was decrypted by libmongocrypt if and only if it was
-    // a BSON Binary object with subtype 6.
-    if (originalValue && originalValue._bsontype === 'Binary' && originalValue.sub_type === 6) {
-      if (!decrypted[kDecoratedKeys]) {
-        Object.defineProperty(decrypted, kDecoratedKeys, {
-          value: [],
-          configurable: true,
-          enumerable: false,
-          writable: false
-        });
-      }
-      // this is defined in the preceding if-statement
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      decrypted[kDecoratedKeys]!.push(k);
-      // Do not recurse into this decrypted value. It could be a sub-document/array,
-      // in which case there is no original value associated with its subfields.
-      continue;
-    }
-
-    decorateDecryptionResult(decrypted[k], originalValue, false);
-  }
-}

src/client-side-encryption/client_encryption.ts

@@ -5,7 +5,7 @@ import type {
   MongoCryptOptions
 } from 'mongodb-client-encryption';
 
-import { type Binary, type Document, type Long, serialize, type UUID } from '../bson';
+import { type Binary, deserialize, type Document, type Long, serialize, type UUID } from '../bson';
 import { type AnyBulkWriteOperation, type BulkWriteResult } from '../bulk/common';
 import { type ProxyOptions } from '../cmap/connection';
 import { type Collection } from '../collection';
@@ -202,7 +202,7 @@ export class ClientEncryption {
       tlsOptions: this._tlsOptions
     });
 
-    const dataKey = await stateMachine.execute<DataKey>(this, context);
+    const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
 
     const { db: dbName, collection: collectionName } = MongoDBCollectionNamespace.fromString(
       this._keyVaultNamespace
@@ -259,7 +259,7 @@ export class ClientEncryption {
       tlsOptions: this._tlsOptions
     });
 
-    const { v: dataKeys } = await stateMachine.execute<{ v: DataKey[] }>(this, context);
+    const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
     if (dataKeys.length === 0) {
       return {};
     }
@@ -640,7 +640,7 @@ export class ClientEncryption {
       tlsOptions: this._tlsOptions
     });
 
-    const { v } = await stateMachine.execute<{ v: T }>(this, context);
+    const { v } = deserialize(await stateMachine.execute(this, context));
 
     return v;
   }
@@ -719,7 +719,7 @@ export class ClientEncryption {
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
-    const result = await stateMachine.execute<{ v: Binary }>(this, context);
+    const result = deserialize(await stateMachine.execute(this, context));
     return result.v;
   }
 }

src/client-side-encryption/state_machine.ts

@@ -112,6 +112,9 @@ export type CSFLEKMSTlsOptions = {
   azure?: ClientEncryptionTlsOptions;
 };
 
+/** `{ v: [] }` */
+const EMPTY_V = Uint8Array.from([13, 0, 0, 0, 4, 118, 0, 5, 0, 0, 0, 0, 0]);
+
 /**
  * @internal
  *
@@ -154,16 +157,13 @@ export class StateMachine {
   /**
    * Executes the state machine according to the specification
    */
-  async execute<T extends Document>(
-    executor: StateMachineExecutable,
-    context: MongoCryptContext
-  ): Promise<T> {
+  async execute(executor: StateMachineExecutable, context: MongoCryptContext): Promise<Uint8Array> {
     const keyVaultNamespace = executor._keyVaultNamespace;
     const keyVaultClient = executor._keyVaultClient;
     const metaDataClient = executor._metaDataClient;
     const mongocryptdClient = executor._mongocryptdClient;
     const mongocryptdManager = executor._mongocryptdManager;
-    let result: T | null = null;
+    let result: Uint8Array | null = null;
 
     while (context.state !== MONGOCRYPT_CTX_DONE && context.state !== MONGOCRYPT_CTX_ERROR) {
       debug(`[context#${context.id}] ${stateToString.get(context.state) || context.state}`);
@@ -220,7 +220,7 @@ export class StateMachine {
             // do not.  We set the result manually here, and let the state machine continue.  `libmongocrypt`
             // will inform us if we need to error by setting the state to `MONGOCRYPT_CTX_ERROR` but
             // otherwise we'll return `{ v: [] }`.
-            result = { v: [] } as any as T;
+            result = EMPTY_V;
           }
           for await (const key of keys) {
             context.addMongoOperationResponse(serialize(key));
@@ -252,7 +252,7 @@ export class StateMachine {
             const message = context.status.message || 'Finalization error';
             throw new MongoCryptError(message);
           }
-          result = deserialize(finalizedContext, this.options) as T;
+          result = finalizedContext;
           break;
         }
 

src/cmap/connection.ts

@@ -1,14 +1,15 @@
 import { type Readable, Transform, type TransformCallback } from 'stream';
 import { clearTimeout, setTimeout } from 'timers';
 
-import type { BSONSerializeOptions, Document, ObjectId } from '../bson';
-import type { AutoEncrypter } from '../client-side-encryption/auto_encrypter';
+import { type BSONSerializeOptions, deserialize, type Document, type ObjectId } from '../bson';
+import { type AutoEncrypter } from '../client-side-encryption/auto_encrypter';
 import {
   CLOSE,
   CLUSTER_TIME_RECEIVED,
   COMMAND_FAILED,
   COMMAND_STARTED,
   COMMAND_SUCCEEDED,
+  kDecorateResult,
   PINNED,
   UNPINNED
 } from '../constants';
@@ -33,6 +34,7 @@ import {
   BufferPool,
   calculateDurationInMs,
   type Callback,
+  decorateDecryptionResult,
   HostAddress,
   maxWireVersion,
   type MongoDBNamespace,
@@ -721,7 +723,7 @@ export class CryptoConnection extends Connection {
     ns: MongoDBNamespace,
     cmd: Document,
     options?: CommandOptions,
-    _responseType?: T | undefined
+    responseType?: T | undefined
   ): Promise<Document> {
     const { autoEncrypter } = this;
     if (!autoEncrypter) {
@@ -735,7 +737,7 @@ export class CryptoConnection extends Connection {
     const serverWireVersion = maxWireVersion(this);
     if (serverWireVersion === 0) {
       // This means the initial handshake hasn't happened yet
-      return await super.command<T>(ns, cmd, options, undefined);
+      return await super.command<T>(ns, cmd, options, responseType);
     }
 
     if (serverWireVersion < 8) {
@@ -769,8 +771,25 @@ export class CryptoConnection extends Connection {
       }
     }
 
-    const response = await super.command<T>(ns, encrypted, options, undefined);
+    const encryptedResponse: MongoDBResponse = (await super.command<T>(
+      ns,
+      encrypted,
+      options,
+      (responseType ?? MongoDBResponse) as any
+    )) as unknown as MongoDBResponse;
 
-    return await autoEncrypter.decrypt(response, options);
+    const result = await autoEncrypter.decrypt(encryptedResponse.toBytes(), options);
+
+    const decryptedResponse = responseType?.make(result) ?? deserialize(result, options);
+
+    if (autoEncrypter[kDecorateResult]) {
+      if (responseType == null) {
+        decorateDecryptionResult(decryptedResponse, encryptedResponse.toObject(), true);
+      } else {
+        decryptedResponse.encryptedResponse = encryptedResponse;
+      }
+    }
+
+    return decryptedResponse;
   }
 }

src/cmap/wire_protocol/responses.ts

@@ -10,7 +10,7 @@ import {
 } from '../../bson';
 import { MongoUnexpectedServerResponseError } from '../../error';
 import { type ClusterTime } from '../../sdam/common';
-import { ns } from '../../utils';
+import { decorateDecryptionResult, ns } from '../../utils';
 import { OnDemandDocument } from './on_demand/document';
 
 // eslint-disable-next-line no-restricted-syntax
@@ -169,6 +169,9 @@ export class MongoDBResponse extends OnDemandDocument {
     }
     return { utf8: { writeErrors: false } };
   }
+
+  // TODO: Supports decorating result
+  encryptedResponse?: MongoDBResponse;
 }
 
 // Here's a litle blast from the past.
@@ -220,6 +223,21 @@ export class CursorResponse extends MongoDBResponse {
     return Math.max(this.batchSize - this.iterated, 0);
   }
 
+  private _encryptedBatch: OnDemandDocument | null = null;
+  get encryptedBatch() {
+    if (this.encryptedResponse == null) return null;
+    if (this._encryptedBatch != null) return this._encryptedBatch;
+
+    const cursor = this.encryptedResponse?.get('cursor', BSONType.object);
+    if (cursor?.has('firstBatch'))
+      this._encryptedBatch = cursor.get('firstBatch', BSONType.array, true);
+    else if (cursor?.has('nextBatch'))
+      this._encryptedBatch = cursor.get('nextBatch', BSONType.array, true);
+    else throw new MongoUnexpectedServerResponseError('Cursor document did not contain a batch');
+
+    return this._encryptedBatch;
+  }
+
   private get batch() {
     if (this._batch != null) return this._batch;
     const cursor = this.cursor;
@@ -249,12 +267,18 @@ export class CursorResponse extends MongoDBResponse {
     }
 
     const result = this.batch.get(this.iterated, BSONType.object, true) ?? null;
+    const encryptedResult = this.encryptedBatch?.get(this.iterated, BSONType.object, true) ?? null;
+
     this.iterated += 1;
 
     if (options?.raw) {
       return result.toBytes();
     } else {
-      return result.toObject(options);
+      const object = result.toObject(options);
+      if (encryptedResult) {
+        decorateDecryptionResult(object, encryptedResult.toObject(options), true);
+      }
+      return object;
     }
   }
 

src/constants.ts

@@ -165,3 +165,12 @@ export const LEGACY_HELLO_COMMAND = 'ismaster';
  * The legacy hello command that was deprecated in MongoDB 5.0.
  */
 export const LEGACY_HELLO_COMMAND_CAMEL_CASE = 'isMaster';
+
+// Typescript errors if we index objects with `Symbol.for(...)`, so
+// to avoid TS errors we pull them out into variables.  Then we can type
+// the objects (and class) that we expect to see them on and prevent TS
+// errors.
+/** @internal */
+export const kDecorateResult = Symbol.for('@@mdb.decorateDecryptionResult');
+/** @internal */
+export const kDecoratedKeys = Symbol.for('@@mdb.decryptedKeys');

src/cursor/abstract_cursor.ts

@@ -653,9 +653,6 @@ export abstract class AbstractCursor<
       const state = await this._initialize(this[kSession]);
       const response = state.response;
       this[kServer] = state.server;
-
-      if (!CursorResponse.is(response)) throw new Error('ah');
-
       this[kId] = response.id;
       this[kNamespace] = response.ns ?? this[kNamespace];
       this[kDocuments] = response;