feat(NODE-3881): require hello command + OP_MSG when 'loadBalanced=Tr…

…ue' (#3907) Co-authored-by: Durran Jordan <[email protected]>
mongodb · Nov 6, 2023 · fd58eec · fd58eec
1 parent bb5fa43
commit fd58eec
Show file tree

Hide file tree

Showing 3 changed files with 157 additions and 24 deletions.
diff --git a/src/cmap/connect.ts b/src/cmap/connect.ts
@@ -214,7 +214,7 @@ export async function prepareHandshakeDocument(
   const { serverApi } = authContext.connection;
 
   const handshakeDoc: HandshakeDocument = {
-    [serverApi?.version ? 'hello' : LEGACY_HELLO_COMMAND]: 1,
+    [serverApi?.version || options.loadBalanced === true ? 'hello' : LEGACY_HELLO_COMMAND]: 1,
     helloOk: true,
     client: options.metadata,
     compression: compressors

diff --git a/test/unit/cmap/connect.test.ts b/test/unit/cmap/connect.test.ts
@@ -234,42 +234,68 @@ describe('Connect Tests', function () {
 
     context('loadBalanced option', () => {
       context('when loadBalanced is not set as an option', () => {
+        const authContext = {
+          connection: {},
+          options: {}
+        };
+
         it('does not set loadBalanced on the handshake document', async () => {
-          const options = {};
-          const authContext = {
-            connection: {},
-            options
-          };
           const handshakeDocument = await prepareHandshakeDocument(authContext);
           expect(handshakeDocument).not.to.have.property('loadBalanced');
         });
+
+        it('does not set hello on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).not.to.have.property('hello');
+        });
+
+        it('sets LEGACY_HELLO_COMMAND on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).to.have.property(LEGACY_HELLO_COMMAND, 1);
+        });
       });
 
       context('when loadBalanced is set to false', () => {
+        const authContext = {
+          connection: {},
+          options: { loadBalanced: false }
+        };
+
         it('does not set loadBalanced on the handshake document', async () => {
-          const options = {
-            loadBalanced: false
-          };
-          const authContext = {
-            connection: {},
-            options
-          };
           const handshakeDocument = await prepareHandshakeDocument(authContext);
           expect(handshakeDocument).not.to.have.property('loadBalanced');
         });
+
+        it('does not set hello on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).not.to.have.property('hello');
+        });
+
+        it('sets LEGACY_HELLO_COMMAND on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).to.have.property(LEGACY_HELLO_COMMAND, 1);
+        });
       });
 
       context('when loadBalanced is set to true', () => {
-        it('does set loadBalanced on the handshake document', async () => {
-          const options = {
-            loadBalanced: true
-          };
-          const authContext = {
-            connection: {},
-            options
-          };
+        const authContext = {
+          connection: {},
+          options: { loadBalanced: true }
+        };
+
+        it('sets loadBalanced on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).to.have.property('loadBalanced');
+        });
+
+        it('sets hello on the handshake document', async () => {
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).to.have.property('hello');
+        });
+
+        it('does not set LEGACY_HELLO_COMMAND on the handshake document', async () => {
           const handshakeDocument = await prepareHandshakeDocument(authContext);
-          expect(handshakeDocument).to.have.property('loadBalanced', true);
+          expect(handshakeDocument).not.have.property(LEGACY_HELLO_COMMAND, 1);
         });
       });
     });

diff --git a/test/unit/cmap/connection.test.ts b/test/unit/cmap/connection.test.ts
@@ -14,12 +14,14 @@ import {
   hasSessionSupport,
   type HostAddress,
   isHello,
-  type MessageStream,
+  MessageStream,
   MongoNetworkError,
   MongoNetworkTimeoutError,
   MongoRuntimeError,
+  Msg,
   ns,
-  type OperationDescription
+  type OperationDescription,
+  Query
 } from '../../mongodb';
 import * as mock from '../../tools/mongodb-mock/index';
 import { generateOpMsgBuffer, getSymbolFrom } from '../../tools/utils';
@@ -1027,4 +1029,109 @@ describe('new Connection()', function () {
       });
     });
   });
+
+  describe('when load-balanced', () => {
+    const CONNECT_DEFAULTS = {
+      id: 1,
+      tls: false,
+      generation: 1,
+      monitorCommands: false,
+      metadata: {} as ClientMetadata
+    };
+    let server;
+    let connectOptions;
+    let connection: Connection;
+    let writeCommandSpy;
+
+    beforeEach(async () => {
+      server = await mock.createServer();
+      server.setMessageHandler(request => {
+        request.reply(mock.HELLO);
+      });
+      writeCommandSpy = sinon.spy(MessageStream.prototype, 'writeCommand');
+    });
+
+    afterEach(async () => {
+      connection?.destroy({ force: true });
+      sinon.restore();
+      await mock.cleanup();
+    });
+
+    it('sends the first command as OP_MSG', async () => {
+      try {
+        connectOptions = {
+          ...CONNECT_DEFAULTS,
+          hostAddress: server.hostAddress() as HostAddress,
+          socketTimeoutMS: 100,
+          loadBalanced: true
+        };
+
+        connection = await promisify<Connection>(callback =>
+          //@ts-expect-error: Callbacks do not have mutual exclusion for error/result existence
+          connect(connectOptions, callback)
+        )();
+
+        await promisify(callback =>
+          connection.command(ns('admin.$cmd'), { hello: 1 }, {}, callback)
+        )();
+      } catch (error) {
+        /** Connection timeouts, but the handshake message is sent. */
+      }
+
+      expect(writeCommandSpy).to.have.been.called;
+      expect(writeCommandSpy.firstCall.args[0] instanceof Msg).to.equal(true);
+    });
+  });
+
+  describe('when not load-balanced', () => {
+    const CONNECT_DEFAULTS = {
+      id: 1,
+      tls: false,
+      generation: 1,
+      monitorCommands: false,
+      metadata: {} as ClientMetadata
+    };
+    let server;
+    let connectOptions;
+    let connection: Connection;
+    let writeCommandSpy;
+
+    beforeEach(async () => {
+      server = await mock.createServer();
+      server.setMessageHandler(request => {
+        request.reply(mock.HELLO);
+      });
+      writeCommandSpy = sinon.spy(MessageStream.prototype, 'writeCommand');
+    });
+
+    afterEach(async () => {
+      connection?.destroy({ force: true });
+      sinon.restore();
+      await mock.cleanup();
+    });
+
+    it('sends the first command as OP_QUERY', async () => {
+      try {
+        connectOptions = {
+          ...CONNECT_DEFAULTS,
+          hostAddress: server.hostAddress() as HostAddress,
+          socketTimeoutMS: 100
+        };
+
+        connection = await promisify<Connection>(callback =>
+          //@ts-expect-error: Callbacks do not have mutual exclusion for error/result existence
+          connect(connectOptions, callback)
+        )();
+
+        await promisify(callback =>
+          connection.command(ns('admin.$cmd'), { hello: 1 }, {}, callback)
+        )();
+      } catch (error) {
+        /** Connection timeouts, but the handshake message is sent. */
+      }
+
+      expect(writeCommandSpy).to.have.been.called;
+      expect(writeCommandSpy.firstCall.args[0] instanceof Query).to.equal(true);
+    });
+  });
 });