Add support for signaling a handshake process in packetizer.

This makes it possible to raise an EOFError if the handshake process
has started but takes too long time to finish.

1 files changed, 45 insertions, 0 deletions

diff --git a/paramiko/packet.py b/paramiko/packet.py
index f516ff9b..b922000c 100644
--- a/paramiko/packet.py
+++ b/paramiko/packet.py
@@ -99,6 +99,10 @@ class Packetizer (object):
         self.__keepalive_last = time.time()
         self.__keepalive_callback = None
 
+        self.__timer = None
+        self.__handshake_complete = False
+        self.__timer_expired = False
+
     def set_log(self, log):
         """
         Set the Python log object to use for logging.
@@ -182,6 +186,45 @@ class Packetizer (object):
         self.__keepalive_callback = callback
         self.__keepalive_last = time.time()
 
+    def read_timer(self):
+        self.__timer_expired = True
+
+    def start_handshake(self, timeout):
+        """
+        Tells `Packetizer` that the handshake process started.
+        Starts a book keeping timer that can signal a timeout in the
+        handshake process.
+
+        :param float timeout: amount of seconds to wait before timing out
+        """
+        if not self.__timer:
+            self.__timer = threading.Timer(float(timeout), self.read_timer)
+            self.__timer.start()
+
+    def handshake_timed_out(self):
+        """
+        Checks if the handshake has timed out.
+        If `start_handshake` wasn't called before the call to this function
+        the return value will always be `False`.
+        If the handshake completed before a time out was reached the return value will be `False`
+
+        :return: handshake time out status, as a `bool`
+        """
+        if not self.__timer:
+            return False
+        if self.__handshake_complete:
+            return False
+        return self.__timer_expired
+
+    def complete_handshake(self):
+        """
+        Tells `Packetizer` that the handshake has completed.
+        """
+        if self.__timer:
+            self.__timer.cancel()
+            self.__timer_expired = False
+            self.__handshake_complete = True
+
     def read_all(self, n, check_rekey=False):
         """
         Read as close to N bytes as possible, blocking as long as necessary.
@@ -200,6 +243,8 @@ class Packetizer (object):
             n -= len(out)
         while n > 0:
             got_timeout = False
+            if self.handshake_timed_out():
+                raise EOFError()
             try:
                 x = self.__socket.recv(n)
                 if len(x) == 0: