From b0f614f82ef2bf2d6484fe21ff3631344530ea20 Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Sun, 11 Nov 2018 22:46:39 -0800
Subject: [PATCH] Tunnel keepalives from client implemented. If client
 dies/exits unexpectedly the server tunnel will disconnect from rport in a
 timely manner.

Signed-off-by: Russ Magee <rmagee@gmail.com>
---
 hkexnet/consts.go  | 15 ++++++++-------
 hkexnet/hkexnet.go |  9 +++++++++
 hkexnet/hkextun.go | 21 +++++++++++++++------
 hkexsh/hkexsh.go   |  9 +++++++++
 4 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/hkexnet/consts.go b/hkexnet/consts.go
index 7f6e3bf..5bdeb9b 100644
--- a/hkexnet/consts.go
+++ b/hkexnet/consts.go
@@ -55,13 +55,14 @@ const (
 	CSOChaff              // Dummy packet, do not pass beyond decryption
 
 	// Tunnel setup/control/status
-	CSOTunSetup    // client -> server tunnel setup request (dstport)
-	CSOTunSetupAck // server -> client tunnel setup ack
-	CSOTunAccept   // client -> server: tunnel client got an Accept()
-	CSOTunRefused  // server -> client: tunnel rport connection refused
-	CSOTunData     // packet contains tunnel data [rport:data]
-	CSOTunDisconn  // server -> client: tunnel rport disconnected
-	CSOTunHangup   // client -> server: tunnel lport hung up
+	CSOTunSetup     // client -> server tunnel setup request (dstport)
+	CSOTunSetupAck  // server -> client tunnel setup ack
+	CSOTunAccept    // client -> server: tunnel client got an Accept()
+	CSOTunRefused   // server -> client: tunnel rport connection refused
+	CSOTunData      // packet contains tunnel data [rport:data]
+	CSOTunKeepAlive // client tunnel heartbeat
+	CSOTunDisconn   // server -> client: tunnel rport disconnected
+	CSOTunHangup    // client -> server: tunnel lport hung up
 )
 
 // TunEndpoint.tunCtl control values - used to control workers for client or server tunnels
diff --git a/hkexnet/hkexnet.go b/hkexnet/hkexnet.go
index 00acdbe..8105a35 100644
--- a/hkexnet/hkexnet.go
+++ b/hkexnet/hkexnet.go
@@ -877,6 +877,15 @@ func (hc Conn) Read(b []byte) (n int, err error) {
 				} else {
 					logger.LogDebug(fmt.Sprintf("[Attempt to write data to closed tun [%d:%d]", lport, rport))
 				}
+			} else if ctrlStatOp == CSOTunKeepAlive {
+					// client side has sent keepalive for tunnels -- if client
+					// dies or exits unexpectedly the absence of this will
+					// let the server know to hang up on Dial()ed server rports.
+				_ = binary.BigEndian.Uint16(payloadBytes[0:2])
+				//logger.LogDebug(fmt.Sprintf("[Server] Got CSOTunKeepAlive"))
+				for _, t := range *hc.tuns {
+					t.KeepAlive = 0
+				}
 			} else if ctrlStatOp == CSONone {
 				hc.dBuf.Write(payloadBytes)
 			} else {
diff --git a/hkexnet/hkextun.go b/hkexnet/hkextun.go
index c95a9a7..078c5f9 100644
--- a/hkexnet/hkextun.go
+++ b/hkexnet/hkextun.go
@@ -42,12 +42,13 @@ type (
 
 	// TunEndpoint [securePort:peer:dataPort]
 	TunEndpoint struct {
-		Rport uint16    // Names are from client's perspective
-		Lport uint16    // ... ie., RPort is on server, LPort is on client
-		Peer  string    //net.Addr
-		Died  bool      // set by client upon receipt of a CSOTunDisconn
-		Ctl   chan rune //See TunCtl_* consts
-		Data  chan []byte
+		Rport     uint16    // Names are from client's perspective
+		Lport     uint16    // ... ie., RPort is on server, LPort is on client
+		Peer      string    //net.Addr
+		Died      bool      // set by client upon receipt of a CSOTunDisconn
+		KeepAlive uint      // must be reset by client to keep server dial() alive
+		Ctl       chan rune //See TunCtl_* consts
+		Data      chan []byte
 	}
 )
 
@@ -328,6 +329,14 @@ func (hc *Conn) StartServerTunnel(lport, rport uint16) {
 								rBuf = append(tunDst.Bytes(), rBuf[:n]...)
 								hc.WritePacket(rBuf[:n+4], CSOTunData)
 							}
+
+							if (*hc.tuns)[rport].KeepAlive > 50 {
+									(*hc.tuns)[rport].Died = true
+									logger.LogDebug("[ServerTun] worker A: Oh dear, client appears to have died. Hanging up rport Dial().")
+							} else {
+								(*hc.tuns)[rport].KeepAlive += 1
+							}
+
 						}
 						logger.LogDebug("[ServerTun] worker A: exiting")
 					}()
diff --git a/hkexsh/hkexsh.go b/hkexsh/hkexsh.go
index 0c89ebd..3fe3106 100755
--- a/hkexsh/hkexsh.go
+++ b/hkexsh/hkexsh.go
@@ -26,6 +26,7 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 
 	hkexsh "blitter.com/go/hkexsh"
 	"blitter.com/go/hkexsh/hkexnet"
@@ -638,6 +639,14 @@ func main() {
 			defer conn.ShutdownChaff()
 		}
 
+		// Keepalive for any tunnels that may exist
+		go func() {
+			for {
+				time.Sleep(time.Duration(2) * time.Second)
+				conn.WritePacket([]byte{0,0}, hkexnet.CSOTunKeepAlive)
+			}
+		}()
+
 		if shellMode {
 			// TESTING - tunnel
 			remAddrs, _ := net.LookupHost(remoteHost)