From 756cd28fb8b8ca16dca2897aef9190a4d0893794 Mon Sep 17 00:00:00 2001 From: Tobias Brunner Date: Mon, 2 May 2016 15:14:40 +0200 Subject: vici: Ensure we read exactly the specified amount of bytes from the socket in Python recv() will return less bytes than specified (as that's the buffer size) if not as many are ready to be read from the socket. --- src/libcharon/plugins/vici/python/vici/protocol.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) (limited to 'src/libcharon/plugins/vici/python') diff --git a/src/libcharon/plugins/vici/python/vici/protocol.py b/src/libcharon/plugins/vici/python/vici/protocol.py index 855a7b2e2..4951817eb 100644 --- a/src/libcharon/plugins/vici/python/vici/protocol.py +++ b/src/libcharon/plugins/vici/python/vici/protocol.py @@ -20,15 +20,22 @@ class Transport(object): self.socket.sendall(struct.pack("!I", len(packet)) + packet) def receive(self): - raw_length = self.socket.recv(self.HEADER_LENGTH) + raw_length = self._recvall(self.HEADER_LENGTH) length, = struct.unpack("!I", raw_length) - payload = self.socket.recv(length) + payload = self._recvall(length) return payload def close(self): self.socket.shutdown(socket.SHUT_RDWR) self.socket.close() + def _recvall(self, count): + """Ensure to read count bytes from the socket""" + data = b"" + while len(data) < count: + data += self.socket.recv(count - len(data)) + return data + class Packet(object): CMD_REQUEST = 0 # Named request message -- cgit v1.2.3