From 8c7e440c19211f8ce42c8fd5f4fb0622917fb9e9 Mon Sep 17 00:00:00 2001
From: Michael Sun <47126816+MichaelSun90@users.noreply.github.com>
Date: Sun, 17 Sep 2023 11:34:36 +0000
Subject: [PATCH] fix: handle timeouts during tls negotiation for `strict`
 encryption (#1564)

Co-authored-by: Arthur Schreiber <schreiber.arthur@gmail.com>
---
 src/connection.ts            | 48 ++++++++++++++++++++++++++++++------
 test/unit/connection-test.ts | 34 ++++++++++++++++++++++++-
 2 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/src/connection.ts b/src/connection.ts
index ec677b047..30c5419ed 100644
--- a/src/connection.ts
+++ b/src/connection.ts
@@ -1931,7 +1931,8 @@ class Connection extends EventEmitter {
         });
       }, (err) => {
         this.clearConnectTimer();
-        if (err.name === 'AbortError') {
+
+        if (signal.aborted) {
           // Ignore the AbortError for now, this is still handled by the connectTimer firing
           return;
         }
@@ -2008,7 +2009,9 @@ class Connection extends EventEmitter {
     this.transitionTo(this.STATE.SENT_PRELOGIN);
   }
 
-  wrapWithTls(socket: net.Socket): Promise<tls.TLSSocket> {
+  wrapWithTls(socket: net.Socket, signal: AbortSignal): Promise<tls.TLSSocket> {
+    signal.throwIfAborted();
+
     return new Promise((resolve, reject) => {
       const secureContext = tls.createSecureContext(this.secureContextOptions);
       // If connect to an ip address directly,
@@ -2023,12 +2026,41 @@ class Connection extends EventEmitter {
         servername: this.config.options.serverName ? this.config.options.serverName : serverName,
       };
 
-      const encryptsocket = tls.connect(encryptOptions, () => {
-        encryptsocket.removeListener('error', reject);
+      const encryptsocket = tls.connect(encryptOptions);
+
+      const onAbort = () => {
+        encryptsocket.removeListener('error', onError);
+        encryptsocket.removeListener('connect', onConnect);
+
+        encryptsocket.destroy();
+
+        reject(signal.reason);
+      };
+
+      const onError = (err: Error) => {
+        signal.removeEventListener('abort', onAbort);
+
+        encryptsocket.removeListener('error', onError);
+        encryptsocket.removeListener('connect', onConnect);
+
+        encryptsocket.destroy();
+
+        reject(err);
+      };
+
+      const onConnect = () => {
+        signal.removeEventListener('abort', onAbort);
+
+        encryptsocket.removeListener('error', onError);
+        encryptsocket.removeListener('connect', onConnect);
+
         resolve(encryptsocket);
-      });
+      };
+
+      signal.addEventListener('abort', onAbort, { once: true });
 
-      encryptsocket.once('error', reject);
+      encryptsocket.on('error', onError);
+      encryptsocket.on('secureConnect', onConnect);
     });
   }
 
@@ -2047,7 +2079,7 @@ class Connection extends EventEmitter {
       if (this.config.options.encrypt === 'strict') {
         try {
           // Wrap the socket with TLS for TDS 8.0
-          socket = await this.wrapWithTls(socket);
+          socket = await this.wrapWithTls(socket, signal);
         } catch (err) {
           socket.end();
 
@@ -2059,7 +2091,7 @@ class Connection extends EventEmitter {
     })().catch((err) => {
       this.clearConnectTimer();
 
-      if (err.name === 'AbortError') {
+      if (signal.aborted) {
         return;
       }
 
diff --git a/test/unit/connection-test.ts b/test/unit/connection-test.ts
index 4067fa3d2..e4a4d8cb6 100644
--- a/test/unit/connection-test.ts
+++ b/test/unit/connection-test.ts
@@ -1,5 +1,5 @@
 import * as net from 'net';
-import { Connection } from '../../src/tedious';
+import { Connection, ConnectionError } from '../../src/tedious';
 import { assert } from 'chai';
 
 describe('Using `strict` encryption', function() {
@@ -42,4 +42,36 @@ describe('Using `strict` encryption', function() {
       done();
     });
   });
+
+  it('handles connection timeout when performing tls handshake', function(done) {
+    server.on('connection', (connection) => {
+      setTimeout(() => {
+        connection.destroy();
+      }, 4000);
+    });
+
+    const addressInfo = server.address() as net.AddressInfo;
+
+    const connection = new Connection({
+      server: addressInfo?.address,
+      options: {
+        port: addressInfo?.port,
+        encrypt: 'strict',
+        connectTimeout: 3000
+      }
+    });
+
+    connection.connect((err) => {
+      assert.instanceOf(err, ConnectionError);
+
+      const message = `Failed to connect to ${addressInfo?.address}:${addressInfo?.port} in 3000ms`;
+      assert.equal(err!.message, message);
+
+      connection.close();
+    });
+
+    connection.on('end', () => {
+      done();
+    });
+  });
 });