From c5b8d78468b5a5e35d3fa28083b01ec087231958 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Wed, 17 Apr 2024 17:01:19 -0400 Subject: [PATCH 01/25] refactor(NODE-6057): implement CursorResponse for lazy document parsing --- src/cmap/wire_protocol/on_demand/document.ts | 15 +++- src/cmap/wire_protocol/responses.ts | 87 +++++++++++++++++++- src/cursor/abstract_cursor.ts | 31 +++++-- src/cursor/find_cursor.ts | 4 +- src/operations/execute_operation.ts | 3 +- src/operations/find.ts | 18 ++-- src/operations/get_more.ts | 10 ++- src/sdam/server.ts | 29 +++++-- 8 files changed, 170 insertions(+), 27 deletions(-) diff --git a/src/cmap/wire_protocol/on_demand/document.ts b/src/cmap/wire_protocol/on_demand/document.ts index ba8804fc6c8..0b851d27d03 100644 --- a/src/cmap/wire_protocol/on_demand/document.ts +++ b/src/cmap/wire_protocol/on_demand/document.ts @@ -197,6 +197,13 @@ export class OnDemandDocument { } } + /** + * Returns the number of elements in this BSON document + */ + public size() { + return this.elements.length; + } + /** * Checks for the existence of an element by name. * @@ -303,12 +310,18 @@ export class OnDemandDocument { }); } + /** Returns this document's bytes only */ + toBytes() { + const size = getInt32LE(this.bson, this.offset); + return this.bson.subarray(this.offset, this.offset + size); + } + /** * Iterates through the elements of a document reviving them using the `as` BSONType. * * @param as - The type to revive all elements as */ - public *valuesAs(as: T): Generator { + public *valuesAs(as: T): Generator { if (!this.isArray) { throw new BSONError('Unexpected conversion of non-array value to array'); } diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index b776a4de568..f93b6390909 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -1,5 +1,13 @@ -import { type BSONSerializeOptions, BSONType, type Document, type Timestamp } from '../../bson'; +import { + type BSONSerializeOptions, + BSONType, + type Document, + Long, + type Timestamp +} from '../../bson'; +import { MongoUnexpectedServerResponseError } from '../../error'; import { type ClusterTime } from '../../sdam/common'; +import { type MongoDBNamespace, ns } from '../../utils'; import { OnDemandDocument } from './on_demand/document'; /** @internal */ @@ -107,3 +115,80 @@ export class MongoDBResponse extends OnDemandDocument { return { utf8: { writeErrors: false } }; } } + +function throwUnsupportedError() { + throw new Error('Unsupported method'); +} + +export class CursorResponse extends MongoDBResponse { + id: Long | null = null; + ns: MongoDBNamespace | null = null; + + documents: any | null = null; + bufferForUnshift: any[] = []; + + private batch: OnDemandDocument | null = null; + private values: Generator | null = null; + private batchSize = 0; + private iterated = 0; + + constructor(b: Uint8Array, o?: number, a?: boolean) { + super(b, o, a); + + if (this.isError) return; + + const cursor = this.get('cursor', BSONType.object, true); + + const id = cursor.get('id', BSONType.long, true); + this.id = new Long(Number(id & 0xffff_ffffn), Number((id >> 32n) & 0xffff_ffffn)); + + const namespace = cursor.get('ns', BSONType.string) ?? ''; + if (namespace) this.ns = ns(namespace); + + if (cursor.has('firstBatch')) this.batch = cursor.get('firstBatch', BSONType.array, true); + else if (cursor.has('nextBatch')) this.batch = cursor.get('nextBatch', BSONType.array, true); + else throw new MongoUnexpectedServerResponseError('Cursor document did not contain a batch'); + + this.values = this.batch.valuesAs(BSONType.object); + this.batchSize = this.batch.size(); + this.iterated = 0; + this.documents = Object.defineProperties(Object.create(null), { + length: { + get: () => { + return this.batchSize - this.iterated; + } + }, + shift: { + value: (options?: BSONSerializeOptions) => { + this.iterated += 1; + if (this.bufferForUnshift.length) return this.bufferForUnshift.pop(); + const r = this.values?.next(); + if (!r || r.done) return null; + if (options.raw) { + return r.value.toBytes(); + } else { + return r.value.toObject(options); + } + } + }, + unshift: { + value: (v: any) => { + this.iterated -= 1; + this.bufferForUnshift.push(v); + } + }, + clear: { + value: () => { + this.iterated = this.batchSize; + this.values?.return(); + } + }, + pushMany: { value: throwUnsupportedError }, + push: { value: throwUnsupportedError } + }); + } + + static isCursorResponse(value: unknown): value is CursorResponse { + return value instanceof CursorResponse; + } +} diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 10aa5eea5f4..0c1a5e1913a 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -1,6 +1,7 @@ import { Readable, Transform } from 'stream'; import { type BSONSerializeOptions, type Document, Long, pluckBSONSerializeOptions } from '../bson'; +import { CursorResponse } from '../cmap/wire_protocol/responses'; import { type AnyError, MongoAPIError, @@ -144,7 +145,14 @@ export abstract class AbstractCursor< /** @internal */ [kNamespace]: MongoDBNamespace; /** @internal */ - [kDocuments]: List; + [kDocuments]: { + length: number; + shift(bsonOptions?: any): TSchema | null; + unshift(doc: TSchema): void; + clear(): void; + pushMany(many: Iterable): void; + push(item: TSchema): void; + }; /** @internal */ [kClient]: MongoClient; /** @internal */ @@ -286,7 +294,7 @@ export abstract class AbstractCursor< const documentsToRead = Math.min(number ?? this[kDocuments].length, this[kDocuments].length); for (let count = 0; count < documentsToRead; count++) { - const document = this[kDocuments].shift(); + const document = this[kDocuments].shift(this[kOptions]); if (document != null) { bufferedDocs.push(document); } @@ -633,12 +641,13 @@ export abstract class AbstractCursor< protected abstract _initialize(session: ClientSession | undefined): Promise; /** @internal */ - async getMore(batchSize: number): Promise { + async getMore(batchSize: number, useCursorResponse = false): Promise { // eslint-disable-next-line @typescript-eslint/no-non-null-assertion const getMoreOperation = new GetMoreOperation(this[kNamespace], this[kId]!, this[kServer]!, { ...this[kOptions], session: this[kSession], - batchSize + batchSize, + useCursorResponse }); return await executeOperation(this[kClient], getMoreOperation); @@ -656,7 +665,11 @@ export abstract class AbstractCursor< const state = await this._initialize(this[kSession]); const response = state.response; this[kServer] = state.server; - if (response.cursor) { + if (CursorResponse.isCursorResponse(response)) { + this[kId] = response.id; + if (response.ns) this[kNamespace] = response.ns; + this[kDocuments] = response.documents; + } else if (response.cursor) { // TODO(NODE-2674): Preserve int64 sent from MongoDB this[kId] = typeof response.cursor.id === 'number' @@ -730,7 +743,7 @@ async function next( } if (cursor[kDocuments].length !== 0) { - const doc = cursor[kDocuments].shift(); + const doc = cursor[kDocuments].shift(cursor[kOptions]); if (doc != null && transform && cursor[kTransform]) { try { @@ -762,8 +775,10 @@ async function next( try { const response = await cursor.getMore(batchSize); - - if (response) { + if (CursorResponse.isCursorResponse(response)) { + cursor[kId] = response.id; + cursor[kDocuments] = response.documents; + } else if (response) { const cursorId = typeof response.cursor.id === 'number' ? Long.fromNumber(response.cursor.id) diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index b76af197e11..1d2fbd5c636 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -111,10 +111,10 @@ export class FindCursor extends AbstractCursor { } } - const response = await super.getMore(batchSize); + const response = await super.getMore(batchSize, true); // TODO: wrap this in some logic to prevent it from happening if we don't need this support if (response) { - this[kNumReturned] = this[kNumReturned] + response.cursor.nextBatch.length; + this[kNumReturned] = this[kNumReturned] + response.batchLength; } return response; diff --git a/src/operations/execute_operation.ts b/src/operations/execute_operation.ts index 6e1b569a7dd..4faf4fd95ad 100644 --- a/src/operations/execute_operation.ts +++ b/src/operations/execute_operation.ts @@ -1,4 +1,5 @@ import type { Document } from '../bson'; +import { type CursorResponse } from '../cmap/wire_protocol/responses'; import { isRetryableReadError, isRetryableWriteError, @@ -44,7 +45,7 @@ export interface ExecutionResult { /** The session used for this operation, may be implicitly created */ session?: ClientSession; /** The raw server response for the operation */ - response: Document; + response: Document | CursorResponse; } /** diff --git a/src/operations/find.ts b/src/operations/find.ts index 3841142c4eb..285720b985b 100644 --- a/src/operations/find.ts +++ b/src/operations/find.ts @@ -1,4 +1,5 @@ import type { Document } from '../bson'; +import { CursorResponse } from '../cmap/wire_protocol/responses'; import type { Collection } from '../collection'; import { MongoInvalidArgumentError } from '../error'; import { ReadConcern } from '../read_concern'; @@ -111,12 +112,17 @@ export class FindOperation extends CommandOperation { findCommand = decorateWithExplain(findCommand, this.explain); } - return await server.command(this.ns, findCommand, { - ...this.options, - ...this.bsonOptions, - documentsReturnedIn: 'firstBatch', - session - }); + return await server.command( + this.ns, + findCommand, + { + ...this.options, + ...this.bsonOptions, + documentsReturnedIn: 'firstBatch', + session + }, + this.explain ? undefined : CursorResponse + ); } } diff --git a/src/operations/get_more.ts b/src/operations/get_more.ts index ada371c956e..05f54b0b57c 100644 --- a/src/operations/get_more.ts +++ b/src/operations/get_more.ts @@ -1,4 +1,5 @@ import type { Document, Long } from '../bson'; +import { CursorResponse } from '../cmap/wire_protocol/responses'; import { MongoRuntimeError } from '../error'; import type { Server } from '../sdam/server'; import type { ClientSession } from '../sessions'; @@ -19,6 +20,8 @@ export interface GetMoreOptions extends OperationOptions { maxTimeMS?: number; /** TODO(NODE-4413): Address bug with maxAwaitTimeMS not being passed in from the cursor correctly */ maxAwaitTimeMS?: number; + + useCursorResponse: boolean; } /** @@ -96,7 +99,12 @@ export class GetMoreOperation extends AbstractOperation { ...this.options }; - return await server.command(this.ns, getMoreCmd, commandOptions); + return await server.command( + this.ns, + getMoreCmd, + commandOptions, + this.options.useCursorResponse ? CursorResponse : undefined + ); } } diff --git a/src/sdam/server.ts b/src/sdam/server.ts index 6dbc31df7d2..8ea91815c60 100644 --- a/src/sdam/server.ts +++ b/src/sdam/server.ts @@ -7,6 +7,7 @@ import { type ConnectionPoolOptions } from '../cmap/connection_pool'; import { PoolClearedError } from '../cmap/errors'; +import { type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses'; import { APM_EVENTS, CLOSED, @@ -262,11 +263,25 @@ export class Server extends TypedEventEmitter { } } - /** - * Execute a command - * @internal - */ - async command(ns: MongoDBNamespace, cmd: Document, options: CommandOptions): Promise { + public async command( + ns: MongoDBNamespace, + command: Document, + options: CommandOptions | undefined, + responseType: T | undefined + ): Promise>; + + public async command( + ns: MongoDBNamespace, + command: Document, + options?: CommandOptions + ): Promise; + + public async command( + ns: MongoDBNamespace, + cmd: Document, + options: CommandOptions, + responseType?: MongoDBResponseConstructor + ): Promise { if (ns.db == null || typeof ns === 'string') { throw new MongoInvalidArgumentError('Namespace must not be a string'); } @@ -308,7 +323,7 @@ export class Server extends TypedEventEmitter { try { try { - return await conn.command(ns, cmd, finalOptions); + return await conn.command(ns, cmd, finalOptions, responseType); } catch (commandError) { throw this.decorateCommandError(conn, cmd, finalOptions, commandError); } @@ -319,7 +334,7 @@ export class Server extends TypedEventEmitter { ) { await this.pool.reauthenticate(conn); try { - return await conn.command(ns, cmd, finalOptions); + return await conn.command(ns, cmd, finalOptions, responseType); } catch (commandError) { throw this.decorateCommandError(conn, cmd, finalOptions, commandError); } From 0e227825331925a24fba840d0e4e5df5a72313af Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Wed, 17 Apr 2024 17:24:39 -0400 Subject: [PATCH 02/25] chore: remove the need for unshift --- src/cmap/wire_protocol/responses.ts | 17 +++------- src/cursor/abstract_cursor.ts | 51 +++++++++++++++++++++-------- 2 files changed, 42 insertions(+), 26 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index f93b6390909..0f7cbcece15 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -121,11 +121,9 @@ function throwUnsupportedError() { } export class CursorResponse extends MongoDBResponse { - id: Long | null = null; - ns: MongoDBNamespace | null = null; - - documents: any | null = null; - bufferForUnshift: any[] = []; + public id: Long | null = null; + public ns: MongoDBNamespace | null = null; + public documents: any | null = null; private batch: OnDemandDocument | null = null; private values: Generator | null = null; @@ -161,22 +159,15 @@ export class CursorResponse extends MongoDBResponse { shift: { value: (options?: BSONSerializeOptions) => { this.iterated += 1; - if (this.bufferForUnshift.length) return this.bufferForUnshift.pop(); const r = this.values?.next(); if (!r || r.done) return null; - if (options.raw) { + if (options?.raw) { return r.value.toBytes(); } else { return r.value.toObject(options); } } }, - unshift: { - value: (v: any) => { - this.iterated -= 1; - this.bufferForUnshift.push(v); - } - }, clear: { value: () => { this.iterated = this.batchSize; diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 0c1a5e1913a..54eb69ccb56 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -148,7 +148,6 @@ export abstract class AbstractCursor< [kDocuments]: { length: number; shift(bsonOptions?: any): TSchema | null; - unshift(doc: TSchema): void; clear(): void; pushMany(many: Iterable): void; push(item: TSchema): void; @@ -390,14 +389,7 @@ export abstract class AbstractCursor< return true; } - const doc = await next(this, { blocking: true, transform: false }); - - if (doc) { - this[kDocuments].unshift(doc); - return true; - } - - return false; + return await next(this, { blocking: true, transform: false, hasNext: true }); } /** Get the next available document from the cursor, returns null if no more documents are available. */ @@ -406,7 +398,7 @@ export abstract class AbstractCursor< throw new MongoCursorExhaustedError(); } - return await next(this, { blocking: true, transform: true }); + return await next(this, { blocking: true, transform: true, hasNext: false }); } /** @@ -417,7 +409,7 @@ export abstract class AbstractCursor< throw new MongoCursorExhaustedError(); } - return await next(this, { blocking: false, transform: true }); + return await next(this, { blocking: false, transform: true, hasNext: false }); } /** @@ -726,13 +718,42 @@ async function next( cursor: AbstractCursor, { blocking, - transform + transform, + hasNext + }: { + blocking: boolean; + transform: boolean; + hasNext: true; + } +): Promise; + +async function next( + cursor: AbstractCursor, + { + blocking, + transform, + hasNext + }: { + blocking: boolean; + transform: boolean; + hasNext: false; + } +): Promise; + +async function next( + cursor: AbstractCursor, + { + blocking, + transform, + hasNext }: { blocking: boolean; transform: boolean; + hasNext: boolean; } -): Promise { +): Promise { if (cursor.closed) { + if (hasNext) return false; return null; } @@ -743,6 +764,7 @@ async function next( } if (cursor[kDocuments].length !== 0) { + if (hasNext) return true; const doc = cursor[kDocuments].shift(cursor[kOptions]); if (doc != null && transform && cursor[kTransform]) { @@ -767,6 +789,7 @@ async function next( // cleanupCursor should never throw, but if it does it indicates a bug in the driver // and we should surface the error await cleanupCursor(cursor, {}); + if (hasNext) return false; return null; } @@ -811,10 +834,12 @@ async function next( } if (cursor[kDocuments].length === 0 && blocking === false) { + if (hasNext) return false; return null; } } while (!cursor.isDead || cursor[kDocuments].length !== 0); + if (hasNext) return false; return null; } From 3144b115a16e2fb6f39ee8d3b75aa2534fd34b48 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Wed, 17 Apr 2024 17:27:02 -0400 Subject: [PATCH 03/25] chore: lint --- src/cmap/wire_protocol/responses.ts | 1 + src/cursor/abstract_cursor.ts | 2 +- src/cursor/run_command_cursor.ts | 3 ++- src/index.ts | 6 +++++- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 0f7cbcece15..17d3f9d37e5 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -120,6 +120,7 @@ function throwUnsupportedError() { throw new Error('Unsupported method'); } +/** @internal */ export class CursorResponse extends MongoDBResponse { public id: Long | null = null; public ns: MongoDBNamespace | null = null; diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 54eb69ccb56..c00dd9705e7 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -961,7 +961,7 @@ class ReadableCursorStream extends Readable { private _readNext() { // eslint-disable-next-line github/no-then - next(this._cursor, { blocking: true, transform: true }).then( + next(this._cursor, { blocking: true, transform: true, hasNext: false }).then( result => { if (result == null) { this.push(null); diff --git a/src/cursor/run_command_cursor.ts b/src/cursor/run_command_cursor.ts index 4f88dc2db51..553041492f4 100644 --- a/src/cursor/run_command_cursor.ts +++ b/src/cursor/run_command_cursor.ts @@ -125,7 +125,8 @@ export class RunCommandCursor extends AbstractCursor { const getMoreOperation = new GetMoreOperation(this.namespace, this.id!, this.server!, { ...this.cursorOptions, session: this.session, - ...this.getMoreOptions + ...this.getMoreOptions, + useCursorResponse: false }); return await executeOperation(this.client, getMoreOperation); diff --git a/src/index.ts b/src/index.ts index bb83a774bf7..812d045ba6a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -290,7 +290,11 @@ export type { ConnectionPoolMetrics } from './cmap/metrics'; export type { StreamDescription, StreamDescriptionOptions } from './cmap/stream_description'; export type { CompressorName } from './cmap/wire_protocol/compression'; export type { JSTypeOf, OnDemandDocument } from './cmap/wire_protocol/on_demand/document'; -export type { MongoDBResponse, MongoDBResponseConstructor } from './cmap/wire_protocol/responses'; +export type { + CursorResponse, + MongoDBResponse, + MongoDBResponseConstructor +} from './cmap/wire_protocol/responses'; export type { CollectionOptions, CollectionPrivate, ModifyResult } from './collection'; export type { COMMAND_FAILED, From 78f9068bc01caadcfc522f76745db1e2804415f3 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Wed, 17 Apr 2024 18:30:53 -0400 Subject: [PATCH 04/25] chore: fix FLE, may revert. --- src/client-side-encryption/auto_encrypter.ts | 15 ++++++++++++--- src/client-side-encryption/state_machine.ts | 13 ++++++++++--- src/cmap/wire_protocol/responses.ts | 14 +++++++++----- 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/src/client-side-encryption/auto_encrypter.ts b/src/client-side-encryption/auto_encrypter.ts index e6334dc57de..898e339a67a 100644 --- a/src/client-side-encryption/auto_encrypter.ts +++ b/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 { - const buffer = Buffer.isBuffer(response) ? response : serialize(response, options); + async decrypt( + response: Uint8Array | Document | MongoDBResponse, + options: CommandOptions = {} + ): Promise { + 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(this, context); + const result = await stateMachine.execute(this, context, response.constructor); if (decorateResult) { decorateDecryptionResult(result, response); } + return result; } diff --git a/src/client-side-encryption/state_machine.ts b/src/client-side-encryption/state_machine.ts index a4b2379fb51..0c6c4da0bc7 100644 --- a/src/client-side-encryption/state_machine.ts +++ b/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( executor: StateMachineExecutable, - context: MongoCryptContext + context: MongoCryptContext, + responseType?: MongoDBResponseConstructor ): Promise { 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; } diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 17d3f9d37e5..d7ca89ab83d 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/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); } } }, From f42584c5015fa34da82babd000c798937d7d1f50 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Wed, 17 Apr 2024 18:37:14 -0400 Subject: [PATCH 05/25] chore: fix ts --- src/client-side-encryption/auto_encrypter.ts | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/client-side-encryption/auto_encrypter.ts b/src/client-side-encryption/auto_encrypter.ts index 898e339a67a..7ced6d74237 100644 --- a/src/client-side-encryption/auto_encrypter.ts +++ b/src/client-side-encryption/auto_encrypter.ts @@ -6,7 +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 { MongoDBResponse, type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses'; import { getMongoDBClientEncryption } from '../deps'; import { MongoRuntimeError } from '../error'; import { MongoClient, type MongoClientOptions } from '../mongo_client'; @@ -484,6 +484,10 @@ export class AutoEncrypter { ? response : serialize(response, options); + const responseType = MongoDBResponse.is(response) + ? (response.constructor as MongoDBResponseConstructor) + : undefined; + const context = this._mongocrypt.makeDecryptionContext(buffer); context.id = this._contextCounter++; @@ -495,7 +499,7 @@ export class AutoEncrypter { }); const decorateResult = this[kDecorateResult]; - const result = await stateMachine.execute(this, context, response.constructor); + const result = await stateMachine.execute(this, context, responseType); if (decorateResult) { decorateDecryptionResult(result, response); } From 8bd2a47fa735e4c19f7778bbe268f2b6d3c0e459 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 18 Apr 2024 15:42:09 -0400 Subject: [PATCH 06/25] fix: old servers --- src/cmap/wire_protocol/responses.ts | 4 ++-- src/cursor/abstract_cursor.ts | 4 ++-- src/cursor/find_cursor.ts | 11 +++++++---- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index d7ca89ab83d..339571e3d6c 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -129,10 +129,10 @@ export class CursorResponse extends MongoDBResponse { public id: Long | null = null; public ns: MongoDBNamespace | null = null; public documents: any | null = null; + public batchSize = 0; private batch: OnDemandDocument | null = null; private values: Generator | null = null; - private batchSize = 0; private iterated = 0; constructor(b: Uint8Array, o?: number, a?: boolean) { @@ -184,7 +184,7 @@ export class CursorResponse extends MongoDBResponse { }); } - static isCursorResponse(value: unknown): value is CursorResponse { + static override is(value: unknown): value is CursorResponse { return value instanceof CursorResponse; } } diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index c00dd9705e7..6d434d48cdf 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -657,7 +657,7 @@ export abstract class AbstractCursor< const state = await this._initialize(this[kSession]); const response = state.response; this[kServer] = state.server; - if (CursorResponse.isCursorResponse(response)) { + if (CursorResponse.is(response)) { this[kId] = response.id; if (response.ns) this[kNamespace] = response.ns; this[kDocuments] = response.documents; @@ -798,7 +798,7 @@ async function next( try { const response = await cursor.getMore(batchSize); - if (CursorResponse.isCursorResponse(response)) { + if (CursorResponse.is(response)) { cursor[kId] = response.id; cursor[kDocuments] = response.documents; } else if (response) { diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index 1d2fbd5c636..9884146f5d9 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -1,4 +1,5 @@ import { type Document, Long } from '../bson'; +import { CursorResponse } from '../cmap/wire_protocol/responses'; import { MongoInvalidArgumentError, MongoTailableCursorError } from '../error'; import { type ExplainVerbosityLike } from '../explain'; import type { MongoClient } from '../mongo_client'; @@ -34,7 +35,7 @@ export class FindCursor extends AbstractCursor { /** @internal */ [kFilter]: Document; /** @internal */ - [kNumReturned]?: number; + [kNumReturned] = 0; /** @internal */ [kBuiltOptions]: FindOptions; @@ -78,7 +79,9 @@ export class FindCursor extends AbstractCursor { const response = await executeOperation(this.client, findOperation); // the response is not a cursor when `explain` is enabled - this[kNumReturned] = response.cursor?.firstBatch?.length; + if (CursorResponse.is(response)) { + this[kNumReturned] = response.batchSize; + } // TODO: NODE-2882 return { server: findOperation.server, session, response }; @@ -113,8 +116,8 @@ export class FindCursor extends AbstractCursor { const response = await super.getMore(batchSize, true); // TODO: wrap this in some logic to prevent it from happening if we don't need this support - if (response) { - this[kNumReturned] = this[kNumReturned] + response.batchLength; + if (CursorResponse.is(response)) { + this[kNumReturned] = this[kNumReturned] + response.batchSize; } return response; From 68077e9c084520540bf7d048876174a9b8722031 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 18 Apr 2024 16:43:13 -0400 Subject: [PATCH 07/25] fix: return consistent type --- src/cmap/wire_protocol/responses.ts | 15 +++++++++++---- src/cursor/find_cursor.ts | 2 +- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 339571e3d6c..f5d8dd8741c 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -126,6 +126,17 @@ function throwUnsupportedError() { /** @internal */ export class CursorResponse extends MongoDBResponse { + static emptyGetMore = new CursorResponse( + Buffer.from( + 'NgAAABBvawABAAAAA2N1cnNvcgAhAAAAEmlkAAAAAAAAAAAABG5leHRCYXRjaAAFAAAAAAAA', + 'base64' + ) + ); + + static override is(value: unknown): value is CursorResponse { + return value instanceof CursorResponse; + } + public id: Long | null = null; public ns: MongoDBNamespace | null = null; public documents: any | null = null; @@ -183,8 +194,4 @@ export class CursorResponse extends MongoDBResponse { push: { value: throwUnsupportedError } }); } - - static override is(value: unknown): value is CursorResponse { - return value instanceof CursorResponse; - } } diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index 9884146f5d9..0b51351fc8f 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -110,7 +110,7 @@ export class FindCursor extends AbstractCursor { // instead, if we determine there are no more documents to request from the server, we preemptively // close the cursor } - return { cursor: { id: Long.ZERO, nextBatch: [] } }; + return CursorResponse.emptyGetMore; } } From 5fb54d33f21a3b1ef2f2b742ec0e8838d46a4a46 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 18 Apr 2024 17:24:07 -0400 Subject: [PATCH 08/25] fix: unit --- src/cursor/find_cursor.ts | 2 +- test/unit/assorted/collations.test.js | 8 ++++---- test/unit/assorted/sessions_collection.test.js | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index 0b51351fc8f..f1b3427e8d6 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -1,4 +1,4 @@ -import { type Document, Long } from '../bson'; +import { type Document } from '../bson'; import { CursorResponse } from '../cmap/wire_protocol/responses'; import { MongoInvalidArgumentError, MongoTailableCursorError } from '../error'; import { type ExplainVerbosityLike } from '../explain'; diff --git a/test/unit/assorted/collations.test.js b/test/unit/assorted/collations.test.js index 0c597a2ce77..124d020afcf 100644 --- a/test/unit/assorted/collations.test.js +++ b/test/unit/assorted/collations.test.js @@ -55,7 +55,7 @@ describe('Collation', function () { request.reply(primary[0]); } else if (doc.aggregate) { commandResult = doc; - request.reply({ ok: 1, cursor: { id: 0, firstBatch: [], ns: 'collation_test' } }); + request.reply({ ok: 1, cursor: { id: 0n, firstBatch: [], ns: 'collation_test' } }); } else if (doc.endSessions) { request.reply({ ok: 1 }); } @@ -183,7 +183,7 @@ describe('Collation', function () { request.reply(primary[0]); } else if (doc.find) { commandResult = doc; - request.reply({ ok: 1, cursor: { id: 0, firstBatch: [] } }); + request.reply({ ok: 1, cursor: { id: 0n, firstBatch: [] } }); } else if (doc.endSessions) { request.reply({ ok: 1 }); } @@ -215,7 +215,7 @@ describe('Collation', function () { request.reply(primary[0]); } else if (doc.find) { commandResult = doc; - request.reply({ ok: 1, cursor: { id: 0, firstBatch: [] } }); + request.reply({ ok: 1, cursor: { id: 0n, firstBatch: [] } }); } else if (doc.endSessions) { request.reply({ ok: 1 }); } @@ -249,7 +249,7 @@ describe('Collation', function () { request.reply(primary[0]); } else if (doc.find) { commandResult = doc; - request.reply({ ok: 1, cursor: { id: 0, firstBatch: [] } }); + request.reply({ ok: 1, cursor: { id: 0n, firstBatch: [] } }); } else if (doc.endSessions) { request.reply({ ok: 1 }); } diff --git a/test/unit/assorted/sessions_collection.test.js b/test/unit/assorted/sessions_collection.test.js index eee1a76ec98..409d818d136 100644 --- a/test/unit/assorted/sessions_collection.test.js +++ b/test/unit/assorted/sessions_collection.test.js @@ -27,7 +27,7 @@ describe('Sessions - unit/sessions', function () { request.reply({ ok: 1, operationTime: insertOperationTime }); } else if (doc.find) { findCommand = doc; - request.reply({ ok: 1, cursor: { id: 0, firstBatch: [] } }); + request.reply({ ok: 1, cursor: { id: 0n, firstBatch: [] } }); } else if (doc.endSessions) { request.reply({ ok: 1 }); } From df5121a8072a2d394541803edfa622a78a0cfdb9 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Fri, 19 Apr 2024 11:50:30 -0400 Subject: [PATCH 09/25] perf: define methods on cursor response --- src/cmap/wire_protocol/responses.ts | 83 ++++++++++++++--------------- src/cursor/abstract_cursor.ts | 4 +- 2 files changed, 41 insertions(+), 46 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index f5d8dd8741c..7b5284dc26e 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -95,35 +95,31 @@ export class MongoDBResponse extends OnDemandDocument { return this.clusterTime ?? null; } - public override toObject(options: BSONSerializeOptions = {}): Record { + public override toObject(options?: BSONSerializeOptions): Record { const exactBSONOptions = { - useBigInt64: options.useBigInt64, - promoteLongs: options.promoteLongs, - promoteValues: options.promoteValues, - promoteBuffers: options.promoteBuffers, - bsonRegExp: options.bsonRegExp, - raw: options.raw ?? false, - fieldsAsRaw: options.fieldsAsRaw ?? {}, + useBigInt64: options?.useBigInt64, + promoteLongs: options?.promoteLongs, + promoteValues: options?.promoteValues, + promoteBuffers: options?.promoteBuffers, + bsonRegExp: options?.bsonRegExp, + raw: options?.raw ?? false, + fieldsAsRaw: options?.fieldsAsRaw ?? {}, validation: this.parseBsonSerializationOptions(options) }; return super.toObject(exactBSONOptions); } - private parseBsonSerializationOptions({ enableUtf8Validation }: BSONSerializeOptions): { + private parseBsonSerializationOptions(options?: { enableUtf8Validation?: boolean }): { utf8: { writeErrors: false } | false; } { + const enableUtf8Validation = options?.enableUtf8Validation; if (enableUtf8Validation === false) { return { utf8: false }; } - return { utf8: { writeErrors: false } }; } } -function throwUnsupportedError() { - throw new Error('Unsupported method'); -} - /** @internal */ export class CursorResponse extends MongoDBResponse { static emptyGetMore = new CursorResponse( @@ -139,7 +135,6 @@ export class CursorResponse extends MongoDBResponse { public id: Long | null = null; public ns: MongoDBNamespace | null = null; - public documents: any | null = null; public batchSize = 0; private batch: OnDemandDocument | null = null; @@ -163,35 +158,35 @@ export class CursorResponse extends MongoDBResponse { else if (cursor.has('nextBatch')) this.batch = cursor.get('nextBatch', BSONType.array, true); else throw new MongoUnexpectedServerResponseError('Cursor document did not contain a batch'); - this.values = this.batch.valuesAs(BSONType.object); this.batchSize = this.batch.size(); - this.iterated = 0; - this.documents = Object.defineProperties(Object.create(null), { - length: { - get: () => { - return Math.max(this.batchSize - this.iterated, 0); - } - }, - shift: { - value: (options?: BSONSerializeOptions) => { - this.iterated += 1; - const result = this.values?.next(); - if (!result || result.done) return null; - if (options?.raw) { - return result.value.toBytes(); - } else { - return result.value.toObject(options); - } - } - }, - clear: { - value: () => { - this.iterated = this.batchSize; - this.values?.return(); - } - }, - pushMany: { value: throwUnsupportedError }, - push: { value: throwUnsupportedError } - }); + } + + get length() { + return Math.max(this.batchSize - this.iterated, 0); + } + + shift(options?: BSONSerializeOptions): any { + this.iterated += 1; + this.values ??= this.batch?.valuesAs(BSONType.object) ?? null; + const result = this.values?.next(); + if (!result || result.done) return null; + if (options?.raw) { + return result.value.toBytes(); + } else { + return result.value.toObject(options); + } + } + + clear() { + this.iterated = this.batchSize; + this.values?.return(); + } + + pushMany() { + throw new Error('pushMany Unsupported method'); + } + + push() { + throw new Error('push Unsupported method'); } } diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 6d434d48cdf..39658c637e4 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -660,7 +660,7 @@ export abstract class AbstractCursor< if (CursorResponse.is(response)) { this[kId] = response.id; if (response.ns) this[kNamespace] = response.ns; - this[kDocuments] = response.documents; + this[kDocuments] = response; } else if (response.cursor) { // TODO(NODE-2674): Preserve int64 sent from MongoDB this[kId] = @@ -800,7 +800,7 @@ async function next( const response = await cursor.getMore(batchSize); if (CursorResponse.is(response)) { cursor[kId] = response.id; - cursor[kDocuments] = response.documents; + cursor[kDocuments] = response; } else if (response) { const cursorId = typeof response.cursor.id === 'number' From 053f01fe7b486cd9e396895fccf1dc1d29389d1c Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 10:01:08 -0400 Subject: [PATCH 10/25] refactor: replace isBuffer check with isUint8Array --- src/client-side-encryption/auto_encrypter.ts | 4 +-- src/utils.ts | 13 ++++++++++ test/unit/utils.test.ts | 26 ++++++++++++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/client-side-encryption/auto_encrypter.ts b/src/client-side-encryption/auto_encrypter.ts index 7ced6d74237..a0a5177cbed 100644 --- a/src/client-side-encryption/auto_encrypter.ts +++ b/src/client-side-encryption/auto_encrypter.ts @@ -10,7 +10,7 @@ import { MongoDBResponse, type MongoDBResponseConstructor } from '../cmap/wire_p import { getMongoDBClientEncryption } from '../deps'; import { MongoRuntimeError } from '../error'; import { MongoClient, type MongoClientOptions } from '../mongo_client'; -import { MongoDBCollectionNamespace } from '../utils'; +import { isUint8Array, MongoDBCollectionNamespace } from '../utils'; import * as cryptoCallbacks from './crypto_callbacks'; import { MongoCryptInvalidArgumentError } from './errors'; import { MongocryptdManager } from './mongocryptd_manager'; @@ -480,7 +480,7 @@ export class AutoEncrypter { ): Promise { const buffer = MongoDBResponse.is(response) ? response.toBytes() - : Buffer.isBuffer(response) + : isUint8Array(response) ? response : serialize(response, options); diff --git a/src/utils.ts b/src/utils.ts index bf34a3d5196..57079b1f639 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -64,6 +64,19 @@ export const ByteUtils = { } }; +/** + * Returns true if value is a Uint8Array or a Buffer + * @param value - any value that may be a Uint8Array + */ +export function isUint8Array(value: unknown): value is Uint8Array { + return ( + value != null && + typeof value === 'object' && + Symbol.toStringTag in value && + value[Symbol.toStringTag] === 'Uint8Array' + ); +} + /** * Determines if a connection's address matches a user provided list * of domain wildcards. diff --git a/test/unit/utils.test.ts b/test/unit/utils.test.ts index 802b9bc5645..0184d44c5b2 100644 --- a/test/unit/utils.test.ts +++ b/test/unit/utils.test.ts @@ -7,6 +7,7 @@ import { HostAddress, hostMatchesWildcards, isHello, + isUint8Array, LEGACY_HELLO_COMMAND, List, matchesParentDomain, @@ -981,4 +982,29 @@ describe('driver utils', function () { }); }); }); + + describe('isUint8Array()', () => { + describe('when given a UintArray', () => + it('returns true', () => expect(isUint8Array(Uint8Array.from([1]))).to.be.true)); + + describe('when given a Buffer', () => + it('returns true', () => expect(isUint8Array(Buffer.from([1]))).to.be.true)); + + describe('when given a value that does not have `Uint8Array` at Symbol.toStringTag', () => { + it('returns false', () => { + const weirdArray = Uint8Array.from([1]); + Object.defineProperty(weirdArray, Symbol.toStringTag, { value: 'blah' }); + expect(isUint8Array(weirdArray)).to.be.false; + }); + }); + + describe('when given null', () => + it('returns false', () => expect(isUint8Array(null)).to.be.false)); + + describe('when given a non object', () => + it('returns false', () => expect(isUint8Array('')).to.be.false)); + + describe('when given an object that does not respond to Symbol.toStringTag', () => + it('returns false', () => expect(isUint8Array(Object.create(null))).to.be.false)); + }); }); From 7edb9480590818409ee9d8855f535d178dcebede Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 10:08:37 -0400 Subject: [PATCH 11/25] fix: add correct ts to execute --- src/client-side-encryption/state_machine.ts | 22 ++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/client-side-encryption/state_machine.ts b/src/client-side-encryption/state_machine.ts index 0c6c4da0bc7..0d40ee07126 100644 --- a/src/client-side-encryption/state_machine.ts +++ b/src/client-side-encryption/state_machine.ts @@ -153,13 +153,29 @@ export class StateMachine { ) {} /** - * Executes the state machine according to the specification + * Executes the state machine according to the specification. + * Will construct the result using `responseType`. + */ + async execute( + executor: StateMachineExecutable, + context: MongoCryptContext, + responseType?: R + ): Promise>; + + /** + * Executes the state machine according to the specification. + * Will return a document from the default BSON deserializer. */ async execute( + executor: StateMachineExecutable, + context: MongoCryptContext + ): Promise; + + async execute( executor: StateMachineExecutable, context: MongoCryptContext, - responseType?: MongoDBResponseConstructor - ): Promise { + responseType?: R + ): Promise> { const keyVaultNamespace = executor._keyVaultNamespace; const keyVaultClient = executor._keyVaultClient; const metaDataClient = executor._metaDataClient; From 4631e34617840fb15fc265a6882c9e94137f8d00 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 12:16:13 -0400 Subject: [PATCH 12/25] docs: add comment for base64 string --- src/cmap/wire_protocol/responses.ts | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 7b5284dc26e..2a5d1fe2e8e 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -122,6 +122,15 @@ export class MongoDBResponse extends OnDemandDocument { /** @internal */ export class CursorResponse extends MongoDBResponse { + /** + * This is a BSON document containing the following: + * ``` + * { ok: 1, cursor: { id: 0n, nextBatch: new Array(0) } } + * ``` + * This is used when the client side findCursor is closed by tracking the number returned and limit + * to avoid an extra round trip. It provides a cursor response that the server _would_ return _if_ + * that round trip were to be made. + */ static emptyGetMore = new CursorResponse( Buffer.from( 'NgAAABBvawABAAAAA2N1cnNvcgAhAAAAEmlkAAAAAAAAAAAABG5leHRCYXRjaAAFAAAAAAAA', From ac5118e71a916e223b5612da432028988b2beac5 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 13:17:57 -0400 Subject: [PATCH 13/25] feat: remove valuesAs just use indexing --- src/cmap/wire_protocol/on_demand/document.ts | 48 +++++++++-------- src/cmap/wire_protocol/responses.ts | 23 ++++---- .../wire_protocol/on_demand/document.test.ts | 53 ++++++------------- 3 files changed, 56 insertions(+), 68 deletions(-) diff --git a/src/cmap/wire_protocol/on_demand/document.ts b/src/cmap/wire_protocol/on_demand/document.ts index 0b851d27d03..638946d647f 100644 --- a/src/cmap/wire_protocol/on_demand/document.ts +++ b/src/cmap/wire_protocol/on_demand/document.ts @@ -58,7 +58,7 @@ export class OnDemandDocument { private readonly indexFound: Record = Object.create(null); /** All bson elements in this document */ - private readonly elements: BSONElement[]; + private readonly elements: ReadonlyArray; constructor( /** BSON bytes, this document begins at offset */ @@ -97,7 +97,7 @@ export class OnDemandDocument { * @param name - a basic latin string name of a BSON element * @returns */ - private getElement(name: string): CachedBSONElement | null { + private getElement(name: string | number): CachedBSONElement | null { const cachedElement = this.cache[name]; if (cachedElement === false) return null; @@ -105,6 +105,22 @@ export class OnDemandDocument { return cachedElement; } + if (typeof name === 'number') { + if (this.isArray) { + if (name < this.elements.length) { + const element = this.elements[name]; + const cachedElement = { element, value: undefined }; + this.cache[name] = cachedElement; + this.indexFound[name] = true; + return cachedElement; + } else { + return null; + } + } else { + return null; + } + } + for (let index = 0; index < this.elements.length; index++) { const element = this.elements[index]; @@ -229,16 +245,20 @@ export class OnDemandDocument { * @param required - whether or not the element is expected to exist, if true this function will throw if it is not present */ public get( - name: string, + name: string | number, as: T, required?: false | undefined ): JSTypeOf[T] | null; /** `required` will make `get` throw if name does not exist or is null/undefined */ - public get(name: string, as: T, required: true): JSTypeOf[T]; + public get( + name: string | number, + as: T, + required: true + ): JSTypeOf[T]; public get( - name: string, + name: string | number, as: T, required?: boolean ): JSTypeOf[T] | null { @@ -315,22 +335,4 @@ export class OnDemandDocument { const size = getInt32LE(this.bson, this.offset); return this.bson.subarray(this.offset, this.offset + size); } - - /** - * Iterates through the elements of a document reviving them using the `as` BSONType. - * - * @param as - The type to revive all elements as - */ - public *valuesAs(as: T): Generator { - if (!this.isArray) { - throw new BSONError('Unexpected conversion of non-array value to array'); - } - let counter = 0; - for (const element of this.elements) { - const value = this.toJSValue(element, as); - this.cache[counter] = { element, value }; - yield value; - counter += 1; - } - } } diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 2a5d1fe2e8e..78694e0efbc 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -147,11 +147,10 @@ export class CursorResponse extends MongoDBResponse { public batchSize = 0; private batch: OnDemandDocument | null = null; - private values: Generator | null = null; private iterated = 0; - constructor(b: Uint8Array, o?: number, a?: boolean) { - super(b, o, a); + constructor(bytes: Uint8Array, offset?: number, isArray?: boolean) { + super(bytes, offset, isArray); if (this.isError) return; @@ -175,20 +174,26 @@ export class CursorResponse extends MongoDBResponse { } shift(options?: BSONSerializeOptions): any { + if (this.iterated >= this.batchSize) { + return null; + } + + const result = this.batch?.get(this.iterated, BSONType.object, true) ?? null; this.iterated += 1; - this.values ??= this.batch?.valuesAs(BSONType.object) ?? null; - const result = this.values?.next(); - if (!result || result.done) return null; + + if (result == null) { + throw new MongoUnexpectedServerResponseError('Cursor batch contains null values'); + } + if (options?.raw) { - return result.value.toBytes(); + return result.toBytes(); } else { - return result.value.toObject(options); + return result.toObject(options); } } clear() { this.iterated = this.batchSize; - this.values?.return(); } pushMany() { diff --git a/test/unit/cmap/wire_protocol/on_demand/document.test.ts b/test/unit/cmap/wire_protocol/on_demand/document.test.ts index 6a7d5bb10cb..82ed4040f64 100644 --- a/test/unit/cmap/wire_protocol/on_demand/document.test.ts +++ b/test/unit/cmap/wire_protocol/on_demand/document.test.ts @@ -73,6 +73,7 @@ describe('class OnDemandDocument', () => { context('get()', () => { let document: OnDemandDocument; + let array: OnDemandDocument; const input = { int: 1, double: 1.2, @@ -86,12 +87,27 @@ describe('class OnDemandDocument', () => { date: new Date(0), object: { a: 1 }, array: [1, 2], - unsupportedType: /abc/ + unsupportedType: /abc/, + [233]: 3 }; beforeEach(async function () { const bytes = BSON.serialize(input); document = new OnDemandDocument(bytes); + array = new OnDemandDocument( + BSON.serialize(Object.fromEntries(Object.values(input).entries())), + 0, + true + ); + }); + + it('supports access by number for arrays', () => { + expect(array.get(1, BSONType.int)).to.equal(1); + }); + + it('does not support access by number for objects', () => { + expect(document.get(233, BSONType.int)).to.be.null; + expect(document.get('233', BSONType.int)).to.equal(3); }); it('returns null if the element does not exist', () => { @@ -277,39 +293,4 @@ describe('class OnDemandDocument', () => { expect(document.getNumber('boolTrue')).to.equal(1); }); }); - - context('*valuesAs()', () => { - let array: OnDemandDocument; - beforeEach(async function () { - const bytes = BSON.serialize( - Object.fromEntries(Array.from({ length: 10 }, () => 1).entries()) - ); - array = new OnDemandDocument(bytes, 0, true); - }); - - it('throws if document is not an array', () => { - const bytes = BSON.serialize( - Object.fromEntries(Array.from({ length: 10 }, () => 1).entries()) - ); - array = new OnDemandDocument(bytes, 0, false); - expect(() => array.valuesAs(BSONType.int).next()).to.throw(); - }); - - it('returns a generator that yields values matching the as BSONType parameter', () => { - let didRun = false; - for (const item of array.valuesAs(BSONType.int)) { - didRun = true; - expect(item).to.equal(1); - } - expect(didRun).to.be.true; - }); - - it('caches the results of array', () => { - const generator = array.valuesAs(BSONType.int); - generator.next(); - generator.next(); - expect(array).to.have.nested.property('cache.0.value', 1); - expect(array).to.have.nested.property('cache.1.value', 1); - }); - }); }); From 9f80b5b4273465d3beb3a11f77cf7254c02853bd Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 13:51:30 -0400 Subject: [PATCH 14/25] feat: required fields by moving error check out of constructor --- src/cmap/connection.ts | 13 +++++-- src/cmap/wire_protocol/responses.ts | 55 ++++++++++++++++++++++++++--- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/src/cmap/connection.ts b/src/cmap/connection.ts index 1213e158ad7..bfdd23bfb8e 100644 --- a/src/cmap/connection.ts +++ b/src/cmap/connection.ts @@ -62,7 +62,11 @@ import type { ClientMetadata } from './handshake/client_metadata'; import { StreamDescription, type StreamDescriptionOptions } from './stream_description'; import { type CompressorName, decompressResponse } from './wire_protocol/compression'; import { onData } from './wire_protocol/on_data'; -import { MongoDBResponse, type MongoDBResponseConstructor } from './wire_protocol/responses'; +import { + isErrorResponse, + MongoDBResponse, + type MongoDBResponseConstructor +} from './wire_protocol/responses'; import { getReadPreference, isSharded } from './wire_protocol/shared'; /** @internal */ @@ -443,7 +447,12 @@ export class Connection extends TypedEventEmitter { this.socket.setTimeout(0); const bson = response.parse(); - const document = new (responseType ?? MongoDBResponse)(bson, 0, false); + const document = + responseType == null + ? new MongoDBResponse(bson) + : isErrorResponse(bson) + ? new MongoDBResponse(bson) + : new responseType(bson); yield document; this.throwIfAborted(); diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 78694e0efbc..c72d3389803 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -3,6 +3,7 @@ import { BSONType, type Document, Long, + parseToElementsToArray, type Timestamp } from '../../bson'; import { MongoUnexpectedServerResponseError } from '../../error'; @@ -10,6 +11,52 @@ import { type ClusterTime } from '../../sdam/common'; import { type MongoDBNamespace, ns } from '../../utils'; import { OnDemandDocument } from './on_demand/document'; +// eslint-disable-next-line no-restricted-syntax +const enum BSONElementOffset { + type = 0, + nameOffset = 1, + nameLength = 2, + offset = 3, + length = 4 +} +/** + * Accepts a BSON payload and checks for na "ok: 0" element. + * This utility is intended to prevent calling response class constructors + * that expect the result to be a success and demand certain properties to exist. + * + * For example, a cursor response always expects a cursor embedded document. + * In order to write the class such that the properties reflect that assertion (non-null) + * we cannot invoke the subclass constructor if the BSON represents an error. + * + * @param bytes - BSON document returned from the server + */ +export function isErrorResponse(bson: Uint8Array): boolean { + const elements = parseToElementsToArray(bson, 0); + for (let eIdx = 0; eIdx < elements.length; eIdx++) { + const element = elements[eIdx]; + + if (element[BSONElementOffset.nameLength] === 2) { + const nameOffset = element[BSONElementOffset.nameOffset]; + + // 111 == "o", 107 == "k" + if (bson[nameOffset] === 111 && bson[nameOffset + 1] === 107) { + const valueOffset = element[BSONElementOffset.offset]; + const valueLength = element[BSONElementOffset.length]; + + // If any byte in the length of the ok number (works for any type) is non zero, + // then it is considered "ok: 1" + for (let i = valueOffset; i < valueOffset + valueLength; i++) { + if (bson[i] !== 0x00) return false; + } + + return true; + } + } + } + + return true; +} + /** @internal */ export type MongoDBResponseConstructor = { new (bson: Uint8Array, offset?: number, isArray?: boolean): MongoDBResponse; @@ -142,18 +189,16 @@ export class CursorResponse extends MongoDBResponse { return value instanceof CursorResponse; } - public id: Long | null = null; - public ns: MongoDBNamespace | null = null; + public id: Long; + public ns: MongoDBNamespace; public batchSize = 0; - private batch: OnDemandDocument | null = null; + private batch: OnDemandDocument; private iterated = 0; constructor(bytes: Uint8Array, offset?: number, isArray?: boolean) { super(bytes, offset, isArray); - if (this.isError) return; - const cursor = this.get('cursor', BSONType.object, true); const id = cursor.get('id', BSONType.long, true); From 0902172725b9770dfaa7b7eae4ee2d4fa92b89d6 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 13:53:13 -0400 Subject: [PATCH 15/25] fix: ns is not a required field --- src/cmap/wire_protocol/responses.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index c72d3389803..81793b7396a 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -190,7 +190,7 @@ export class CursorResponse extends MongoDBResponse { } public id: Long; - public ns: MongoDBNamespace; + public ns: MongoDBNamespace | null = null; public batchSize = 0; private batch: OnDemandDocument; @@ -204,8 +204,8 @@ export class CursorResponse extends MongoDBResponse { const id = cursor.get('id', BSONType.long, true); this.id = new Long(Number(id & 0xffff_ffffn), Number((id >> 32n) & 0xffff_ffffn)); - const namespace = cursor.get('ns', BSONType.string) ?? ''; - if (namespace) this.ns = ns(namespace); + const namespace = cursor.get('ns', BSONType.string); + if (namespace != null) this.ns = ns(namespace); if (cursor.has('firstBatch')) this.batch = cursor.get('firstBatch', BSONType.array, true); else if (cursor.has('nextBatch')) this.batch = cursor.get('nextBatch', BSONType.array, true); From ea365c11fc84bfcd3f078194aea7031063369674 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 15:01:53 -0400 Subject: [PATCH 16/25] fix: rename hasNext --- src/cursor/abstract_cursor.ts | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 39658c637e4..1e92dace1eb 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -389,7 +389,7 @@ export abstract class AbstractCursor< return true; } - return await next(this, { blocking: true, transform: false, hasNext: true }); + return await next(this, { blocking: true, transform: false, shift: false }); } /** Get the next available document from the cursor, returns null if no more documents are available. */ @@ -398,7 +398,7 @@ export abstract class AbstractCursor< throw new MongoCursorExhaustedError(); } - return await next(this, { blocking: true, transform: true, hasNext: false }); + return await next(this, { blocking: true, transform: true, shift: true }); } /** @@ -409,7 +409,7 @@ export abstract class AbstractCursor< throw new MongoCursorExhaustedError(); } - return await next(this, { blocking: false, transform: true, hasNext: false }); + return await next(this, { blocking: false, transform: true, shift: true }); } /** @@ -719,11 +719,11 @@ async function next( { blocking, transform, - hasNext + shift }: { blocking: boolean; transform: boolean; - hasNext: true; + shift: false; } ): Promise; @@ -732,11 +732,11 @@ async function next( { blocking, transform, - hasNext + shift }: { blocking: boolean; transform: boolean; - hasNext: false; + shift: true; } ): Promise; @@ -745,15 +745,15 @@ async function next( { blocking, transform, - hasNext + shift }: { blocking: boolean; transform: boolean; - hasNext: boolean; + shift: boolean; } ): Promise { if (cursor.closed) { - if (hasNext) return false; + if (!shift) return false; return null; } @@ -764,7 +764,7 @@ async function next( } if (cursor[kDocuments].length !== 0) { - if (hasNext) return true; + if (!shift) return true; const doc = cursor[kDocuments].shift(cursor[kOptions]); if (doc != null && transform && cursor[kTransform]) { @@ -789,7 +789,7 @@ async function next( // cleanupCursor should never throw, but if it does it indicates a bug in the driver // and we should surface the error await cleanupCursor(cursor, {}); - if (hasNext) return false; + if (!shift) return false; return null; } @@ -834,12 +834,12 @@ async function next( } if (cursor[kDocuments].length === 0 && blocking === false) { - if (hasNext) return false; + if (!shift) return false; return null; } } while (!cursor.isDead || cursor[kDocuments].length !== 0); - if (hasNext) return false; + if (!shift) return false; return null; } @@ -961,7 +961,7 @@ class ReadableCursorStream extends Readable { private _readNext() { // eslint-disable-next-line github/no-then - next(this._cursor, { blocking: true, transform: true, hasNext: false }).then( + next(this._cursor, { blocking: true, transform: true, shift: false }).then( result => { if (result == null) { this.push(null); From 24b7678d42224f0e4cae1282c4390a945e093a26 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 16:20:23 -0400 Subject: [PATCH 17/25] fix: shift: false wrong for stream, stream uses same settings as `cursor.next` --- src/cursor/abstract_cursor.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 1e92dace1eb..7f0a548974d 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -961,7 +961,7 @@ class ReadableCursorStream extends Readable { private _readNext() { // eslint-disable-next-line github/no-then - next(this._cursor, { blocking: true, transform: true, shift: false }).then( + this._cursor.next().then( result => { if (result == null) { this.push(null); From 195bd62321fb39d61e0bc369cf07529fe2343b8b Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 16:51:07 -0400 Subject: [PATCH 18/25] fix: stream cannot use cursor's next because it reaches exhaustion error in changestream prose test --- src/cursor/abstract_cursor.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts index 7f0a548974d..c4f349500a1 100644 --- a/src/cursor/abstract_cursor.ts +++ b/src/cursor/abstract_cursor.ts @@ -961,7 +961,7 @@ class ReadableCursorStream extends Readable { private _readNext() { // eslint-disable-next-line github/no-then - this._cursor.next().then( + next(this._cursor, { blocking: true, transform: true, shift: true }).then( result => { if (result == null) { this.push(null); From 2f10f0e8d1e80a9347d20472c4e9d28e28934731 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 17:36:04 -0400 Subject: [PATCH 19/25] test: add unit tests for cursor response ctor --- .../unit/cmap/wire_protocol/responses.test.ts | 97 ++++++++++++++++++- 1 file changed, 96 insertions(+), 1 deletion(-) diff --git a/test/unit/cmap/wire_protocol/responses.test.ts b/test/unit/cmap/wire_protocol/responses.test.ts index fc5ee88ae16..91e052da84f 100644 --- a/test/unit/cmap/wire_protocol/responses.test.ts +++ b/test/unit/cmap/wire_protocol/responses.test.ts @@ -1,7 +1,15 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; -import { BSON, MongoDBResponse, OnDemandDocument } from '../../../mongodb'; +import { + BSON, + BSONError, + CursorResponse, + Int32, + MongoDBResponse, + MongoUnexpectedServerResponseError, + OnDemandDocument +} from '../../../mongodb'; describe('class MongoDBResponse', () => { it('is a subclass of OnDemandDocument', () => { @@ -76,3 +84,90 @@ describe('class MongoDBResponse', () => { }); }); }); + +describe('class CursorResponse', () => { + describe('constructor()', () => { + it('throws if input does not contain cursor embedded document', () => { + expect(() => new CursorResponse(BSON.serialize({ ok: 1 }))).to.throw(BSONError); + }); + + it('throws if input does not contain cursor.id int64', () => { + expect(() => new CursorResponse(BSON.serialize({ ok: 1, cursor: {} }))).to.throw(BSONError); + }); + + it('sets namespace to null if input does not contain cursor.ns', () => { + expect(new CursorResponse(BSON.serialize({ ok: 1, cursor: { id: 0n, firstBatch: [] } })).ns) + .to.be.null; + }); + + it('throws if input does not contain firstBatch nor nextBatch', () => { + expect( + () => new CursorResponse(BSON.serialize({ ok: 1, cursor: { id: 0n, batch: [] } })) + ).to.throw(MongoUnexpectedServerResponseError); + }); + + it('reports a length equal to the batch', () => { + expect( + new CursorResponse(BSON.serialize({ ok: 1, cursor: { id: 0n, nextBatch: [1, 2, 3] } })) + ).to.have.lengthOf(3); + }); + }); + + describe('shift()', () => { + let response; + + beforeEach(async function () { + response = new CursorResponse( + BSON.serialize({ + ok: 1, + cursor: { id: 0n, nextBatch: [{ _id: 1 }, { _id: 2 }, { _id: 3 }] } + }) + ); + }); + + it('returns a document from the batch', () => { + expect(response.shift()).to.deep.equal({ _id: 1 }); + expect(response.shift()).to.deep.equal({ _id: 2 }); + expect(response.shift()).to.deep.equal({ _id: 3 }); + expect(response.shift()).to.deep.equal(null); + }); + + it('passes BSON options to deserialization', () => { + expect(response.shift({ promoteValues: false })).to.deep.equal({ _id: new Int32(1) }); + expect(response.shift({ promoteValues: true })).to.deep.equal({ _id: 2 }); + expect(response.shift({ promoteValues: false })).to.deep.equal({ _id: new Int32(3) }); + expect(response.shift()).to.deep.equal(null); + }); + }); + + describe('clear()', () => { + let response; + + beforeEach(async function () { + response = new CursorResponse( + BSON.serialize({ + ok: 1, + cursor: { id: 0n, nextBatch: [{ _id: 1 }, { _id: 2 }, { _id: 3 }] } + }) + ); + }); + + it('makes length equal to 0', () => { + expect(response.clear()).to.be.undefined; + expect(response).to.have.lengthOf(0); + }); + + it('makes shift return null', () => { + expect(response.clear()).to.be.undefined; + expect(response.shift()).to.be.null; + }); + }); + + describe('pushMany()', () => + it('throws unsupported error', () => + expect(CursorResponse.prototype.pushMany).to.throw(/Unsupported/i))); + + describe('push()', () => + it('throws unsupported error', () => + expect(CursorResponse.prototype.push).to.throw(/Unsupported/i))); +}); From ca8816848eea788100575bd4b0737f96c243d64c Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Tue, 23 Apr 2024 17:36:29 -0400 Subject: [PATCH 20/25] fix: batch always exists and emptyGetMore made simpler --- src/cmap/wire_protocol/responses.ts | 24 +++++------------------- 1 file changed, 5 insertions(+), 19 deletions(-) diff --git a/src/cmap/wire_protocol/responses.ts b/src/cmap/wire_protocol/responses.ts index 81793b7396a..65515cbb316 100644 --- a/src/cmap/wire_protocol/responses.ts +++ b/src/cmap/wire_protocol/responses.ts @@ -170,23 +170,13 @@ export class MongoDBResponse extends OnDemandDocument { /** @internal */ export class CursorResponse extends MongoDBResponse { /** - * This is a BSON document containing the following: - * ``` - * { ok: 1, cursor: { id: 0n, nextBatch: new Array(0) } } - * ``` - * This is used when the client side findCursor is closed by tracking the number returned and limit - * to avoid an extra round trip. It provides a cursor response that the server _would_ return _if_ - * that round trip were to be made. + * This supports a feature of the FindCursor. + * It is an optimization to avoid an extra getMore when the limit has been reached */ - static emptyGetMore = new CursorResponse( - Buffer.from( - 'NgAAABBvawABAAAAA2N1cnNvcgAhAAAAEmlkAAAAAAAAAAAABG5leHRCYXRjaAAFAAAAAAAA', - 'base64' - ) - ); + static emptyGetMore = { id: new Long(0), length: 0, shift: () => null }; static override is(value: unknown): value is CursorResponse { - return value instanceof CursorResponse; + return value instanceof CursorResponse || value === CursorResponse.emptyGetMore; } public id: Long; @@ -223,13 +213,9 @@ export class CursorResponse extends MongoDBResponse { return null; } - const result = this.batch?.get(this.iterated, BSONType.object, true) ?? null; + const result = this.batch.get(this.iterated, BSONType.object, true) ?? null; this.iterated += 1; - if (result == null) { - throw new MongoUnexpectedServerResponseError('Cursor batch contains null values'); - } - if (options?.raw) { return result.toBytes(); } else { From 05334145681c7a8021c215b1d41d4719dee9227d Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 25 Apr 2024 17:32:39 -0400 Subject: [PATCH 21/25] test: demonstrate internal decorate feature breakage --- .../client-side-encryption/driver.test.ts | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/test/integration/client-side-encryption/driver.test.ts b/test/integration/client-side-encryption/driver.test.ts index d22033c890a..35eff329256 100644 --- a/test/integration/client-side-encryption/driver.test.ts +++ b/test/integration/client-side-encryption/driver.test.ts @@ -297,4 +297,108 @@ describe('Client Side Encryption Functional', function () { }); }); }); + + describe('when @@mdb.decorateDecryptionResult is set on autoEncrypter', () => { + let client: MongoClient; + let encryptedClient: MongoClient; + + beforeEach(async function () { + client = this.configuration.newClient(); + + const encryptSchema = (keyId: unknown, bsonType: string) => ({ + encrypt: { + bsonType, + algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Random', + keyId: [keyId] + } + }); + + const kmsProviders = this.configuration.kmsProviders(crypto.randomBytes(96)); + + await client.connect(); + + const encryption = new ClientEncryption(client, { + keyVaultNamespace, + kmsProviders, + extraOptions: getEncryptExtraOptions() + }); + + const dataDb = client.db(dataDbName); + const keyVaultDb = client.db(keyVaultDbName); + + await dataDb.dropCollection(dataCollName).catch(() => null); + await keyVaultDb.dropCollection(keyVaultCollName).catch(() => null); + await keyVaultDb.createCollection(keyVaultCollName); + const dataKey = await encryption.createDataKey('local'); + + const $jsonSchema = { + bsonType: 'object', + properties: { + a: encryptSchema(dataKey, 'int'), + b: encryptSchema(dataKey, 'string'), + c: { + bsonType: 'object', + properties: { + d: { + encrypt: { + keyId: [dataKey], + algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic', + bsonType: 'string' + } + } + } + } + } + }; + + await dataDb.createCollection(dataCollName, { + validator: { $jsonSchema } + }); + + encryptedClient = this.configuration.newClient( + {}, + { + autoEncryption: { + keyVaultNamespace, + kmsProviders, + extraOptions: getEncryptExtraOptions() + } + } + ); + + encryptedClient.autoEncrypter[Symbol.for('@@mdb.decorateDecryptionResult')] = true; + await encryptedClient.connect(); + }); + + afterEach(function () { + return Promise.resolve() + .then(() => encryptedClient?.close()) + .then(() => client?.close()); + }); + + it('adds decrypted keys to result at @@mdb.decryptedKeys', async function () { + const coll = encryptedClient.db(dataDbName).collection(dataCollName); + + const data = { + _id: new BSON.ObjectId(), + a: 1, + b: 'abc', + c: { d: 'def' } + }; + + const result = await coll.insertOne(data); + const decrypted = await coll.findOne({ _id: result.insertedId }); + + expect(decrypted).to.deep.equal(data); + expect(decrypted) + .to.have.property(Symbol.for('@@mdb.decryptedKeys')) + .that.deep.equals(['a', 'b']); + + // Nested + expect(decrypted).to.have.property('c'); + expect(decrypted.c) + .to.have.property(Symbol.for('@@mdb.decryptedKeys')) + .that.deep.equals(['d']); + }); + }); }); From 3ee316d629839e03d38c0a2beeff0fcebeebd0d1 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 25 Apr 2024 17:37:18 -0400 Subject: [PATCH 22/25] fix: revert changes to FLE and exempt from using responseType --- src/client-side-encryption/auto_encrypter.ts | 21 +++---------- src/client-side-encryption/state_machine.ts | 31 +++---------------- src/cmap/connection.ts | 6 ++-- src/cursor/find_cursor.ts | 8 +++-- src/operations/find.ts | 13 +++----- .../crud/abstract_operation.test.ts | 2 +- test/unit/operations/find.test.ts | 6 ++-- 7 files changed, 25 insertions(+), 62 deletions(-) diff --git a/src/client-side-encryption/auto_encrypter.ts b/src/client-side-encryption/auto_encrypter.ts index a0a5177cbed..e6334dc57de 100644 --- a/src/client-side-encryption/auto_encrypter.ts +++ b/src/client-side-encryption/auto_encrypter.ts @@ -6,11 +6,10 @@ import { import { deserialize, type Document, serialize } from '../bson'; import { type CommandOptions, type ProxyOptions } from '../cmap/connection'; -import { MongoDBResponse, type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses'; import { getMongoDBClientEncryption } from '../deps'; import { MongoRuntimeError } from '../error'; import { MongoClient, type MongoClientOptions } from '../mongo_client'; -import { isUint8Array, MongoDBCollectionNamespace } from '../utils'; +import { MongoDBCollectionNamespace } from '../utils'; import * as cryptoCallbacks from './crypto_callbacks'; import { MongoCryptInvalidArgumentError } from './errors'; import { MongocryptdManager } from './mongocryptd_manager'; @@ -474,19 +473,8 @@ export class AutoEncrypter { /** * Decrypt a command response */ - async decrypt( - response: Uint8Array | Document | MongoDBResponse, - options: CommandOptions = {} - ): Promise { - const buffer = MongoDBResponse.is(response) - ? response.toBytes() - : isUint8Array(response) - ? response - : serialize(response, options); - - const responseType = MongoDBResponse.is(response) - ? (response.constructor as MongoDBResponseConstructor) - : undefined; + async decrypt(response: Uint8Array | Document, options: CommandOptions = {}): Promise { + const buffer = Buffer.isBuffer(response) ? response : serialize(response, options); const context = this._mongocrypt.makeDecryptionContext(buffer); @@ -499,11 +487,10 @@ export class AutoEncrypter { }); const decorateResult = this[kDecorateResult]; - const result = await stateMachine.execute(this, context, responseType); + const result = await stateMachine.execute(this, context); if (decorateResult) { decorateDecryptionResult(result, response); } - return result; } diff --git a/src/client-side-encryption/state_machine.ts b/src/client-side-encryption/state_machine.ts index 0d40ee07126..a4b2379fb51 100644 --- a/src/client-side-encryption/state_machine.ts +++ b/src/client-side-encryption/state_machine.ts @@ -11,7 +11,6 @@ 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'; @@ -153,35 +152,18 @@ export class StateMachine { ) {} /** - * Executes the state machine according to the specification. - * Will construct the result using `responseType`. - */ - async execute( - executor: StateMachineExecutable, - context: MongoCryptContext, - responseType?: R - ): Promise>; - - /** - * Executes the state machine according to the specification. - * Will return a document from the default BSON deserializer. + * Executes the state machine according to the specification */ async execute( executor: StateMachineExecutable, context: MongoCryptContext - ): Promise; - - async execute( - executor: StateMachineExecutable, - context: MongoCryptContext, - responseType?: R - ): Promise> { + ): Promise { const keyVaultNamespace = executor._keyVaultNamespace; const keyVaultClient = executor._keyVaultClient; const metaDataClient = executor._metaDataClient; const mongocryptdClient = executor._mongocryptdClient; const mongocryptdManager = executor._mongocryptdManager; - let result: any | null = null; + let result: T | null = null; while (context.state !== MONGOCRYPT_CTX_DONE && context.state !== MONGOCRYPT_CTX_ERROR) { debug(`[context#${context.id}] ${stateToString.get(context.state) || context.state}`); @@ -270,12 +252,7 @@ export class StateMachine { const message = context.status.message || 'Finalization error'; throw new MongoCryptError(message); } - - result = - responseType != null - ? new responseType(finalizedContext) - : (result = deserialize(finalizedContext, this.options) as T); - + result = deserialize(finalizedContext, this.options) as T; break; } diff --git a/src/cmap/connection.ts b/src/cmap/connection.ts index bfdd23bfb8e..e1ad9a02935 100644 --- a/src/cmap/connection.ts +++ b/src/cmap/connection.ts @@ -748,7 +748,7 @@ export class CryptoConnection extends Connection { ns: MongoDBNamespace, cmd: Document, options?: CommandOptions, - responseType?: T | undefined + _responseType?: T | undefined ): Promise { const { autoEncrypter } = this; if (!autoEncrypter) { @@ -762,7 +762,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(ns, cmd, options, responseType); + return await super.command(ns, cmd, options, undefined); } if (serverWireVersion < 8) { @@ -796,7 +796,7 @@ export class CryptoConnection extends Connection { } } - const response = await super.command(ns, encrypted, options, responseType); + const response = await super.command(ns, encrypted, options, undefined); return await autoEncrypter.decrypt(response, options); } diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index f1b3427e8d6..57d12acc300 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -70,12 +70,14 @@ export class FindCursor extends AbstractCursor { /** @internal */ async _initialize(session: ClientSession): Promise { - const findOperation = new FindOperation(undefined, this.namespace, this[kFilter], { + const findOperation = new FindOperation(this.namespace, this[kFilter], { ...this[kBuiltOptions], // NOTE: order matters here, we may need to refine this ...this.cursorOptions, session }); + findOperation.encryptionEnabled = !!this.client.autoEncrypter; + const response = await executeOperation(this.client, findOperation); // the response is not a cursor when `explain` is enabled @@ -114,7 +116,7 @@ export class FindCursor extends AbstractCursor { } } - const response = await super.getMore(batchSize, true); + const response = await super.getMore(batchSize, this.client.autoEncrypter ? false : true); // TODO: wrap this in some logic to prevent it from happening if we don't need this support if (CursorResponse.is(response)) { this[kNumReturned] = this[kNumReturned] + response.batchSize; @@ -148,7 +150,7 @@ export class FindCursor extends AbstractCursor { async explain(verbosity?: ExplainVerbosityLike): Promise { return await executeOperation( this.client, - new FindOperation(undefined, this.namespace, this[kFilter], { + new FindOperation(this.namespace, this[kFilter], { ...this[kBuiltOptions], // NOTE: order matters here, we may need to refine this ...this.cursorOptions, explain: verbosity ?? true diff --git a/src/operations/find.ts b/src/operations/find.ts index 285720b985b..c776bcaf434 100644 --- a/src/operations/find.ts +++ b/src/operations/find.ts @@ -78,13 +78,10 @@ export class FindOperation extends CommandOperation { override options: FindOptions & { writeConcern?: never }; filter: Document; - constructor( - collection: Collection | undefined, - ns: MongoDBNamespace, - filter: Document = {}, - options: FindOptions = {} - ) { - super(collection, options); + public encryptionEnabled = false; + + constructor(ns: MongoDBNamespace, filter: Document = {}, options: FindOptions = {}) { + super(undefined, options); this.options = { ...options }; delete this.options.writeConcern; @@ -121,7 +118,7 @@ export class FindOperation extends CommandOperation { documentsReturnedIn: 'firstBatch', session }, - this.explain ? undefined : CursorResponse + this.explain || this.encryptionEnabled ? undefined : CursorResponse ); } } diff --git a/test/integration/crud/abstract_operation.test.ts b/test/integration/crud/abstract_operation.test.ts index fcac3e6ffe4..54eddb5e689 100644 --- a/test/integration/crud/abstract_operation.test.ts +++ b/test/integration/crud/abstract_operation.test.ts @@ -102,7 +102,7 @@ describe('abstract operation', function () { correctCommandName: 'count' }, { - subclassCreator: () => new mongodb.FindOperation(collection, collection.fullNamespace), + subclassCreator: () => new mongodb.FindOperation(undefined, collection.fullNamespace), subclassType: mongodb.FindOperation, correctCommandName: 'find' }, diff --git a/test/unit/operations/find.test.ts b/test/unit/operations/find.test.ts index bfb67d8e818..f208636d238 100644 --- a/test/unit/operations/find.test.ts +++ b/test/unit/operations/find.test.ts @@ -18,7 +18,7 @@ describe('FindOperation', function () { }); describe('#constructor', function () { - const operation = new FindOperation(undefined, namespace, filter, options); + const operation = new FindOperation(namespace, filter, options); it('sets the namespace', function () { expect(operation.ns).to.deep.equal(namespace); @@ -40,7 +40,7 @@ describe('FindOperation', function () { const server = new Server(topology, new ServerDescription('a:1'), {} as any); it('should build basic find command with filter', async () => { - const findOperation = new FindOperation(undefined, namespace, filter); + const findOperation = new FindOperation(namespace, filter); const stub = sinon.stub(server, 'command').resolves({}); await findOperation.execute(server, undefined); expect(stub).to.have.been.calledOnceWith(namespace, { @@ -53,7 +53,7 @@ describe('FindOperation', function () { const options = { oplogReplay: true }; - const findOperation = new FindOperation(undefined, namespace, {}, options); + const findOperation = new FindOperation(namespace, {}, options); const stub = sinon.stub(server, 'command').resolves({}); await findOperation.execute(server, undefined); expect(stub).to.have.been.calledOnceWith( From c0b0a8689178fdebdf3a330170aff67db361c82a Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Thu, 25 Apr 2024 18:03:02 -0400 Subject: [PATCH 23/25] test: fixes --- src/operations/find.ts | 1 - .../client-side-encryption/driver.test.ts | 173 +++++++++--------- .../crud/abstract_operation.test.ts | 2 +- 3 files changed, 89 insertions(+), 87 deletions(-) diff --git a/src/operations/find.ts b/src/operations/find.ts index c776bcaf434..d51e762e8d3 100644 --- a/src/operations/find.ts +++ b/src/operations/find.ts @@ -1,6 +1,5 @@ import type { Document } from '../bson'; import { CursorResponse } from '../cmap/wire_protocol/responses'; -import type { Collection } from '../collection'; import { MongoInvalidArgumentError } from '../error'; import { ReadConcern } from '../read_concern'; import type { Server } from '../sdam/server'; diff --git a/test/integration/client-side-encryption/driver.test.ts b/test/integration/client-side-encryption/driver.test.ts index 35eff329256..d472271e83e 100644 --- a/test/integration/client-side-encryption/driver.test.ts +++ b/test/integration/client-side-encryption/driver.test.ts @@ -298,107 +298,110 @@ describe('Client Side Encryption Functional', function () { }); }); - describe('when @@mdb.decorateDecryptionResult is set on autoEncrypter', () => { - let client: MongoClient; - let encryptedClient: MongoClient; - - beforeEach(async function () { - client = this.configuration.newClient(); - - const encryptSchema = (keyId: unknown, bsonType: string) => ({ - encrypt: { - bsonType, - algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Random', - keyId: [keyId] - } - }); + describe( + 'when @@mdb.decorateDecryptionResult is set on autoEncrypter', + { requires: { mongodb: '>=7.0' } }, + () => { + let client: MongoClient; + let encryptedClient: MongoClient; + + beforeEach(async function () { + client = this.configuration.newClient(); + + const encryptSchema = (keyId: unknown, bsonType: string) => ({ + encrypt: { + bsonType, + algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Random', + keyId: [keyId] + } + }); - const kmsProviders = this.configuration.kmsProviders(crypto.randomBytes(96)); + const kmsProviders = this.configuration.kmsProviders(crypto.randomBytes(96)); - await client.connect(); + await client.connect(); - const encryption = new ClientEncryption(client, { - keyVaultNamespace, - kmsProviders, - extraOptions: getEncryptExtraOptions() - }); - - const dataDb = client.db(dataDbName); - const keyVaultDb = client.db(keyVaultDbName); - - await dataDb.dropCollection(dataCollName).catch(() => null); - await keyVaultDb.dropCollection(keyVaultCollName).catch(() => null); - await keyVaultDb.createCollection(keyVaultCollName); - const dataKey = await encryption.createDataKey('local'); + const encryption = new ClientEncryption(client, { + keyVaultNamespace, + kmsProviders, + extraOptions: getEncryptExtraOptions() + }); - const $jsonSchema = { - bsonType: 'object', - properties: { - a: encryptSchema(dataKey, 'int'), - b: encryptSchema(dataKey, 'string'), - c: { - bsonType: 'object', - properties: { - d: { - encrypt: { - keyId: [dataKey], - algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic', - bsonType: 'string' + const dataDb = client.db(dataDbName); + const keyVaultDb = client.db(keyVaultDbName); + + await dataDb.dropCollection(dataCollName).catch(() => null); + await keyVaultDb.dropCollection(keyVaultCollName).catch(() => null); + await keyVaultDb.createCollection(keyVaultCollName); + const dataKey = await encryption.createDataKey('local'); + + const $jsonSchema = { + bsonType: 'object', + properties: { + a: encryptSchema(dataKey, 'int'), + b: encryptSchema(dataKey, 'string'), + c: { + bsonType: 'object', + properties: { + d: { + encrypt: { + keyId: [dataKey], + algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic', + bsonType: 'string' + } } } } } - } - }; + }; - await dataDb.createCollection(dataCollName, { - validator: { $jsonSchema } - }); + await dataDb.createCollection(dataCollName, { + validator: { $jsonSchema } + }); - encryptedClient = this.configuration.newClient( - {}, - { - autoEncryption: { - keyVaultNamespace, - kmsProviders, - extraOptions: getEncryptExtraOptions() + encryptedClient = this.configuration.newClient( + {}, + { + autoEncryption: { + keyVaultNamespace, + kmsProviders, + extraOptions: getEncryptExtraOptions() + } } - } - ); + ); - encryptedClient.autoEncrypter[Symbol.for('@@mdb.decorateDecryptionResult')] = true; - await encryptedClient.connect(); - }); + encryptedClient.autoEncrypter[Symbol.for('@@mdb.decorateDecryptionResult')] = true; + await encryptedClient.connect(); + }); - afterEach(function () { - return Promise.resolve() - .then(() => encryptedClient?.close()) - .then(() => client?.close()); - }); + afterEach(async function () { + await encryptedClient?.close(); + await client?.close(); + }); - it('adds decrypted keys to result at @@mdb.decryptedKeys', async function () { - const coll = encryptedClient.db(dataDbName).collection(dataCollName); + it('adds decrypted keys to result at @@mdb.decryptedKeys', async function () { + const coll = encryptedClient.db(dataDbName).collection(dataCollName); - const data = { - _id: new BSON.ObjectId(), - a: 1, - b: 'abc', - c: { d: 'def' } - }; + const data = { + _id: new BSON.ObjectId(), + a: 1, + b: 'abc', + c: { d: 'def' } + }; - const result = await coll.insertOne(data); - const decrypted = await coll.findOne({ _id: result.insertedId }); + const result = await coll.insertOne(data); + const decrypted = await coll.findOne({ _id: result.insertedId }); - expect(decrypted).to.deep.equal(data); - expect(decrypted) - .to.have.property(Symbol.for('@@mdb.decryptedKeys')) - .that.deep.equals(['a', 'b']); + expect(decrypted).to.deep.equal(data); + expect(decrypted) + .to.have.property(Symbol.for('@@mdb.decryptedKeys')) + .that.deep.equals(['a', 'b']); - // Nested - expect(decrypted).to.have.property('c'); - expect(decrypted.c) - .to.have.property(Symbol.for('@@mdb.decryptedKeys')) - .that.deep.equals(['d']); - }); - }); + // Nested + expect(decrypted).to.have.property('c'); + expect(decrypted.c) + .to.have.property(Symbol.for('@@mdb.decryptedKeys')) + .that.deep.equals(['d']); + }); + } + ); }); diff --git a/test/integration/crud/abstract_operation.test.ts b/test/integration/crud/abstract_operation.test.ts index 54eddb5e689..77c45890e26 100644 --- a/test/integration/crud/abstract_operation.test.ts +++ b/test/integration/crud/abstract_operation.test.ts @@ -102,7 +102,7 @@ describe('abstract operation', function () { correctCommandName: 'count' }, { - subclassCreator: () => new mongodb.FindOperation(undefined, collection.fullNamespace), + subclassCreator: () => new mongodb.FindOperation(collection.fullNamespace), subclassType: mongodb.FindOperation, correctCommandName: 'find' }, From 26be067c85f916fa8794420a627d4c8bc884bcbe Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Mon, 29 Apr 2024 13:22:17 -0400 Subject: [PATCH 24/25] test: lower server version for decorateDecryptionResult test --- test/integration/client-side-encryption/driver.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/client-side-encryption/driver.test.ts b/test/integration/client-side-encryption/driver.test.ts index d472271e83e..342f5adf1ca 100644 --- a/test/integration/client-side-encryption/driver.test.ts +++ b/test/integration/client-side-encryption/driver.test.ts @@ -300,7 +300,7 @@ describe('Client Side Encryption Functional', function () { describe( 'when @@mdb.decorateDecryptionResult is set on autoEncrypter', - { requires: { mongodb: '>=7.0' } }, + { requires: { clientSideEncryption: true, mongodb: '>=4.4' } }, () => { let client: MongoClient; let encryptedClient: MongoClient; From 0d801a68b58d6b7ea6feff9d477256ca7da9c304 Mon Sep 17 00:00:00 2001 From: Neal Beeken Date: Mon, 29 Apr 2024 13:32:01 -0400 Subject: [PATCH 25/25] fix: remove encryption flag from find --- src/cursor/find_cursor.ts | 2 -- src/operations/find.ts | 4 +--- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/cursor/find_cursor.ts b/src/cursor/find_cursor.ts index 57d12acc300..19c599fe8c1 100644 --- a/src/cursor/find_cursor.ts +++ b/src/cursor/find_cursor.ts @@ -76,8 +76,6 @@ export class FindCursor extends AbstractCursor { session }); - findOperation.encryptionEnabled = !!this.client.autoEncrypter; - const response = await executeOperation(this.client, findOperation); // the response is not a cursor when `explain` is enabled diff --git a/src/operations/find.ts b/src/operations/find.ts index d51e762e8d3..1c2ccdb1cac 100644 --- a/src/operations/find.ts +++ b/src/operations/find.ts @@ -77,8 +77,6 @@ export class FindOperation extends CommandOperation { override options: FindOptions & { writeConcern?: never }; filter: Document; - public encryptionEnabled = false; - constructor(ns: MongoDBNamespace, filter: Document = {}, options: FindOptions = {}) { super(undefined, options); @@ -117,7 +115,7 @@ export class FindOperation extends CommandOperation { documentsReturnedIn: 'firstBatch', session }, - this.explain || this.encryptionEnabled ? undefined : CursorResponse + this.explain ? undefined : CursorResponse ); } }