diff --git a/test/functional/test_framework/p2p.py b/test/functional/test_framework/p2p.py index e33f6e91e71462..11c65fd0aa345c 100755 --- a/test/functional/test_framework/p2p.py +++ b/test/functional/test_framework/p2p.py @@ -303,6 +303,9 @@ def _on_data(self): msg_len, msg = self.v2_conn.v2_dec_msg(self.recvbuf) if msg is None: return + self.recvbuf = self.recvbuf[3+msg_len+16:] + if msg == b'': # reject decoy messages and messages longer than permitted payload length + return size_or_shortid = msg[0] if 0 < size_or_shortid <= 12: # TODO: check if this works # a string command @@ -316,7 +319,6 @@ def _on_data(self): else: msgtype = "unknown-" + str(msg[1]) # TODO: check if this works f = BytesIO(payload) - self.recvbuf = self.recvbuf[3+msg_len+16:] # TODO: Remove later, depending on disconnection logic. # not using an "else" here so that if v2 doesn't work out, normal v1 handling can be done. @@ -356,12 +358,12 @@ def on_message(self, message): # Socket write methods - def send_message(self, message): + def send_message(self, message, is_decoy=False): """Send a P2P message over the socket. This method takes a P2P payload, builds the P2P header and adds the message to the send buffer to be sent over the socket.""" - tmsg = self.build_message(message) + tmsg = self.build_message(message, is_decoy) self._log_message("send", message) return self.send_raw_message(tmsg) @@ -379,7 +381,7 @@ def maybe_write(): # Class utility methods - def build_message(self, message): + def build_message(self, message, is_decoy=False): """Build a serialized P2P message""" if self.support_v2_p2p: msgtype = message.msgtype @@ -391,7 +393,7 @@ def build_message(self, message): tmsg += bytes(msgtype, 'ascii') tmsg += data # Need to return (3 bytes length + encrypted payload + MAC tag) - return self.v2_conn.v2_enc_msg(tmsg) + return self.v2_conn.v2_enc_msg(tmsg, is_decoy) else: msgtype = message.msgtype data = message.serialize()