From bdc8111efcb6e17a0ee0866ef88e572907b74d52 Mon Sep 17 00:00:00 2001
From: tcely <tcely@users.noreply.github.com>
Date: Fri, 17 Aug 2018 01:37:50 -0400
Subject: [PATCH] krssh: call tryParse on fewer packets

We're looking for specific packets that happen early in a SSH connection.

Once we send the host auth packet to krd, we should not continue trying to
parse every packet coming back from the server.

I agree to license all rights to my contributions in each modified file
exclusively to KryptCo, Inc.

Signed-off-by: tcely <tcely@users.noreply.github.com>
---
 krssh/krssh.go | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/krssh/krssh.go b/krssh/krssh.go
index e028dabe..b621e3b0 100644
--- a/krssh/krssh.go
+++ b/krssh/krssh.go
@@ -48,8 +48,9 @@ type kexECDHReplyMsg struct {
 	Signature       []byte
 }
 
-func sendHostAuth(hostAuth kr.HostAuth) {
+func sendHostAuth(hostAuth kr.HostAuth, sentHostAuth chan bool) {
 	conn, err := kr.HostAuthDial()
+	sentHostAuth <- nil == err
 	if err != nil {
 		os.Stderr.WriteString(kr.Red("Krypton ▶ Could not connect to Krypton daemon. Make sure it is running by typing \"kr restart\".\r\n"))
 		return
@@ -58,7 +59,7 @@ func sendHostAuth(hostAuth kr.HostAuth) {
 	json.NewEncoder(conn).Encode(hostAuth)
 }
 
-func tryParse(hostname string, onHostPrefix chan string, buf []byte) (err error) {
+func tryParse(hostname string, onHostPrefix chan string, sentHostAuth chan bool, buf []byte) (err error) {
 	kexECDHReplyTemplate := kexECDHReplyMsg{}
 	err = ssh.Unmarshal(buf, &kexECDHReplyTemplate)
 
@@ -78,7 +79,7 @@ func tryParse(hostname string, onHostPrefix chan string, buf []byte) (err error)
 		case onHostPrefix <- "[" + basex.Base62StdEncoding.EncodeToString(sigHash[:]) + "]":
 		default:
 		}
-		sendHostAuth(hostAuth)
+		sendHostAuth(hostAuth, sentHostAuth)
 	}
 	return
 }
@@ -174,6 +175,8 @@ func startRemoteOutputParsing(remoteConn io.Reader, doneChan chan bool, notifyPr
 		func() {
 			buf := make([]byte, 1<<18)
 			packetNum := 0
+			doneHostAuth := false
+			sentHostAuth := make(chan bool, 1)
 			for {
 				n, err := remoteConn.Read(buf)
 				if err != nil {
@@ -181,14 +184,18 @@ func startRemoteOutputParsing(remoteConn io.Reader, doneChan chan bool, notifyPr
 				}
 				if n > 0 {
 					packetNum++
-					if packetNum > 1 {
-						sshPacket := parseSSHPacket(buf)
-						tryParse(host, notifyPrefixChan, sshPacket)
-					}
 					byteBuf := bytes.NewBuffer(buf[:n])
 					_, err := byteBuf.WriteTo(os.Stdout)
-					if err != nil {
+					switch {
+					case err != nil:
 						return
+					case doneHostAuth:
+						continue
+					case packetNum > 1:
+						sshPacket := parseSSHPacket(buf)
+						if err := tryParse(host, notifyPrefixChan, sentHostAuth, sshPacket); err == nil {
+							doneHostAuth = <-sentHostAuth
+						}
 					}
 				}
 			}