def send(self, header: bytes, body: bytes) -> None:
cmd_id = get_devp2p_cmd_id(body)
self.logger.debug2("Sending msg with cmd id %d to %s", cmd_id, self)
if self.is_closing:
raise PeerConnectionLost(
"Attempted to send msg with cmd id %d to disconnected peer %s",
cmd_id, self)
self.write(self._encrypt(header, body))
def decode(self, data: bytes) -> Payload:
packet_type = get_devp2p_cmd_id(data)
if packet_type != self.cmd_id:
raise MalformedMessage(f"Wrong packet type: {packet_type}, expected {self.cmd_id}")
compressed_payload = data[1:]
encoded_payload = self.decompress_payload(compressed_payload)
return self.decode_payload(encoded_payload)
def send(self, header: bytes, body: bytes) -> None:
(size, ) = struct.unpack(b'>I', b'\x00' + header[:3])
if self.is_closing:
cmd_id = get_devp2p_cmd_id(body)
self.logger.error(
"Attempted to send msg with cmd id %d to disconnected peer %s",
cmd_id, self)
return
self.write(header[:3] + body[:size])
def get_protocol_command_for(self, msg: bytes) -> Command:
"""Return the Command corresponding to the cmd_id encoded in the given msg."""
cmd_id = get_devp2p_cmd_id(msg)
self.logger.debug2("Got msg with cmd_id: %s", cmd_id)
if cmd_id < self.base_protocol.cmd_length:
return self.base_protocol.cmd_by_id[cmd_id]
elif cmd_id < self.sub_proto.cmd_id_offset + self.sub_proto.cmd_length:
return self.sub_proto.cmd_by_id[cmd_id]
else:
raise UnknownProtocolCommand(f"No protocol found for cmd_id {cmd_id}")
async def stream_transport_messages(transport: TransportAPI,
base_protocol: BaseP2PProtocol,
*protocols: ProtocolAPI,
token: CancelToken = None,
) -> AsyncIterator[Tuple[ProtocolAPI, CommandAPI, Payload]]:
"""
Streams 3-tuples of (Protocol, Command, Payload) over the provided `Transport`
"""
# A cache for looking up the proper protocol instance for a given command
# id.
cmd_id_cache: Dict[int, ProtocolAPI] = {}
while not transport.is_closing:
raw_msg = await transport.recv(token)
cmd_id = get_devp2p_cmd_id(raw_msg)
if cmd_id not in cmd_id_cache:
if cmd_id < base_protocol.cmd_length:
cmd_id_cache[cmd_id] = base_protocol
else:
for protocol in protocols:
if cmd_id < protocol.cmd_id_offset + protocol.cmd_length:
cmd_id_cache[cmd_id] = protocol
break
else:
protocol_infos = ' '.join(tuple(
f"{proto.name}@{proto.version}[offset={proto.cmd_id_offset},cmd_length={proto.cmd_length}]" # noqa: E501
for proto in cons(base_protocol, protocols)
))
raise UnknownProtocolCommand(
f"No protocol found for cmd_id {cmd_id}: Available "
f"protocol/offsets are: {protocol_infos}"
)
msg_proto = cmd_id_cache[cmd_id]
cmd = msg_proto.cmd_by_id[cmd_id]
msg = cmd.decode(raw_msg)
yield msg_proto, cmd, msg
# yield to the event loop for a moment to allow `transport.is_closing`
# a chance to update.
await asyncio.sleep(0)
