Bitmessage
/
MiNode
3
0
Fork 1
Code Pull Requests 2 Actions Activity

Improve malformed message handling

This commit is contained in:
TheKysek 2016-10-15 17:00:18 +02:00
parent 90a9cee834
commit c525d4f2e1
2 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/connection.py
View File

@ -197,11 +197,20 @@ class Connection(threading.Thread):
while len(self.buffer_receive) >= self.next_message_size:
if self.next_header:
self.next_header = False
h = message.Header.from_bytes(self.buffer_receive[:shared.header_length])
try:
h = message.Header.from_bytes(self.buffer_receive[:shared.header_length])
except ValueError as e:
self.status = 'disconnecting'
logging.warning('Received malformed message from {}:{}: {}'.format(self.host, self.port, e))
break
self.next_message_size += h.payload_length
else:
m = message.Message.from_bytes(self.buffer_receive[:self.next_message_size])
try:
m = message.Message.from_bytes(self.buffer_receive[:self.next_message_size])
except ValueError as e:
self.status = 'disconnecting'
logging.warning('Received malformed message from {}:{}, {}'.format(self.host, self.port, e))
break
self.next_header = True
self.buffer_receive = self.buffer_receive[self.next_message_size:]
self.next_message_size = shared.header_length

6
src/message.py
View File

@ -31,7 +31,7 @@ class Header(object):
magic_bytes, command, payload_length, payload_checksum = struct.unpack('>4s12sL4s', b)
if magic_bytes != shared.magic_bytes:
raise IOError('magic_bytes do not match')
raise ValueError('magic_bytes do not match')
command = command.rstrip(b'\x00')
@ -63,12 +63,12 @@ class Message(object):
payload_length = len(payload)
if payload_length != h.payload_length:
raise Exception('wrong payload length, expected {}, got {}'.format(h.payload_length, payload_length))
raise ValueError('wrong payload length, expected {}, got {}'.format(h.payload_length, payload_length))
payload_checksum = hashlib.sha512(payload).digest()[:4]
if payload_checksum != h.payload_checksum:
raise Exception('wrong payload checksum, expected {}, got {}'.format(h.payload_checksum, payload_checksum))
raise ValueError('wrong payload checksum, expected {}, got {}'.format(h.payload_checksum, payload_checksum))
return cls(h.command, payload)