From 63e55c4260ed3837885ba542d01df733492d8036 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tarjei=20Hus=C3=B8y?= Date: Thu, 2 Jun 2016 08:57:34 -0700 Subject: [PATCH] Remove unread bytes shenanigans Read directly into buffer instead, should not read more than wait we're waiting for. --- postgres_mitm.py | 30 +++++++++++------------------- 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/postgres_mitm.py b/postgres_mitm.py index 6504f1d..2d016eb 100644 --- a/postgres_mitm.py +++ b/postgres_mitm.py @@ -133,8 +133,6 @@ class ClientConnection(threading.Thread): self.target_backend = target_backend self.server_socket = None self._stop = threading.Event() - self.unread_client_data = '' - self.buffer_size = 4096 def stop(self): @@ -236,25 +234,9 @@ class ClientConnection(threading.Thread): def read_n_bytes_from_client(self, n): - chunks = [] - bytes_read = 0 + return read_n_bytes_from_socket(self.socket, n) - if self.unread_client_data: - bytes_read += len(self.unread_client_data) - _logger.debug('Already had %d bytes from client', bytes_read) - chunks.append(self.unread_client_data) - self.unread_client_data = '' - while bytes_read < n: - chunk = self.socket.recv(self.buffer_size) - if not chunk: - raise Exception('Not enough data read') - chunks.append(chunk) - bytes_read += len(chunk) - - received_so_far = ''.join(chunks) - self.unread_client_data = received_so_far[n:] - return received_so_far[:n] def send_to_client(self, msg): @@ -362,6 +344,16 @@ def socket_is_closed(sock): return isinstance(sock._sock, socket._closedsocket) +def read_n_bytes_from_socket(sock, n): + buf = bytearray(n) + view = memoryview(buf) + while n: + nbytes = sock.recv_into(view, n) + view = view[nbytes:] # slicing views is cheap + n -= nbytes + return str(buf) + + def create_md5_auth_packet(username, password, salt): pw_and_username = password + username pw_hash = hashlib.md5(pw_and_username).hexdigest()