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)