[project @ robey@master-shake.local-20060120182320-d569b04adc2bd622]

some performance improvements: be a LOT less aggressive about stirring the randpool; use buffering when reading the banner; add a hook for using a native-compiled hmac (which gives the biggest boost, but should probably be done in pycrypto)

1 files changed, 34 insertions, 14 deletions

diff --git a/paramiko/packet.py b/paramiko/packet.py
index 7e9afc78..181f3eb5 100644
--- a/paramiko/packet.py
+++ b/paramiko/packet.py
@@ -25,7 +25,6 @@ import socket
 import struct
 import threading
 import time
-from Crypto.Hash import HMAC
 
 from paramiko.common import *
 from paramiko import util
@@ -33,6 +32,19 @@ from paramiko.ssh_exception import SSHException
 from paramiko.message import Message
 
 
+got_r_hmac = False
+try:
+    import r_hmac
+    got_r_hmac = True
+except ImportError:
+    pass
+def compute_hmac(key, message, digest_class):
+    if got_r_hmac:
+        return r_hmac.HMAC(key, message, digest_class).digest()
+    from Crypto.Hash import HMAC
+    return HMAC.HMAC(key, message, digest_class).digest()
+
+
 class NeedRekeyException (Exception):
     pass
 
@@ -54,6 +66,7 @@ class Packetizer (object):
         self.__dump_packets = False
         self.__need_rekey = False
         self.__init_count = 0
+        self.__remainder = ''
         
         # used for noticing when to re-key:
         self.__sent_bytes = 0
@@ -186,9 +199,14 @@ class Packetizer (object):
         @raise EOFError: if the socket was closed before all the bytes could
             be read
         """
-        if PY22:
-            return self._py22_read_all(n)
         out = ''
+        # handle over-reading from reading the banner line
+        if len(self.__remainder) > 0:
+            out = self.__remainder[:n]
+            self.__remainder = self.__remainder[n:]
+            n -= len(out)
+        if PY22:
+            return self._py22_read_all(n, out)
         while n > 0:
             try:
                 x = self.__socket.recv(n)
@@ -225,14 +243,15 @@ class Packetizer (object):
         
     def readline(self, timeout):
         """
-        Read a line from the socket.  This is done in a fairly inefficient
-        way, but is only used for initial banner negotiation so it's not worth
-        optimising.
+        Read a line from the socket.  We assume no data is pending after the
+        line, so it's okay to attempt large reads.
         """
         buf = ''
         while not '\n' in buf:
             buf += self._read_timeout(timeout)
-        buf = buf[:-1]
+        n = buf.index('\n')
+        self.__remainder += buf[n+1:]
+        buf = buf[:n]
         if (len(buf) > 0) and (buf[-1] == '\r'):
             buf = buf[:-1]
         return buf
@@ -242,7 +261,6 @@ class Packetizer (object):
         Write a block of data using the current cipher, as an SSH block.
         """
         # encrypt this sucka
-        randpool.stir()
         data = str(data)
         cmd = ord(data[0])
         if cmd in MSG_NAMES:
@@ -264,12 +282,15 @@ class Packetizer (object):
             # + mac
             if self.__block_engine_out != None:
                 payload = struct.pack('>I', self.__sequence_number_out) + packet
-                out += HMAC.HMAC(self.__mac_key_out, payload, self.__mac_engine_out).digest()[:self.__mac_size_out]
+                out += compute_hmac(self.__mac_key_out, payload, self.__mac_engine_out)[:self.__mac_size_out]
             self.__sequence_number_out = (self.__sequence_number_out + 1) & 0xffffffffL
             self.write_all(out)
 
             self.__sent_bytes += len(out)
             self.__sent_packets += 1
+            if (self.__sent_packets % 100) == 0:
+                # stirring the randpool takes 30ms on my ibook!!
+                randpool.stir()
             if ((self.__sent_packets >= self.REKEY_PACKETS) or (self.__sent_bytes >= self.REKEY_BYTES)) \
                    and not self.__need_rekey:
                 # only ask once for rekeying
@@ -310,12 +331,12 @@ class Packetizer (object):
         if self.__mac_size_in > 0:
             mac = post_packet[:self.__mac_size_in]
             mac_payload = struct.pack('>II', self.__sequence_number_in, packet_size) + packet
-            my_mac = HMAC.HMAC(self.__mac_key_in, mac_payload, self.__mac_engine_in).digest()[:self.__mac_size_in]
+            my_mac = compute_hmac(self.__mac_key_in, mac_payload, self.__mac_engine_in)[:self.__mac_size_in]
             if my_mac != mac:
                 raise SSHException('Mismatched MAC')
         padding = ord(packet[0])
         payload = packet[1:packet_size - padding]
-        randpool.add_event(packet[packet_size - padding])
+        randpool.add_event()
         if self.__dump_packets:
             self._log(DEBUG, 'Got payload (%d bytes, %d padding)' % (packet_size, padding))
 
@@ -374,8 +395,7 @@ class Packetizer (object):
             self.__keepalive_callback()
             self.__keepalive_last = now
     
-    def _py22_read_all(self, n):
-        out = ''
+    def _py22_read_all(self, n, out):
         while n > 0:
             r, w, e = select.select([self.__socket], [], [], 0.1)
             if self.__socket not in r:
@@ -412,7 +432,7 @@ class Packetizer (object):
         start = time.time()
         while True:
             try:
-                x = self.__socket.recv(1)
+                x = self.__socket.recv(128)
                 if len(x) == 0:
                     raise EOFError()
                 break