Ensuring disconnect happens after close_channel in TunnelServer (#307)

* Fix issue with TunnelServer disconnect while channel still open
* Moving client disconnect() into TunnelServer
* Added test for proxy reconnection
* Resolves #304

Co-authored-by: Max Ballenger <max@aeva.ai>

Author: maxballenger

pssh/clients/native/single.py

@@ -173,9 +173,18 @@ def disconnect(self):
             except Exception:
                 pass
             self.session = None
-        self.sock = None
         if isinstance(self._proxy_client, SSHClient):
-            self._proxy_client.disconnect()
+            # Don't disconnect proxy client here - let the TunnelServer do it at the time that
+            # _wait_send_receive_lets ends. The cleanup_server call here triggers the TunnelServer
+            # to stop.
+            FORWARDER.cleanup_server(self._proxy_client)
+
+            # I wanted to clean up all the sockets here to avoid a ResourceWarning from unittest,
+            # but unfortunately closing this socket here causes a segfault, not sure why yet.
+            # self.sock.close()
+        else:
+            self.sock.close()
+        self.sock = None
 
     def spawn_send_keepalive(self):
         """Spawns a new greenlet that sends keep alive messages every

pssh/clients/native/tunnel.py

@@ -90,10 +90,8 @@ def _cleanup_servers_let(self):
 
     def _cleanup_servers(self):
         for client in list(self._servers.keys()):
-            server = self._servers[client]
             if client.sock is None or client.sock.closed:
-                server.stop()
-                del self._servers[client]
+                self.cleanup_server(client)
 
     def run(self):
         """Thread runner ensures a non main hub has been created for all subsequent
@@ -118,6 +116,13 @@ def run(self):
                          exc_info=1)
             self.shutdown()
 
+    def cleanup_server(self, client):
+        """The purpose of this function is for a proxied client to notify the LocalForwarder that it
+         is shutting down and its corresponding server can also be shut down."""
+        server = self._servers[client]
+        server.stop()
+        del self._servers[client]
+
 
 class TunnelServer(StreamServer):
     """Local port forwarding server for tunneling connections from remote SSH server.
@@ -165,9 +170,12 @@ def _wait_send_receive_lets(self, source, dest, channel, forward_sock):
         try:
             joinall((source, dest), raise_error=True)
         finally:
-            logger.debug("Closing channel and forward socket")
+            # Forward socket does not need to be closed here; StreamServer does it in do_close
+            logger.debug("Closing channel")
             self._client.close_channel(channel)
-            forward_sock.close()
+
+            # Disconnect client here to make sure it happens AFTER close_channel
+            self._client.disconnect()
 
     def _read_forward_sock(self, forward_sock, channel):
         while True:

tests/native/test_tunnel.py

@@ -23,6 +23,7 @@
 import string
 import random
 import time
+import gc
 
 from datetime import datetime
 from socket import timeout as socket_timeout
@@ -32,7 +33,7 @@
 
 from pssh.config import HostConfig
 from pssh.clients.native import SSHClient, ParallelSSHClient
-from pssh.clients.native.tunnel import LocalForwarder, TunnelServer
+from pssh.clients.native.tunnel import LocalForwarder, TunnelServer, FORWARDER
 from pssh.exceptions import UnknownHostException, \
     AuthenticationException, ConnectionErrorException, SessionError, \
     HostArgumentException, SFTPError, SFTPIOError, Timeout, SCPError, \
@@ -95,6 +96,36 @@ def test_tunnel_server(self):
             self.assertEqual(self.port, client.port)
         finally:
             remote_server.stop()
+    
+    # The purpose of this test is to exercise 
+    # https://github.com/ParallelSSH/parallel-ssh/issues/304 
+    def test_tunnel_server_reconn(self):
+        remote_host = '127.0.0.8'
+        remote_server = OpenSSHServer(listen_ip=remote_host, port=self.port)
+        remote_server.start_server()
+
+        reconn_n = 20       # Number of reconnect attempts
+        reconn_delay = 1    # Number of seconds to delay betwen reconnects
+        try:
+            for _ in range(reconn_n):
+                client = SSHClient(
+                    remote_host, port=self.port, pkey=self.user_key,
+                    num_retries=1,
+                    proxy_host=self.proxy_host,
+                    proxy_pkey=self.user_key,
+                    proxy_port=self.proxy_port,
+                )
+                output = client.run_command(self.cmd)
+                _stdout = list(output.stdout)
+                self.assertListEqual(_stdout, [self.resp])
+                self.assertEqual(remote_host, client.host)
+                self.assertEqual(self.port, client.port)
+                client.disconnect()
+                FORWARDER._cleanup_servers()
+                time.sleep(reconn_delay)
+                gc.collect()
+        finally:
+            remote_server.stop()
 
     def test_tunnel_server_same_port(self):
         remote_host = '127.0.0.7'