using bson for buffer calls

Author: PavelSafronov

src/bson.ts

@@ -32,6 +32,13 @@ export {
 /** @internal */
 export type BSONElement = BSON.OnDemand['BSONElement'];
 
+/** @internal */
+export function toLocalBufferType(this: void, buffer: Buffer | Uint8Array): Buffer {
+  return Buffer.isBuffer(buffer)
+    ? buffer
+    : Buffer.from(buffer.buffer, buffer.byteOffset, buffer.byteLength);
+}
+
 export function parseToElementsToArray(bytes: Uint8Array, offset?: number): BSONElement[] {
   const res = BSON.onDemand.parseToElements(bytes, offset);
   return Array.isArray(res) ? res : [...res];
@@ -41,6 +48,34 @@ export const getInt32LE = BSON.onDemand.NumberUtils.getInt32LE;
 export const getFloat64LE = BSON.onDemand.NumberUtils.getFloat64LE;
 export const getBigInt64LE = BSON.onDemand.NumberUtils.getBigInt64LE;
 export const toUTF8 = BSON.onDemand.ByteUtils.toUTF8;
+export const writeInt32LE = BSON.onDemand.NumberUtils.setInt32LE;
+
+export const fromUTF8 = (text: string) => toLocalBufferType(BSON.onDemand.ByteUtils.fromUTF8(text));
+
+export const concatBuffers = (list: Buffer[] | Uint8Array[]) => {
+  return toLocalBufferType(BSON.onDemand.ByteUtils.concat(list));
+};
+export const allocateBuffer = (size: number) =>
+  toLocalBufferType(BSON.onDemand.ByteUtils.allocate(size));
+export const allocateUnsafeBuffer = (size: number) =>
+  toLocalBufferType(BSON.onDemand.ByteUtils.allocateUnsafe(size));
+
+export const utf8ByteLength = BSON.onDemand.ByteUtils.utf8ByteLength;
+export const toBase64 = BSON.onDemand.ByteUtils.toBase64;
+export const encodeUTF8Into = BSON.onDemand.ByteUtils.encodeUTF8Into;
+
+const validateBufferInputs = (buffer: Uint8Array, offset: number, length: number) => {
+  if (offset < 0 || offset + length > buffer.length) {
+    throw new RangeError(
+      `Attempt to access memory outside buffer bounds: buffer length: ${buffer.length}, offset: ${offset}, length: ${length}`
+    );
+  }
+};
+
+export const readInt32LE = (buffer: Uint8Array, offset: number): number => {
+  validateBufferInputs(buffer, offset, 4);
+  return getInt32LE(buffer, offset);
+};
 
 /**
  * BSON Serialization options.

src/cmap/auth/scram.ts

@@ -1,7 +1,14 @@
 import { saslprep } from '@mongodb-js/saslprep';
 import * as crypto from 'crypto';
 
-import { Binary, type Document } from '../../bson';
+import {
+  allocateBuffer,
+  Binary,
+  concatBuffers,
+  type Document,
+  fromUTF8,
+  toBase64
+} from '../../bson';
 import {
   MongoInvalidArgumentError,
   MongoMissingCredentialsError,
@@ -68,11 +75,11 @@ function cleanUsername(username: string) {
 function clientFirstMessageBare(username: string, nonce: Buffer) {
   // NOTE: This is done b/c Javascript uses UTF-16, but the server is hashing in UTF-8.
   // Since the username is not sasl-prep-d, we need to do this here.
-  return Buffer.concat([
-    Buffer.from('n=', 'utf8'),
-    Buffer.from(username, 'utf8'),
-    Buffer.from(',r=', 'utf8'),
-    Buffer.from(nonce.toString('base64'), 'utf8')
+  return concatBuffers([
+    fromUTF8('n='),
+    fromUTF8(username),
+    fromUTF8(',r='),
+    fromUTF8(toBase64(nonce))
   ]);
 }
 
@@ -91,7 +98,7 @@ function makeFirstMessage(
     saslStart: 1,
     mechanism,
     payload: new Binary(
-      Buffer.concat([Buffer.from('n,,', 'utf8'), clientFirstMessageBare(username, nonce)])
+      concatBuffers([Buffer.from('n,,', 'utf8'), clientFirstMessageBare(username, nonce)])
     ),
     autoAuthorize: 1,
     options: { skipEmptyExchange: true }
@@ -199,7 +206,7 @@ async function continueScramConversation(
   const retrySaslContinueCmd = {
     saslContinue: 1,
     conversationId: r.conversationId,
-    payload: Buffer.alloc(0)
+    payload: allocateBuffer(0)
   };
 
   await connection.command(ns(`${db}.$cmd`), retrySaslContinueCmd, undefined);

src/cmap/commands.ts

@@ -1,5 +1,15 @@
-import type { BSONSerializeOptions, Document, Long } from '../bson';
-import * as BSON from '../bson';
+import {
+  allocateBuffer,
+  allocateUnsafeBuffer,
+  BSON,
+  type BSONSerializeOptions,
+  concatBuffers,
+  type Document,
+  type Long,
+  readInt32LE,
+  utf8ByteLength,
+  writeInt32LE
+} from '../bson';
 import { MongoInvalidArgumentError, MongoRuntimeError } from '../error';
 import { type ReadPreference } from '../read_preference';
 import type { ClientSession } from '../sessions';
@@ -30,7 +40,7 @@ const QUERY_FAILURE = 2;
 const SHARD_CONFIG_STALE = 4;
 const AWAIT_CAPABLE = 8;
 
-const encodeUTF8Into = BSON.BSON.onDemand.ByteUtils.encodeUTF8Into;
+const encodeUTF8Into = BSON.onDemand.ByteUtils.encodeUTF8Into;
 
 /** @internal */
 export type WriteProtocolMessageType = OpQueryRequest | OpMsgRequest;
@@ -182,10 +192,10 @@ export class OpQueryRequest {
     if (this.batchSize !== this.numberToReturn) this.numberToReturn = this.batchSize;
 
     // Allocate write protocol header buffer
-    const header = Buffer.alloc(
+    const header = allocateBuffer(
       4 * 4 + // Header
         4 + // Flags
-        Buffer.byteLength(this.ns) +
+        utf8ByteLength(this.ns) +
         1 + // namespace
         4 + // numberToSkip
         4 // numberToReturn
@@ -256,7 +266,7 @@ export class OpQueryRequest {
     index = index + 4;
 
     // Write collection name
-    index = index + header.write(this.ns, index, 'utf8') + 1;
+    index = index + encodeUTF8Into(header, this.ns, index) + 1;
     header[index - 1] = 0;
 
     // Write header information flags numberToSkip
@@ -364,10 +374,10 @@ export class OpReply {
     this.index = 20;
 
     // Read the message body
-    this.responseFlags = this.data.readInt32LE(0);
-    this.cursorId = new BSON.Long(this.data.readInt32LE(4), this.data.readInt32LE(8));
-    this.startingFrom = this.data.readInt32LE(12);
-    this.numberReturned = this.data.readInt32LE(16);
+    this.responseFlags = readInt32LE(this.data, 0);
+    this.cursorId = new BSON.Long(readInt32LE(this.data, 4), readInt32LE(this.data, 8));
+    this.startingFrom = readInt32LE(this.data, 12);
+    this.numberReturned = readInt32LE(this.data, 16);
 
     if (this.numberReturned < 0 || this.numberReturned > 2 ** 32 - 1) {
       throw new RangeError(
@@ -446,7 +456,7 @@ export class DocumentSequence {
     this.serializedDocumentsLength = 0;
     // Document sequences starts with type 1 at the first byte.
     // Field strings must always be UTF-8.
-    const buffer = Buffer.allocUnsafe(1 + 4 + this.field.length + 1);
+    const buffer = allocateUnsafeBuffer(1 + 4 + this.field.length + 1);
     buffer[0] = 1;
     // Third part is the field name at offset 5 with trailing null byte.
     encodeUTF8Into(buffer, `${this.field}\0`, 5);
@@ -482,7 +492,7 @@ export class DocumentSequence {
    * @returns The section bytes.
    */
   toBin(): Uint8Array {
-    return Buffer.concat(this.chunks);
+    return concatBuffers(this.chunks);
   }
 }
 
@@ -547,7 +557,7 @@ export class OpMsgRequest {
       flags |= OPTS_EXHAUST_ALLOWED;
     }
 
-    const header = Buffer.alloc(
+    const header = allocateBuffer(
       4 * 4 + // Header
         4 // Flags
     );
@@ -558,11 +568,11 @@ export class OpMsgRequest {
     const command = this.command;
     totalLength += this.makeSections(buffers, command);
 
-    header.writeInt32LE(totalLength, 0); // messageLength
-    header.writeInt32LE(this.requestId, 4); // requestID
-    header.writeInt32LE(0, 8); // responseTo
-    header.writeInt32LE(OP_MSG, 12); // opCode
-    header.writeUInt32LE(flags, 16); // flags
+    writeInt32LE(header, 0, totalLength); // messageLength
+    writeInt32LE(header, 4, this.requestId); // requestID
+    writeInt32LE(header, 8, 0); // responseTo
+    writeInt32LE(header, 12, OP_MSG); // opCode
+    writeInt32LE(header, 16, flags); // flags
     return buffers;
   }
 
@@ -571,7 +581,7 @@ export class OpMsgRequest {
    */
   makeSections(buffers: Uint8Array[], document: Document): number {
     const sequencesBuffer = this.extractDocumentSequences(document);
-    const payloadTypeBuffer = Buffer.allocUnsafe(1);
+    const payloadTypeBuffer = allocateUnsafeBuffer(1);
     payloadTypeBuffer[0] = 0;
 
     const documentBuffer = this.serializeBson(document);
@@ -606,11 +616,11 @@ export class OpMsgRequest {
       }
     }
     if (chunks.length > 0) {
-      return Buffer.concat(chunks);
+      return concatBuffers(chunks);
     }
     // If we have no document sequences we return an empty buffer for nothing to add
     // to the payload.
-    return Buffer.alloc(0);
+    return allocateBuffer(0);
   }
 
   serializeBson(document: Document): Uint8Array {
@@ -676,7 +686,7 @@ export class OpMsgResponse {
     this.fromCompressed = msgHeader.fromCompressed;
 
     // Read response flags
-    this.responseFlags = msgBody.readInt32LE(0);
+    this.responseFlags = readInt32LE(msgBody, 0);
     this.checksumPresent = (this.responseFlags & OPTS_CHECKSUM_PRESENT) !== 0;
     this.moreToCome = (this.responseFlags & OPTS_MORE_TO_COME) !== 0;
     this.exhaustAllowed = (this.responseFlags & OPTS_EXHAUST_ALLOWED) !== 0;
@@ -700,9 +710,9 @@ export class OpMsgResponse {
     this.index = 4;
 
     while (this.index < this.data.length) {
-      const payloadType = this.data.readUInt8(this.index++);
+      const payloadType = this.data[this.index++];
       if (payloadType === 0) {
-        const bsonSize = this.data.readUInt32LE(this.index);
+        const bsonSize = readInt32LE(this.data, this.index);
         const bin = this.data.subarray(this.index, this.index + bsonSize);
 
         this.sections.push(bin);
@@ -759,30 +769,31 @@ export class OpCompressedRequest {
   }
 
   async toBin(): Promise<Buffer[]> {
-    const concatenatedOriginalCommandBuffer = Buffer.concat(this.command.toBin());
+    const concatenatedOriginalCommandBuffer = concatBuffers(this.command.toBin());
     // otherwise, compress the message
     const messageToBeCompressed = concatenatedOriginalCommandBuffer.slice(MESSAGE_HEADER_SIZE);
 
     // Extract information needed for OP_COMPRESSED from the uncompressed message
-    const originalCommandOpCode = concatenatedOriginalCommandBuffer.readInt32LE(12);
+    const originalCommandOpCode = readInt32LE(concatenatedOriginalCommandBuffer, 12);
 
     // Compress the message body
     const compressedMessage = await compress(this.options, messageToBeCompressed);
     // Create the msgHeader of OP_COMPRESSED
-    const msgHeader = Buffer.alloc(MESSAGE_HEADER_SIZE);
-    msgHeader.writeInt32LE(
-      MESSAGE_HEADER_SIZE + COMPRESSION_DETAILS_SIZE + compressedMessage.length,
-      0
+    const msgHeader = allocateBuffer(MESSAGE_HEADER_SIZE);
+    writeInt32LE(
+      msgHeader,
+      0,
+      MESSAGE_HEADER_SIZE + COMPRESSION_DETAILS_SIZE + compressedMessage.length
     ); // messageLength
-    msgHeader.writeInt32LE(this.command.requestId, 4); // requestID
-    msgHeader.writeInt32LE(0, 8); // responseTo (zero)
-    msgHeader.writeInt32LE(OP_COMPRESSED, 12); // opCode
+    writeInt32LE(msgHeader, 4, this.command.requestId); // requestID
+    writeInt32LE(msgHeader, 8, 0); // responseTo (zero)
+    writeInt32LE(msgHeader, 12, OP_COMPRESSED); // opCode
 
     // Create the compression details of OP_COMPRESSED
-    const compressionDetails = Buffer.alloc(COMPRESSION_DETAILS_SIZE);
-    compressionDetails.writeInt32LE(originalCommandOpCode, 0); // originalOpcode
-    compressionDetails.writeInt32LE(messageToBeCompressed.length, 4); // Size of the uncompressed compressedMessage, excluding the MsgHeader
-    compressionDetails.writeUInt8(Compressor[this.options.agreedCompressor], 8); // compressorID
+    const compressionDetails = allocateBuffer(COMPRESSION_DETAILS_SIZE);
+    writeInt32LE(compressionDetails, 0, originalCommandOpCode); // originalOpcode
+    writeInt32LE(compressionDetails, 4, messageToBeCompressed.length); // Size of the uncompressed compressedMessage, excluding the MsgHeader
+    writeInt32LE(compressionDetails, 8, Compressor[this.options.agreedCompressor]); // compressorID
     return [msgHeader, compressionDetails, compressedMessage];
   }
 }

src/cmap/connection.ts

@@ -2,7 +2,9 @@ import { type Readable, Transform, type TransformCallback } from 'stream';
 import { clearTimeout, setTimeout } from 'timers';
 
 import {
+  BSON,
   type BSONSerializeOptions,
+  concatBuffers,
   deserialize,
   type DeserializeOptions,
   type Document,
@@ -174,7 +176,7 @@ function streamIdentifier(stream: Stream, options: ConnectionOptions): string {
     return HostAddress.fromHostPort(remoteAddress, remotePort).toString();
   }
 
-  return uuidV4().toString('hex');
+  return BSON.onDemand.ByteUtils.toHex(uuidV4());
 }
 
 /** @internal */
@@ -204,7 +206,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   private lastUseTime: number;
   private clusterTime: Document | null = null;
   private error: Error | null = null;
-  private dataEvents: AsyncGenerator<Buffer, void, void> | null = null;
+  private dataEvents: AsyncGenerator<Uint8Array, void, void> | null = null;
 
   private readonly socketTimeoutMS: number;
   private readonly monitorCommands: boolean;
@@ -696,7 +698,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
             zlibCompressionLevel: options.zlibCompressionLevel ?? 0
           });
 
-    const buffer = Buffer.concat(await finalCommand.toBin());
+    const buffer = concatBuffers(await finalCommand.toBin());
 
     if (options.timeoutContext?.csotEnabled()) {
       if (
@@ -794,7 +796,7 @@ export class SizedMessageTransform extends Transform {
     this.connection = connection;
   }
 
-  override _transform(chunk: Buffer, encoding: unknown, callback: TransformCallback): void {
+  override _transform(chunk: Uint8Array, encoding: unknown, callback: TransformCallback): void {
     if (this.connection.delayedTimeoutId != null) {
       clearTimeout(this.connection.delayedTimeoutId);
       this.connection.delayedTimeoutId = null;

src/cmap/wire_protocol/compression.ts

@@ -1,5 +1,6 @@
 import * as zlib from 'zlib';
 
+import { concatBuffers, readInt32LE } from '../../bson';
 import { LEGACY_HELLO_COMMAND } from '../../constants';
 import { getSnappy, getZstdLibrary, type SnappyLib, type ZStandard } from '../../deps';
 import { MongoDecompressionError, MongoInvalidArgumentError } from '../../error';
@@ -168,7 +169,7 @@ export async function compressCommand(
           zlibCompressionLevel: description.zlibCompressionLevel ?? 0
         });
   const data = await finalCommand.toBin();
-  return Buffer.concat(data);
+  return concatBuffers(data);
 }
 
 /**
@@ -180,10 +181,10 @@ export async function compressCommand(
  */
 export async function decompressResponse(message: Buffer): Promise<OpMsgResponse | OpReply> {
   const messageHeader: MessageHeader = {
-    length: message.readInt32LE(0),
-    requestId: message.readInt32LE(4),
-    responseTo: message.readInt32LE(8),
-    opCode: message.readInt32LE(12)
+    length: readInt32LE(message, 0),
+    requestId: readInt32LE(message, 4),
+    responseTo: readInt32LE(message, 8),
+    opCode: readInt32LE(message, 12)
   };
 
   if (messageHeader.opCode !== OP_COMPRESSED) {
@@ -195,8 +196,8 @@ export async function decompressResponse(message: Buffer): Promise<OpMsgResponse
   const header: MessageHeader = {
     ...messageHeader,
     fromCompressed: true,
-    opCode: message.readInt32LE(MESSAGE_HEADER_SIZE),
-    length: message.readInt32LE(MESSAGE_HEADER_SIZE + 4)
+    opCode: readInt32LE(message, MESSAGE_HEADER_SIZE),
+    length: readInt32LE(message, MESSAGE_HEADER_SIZE + 4)
   };
   const compressorID = message[MESSAGE_HEADER_SIZE + 8];
   const compressedBuffer = message.slice(MESSAGE_HEADER_SIZE + 9);