[test] Build v2 P2P messages

This commit is contained in:
stratospher 2022-02-05 23:28:02 +05:30
parent bb7bffed79
commit a94e350ac0

View file

@ -82,6 +82,7 @@ from test_framework.util import (
) )
from test_framework.v2_p2p import ( from test_framework.v2_p2p import (
EncryptedP2PState, EncryptedP2PState,
MSGTYPE_TO_SHORTID,
SHORTID, SHORTID,
) )
@ -386,15 +387,25 @@ class P2PConnection(asyncio.Protocol):
"""Build a serialized P2P message""" """Build a serialized P2P message"""
msgtype = message.msgtype msgtype = message.msgtype
data = message.serialize() data = message.serialize()
tmsg = self.magic_bytes if self.supports_v2_p2p:
tmsg += msgtype if msgtype in SHORTID.values():
tmsg += b"\x00" * (12 - len(msgtype)) tmsg = MSGTYPE_TO_SHORTID.get(msgtype).to_bytes(1, 'big')
tmsg += struct.pack("<I", len(data)) else:
th = sha256(data) tmsg = b"\x00"
h = sha256(th) tmsg += msgtype
tmsg += h[:4] tmsg += b"\x00" * (12 - len(msgtype))
tmsg += data tmsg += data
return tmsg return self.v2_state.v2_enc_packet(tmsg)
else:
tmsg = self.magic_bytes
tmsg += msgtype
tmsg += b"\x00" * (12 - len(msgtype))
tmsg += struct.pack("<I", len(data))
th = sha256(data)
h = sha256(th)
tmsg += h[:4]
tmsg += data
return tmsg
def _log_message(self, direction, msg): def _log_message(self, direction, msg):
"""Logs a message being sent or received over the connection.""" """Logs a message being sent or received over the connection."""