WIP - try and fix event loop block on embedded server proxy connections. Added invalid host args when using dict

Author: Dan

embedded_server/embedded_server.py

@@ -38,12 +38,13 @@
 *Warning* - Note that commands, with or without a shell, are actually run on the system running this server. Destructive commands will affect the system as permissions of user running the server allow. **Use at your own risk**.
 """
 
+from gipc import start_process
+from multiprocessing import Process
 import sys
 if 'threading' in sys.modules:
     del sys.modules['threading']
 from gevent import monkey
 monkey.patch_all()
-from multiprocessing import Process
 import os
 import gevent
 from gevent import socket
@@ -53,9 +54,11 @@
 import logging
 import paramiko
 import time
+import gevent.subprocess
+import gevent.hub
+
 from .stub_sftp import StubSFTPServer
 from .tunnel import Tunneler
-import gevent.subprocess
 
 logger = logging.getLogger("embedded_server")
 paramiko_logger = logging.getLogger('paramiko.transport')
@@ -86,16 +89,16 @@ def __init__(self, host_key, fail_auth=False,
                  ssh_exception=False,
                  socket=None,
                  port=0,
-                 listen_ip='127.0.0.1',
+                 host='127.0.0.1',
                  timeout=None):
         if not socket:
-            self.socket = make_socket(listen_ip, port)
+            self.socket = make_socket(host, port)
         if not self.socket:
             msg = "Could not establish listening connection on %s:%s"
-            logger.error(msg, listen_ip, port)
-            raise Exception(msg, listen_ip, port)
-        self.listen_ip = listen_ip
-        self.listen_port = self.socket.getsockname()[1]
+            logger.error(msg, host, port)
+            raise Exception(msg, host, port)
+        self.host = host
+        self.port = self.socket.getsockname()[1]
         self.event = Event()
         self.fail_auth = fail_auth
         self.ssh_exception = ssh_exception
@@ -105,15 +108,20 @@ def __init__(self, host_key, fail_auth=False,
 
     def start_listening(self):
         try:
+            gevent.sleep(0)
             self.socket.listen(100)
-            logger.info('Listening for connection on %s:%s..', self.listen_ip,
-                        self.listen_port)
+            logger.info('Listening for connection on %s:%s..', self.host,
+                        self.port)
         except Exception as e:
             logger.error('*** Listen failed: %s' % (str(e),))
             traceback.print_exc()
             raise
+        gevent.sleep()
         conn, addr = self.socket.accept()
+        gevent.sleep(.2)
         logger.info('Got connection..')
+        # import ipdb; ipdb.set_trace()
+        gevent.sleep(.2)
         if self.timeout:
             logger.debug("SSH server sleeping for %s then raising socket.timeout",
                          self.timeout)
@@ -123,21 +131,26 @@ def start_listening(self):
         self.transport.add_server_key(self.host_key)
         self.transport.set_subsystem_handler('sftp', paramiko.SFTPServer,
                                              StubSFTPServer)
+        gevent.sleep()
         try:
             self.transport.start_server(server=self)
         except paramiko.SSHException as e:
             logger.exception('SSH negotiation failed')
             raise
+        gevent.sleep(0)
 
     def run(self):
         while True:
             try:
                 self.start_listening()
+                gevent.sleep(0)
             except Exception:
                 logger.exception("Error occured starting server")
                 continue
+            gevent.sleep(0)
             try:
                 self.accept_connections()
+                gevent.sleep(0)
             except Exception as e:
                 logger.error('*** Caught exception: %s: %s' % (str(e.__class__), str(e),))
                 traceback.print_exc()
@@ -152,8 +165,10 @@ def accept_connections(self):
             gevent.sleep(0)
             channel = self.transport.accept(20)
             if not channel:
-                logger.error("Could not establish channel")
-                return
+                logger.error("Could not establish channel on %s:%s",
+                             self.host, self.port)
+                gevent.sleep(0)
+                continue
             while self.transport.is_active():
                 logger.debug("Transport active, waiting..")
                 gevent.sleep(1)
@@ -190,21 +205,28 @@ def check_channel_pty_request(self, channel, term, width, height, pixelwidth,
         return True
 
     def check_channel_direct_tcpip_request(self, chanid, origin, destination):
+        # import ipdb; ipdb.set_trace()
         logger.debug("Proxy connection %s -> %s requested", origin, destination,)
         extra = {'username' : self.transport.get_username()}
         logger.debug("Starting proxy connection %s -> %s",
                      origin, destination, extra=extra)
+        self.event.set()
         try:
-            tunnel = Tunneler(destination, self.transport, chanid)
+            gevent.sleep(.2)
+            tunnel = Process(target=Tunneler, args=(destination, self.transport, chanid,))
+            tunnel.daemon = True
             tunnel.start()
+            gevent.sleep(.2)
         except Exception as ex:
             logger.error("Error creating proxy connection to %s - %s",
                          destination, ex,)
             return paramiko.OPEN_FAILED_CONNECT_FAILED
+        gevent.sleep(2)
         return paramiko.OPEN_SUCCEEDED
 
     def check_channel_forward_agent_request(self, channel):
         logger.debug("Forward agent key request for channel %s" % (channel,))
+        gevent.sleep(0)
         return True
 
     def check_channel_exec_request(self, channel, cmd,
@@ -217,6 +239,7 @@ def check_channel_exec_request(self, channel, cmd,
                                           stdin=gevent.subprocess.PIPE,
                                           shell=True, env=_env)
         gevent.spawn(self._read_response, channel, process)
+        gevent.sleep(0)
         return True
 
     def _read_response(self, channel, process):
@@ -230,6 +253,7 @@ def _read_response(self, channel, process):
         # Let clients consume output from channel before closing
         gevent.sleep(.1)
         channel.close()
+        gevent.sleep(0)
 
 def make_socket(listen_ip, port=0):
     """Make socket on given address and available port chosen by OS"""
@@ -246,21 +270,33 @@ def make_socket(listen_ip, port=0):
 def start_server(listen_ip, fail_auth=False, ssh_exception=False,
                  timeout=None,
                  listen_port=0):
-    server = Server(host_key, listen_ip=listen_ip, port=listen_port,
+    # gevent.reinit()
+    gevent.hub.reinit()
+    # h.destroy(destroy_loop=True)
+    # h = gevent.hub.Hub()
+    # h.NOT_ERROR = (Exception,)
+    # gevent.hub.set_hub(h)
+    server = Server(host_key, host=listen_ip, port=listen_port,
                     fail_auth=fail_auth, ssh_exception=ssh_exception,
                     timeout=timeout)
     try:
         server.run()
     except KeyboardInterrupt:
         sys.exit(0)
+    # listen_process = Process(target=server.run)
+    # listen_process.start()
+    # listen_process.join()
 
 def start_server_process(listen_ip, fail_auth=False, ssh_exception=False,
                          timeout=None, listen_port=0):
+    gevent.reinit()
     server = Process(target=start_server, args=(listen_ip,),
                      kwargs={
                          'listen_port': listen_port,
                          'fail_auth': fail_auth,
                          'ssh_exception': ssh_exception,
                          'timeout': timeout,
                          })
+    server.start()
+    gevent.sleep(.2)
     return server

embedded_server/tunnel.py

@@ -29,9 +29,11 @@
 class Tunneler(gevent.Greenlet):    
     def __init__(self, address, transport, chanid):
         gevent.Greenlet.__init__(self)
+        gevent.sleep(.2)
         self.socket = socket.create_connection(address)
         self.transport = transport
         self.chanid = chanid
+        gevent.sleep(0)
 
     def close(self):
         try:
@@ -50,13 +52,16 @@ def tunnel(self, dest_socket, source_chan):
                 response_data = dest_socket.recv(1024)
                 source_chan.sendall(response_data)
                 logger.debug("Tunnel sent data..")
-                gevent.sleep(0)
+                gevent.sleep(.1)
         finally:
             source_chan.close()
             dest_socket.close()
+        gevent.sleep(0)
 
     def run(self):
+        gevent.sleep(.2)
         channel = self.transport.accept(20)
+        gevent.sleep(0)
         if not channel:
             return
         if not channel.get_id() == self.chanid:
@@ -69,3 +74,4 @@ def run(self):
         except Exception as ex:
             logger.exception("Got exception creating tunnel - %s", ex,)
         logger.debug("Finished tunneling")
+        gevent.sleep(0)

pssh/pssh_client.py

@@ -369,13 +369,19 @@ def run_command(self, *args, **kwargs):
         :rtype: Dictionary with host as key as per \
           :mod:`pssh.pssh_client.ParallelSSHClient.get_output`
         
-        :raises: :mod:`pssh.exceptions.AuthenticationException` on authentication error
-        :raises: :mod:`pssh.exceptions.UnknownHostException` on DNS resolution error
-        :raises: :mod:`pssh.exceptions.ConnectionErrorException` on error connecting
-        :raises: :mod:`pssh.exceptions.SSHException` on other undefined SSH errors
-        :raises: :mod:`pssh.exceptions.HostArgumentException` on number of host \
-        arguments not equal to number of hosts
+        :raises: :mod:`pssh.exceptions.AuthenticationException` on \
+        authentication error
+        :raises: :mod:`pssh.exceptions.UnknownHostException` on DNS resolution \
+        error
+        :raises: :mod:`pssh.exceptions.ConnectionErrorException` on error \
+        connecting
+        :raises: :mod:`pssh.exceptions.SSHException` on other undefined SSH \
+        errors
+        :raises: :mod:`pssh.exceptions.HostArgumentException` on number of \
+        host arguments not equal to number of hosts
         :raises: `TypeError` on not enough host arguments for cmd string format
+        :raises: `KeyError` on no host argument key in arguments dict for cmd \
+        string format
         
         **Example Usage**
         

tests/test_pssh_client.py

@@ -27,18 +27,17 @@
 import warnings
 import shutil
 import sys
-from multiprocessing import Process
 
+from embedded_server.embedded_server import start_server, make_socket, \
+     logger as server_logger, paramiko_logger, start_server_process
+from embedded_server.fake_agent import FakeAgent
+from paramiko import RSAKey
 import gevent
 from pssh import ParallelSSHClient, UnknownHostException, \
      AuthenticationException, ConnectionErrorException, SSHException, \
      logger as pssh_logger
 from pssh.exceptions import HostArgumentException
 from pssh.utils import load_private_key
-from embedded_server.embedded_server import start_server, make_socket, \
-     logger as server_logger, paramiko_logger, start_server_process
-from embedded_server.fake_agent import FakeAgent
-from paramiko import RSAKey
 
 PKEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
 USER_KEY = RSAKey.from_private_key_file(PKEY_FILENAME)
@@ -538,22 +537,24 @@ def test_ssh_proxy(self):
         proxy_server_port = self.make_random_port(proxy_host)
         proxy_server = start_server_process(proxy_host,
                                             listen_port=proxy_server_port)
-        client = ParallelSSHClient([self.host], port=self.listen_port,
+        client = ParallelSSHClient([self.host], port=39783,
                                    pkey=self.user_key,
                                    proxy_host=proxy_host,
-                                   proxy_port=proxy_server_port
+                                   proxy_port=proxy_server_port,
                                    )
         gevent.sleep(2)
-        output = client.run_command(self.fake_cmd)
-        gevent.sleep(.2)
-        stdout = list(output[self.host]['stdout'])
-        expected_stdout = [self.fake_resp]
-        self.assertEqual(expected_stdout, stdout,
-                         msg="Got unexpected stdout - %s, expected %s" % 
-                         (stdout,
-                          expected_stdout,))
-        del client
-        proxy_server.terminate()
+        try:
+            output = client.run_command(self.fake_cmd)
+            gevent.sleep(1)
+            stdout = list(output[self.host]['stdout'])
+            expected_stdout = [self.fake_resp]
+            self.assertEqual(expected_stdout, stdout,
+                             msg="Got unexpected stdout - %s, expected %s" % 
+                             (stdout,
+                              expected_stdout,))
+        finally:
+            del client
+            proxy_server.terminate()
 
     def test_ssh_proxy_auth(self):
         """Test connecting to remote destination via SSH proxy
@@ -589,18 +590,19 @@ def test_ssh_proxy_auth(self):
                              proxy_password)
             self.assertTrue(client.host_clients[self.host].proxy_pkey)
         finally:
+            del client
             proxy_server.terminate()
 
     def test_ssh_proxy_auth_fail(self):
         """Test failures while connecting via proxy"""
-        listen_socket = make_socket(self.host)
-        listen_port = listen_socket.getsockname()[1]
-        self.server.kill()
-        server = start_server(listen_socket,
-                              fail_auth=True)
-        proxy_server_socket = make_socket('127.0.0.2')
-        proxy_server_port = proxy_server_socket.getsockname()[1]
-        proxy_server = start_server(proxy_server_socket)
+        # listen_socket = make_socket(self.host)
+        proxy_host = '127.0.0.2'
+        listen_port = self.make_random_port()
+        server = start_server_process(self.host, listen_port=listen_port,
+                                      fail_auth=True)
+        proxy_server_port = self.make_random_port(host=proxy_host)
+        proxy_server = start_server_process(proxy_host,
+                                            listen_port=proxy_server_port)
         proxy_user = 'proxy_user'
         proxy_password = 'fake'
         gevent.sleep(2)
@@ -611,12 +613,15 @@ def test_ssh_proxy_auth_fail(self):
                                    proxy_user=proxy_user,
                                    proxy_password='fake',
                                    proxy_pkey=self.user_key,
+                                   num_retries=1,
                                    )
         gevent.sleep(2)
-        self.assertRaises(AuthenticationException, client.run_command, self.fake_cmd)
-        del client
-        server.kill()
-        proxy_server.kill()
+        try:
+            self.assertRaises(AuthenticationException, client.run_command, self.fake_cmd)
+        finally:
+            del client
+            server.terminate()
+            proxy_server.terminate()
 
     def test_bash_variable_substitution(self):
         """Test bash variables work correctly"""