fix: Use latest connection-model and keep reference to tunnel to be able to close it when disconnecting native client COMPASS-4474 (#308)

* fix: Use latest connection-model and keep reference to tunnel to be able to close it when disconnecting native client COMPASS-4474

* feat: Do not allow to call connect multiple times without disconnecting first

* test: Test that we are closing the tunnel before resolving the disconnect callback

Author: gribnoysup

lib/native-client.js

@@ -92,7 +92,13 @@ class NativeClient extends EventEmitter {
   constructor(model) {
     super();
     this.model = model;
+
     this.connectionOptions = null;
+    this.tunnel = null;
+
+    this.isWritable = false;
+    this.isMongos = false;
+    this._isConnected = false;
   }
 
   /**
@@ -104,19 +110,33 @@ class NativeClient extends EventEmitter {
   connect(done) {
     debug('connecting...');
 
-    this.connectionOptions = null;
-    this.isWritable = false;
-    this.isMongos = false;
+    if (this._isConnected) {
+      setImmediate(() => {
+        done(
+          new Error(
+            'Connect method has been called more than once without disconnecting.'
+          )
+        );
+      });
+
+      return this;
+    }
+
+    // Not really true at that point, we are doing it just so we don't allow
+    // simultaneous syncronous calls to the connect method
+    this._isConnected = true;
 
     connect(
       this.model,
       this.setupListeners.bind(this),
-      (err, _client, connectionOptions) => {
+      (err, _client, tunnel, connectionOptions) => {
         if (err) {
+          this._isConnected = false;
           return done(this._translateMessage(err));
         }
 
         this.connectionOptions = connectionOptions;
+        this.tunnel = tunnel;
 
         this.isWritable = this.client.isWritable;
         this.isMongos = this.client.isMongos;
@@ -128,6 +148,7 @@ class NativeClient extends EventEmitter {
 
         this.client.on('status', (evt) => this.emit('status', evt));
         this.database = this.client.db(this.model.ns || ADMIN);
+
         done(null, this);
       }
     );
@@ -298,6 +319,8 @@ class NativeClient extends EventEmitter {
    *
    * @param {String} databaseName - The database name.
    * @param {Function} callback - The callback.
+   *
+   * @returns {void}
    */
   collections(databaseName, callback) {
     if (databaseName === SYSTEM) {
@@ -521,9 +544,22 @@ class NativeClient extends EventEmitter {
 
   /**
    * Disconnect the client.
+   * @param {Function} callback
    */
   disconnect(callback) {
-    this.client.close(true, callback);
+    this.client.close(true, err => {
+      if (this.tunnel) {
+        debug('mongo client closed. shutting down ssh tunnel');
+        this.tunnel.close().finally(() => {
+          this._cleanup();
+          debug('ssh tunnel stopped');
+          callback(err);
+        });
+      } else {
+        this._cleanup();
+        return callback(err);
+      }
+    });
   }
 
   /**
@@ -1249,6 +1285,15 @@ class NativeClient extends EventEmitter {
     }
     return error;
   }
+
+  _cleanup() {
+    this.client = null;
+    this.connectionOptions = null;
+    this.tunnel = null;
+    this.isWritable = false;
+    this.isMongos = false;
+    this._isConnected = false;
+  }
 }
 
 function addDebugToClass(cls) {

package-lock.json

@@ -29,7 +29,7 @@
   ],
   "peerDependencies": {
     "mongodb": "3.x",
-    "mongodb-connection-model": "*"
+    "mongodb-connection-model": ">=19.1.0"
   },
   "dependencies": {
     "async": "^3.2.0",
@@ -52,7 +52,7 @@
     "mocha": "^8.2.1",
     "mock-require": "^3.0.3",
     "mongodb": "^3.6.3",
-    "mongodb-connection-model": "^18.0.0",
+    "mongodb-connection-model": "^19.1.0",
     "mongodb-runner": "^4.8.0",
     "pre-commit": "^1.2.2",
     "sinon": "^9.2.3",

package.json

@@ -8,6 +8,58 @@ const EventEmitter = require('events');
 
 var NativeClient = require('../lib/native-client');
 
+function mockedTopologyDescription(
+  topologyType = 'Standalone',
+  serverType = 'Single'
+) {
+  return {
+    type: topologyType,
+    servers: new Map([['127.0.0.1:27017', { type: serverType }]])
+  };
+}
+
+/*
+ * pretends to be a connection-model providing every function call
+ * required in NativeClient#connect, but returns topology and connection
+ * params of our choice
+ */
+function mockedConnectionModel(topologyDescription, connectionOptions) {
+  const _topologyDescription =
+    topologyDescription || mockedTopologyDescription();
+
+  const _connectionOptions = connectionOptions || {
+    url:
+      'mongodb://127.0.0.1:27018/data-service?readPreference=primary&ssl=false',
+    options: {
+      connectWithNoPrimary: true,
+      readPreference: 'primary',
+      useNewUrlParser: true,
+      useUnifiedTopology: true
+    }
+  };
+
+  const mockedTunnel = {
+    close() {
+      return Promise.resolve();
+    }
+  };
+
+  return {
+    connect(_model, setupListeners, cb) {
+      const mockedClient = new EventEmitter();
+      mockedClient.db = () => {};
+      mockedClient.close = (_force, closeCb) => {
+        closeCb();
+      };
+      setupListeners(mockedClient);
+      mockedClient.emit('topologyDescriptionChanged', {
+        newDescription: _topologyDescription
+      });
+      cb(null, mockedClient, mockedTunnel, _connectionOptions);
+    }
+  };
+}
+
 describe('NativeClient', function() {
   this.slow(10000);
   this.timeout(20000);
@@ -34,52 +86,32 @@ describe('NativeClient', function() {
 
   describe('#connect', function() {
     context('when mocking connection-model', function() {
-      function mockedTopologyDescription(
-        topologyType = 'Standalone',
-        serverType = 'Single'
-      ) {
-        return {
-          type: topologyType,
-          servers: new Map([['127.0.0.1:27017', { type: serverType }]])
-        };
-      }
-
-      /*
-       * pretends to be a connection-model providing every function call
-       * required in NativeClient#connect, but returns topology and connection
-       * params of our choice
-       */
-      function mockedConnectionModel(topologyDescription, connectionOptions) {
-        const _topologyDescription =
-          topologyDescription || mockedTopologyDescription();
-
-        const _connectionOptions = connectionOptions || {
-          url: 'mongodb://127.0.0.1:27018/data-service?readPreference=primary&ssl=false',
-          options: {
-            connectWithNoPrimary: true,
-            readPreference: 'primary',
-            useNewUrlParser: true,
-            useUnifiedTopology: true
-          }
-        };
-
-        return {
-          connect(_model, setupListeners, cb) {
-            const mockedClient = new EventEmitter();
-            mockedClient.db = () => {};
-            setupListeners(mockedClient);
-            mockedClient.emit('topologyDescriptionChanged', {
-              newDescription: _topologyDescription
-            });
-            cb(null, mockedClient, _connectionOptions);
-          }
-        };
-      }
-
       after(function() {
         mock.stop('mongodb-connection-model');
       });
 
+      it('does not allow to connect twice without disonnecting first', (done) => {
+        mock(
+          'mongodb-connection-model',
+          mockedConnectionModel()
+        );
+
+        const MockedNativeClient = mock.reRequire('../lib/native-client');
+        const mockedClient = new MockedNativeClient(helper.connection);
+
+        mockedClient.connect(() => {});
+        mockedClient.connect(err => {
+          expect(err).to.be.instanceOf(Error);
+          expect(err)
+            .to.have.property('message')
+            .match(
+              /Connect method has been called more than once without disconnecting/
+            );
+
+          done();
+        });
+      });
+
       it('sets .connectionOptions after successful connection', function(done) {
         mock(
           'mongodb-connection-model',
@@ -183,6 +215,40 @@ describe('NativeClient', function() {
     });
   });
 
+  describe('#disconnect', () => {
+    context('when mocking connection-model', () => {
+      after(() => {
+        mock.stop('mongodb-connection-model');
+      });
+
+      it('should close tunnel before calling disconnect callback', (done) => {
+        mock(
+          'mongodb-connection-model',
+          mockedConnectionModel()
+        );
+
+        const MockedNativeClient = mock.reRequire('../lib/native-client');
+        const mockedClient = new MockedNativeClient(helper.connection);
+
+        mockedClient.connect(() => {
+          const closeSpy = sandbox.spy(mockedClient.tunnel, 'close');
+
+          const disconnectCallbackSpy = sandbox.spy(() => {
+            try {
+              expect(closeSpy).to.have.been.calledOnce;
+              expect(closeSpy).to.have.been.calledBefore(disconnectCallbackSpy);
+              done();
+            } catch (err) {
+              done(err);
+            }
+          });
+
+          mockedClient.disconnect(disconnectCallbackSpy);
+        });
+      });
+    });
+  });
+
   describe('#command', function() {
     it('executes the command', function(done) {
       client.command('data-service', { ping: 1 }, function(error, result) {