From 22da88af7d8d80fa39a8f168e3516b9314f722b5 Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Tue, 26 Jun 2018 20:14:43 -0700
Subject: [PATCH] Attempts to handle disconnects better.. TODO: torture tests
 and implement exit status for -x commands

---
 hkexnet.go         | 24 ++++++++++++++++++++++--
 hkexsh/hkexsh.go   |  6 ++++--
 hkexshd/hkexshd.go | 16 +++++++++++++---
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/hkexnet.go b/hkexnet.go
index e62ae7c..1db5787 100644
--- a/hkexnet.go
+++ b/hkexnet.go
@@ -34,6 +34,7 @@ const (
 	CSONone        = iota // No error, normal packet
 	CSOHmacInvalid        // HMAC mismatch detected on remote end
 	CSOTermSize           // set term size (rows:cols)
+	CSOExitStatus         // Remote cmd exit status (TODO)
 	CSOChaff              // Dummy packet, do not pass beyond decryption
 )
 
@@ -45,6 +46,7 @@ type WinSize struct {
 }
 
 type ChaffConfig struct {
+	shutdown bool //set to inform chaffHelper to shut down
 	enabled  bool
 	msecsMin uint //msecs min interval
 	msecsMax uint //msecs max interval
@@ -184,6 +186,7 @@ func Dial(protocol string, ipport string, extensions ...string) (hc *Conn, err e
 
 // Close a hkex.Conn
 func (c Conn) Close() (err error) {
+	c.DisableChaff()
 	err = c.c.Close()
 	log.Println("[Conn Closing]")
 	return
@@ -358,7 +361,8 @@ func (c Conn) Read(b []byte) (n int, err error) {
 		err = binary.Read(c.c, binary.BigEndian, &payloadLen)
 		if err != nil {
 			if err.Error() != "EOF" {
-				panic(err)
+				log.Println("unexpected Read() err:", err)
+				//panic(err)
 				// Cannot just return 0, err here - client won't hang up properly
 				// when 'exit' from shell. TODO: try server sending ctrlStatOp to
 				// indicate to Reader? -rlm 20180428
@@ -504,12 +508,23 @@ func (c Conn) WritePacket(b []byte, op byte) (n int, err error) {
 }
 
 func (c *Conn) EnableChaff() {
+	c.chaff.shutdown = false
 	c.chaff.enabled = true
 	log.Println("Chaffing ENABLED")
 	c.chaffHelper()
 }
 
-func (c *Conn) Chaff(msecsMin uint, msecsMax uint, szMax uint) {
+func (c *Conn) DisableChaff() {
+	c.chaff.enabled = false
+	log.Println("Chaffing DISABLED")
+}
+
+func (c *Conn) ShutdownChaff() {
+	c.chaff.shutdown = true
+	log.Println("Chaffing SHUTDOWN")
+}
+
+func (c *Conn) SetupChaff(msecsMin uint, msecsMax uint, szMax uint) {
 	c.chaff.msecsMin = msecsMin //move these to params of chaffHelper() ?
 	c.chaff.msecsMax = msecsMax
 	c.chaff.szMax = szMax
@@ -533,6 +548,11 @@ func (c *Conn) chaffHelper() {
 				}
 			}
 			time.Sleep(time.Duration(nextDuration) * time.Millisecond)
+			if c.chaff.shutdown {
+				log.Println("*** chaffHelper shutting down")
+				break
+			}
+
 		}
 	}()
 }
diff --git a/hkexsh/hkexsh.go b/hkexsh/hkexsh.go
index b713fd9..40ec2f1 100644
--- a/hkexsh/hkexsh.go
+++ b/hkexsh/hkexsh.go
@@ -184,11 +184,13 @@ func main() {
 	_, err = conn.Write(rec.authCookie)
 
 	// Set up chaffing to server
-	conn.Chaff(chaffFreqMin, chaffFreqMax, chaffBytesMax) // enable client->server chaffing
+	conn.SetupChaff(chaffFreqMin, chaffFreqMax, chaffBytesMax) // enable client->server chaffing
 	if chaffEnabled {
 		conn.EnableChaff()
 	}
-
+	defer conn.DisableChaff()
+	defer conn.ShutdownChaff()
+	
 	//client reader (from server) goroutine
 	wg.Add(1)
 	go func() {
diff --git a/hkexshd/hkexshd.go b/hkexshd/hkexshd.go
index 8bb1145..e737989 100644
--- a/hkexshd/hkexshd.go
+++ b/hkexshd/hkexshd.go
@@ -132,15 +132,25 @@ func runShellAs(who string, cmd string, interactive bool, conn hkexsh.Conn, chaf
 
 	// Copy stdin to the pty.. (bgnd goroutine)
 	go func() {
-		_, _ = io.Copy(ptmx, conn)
+		_, e := io.Copy(ptmx, conn)
+		if e != nil {
+			log.Printf("** std->pty ended **\n")
+			return
+		}
 	}()
 
 	if chaffing {
 		conn.EnableChaff()
 	}
+	defer conn.DisableChaff()
+	defer conn.ShutdownChaff()
 
 	// ..and the pty to stdout.
-	_, _ = io.Copy(conn, ptmx)
+	_, e := io.Copy(conn, ptmx)
+	if e != nil {
+		log.Printf("** pty->stdout ended **\n")
+		return
+	}
 
 	//err = c.Run()  // returns when c finishes.
 
@@ -212,7 +222,7 @@ func main() {
 			// Set up chaffing to client
 			// Will only start when runShellAs() is called
 			// after stdin/stdout are hooked up
-			conn.Chaff(chaffFreqMin, chaffFreqMax, chaffBytesMax) // configure server->client chaffing
+			conn.SetupChaff(chaffFreqMin, chaffFreqMax, chaffBytesMax) // configure server->client chaffing
 
 			// Handle the connection in a new goroutine.
 			// The loop then returns to accepting, so that