From 316b4a6c54ea34901d2fbdd8795c39ae2e14e03a Mon Sep 17 00:00:00 2001
From: Thorsten Sommer <service@sommer-engineering.com>
Date: Fri, 2 Jan 2015 18:45:38 +0100
Subject: [PATCH] Added gracefully shutdown for clients

Added the gracefully shutdown for all client connections to release all
these resources.
---
 Forward.go  | 25 +++++++++++++++++++++++--
 Transfer.go |  4 +++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/Forward.go b/Forward.go
index a0b437a..4f24e8d 100644
--- a/Forward.go
+++ b/Forward.go
@@ -68,9 +68,30 @@ func forward(localConn net.Conn, config *ssh.ClientConfig) {
 			// Fine, the connections are up and ready :-)
 			log.Printf("The remote end-point %s is connected.\n", remoteAddrString)
 
+			// To be able to close down both transfer threads, we create a channel:
+			quit := make(chan bool)
+
 			// Create the transfers to/from both sides (two new threads are created for this):
-			go transfer(localConn, sshConn, `Local => Remote`)
-			go transfer(sshConn, localConn, `Remote => Local`)
+			go transfer(localConn, sshConn, `Local => Remote`, quit)
+			go transfer(sshConn, localConn, `Remote => Local`, quit)
+
+			// Wait and look if any of the two transfer theads are down:
+			isRunning := true
+			for isRunning {
+				select {
+				case <-quit:
+					log.Println(`At least one transfer was stopped.`)
+					isRunning = false
+					break
+				}
+			}
+
+			// Now, close all the channels and therefore, force the other / second thread to go down:
+			log.Println(`Close now all connections.`)
+			sshConn.Close()
+			localConn.Close()
+			sshClientConnection.Close()
+
 			return
 		}
 	}
diff --git a/Transfer.go b/Transfer.go
index 119c40a..b5d3fac 100644
--- a/Transfer.go
+++ b/Transfer.go
@@ -6,7 +6,7 @@ import (
 )
 
 // The transfer function.
-func transfer(fromReader io.Reader, toWriter io.Writer, name string) {
+func transfer(fromReader io.Reader, toWriter io.Writer, name string, quit chan bool) {
 
 	log.Printf("%s transfer started.", name)
 
@@ -22,4 +22,6 @@ func transfer(fromReader io.Reader, toWriter io.Writer, name string) {
 	} else {
 		log.Printf("%s transfer closed.\n", name)
 	}
+
+	quit <- true
 }