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.
This commit is contained in:
Tobias Brunner 2016-05-02 15:14:40 +02:00
parent 26aa9c638d
commit 756cd28fb8
1 changed files with 9 additions and 2 deletions

View File

@ -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