From 555857319605d72ce57c639166ce25f57739bf74 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Wed, 4 Dec 2024 18:23:08 +0100
Subject: [PATCH] fix(NODE-6600): set object mode correctly for message
 chunking in SizedMessageTransform (#4345)

Co-authored-by: Bailey Pearson <bailey.pearson@mongodb.com>
---
 src/cmap/connection.ts            |  2 +-
 test/unit/cmap/connection.test.ts | 18 +++++++++++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/cmap/connection.ts b/src/cmap/connection.ts
index b92adb50511..6df81b34d94 100644
--- a/src/cmap/connection.ts
+++ b/src/cmap/connection.ts
@@ -769,7 +769,7 @@ export class SizedMessageTransform extends Transform {
   connection: Connection;
 
   constructor({ connection }: { connection: Connection }) {
-    super({ objectMode: false });
+    super({ writableObjectMode: false, readableObjectMode: true });
     this.bufferPool = new BufferPool();
     this.connection = connection;
   }
diff --git a/test/unit/cmap/connection.test.ts b/test/unit/cmap/connection.test.ts
index 05e66f3dcfc..df3ea0f53ee 100644
--- a/test/unit/cmap/connection.test.ts
+++ b/test/unit/cmap/connection.test.ts
@@ -11,7 +11,8 @@ import {
   MongoClientAuthProviders,
   MongoDBCollectionNamespace,
   MongoNetworkTimeoutError,
-  ns
+  ns,
+  SizedMessageTransform
 } from '../../mongodb';
 import * as mock from '../../tools/mongodb-mock/index';
 import { getSymbolFrom } from '../../tools/utils';
@@ -323,4 +324,19 @@ describe('new Connection()', function () {
       });
     });
   });
+
+  describe('SizedMessageTransform', function () {
+    it('parses chunks of wire messages', function () {
+      const stream = new SizedMessageTransform({ connection: {} as any });
+      // Message of length 4 + 4 = 8
+      stream.write(Buffer.from([8, 0, 0, 0]));
+      stream.write(Buffer.from([1, 2, 3, 4]));
+      // Message of length 4 + 2 = 6, chunked differently
+      stream.write(Buffer.from([6, 0, 0]));
+      stream.write(Buffer.from([0, 5, 6]));
+      expect(stream.read(1)).to.deep.equal(Buffer.from([8, 0, 0, 0, 1, 2, 3, 4]));
+      expect(stream.read(1)).to.deep.equal(Buffer.from([6, 0, 0, 0, 5, 6]));
+      expect(stream.read(1)).to.equal(null);
+    });
+  });
 });