From 7ebd2366d5d2ee0df4868520bb042b43f83264b8 Mon Sep 17 00:00:00 2001
From: mShan0 <96149598+mShan0@users.noreply.github.com>
Date: Sat, 18 Jun 2022 02:05:57 -0700
Subject: [PATCH] fix: fix ntlm connection state bug (#1446)

---
 src/connection.ts                             | 12 ++++----
 .../child-processes/ntlm-connect-node17.js    | 14 +++++++--
 test/integration/connection-test.js           | 29 ++++++++++++++++---
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/src/connection.ts b/src/connection.ts
index cfd937f52..dd24405d9 100644
--- a/src/connection.ts
+++ b/src/connection.ts
@@ -3411,9 +3411,9 @@ Connection.prototype.STATE = {
           if (handler.loginAckReceived) {
             if (handler.routingData) {
               this.routingData = handler.routingData;
-              this.transitionTo(this.STATE.REROUTING);
+              return this.transitionTo(this.STATE.REROUTING);
             } else {
-              this.transitionTo(this.STATE.LOGGED_IN_SENDING_INITIAL_SQL);
+              return this.transitionTo(this.STATE.LOGGED_IN_SENDING_INITIAL_SQL);
             }
           } else if (this.ntlmpacket) {
             try {
@@ -3442,19 +3442,19 @@ Connection.prototype.STATE = {
               } else {
                 throw error;
               }
-              this.transitionTo(this.STATE.FINAL);
+              return this.transitionTo(this.STATE.FINAL);
             }
           } else if (this.loginError) {
             if (isTransientError(this.loginError)) {
               this.debug.log('Initiating retry on transient error');
-              this.transitionTo(this.STATE.TRANSIENT_FAILURE_RETRY);
+              return this.transitionTo(this.STATE.TRANSIENT_FAILURE_RETRY);
             } else {
               this.emit('connect', this.loginError);
-              this.transitionTo(this.STATE.FINAL);
+              return this.transitionTo(this.STATE.FINAL);
             }
           } else {
             this.emit('connect', new ConnectionError('Login failed.', 'ELOGIN'));
-            this.transitionTo(this.STATE.FINAL);
+            return this.transitionTo(this.STATE.FINAL);
           }
         }
 
diff --git a/test/integration/child-processes/ntlm-connect-node17.js b/test/integration/child-processes/ntlm-connect-node17.js
index c49bb5dfc..869428ade 100644
--- a/test/integration/child-processes/ntlm-connect-node17.js
+++ b/test/integration/child-processes/ntlm-connect-node17.js
@@ -1,4 +1,4 @@
-const { Connection } = require('../../../');
+const { Connection, Request } = require('../../../');
 const fs = require('fs');
 const homedir = require('os').homedir();
 
@@ -7,6 +7,15 @@ function getNtlmConfig() {
     fs.readFileSync(homedir + '/.tedious/test-connection.json', 'utf8')
   ).ntlm;
 }
+
+const request = new Request('select 1; select 2;', function(err, rowCount) {
+  if (err) {
+    throw err;
+  }
+  connection.close();
+});
+
+
 const ntlmConfig = getNtlmConfig();
 
 const connection = new Connection(ntlmConfig);
@@ -19,5 +28,6 @@ connection.connect(function(err) {
       throw err;
     }
   }
-  connection.close();
+
+  connection.execSql(request);
 });
diff --git a/test/integration/connection-test.js b/test/integration/connection-test.js
index 54b745e1c..b88453d3b 100644
--- a/test/integration/connection-test.js
+++ b/test/integration/connection-test.js
@@ -435,14 +435,34 @@ describe('Ntlm Test', function() {
         assert.ok(false, 'Unexpected value for domainCase: ' + domainCase);
     }
 
-    const connection = new Connection(ntlmConfig);
+    let row = 0;
 
-    connection.connect(function(err) {
+    const request = new Request('select 1; select 2;', function(err, rowCount) {
       assert.ifError(err);
+      assert.strictEqual(rowCount, 2);
 
       connection.close();
     });
 
+    request.on('doneInProc', function(rowCount, more) {
+      assert.strictEqual(rowCount, 1);
+    });
+
+    request.on('columnMetadata', function(columnsMetadata) {
+      assert.strictEqual(columnsMetadata.length, 1);
+    });
+
+    request.on('row', function(columns) {
+      assert.strictEqual(columns[0].value, ++row);
+    });
+
+    const connection = new Connection(ntlmConfig);
+
+    connection.connect(function(err) {
+      assert.ifError(err);
+      connection.execSql(request);
+    });
+
     connection.on('end', function() {
       done();
     });
@@ -471,8 +491,9 @@ describe('Ntlm Test', function() {
     });
 
   } else {
+    const ntlmConfig = getNtlmConfig();
 
-    it('should throw an aggregate error with node 17', () => {
+    (ntlmConfig ? it : it.skip)('should throw an aggregate error with node 17', () => {
       const child = childProcess.spawnSync(process.execPath,
                                            ['./test/integration/child-processes/ntlm-connect-node17.js'],
                                            { encoding: 'utf8' });
@@ -483,7 +504,7 @@ describe('Ntlm Test', function() {
       assert.strictEqual(child.status, 1);
     });
 
-    it('should ntlm with node 17 when `--openssl-legacy-provider` flag enabled', () => {
+    (ntlmConfig ? it : it.skip)('should ntlm with node 17 when `--openssl-legacy-provider` flag enabled', () => {
       const child = childProcess.spawnSync(process.execPath,
                                            ['--openssl-legacy-provider',
                                             './test/integration/child-processes/ntlm-connect-node17.js'],