From 94842ce0efb0d40ef6fe99ac2e735628b035d304 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adolfo=20G=C3=B3mez=20Garc=C3=ADa?= <dkmaster@dkmon.com>
Date: Thu, 28 Apr 2016 09:34:03 +0200
Subject: [PATCH] A better implementation for shared connections

---
 client/src/uds/forward.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/client/src/uds/forward.py b/client/src/uds/forward.py
index 98757f81d..dd43d6931 100644
--- a/client/src/uds/forward.py
+++ b/client/src/uds/forward.py
@@ -73,7 +73,6 @@ class Handler(SocketServer.BaseRequestHandler):
 
 class ForwardThread(threading.Thread):
     status = 0  # Connecting
-    clientUseCounter = 0
 
     def __init__(self, server, port, username, password, localPort, redirectHost, redirectPort, waitTime):
         threading.Thread.__init__(self)
@@ -104,6 +103,7 @@ class ForwardThread(threading.Thread):
 
         ft = ForwardThread(self.server, self.port, self.username, self.password, localPort, redirectHost, redirectPort. self.waitTime)
         ft.client = self.client
+        self.client.useCount += 1  # One more using this client
         ft.start()
 
         while ft.status == 0:
@@ -122,6 +122,7 @@ class ForwardThread(threading.Thread):
     def run(self):
         if self.client is None:
             self.client = paramiko.SSHClient()
+            self.client.useCount = 1  # Custom added variable, to keep track on when to close tunnel
             self.client.load_system_host_keys()
             self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
@@ -134,6 +135,8 @@ class ForwardThread(threading.Thread):
                 self.status = 2  # Error
                 return
 
+        self.clientUseCount += 1
+
         class SubHandler(Handler):
             chain_host = self.redirectHost
             chain_port = self.redirectPort
@@ -147,8 +150,6 @@ class ForwardThread(threading.Thread):
 
         self.status = 1  # Ok, listening
 
-        ForwardThread.clientUseCounter += 1
-
         self.fs = ForwardServer(('', self.localPort), SubHandler)
         self.fs.serve_forever()
 
@@ -161,9 +162,10 @@ class ForwardThread(threading.Thread):
             self.fs.shutdown()
 
             if self.client is not None:
-                ForwardThread.clientUseCounter -= 1
-                if ForwardThread.clientUseCounter == 0:
+                self.client.useCount -= 1
+                if self.client.useCount == 0:
                     self.client.close()
+                self.client = None  # Clean up
         except Exception:
             logger.exception('Exception stopping')
             pass