From 0af296dbebdefc9d761d23395ec27b3d716fd3b3 Mon Sep 17 00:00:00 2001
From: Brian White <mscdex@mscdex.net>
Date: Thu, 12 Aug 2021 12:28:12 -0400
Subject: [PATCH] protocol/kex: do not wait to send NEWKEYS

---
 lib/protocol/kex.js | 31 +++++++++++++++++++------------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/lib/protocol/kex.js b/lib/protocol/kex.js
index 59fb852f..507c88af 100644
--- a/lib/protocol/kex.js
+++ b/lib/protocol/kex.js
@@ -778,18 +778,7 @@ const createKeyExchange = (() => {
           this._protocol._packetRW.write.finalize(packet, true)
         );
       }
-      if (!this._sentNEWKEYS) {
-        this._protocol._debug && this._protocol._debug(
-          'Outbound: Sending NEWKEYS'
-        );
-        const p = this._protocol._packetRW.write.allocStartKEX;
-        const packet = this._protocol._packetRW.write.alloc(1, true);
-        packet[p] = MESSAGE.NEWKEYS;
-        this._protocol._cipher.encrypt(
-          this._protocol._packetRW.write.finalize(packet, true)
-        );
-        this._sentNEWKEYS = true;
-      }
+      trySendNEWKEYS(this);
 
       const completeHandshake = () => {
         if (!this.sessionID)
@@ -1180,6 +1169,8 @@ const createKeyExchange = (() => {
                 this._hostVerified = true;
                 if (this._receivedNEWKEYS)
                   this.finish();
+                else
+                  trySendNEWKEYS(this);
               });
             }
             if (ret === undefined) {
@@ -1203,6 +1194,7 @@ const createKeyExchange = (() => {
               'Host accepted (verified)'
             );
             this._hostVerified = true;
+            trySendNEWKEYS(this);
           }
           ++this._step;
           break;
@@ -1798,6 +1790,21 @@ function dhEstimate(neg) {
   return 8192;
 }
 
+function trySendNEWKEYS(kex) {
+  if (!kex._sentNEWKEYS) {
+    kex._protocol._debug && kex._protocol._debug(
+      'Outbound: Sending NEWKEYS'
+    );
+    const p = kex._protocol._packetRW.write.allocStartKEX;
+    const packet = kex._protocol._packetRW.write.alloc(1, true);
+    packet[p] = MESSAGE.NEWKEYS;
+    kex._protocol._cipher.encrypt(
+      kex._protocol._packetRW.write.finalize(packet, true)
+    );
+    kex._sentNEWKEYS = true;
+  }
+}
+
 module.exports = {
   KexInit,
   kexinit,