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);
+    });
   });
 });