From 09437973006ce893f59c47ea14596557d9772fbc Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Sun, 11 Nov 2018 21:05:25 -0800
Subject: [PATCH] Minor cleanup prior to tackling os.Exit() in doShellMode()
 and server-side tun hangs

---
 hkexnet/hkextun.go | 9 +++++----
 hkexsh/hkexsh.go   | 2 ++
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/hkexnet/hkextun.go b/hkexnet/hkextun.go
index 4f85d65..da2e3ce 100644
--- a/hkexnet/hkextun.go
+++ b/hkexnet/hkextun.go
@@ -96,19 +96,17 @@ func (hc *Conn) StartClientTunnel(lport, rport uint16) {
 
 	go func() {
 		var wg sync.WaitGroup
-		weAreListening := false
 
 		for cmd := range (*hc.tuns)[rport].Ctl {
-			if cmd == 'a' && !weAreListening {
+			if cmd == 'a' {
 				l, e := net.Listen("tcp4", fmt.Sprintf(":%d", lport))
 				if e != nil {
 					logger.LogDebug(fmt.Sprintf("[ClientTun] Could not get lport %d! (%s)", lport, e))
 				} else {
-					weAreListening = true
 					logger.LogDebug(fmt.Sprintf("[ClientTun] Listening for client tunnel port %d", lport))
 
 					for {
-						c, e := l.Accept()
+						c, e := l.Accept() // blocks until new conn
 						// If tunnel is being re-used, re-init it
 						if (*hc.tuns)[rport] == nil {
 							hc.InitTunEndpoint(lport, "", rport)
@@ -226,6 +224,9 @@ func (hc *Conn) StartClientTunnel(lport, rport uint16) {
 
 						} // end Accept() worker block
 						wg.Wait()
+
+						// When both workers have exited due to a disconnect or other
+						// condition, it's safe to remove the tunnel descriptor.
 						logger.LogDebug("[ClientTun] workers exited")
 						delete((*hc.tuns), rport)
 					} // end for-accept
diff --git a/hkexsh/hkexsh.go b/hkexsh/hkexsh.go
index 0e7718a..0c89ebd 100755
--- a/hkexsh/hkexsh.go
+++ b/hkexsh/hkexsh.go
@@ -256,6 +256,7 @@ func doCopyMode(conn *hkexnet.Conn, remoteDest bool, files string, rec *hkexsh.S
 func doShellMode(isInteractive bool, conn *hkexnet.Conn, oldState *hkexsh.State, rec *hkexsh.Session) {
 	//client reader (from server) goroutine
 	//Read remote end's stdout
+
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
@@ -318,6 +319,7 @@ func doShellMode(isInteractive bool, conn *hkexnet.Conn, oldState *hkexsh.State,
 	// Wait until both stdin and stdout goroutines finish before returning
 	// (ensure client gets all data from server before closing)
 	wg.Wait()
+	return
 }
 
 func UsageShell() {