diff --git a/hkexnet/hkexnet.go b/hkexnet/hkexnet.go
index 28609d8..e17bcb7 100644
--- a/hkexnet/hkexnet.go
+++ b/hkexnet/hkexnet.go
@@ -443,7 +443,7 @@ func (hc Conn) Read(b []byte) (n int, err error) {
 				*hc.closeStat = uint8(payloadBytes[0])
 			} else {
 				log.Println("[truncated payload, cannot determine CSOExitStatus]")
-				*hc.closeStat = 99
+				*hc.closeStat = 98
 			}
 		} else {
 			hc.dBuf.Write(payloadBytes)
@@ -455,7 +455,7 @@ func (hc Conn) Read(b []byte) (n int, err error) {
 		hTmp := hc.rm.Sum(nil)[0:4]
 		log.Printf("<%04x) HMAC:(i)%s (c)%02x\r\n", decryptN, hex.EncodeToString([]byte(hmacIn[0:])), hTmp)
 
-		if *hc.closeStat == 99 {
+		if *hc.closeStat >90 {
 			log.Println("[cannot verify HMAC]")
 		} else {
 			// Log alert if hmac didn't match, corrupted channel
diff --git a/hkexsh/hkexsh.go b/hkexsh/hkexsh.go
index 772a3dc..34db2b3 100755
--- a/hkexsh/hkexsh.go
+++ b/hkexsh/hkexsh.go
@@ -302,14 +302,15 @@ func main() {
 	conn.SetupChaff(chaffFreqMin, chaffFreqMax, chaffBytesMax) // enable client->server chaffing
 	if chaffEnabled {
 		conn.EnableChaff()
-		//defer conn.DisableChaff()
-		//defer conn.ShutdownChaff()
+		defer conn.DisableChaff()
+		defer conn.ShutdownChaff()
 	}
 
 	//client reader (from server) goroutine
 	//Read remote end's stdout
 	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		// By deferring a call to wg.Done(),
 		// each goroutine guarantees that it marks
 		// its direction's stream as finished.
@@ -330,7 +331,6 @@ func main() {
 			log.Println("[* Got EOF *]")
 			_ = hkexsh.Restore(int(os.Stdin.Fd()), oldState) // Best effort.
 		}
-		wg.Done()
 	}()
 
 	// Only look for data from stdin to send to remote end
@@ -343,6 +343,7 @@ func main() {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
+			//!defer wg.Done()
 			// Copy() expects EOF so this will
 			// exit with outerr == nil
 			//!_, outerr := io.Copy(conn, os.Stdin)
@@ -358,15 +359,15 @@ func main() {
 				os.Exit(255)
 			}
 			log.Println("[Sent EOF]")
-			wg.Done()
 		}()
 	}
 
 	// Wait until both stdin and stdout goroutines finish
+	// ** IMPORTANT! This must come before the Restore() tty call below
+	// in order to maintain raw mode for interactive sessions. -rlm 20180805
 	wg.Wait()
-	conn.DisableChaff()
-	conn.ShutdownChaff()
-
+	
 	_ = hkexsh.Restore(int(os.Stdin.Fd()), oldState) // Best effort.
+
 	os.Exit(rec.status)
 }