chore: fix FLE, may revert.

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 { MongoDBResponse } from '../cmap/wire_protocol/responses';
 import { getMongoDBClientEncryption } from '../deps';
 import { MongoRuntimeError } from '../error';
 import { MongoClient, type MongoClientOptions } from '../mongo_client';
@@ -473,8 +474,15 @@ export class AutoEncrypter {
   /**
    * Decrypt a command response
    */
-  async decrypt(response: Uint8Array | Document, options: CommandOptions = {}): Promise<Document> {
-    const buffer = Buffer.isBuffer(response) ? response : serialize(response, options);
+  async decrypt(
+    response: Uint8Array | Document | MongoDBResponse,
+    options: CommandOptions = {}
+  ): Promise<Document> {
+    const buffer = MongoDBResponse.is(response)
+      ? response.toBytes()
+      : Buffer.isBuffer(response)
+      ? response
+      : serialize(response, options);
 
     const context = this._mongocrypt.makeDecryptionContext(buffer);
 
@@ -487,10 +495,11 @@ export class AutoEncrypter {
     });
 
     const decorateResult = this[kDecorateResult];
-    const result = await stateMachine.execute<Document>(this, context);
+    const result = await stateMachine.execute(this, context, response.constructor);
     if (decorateResult) {
       decorateDecryptionResult(result, response);
     }
+
     return result;
   }
 

src/client-side-encryption/state_machine.ts

@@ -11,6 +11,7 @@ import {
   serialize
 } from '../bson';
 import { type ProxyOptions } from '../cmap/connection';
+import { type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses';
 import { getSocks, type SocksLib } from '../deps';
 import { type MongoClient, type MongoClientOptions } from '../mongo_client';
 import { BufferPool, MongoDBCollectionNamespace, promiseWithResolvers } from '../utils';
@@ -156,14 +157,15 @@ export class StateMachine {
    */
   async execute<T extends Document>(
     executor: StateMachineExecutable,
-    context: MongoCryptContext
+    context: MongoCryptContext,
+    responseType?: MongoDBResponseConstructor
   ): Promise<T> {
     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: any | null = null;
 
     while (context.state !== MONGOCRYPT_CTX_DONE && context.state !== MONGOCRYPT_CTX_ERROR) {
       debug(`[context#${context.id}] ${stateToString.get(context.state) || context.state}`);
@@ -252,7 +254,12 @@ export class StateMachine {
             const message = context.status.message || 'Finalization error';
             throw new MongoCryptError(message);
           }
-          result = deserialize(finalizedContext, this.options) as T;
+
+          result =
+            responseType != null
+              ? new responseType(finalizedContext)
+              : (result = deserialize(finalizedContext, this.options) as T);
+
           break;
         }
 

src/cmap/wire_protocol/responses.ts

@@ -17,6 +17,10 @@ export type MongoDBResponseConstructor = {
 
 /** @internal */
 export class MongoDBResponse extends OnDemandDocument {
+  static is(value: unknown): value is MongoDBResponse {
+    return value instanceof MongoDBResponse;
+  }
+
   // {ok:1}
   static empty = new MongoDBResponse(new Uint8Array([13, 0, 0, 0, 16, 111, 107, 0, 1, 0, 0, 0, 0]));
 
@@ -154,18 +158,18 @@ export class CursorResponse extends MongoDBResponse {
     this.documents = Object.defineProperties(Object.create(null), {
       length: {
         get: () => {
-          return this.batchSize - this.iterated;
+          return Math.max(this.batchSize - this.iterated, 0);
         }
       },
       shift: {
         value: (options?: BSONSerializeOptions) => {
           this.iterated += 1;
-          const r = this.values?.next();
-          if (!r || r.done) return null;
+          const result = this.values?.next();
+          if (!result || result.done) return null;
           if (options?.raw) {
-            return r.value.toBytes();
+            return result.value.toBytes();
           } else {
-            return r.value.toObject(options);
+            return result.value.toObject(options);
           }
         }
       },