fix(NODE-6630): read all messages in buffer when chunk arrives

Author: nbbeeken

src/cmap/connection.ts

@@ -27,6 +27,7 @@ import {
  MongoNetworkTimeoutError,
  MongoOperationTimeoutError,
  MongoParseError,
+ MongoRuntimeError,
  MongoServerError,
  MongoUnexpectedServerResponseError
 } from '../error';
@@ -791,22 +792,46 @@ export class SizedMessageTransform extends Transform {
  }
 
  this.bufferPool.append(chunk);
- const sizeOfMessage = this.bufferPool.getInt32();
 
- if (sizeOfMessage == null) {
- return callback();
- }
+ while (this.bufferPool.length) {
+ // While there are any bytes in the buffer
 
- if (sizeOfMessage < 0) {
- return callback(new MongoParseError(`Invalid message size: ${sizeOfMessage}, too small`));
- }
+ // Try to fetch a size from the top 4 bytes
+ const sizeOfMessage = this.bufferPool.getInt32();
+
+ if (sizeOfMessage == null) {
+ // Not even an int32 worth of data. Stop the loop, we need more chunks.
+ break;
+ }
 
- if (sizeOfMessage > this.bufferPool.length) {
- return callback();
+ if (sizeOfMessage < 0) {
+ // The size in the message has a negative value, this is probably corruption, throw:
+ return callback(new MongoParseError(`Message size cannot be negative: ${sizeOfMessage}`));
+ }
+
+ if (sizeOfMessage > this.bufferPool.length) {
+ // We do not have enough bytes to make a sizeOfMessage chunk
+ break;
+ }
+
+ // Add a message to the stream
+ const message = this.bufferPool.read(sizeOfMessage);
+
+ if (!this.push(message)) {
+ // We only subscribe to data events so we should never get backpressure
+ // if we do, we do not have the handling for it.
+ return callback(
+ new MongoRuntimeError(`SizedMessageTransform does not support backpressure`)
+ );
+ }
  }
 
- const message = this.bufferPool.read(sizeOfMessage);
- return callback(null, message);
+ callback();
+ }
+
+ override pipe<T extends NodeJS.WritableStream>(destination: T, options?: { end?: boolean }): T {
+ destination.on('drain', this.emit.bind('drain'));
+ return super.pipe(destination, options);
  }
 }
 

test/unit/cmap/connection.test.ts

@@ -1,4 +1,5 @@
 import { Socket } from 'node:net';
+import { Writable } from 'node:stream';
 
 import { expect } from 'chai';
 import * as sinon from 'sinon';
@@ -11,7 +12,9 @@ import {
  MongoClientAuthProviders,
  MongoDBCollectionNamespace,
  MongoNetworkTimeoutError,
+ MongoRuntimeError,
  ns,
+ promiseWithResolvers,
  SizedMessageTransform
 } from '../../mongodb';
 import * as mock from '../../tools/mongodb-mock/index';
@@ -333,5 +336,49 @@ describe('new Connection()', function () {
  expect(stream.read(1)).to.deep.equal(Buffer.from([6, 0, 0, 0, 5, 6]));
  expect(stream.read(1)).to.equal(null);
  });
+
+ it('parses many wire messages when chunk arrives', function () {
+ const stream = new SizedMessageTransform({ connection: {} as any });
+
+ let dataCount = 0;
+ stream.on('data', () => {
+ dataCount += 1;
+ });
+
+ // 3 messages of size 8
+ stream.write(
+ Buffer.from([
+ ...[8, 0, 0, 0, 0, 0, 0, 0],
+ ...[8, 0, 0, 0, 0, 0, 0, 0],
+ ...[8, 0, 0, 0, 0, 0, 0, 0]
+ ])
+ );
+
+ expect(dataCount).to.equal(3);
+ });
+
+ it('waits for a drain event when destination needs backpressure', async function () {
+ const stream = new SizedMessageTransform({ connection: {} as any });
+ const destination = new Writable({
+ highWaterMark: 1,
+ objectMode: true,
+ write: (chunk, encoding, callback) => {
+ void stream;
+ setTimeout(1).then(() => callback());
+ }
+ });
+
+ // 1000 messages of size 8
+ stream.write(
+ Buffer.from(Array.from({ length: 1000 }, () => [8, 0, 0, 0, 0, 0, 0, 0]).flat(1))
+ );
+
+ const { promise, resolve, reject } = promiseWithResolvers();
+
+ stream.on('error', reject).pipe(destination).on('error', reject).on('finish', resolve);
+
+ const error = await promise.catch(error => error);
+ expect(error).to.be.instanceOf(MongoRuntimeError);
+ });
  });
 });