fix(ssh-tunnel): reconnect the ssh tunnel when it gets disconnected COMPASS-5454

.evergreen/connectivity-tests/run.sh

@@ -14,6 +14,7 @@ echo running connectivity tests image
 
 docker run \
   --rm \
+  -e DEBUG="${DEBUG}" \
   -e E2E_TESTS_ATLAS_HOST="${E2E_TESTS_ATLAS_HOST}" \
   -e E2E_TESTS_DATA_LAKE_HOST="${E2E_TESTS_DATA_LAKE_HOST}" \
   -e E2E_TESTS_ANALYTICS_NODE_HOST="${E2E_TESTS_ANALYTICS_NODE_HOST}" \

packages/ssh-tunnel/package.json

@@ -54,12 +54,14 @@
     "@mongodb-js/prettier-config-compass": "^0.5.0",
     "@mongodb-js/tsconfig-compass": "^0.6.0",
     "@types/chai": "^4.2.21",
+    "@types/chai-as-promised": "^7.1.4",
     "@types/debug": "^4.1.7",
     "@types/mocha": "^9.0.0",
     "@types/node-fetch": "^2.5.8",
     "@types/sinon-chai": "^3.2.5",
     "@types/ssh2": "^0.5.46",
     "chai": "^4.3.4",
+    "chai-as-promised": "*",
     "depcheck": "^1.4.1",
     "eslint": "^7.25.0",
     "gen-esm-wrapper": "^1.1.0",

packages/ssh-tunnel/src/index.spec.ts

@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/restrict-template-expressions */
+import { once } from 'events';
 import type { ServerConfig } from 'ssh2';
 import { Server as SSHServer } from 'ssh2';
 import type { Server as HttpServer } from 'http';
@@ -7,12 +8,16 @@ import { promisify } from 'util';
 import { readFileSync } from 'fs';
 import path from 'path';
 import { Socket } from 'net';
-import fetch from 'node-fetch';
+import fetch, { FetchError } from 'node-fetch';
 import { expect } from 'chai';
 import { SocksClient } from 'socks';
-
+import chai from 'chai';
+import chaiAsPromised from 'chai-as-promised';
 import type { SshTunnelConfig } from './index';
 import SSHTunnel from './index';
+import sinon from 'sinon';
+
+chai.use(chaiAsPromised);
 
 function sleep(ms: number) {
   return new Promise((resolve) => setTimeout(resolve, ms));
@@ -94,6 +99,7 @@ async function createTestSshTunnel(config: Partial<SshTunnelConfig> = {}) {
     localPort: 0,
     ...config,
   });
+  sinon.spy(sshTunnel.sshClient, 'connect');
   await sshTunnel.listen();
 }
 
@@ -106,6 +112,12 @@ async function stopTestSshTunnel() {
   }
 }
 
+function breakSshTunnelConnection() {
+  const promise = once(sshTunnel.sshClient, 'close');
+  sshTunnel.sshClient.end();
+  return promise;
+}
+
 interface Socks5ProxyOptions {
   proxyHost: string;
   proxyPort: number;
@@ -273,4 +285,129 @@ describe('SSHTunnel', function () {
       expect(sshTunnel.server.address()).to.equal(null);
     }
   });
+
+  it('does not reconnect if the tunnel is already connected', async function () {
+    await createTestSshTunnel();
+
+    const address = `http://localhost:${httpServer.address().port}/`;
+    const options = {
+      proxyHost: sshTunnel.config.localAddr,
+      proxyPort: sshTunnel.config.localPort,
+    };
+    const expected = 'Hello from http server';
+
+    const res1 = await httpFetchWithSocks5(address, options);
+    expect(await res1.text()).to.equal(expected);
+
+    const res2 = await httpFetchWithSocks5(address, options);
+    expect(await res2.text()).to.equal(expected);
+
+    expect(sshTunnel.sshClient.connect.callCount).to.equal(1);
+  });
+
+  it('reconnects tunnel if it got accidentally disconnected', async function () {
+    await createTestSshTunnel();
+
+    await breakSshTunnelConnection();
+
+    const address = `http://localhost:${httpServer.address().port}/`;
+    const options = {
+      proxyHost: sshTunnel.config.localAddr,
+      proxyPort: sshTunnel.config.localPort,
+    };
+    const expected = 'Hello from http server';
+
+    const res1 = await httpFetchWithSocks5(address, options);
+    expect(await res1.text()).to.equal(expected);
+
+    const res2 = await httpFetchWithSocks5(address, options);
+    expect(await res2.text()).to.equal(expected);
+
+    expect(sshTunnel.sshClient.connect.callCount).to.equal(2);
+  });
+
+  it('reuses the connection promise if a request comes in before the tunnel connects', async function () {
+    await createTestSshTunnel();
+
+    await breakSshTunnelConnection();
+
+    const address = `http://localhost:${httpServer.address().port}/`;
+    const options = {
+      proxyHost: sshTunnel.config.localAddr,
+      proxyPort: sshTunnel.config.localPort,
+    };
+    const expected = 'Hello from http server';
+
+    const [res1, res2] = await Promise.all([
+      httpFetchWithSocks5(address, options),
+      httpFetchWithSocks5(address, options),
+    ]);
+    expect(await res1.text()).to.equal(expected);
+    expect(await res2.text()).to.equal(expected);
+
+    expect(sshTunnel.sshClient.connect.callCount).to.equal(2);
+  });
+
+  it('does not reconnect the tunnel after it was deliberately closed', async function () {
+    await createTestSshTunnel();
+
+    // NOTE: normally you'd call close(), but that also closes the server so
+    // you'd get a different error first. Trying to trigger the race condition
+    // where the request made it to the socks5 server in time.
+    await sshTunnel.closeSshClient();
+
+    const remotePort = httpServer.address().port;
+    const address = `http://localhost:${remotePort}/`;
+
+    const options = {
+      proxyHost: sshTunnel.config.localAddr,
+      proxyPort: sshTunnel.config.localPort,
+    };
+
+    const promise = httpFetchWithSocks5(address, options);
+
+    await expect(promise).to.be.rejectedWith(
+      FetchError,
+      `request to http://localhost:${remotePort}/ failed, reason: Socket closed`
+    );
+
+    expect(sshTunnel.sshClient.connect.callCount).to.equal(1);
+  });
+
+  it('reconnects if the ssh connection times out while we try and open the channel', async function () {
+    await createTestSshTunnel();
+
+    const forwardOut = sshTunnel.forwardOut;
+    sinon
+      .stub(sshTunnel, 'forwardOut')
+      .callsFake(async function (
+        srcAddr: string,
+        srcPort: number,
+        dstAddr: string,
+        dstPort: number
+      ) {
+        await breakSshTunnelConnection();
+        const promise = forwardOut.call(
+          this,
+          srcAddr,
+          srcPort,
+          dstAddr,
+          dstPort
+        );
+        sshTunnel.forwardOut.restore();
+        return promise;
+      });
+
+    const address = `http://localhost:${httpServer.address().port}/`;
+    const options = {
+      proxyHost: sshTunnel.config.localAddr,
+      proxyPort: sshTunnel.config.localPort,
+    };
+    const expected = 'Hello from http server';
+
+    const res = await httpFetchWithSocks5(address, options);
+    expect(await res.text()).to.equal(expected);
+
+    expect(sshTunnel.sshClient.connect.callCount).to.equal(2);
+  });
 });