def send(self, message: MessageAPI) -> None: header = message.header.ljust(roundup_16(len(message.header)), b'\x00') body = message.body.ljust(roundup_16(len(message.body)), b'\x00') if self.is_closing: raise PeerConnectionLost( f"Attempted to send msg with cmd id {message.command_id} to " f"disconnected peer {self.remote}") self.write(self._encrypt(header, body))
async def recv(self, token: CancelToken) -> bytes: header_data = await self.read(HEADER_LEN + MAC_LEN, token) try: header = self._decrypt_header(header_data) except DecryptionError as err: self.logger.debug( "Bad message header from peer %s: Error: %r", self, err, ) raise MalformedMessage from err frame_size = self._get_frame_size(header) # The frame_size specified in the header does not include the padding to 16-byte boundary, # so need to do this here to ensure we read all the frame's data. read_size = roundup_16(frame_size) frame_data = await self.read(read_size + MAC_LEN, token) try: msg = self._decrypt_body(frame_data, frame_size) except DecryptionError as err: self.logger.debug( "Bad message body from peer %s: Error: %r", self, err, ) raise MalformedMessage from err return msg
async def recv(self) -> MessageAPI: header_bytes = await self.read(HEADER_LEN + MAC_LEN) try: padded_header = self._decrypt_header(header_bytes) except (ValueError, DecryptionError) as err: self.logger.info("Bad message header from peer %s: Error: %r", self, err) raise MalformedMessage(*err.args) from err # TODO: use `int.from_bytes(...)` frame_size = self._get_frame_size(padded_header) # The frame_size specified in the header does not include the padding to 16-byte boundary, # so need to do this here to ensure we read all the frame's data. read_size = roundup_16(frame_size) frame_data = await self.read(read_size + MAC_LEN) try: body = self._decrypt_body(frame_data, frame_size) except (ValueError, DecryptionError) as err: self.logger.info("Bad message body from peer %s: Error: %r", self, err) raise MalformedMessage(*err.args) from err # Decode the header data and re-encode to recover the unpadded header size. try: header_data = _decode_header_data(padded_header[3:]) except rlp.exceptions.DeserializationError as err: raise MalformedMessage(*err.args) from err header = padded_header[:3] + rlp.encode(header_data) return Message(header, body)
async def recv(self) -> MessageAPI: # Check that Transport read state is IDLE. if self.read_state is not TransportState.IDLE: # This is logged at INFO level because it indicates we are not # properly managing the Transport and are interrupting it mid-read # somewhere. self.logger.info( 'Corrupted transport: %s - state=%s', self, self.read_state.name, ) raise CorruptTransport(f"Corrupted transport: {self} - state={self.read_state.name}") # Set status to indicate we are waiting to read the message header self.read_state = TransportState.HEADER try: header_bytes = await self.read(HEADER_LEN + MAC_LEN) except asyncio.CancelledError: self.logger.debug('Transport cancelled during header read. resetting to IDLE state') self.read_state = TransportState.IDLE raise # Set status to indicate we are waiting to read the message body self.read_state = TransportState.BODY try: padded_header = self._decrypt_header(header_bytes) except DecryptionError as err: self.logger.debug( "Bad message header from peer %s: Error: %r", self, err, ) raise MalformedMessage from err # TODO: use `int.from_bytes(...)` frame_size = self._get_frame_size(padded_header) # The frame_size specified in the header does not include the padding to 16-byte boundary, # so need to do this here to ensure we read all the frame's data. read_size = roundup_16(frame_size) frame_data = await self.read(read_size + MAC_LEN) try: body = self._decrypt_body(frame_data, frame_size) except DecryptionError as err: self.logger.debug( "Bad message body from peer %s: Error: %r", self, err, ) raise MalformedMessage from err # Reset status back to IDLE self.read_state = TransportState.IDLE # Decode the header data and re-encode to recover the unpadded header size. try: header_data = _decode_header_data(padded_header[3:]) except rlp.exceptions.DeserializationError as err: raise MalformedMessage from err header = padded_header[:3] + rlp.encode(header_data) return Message(header, body)
async def recv(self, token: CancelToken) -> bytes: # Check that Transport read state is IDLE. if self.read_state is not TransportState.IDLE: # This is logged at INFO level because it indicates we are not # properly managing the Transport and are interrupting it mid-read # somewhere. self.logger.info( 'Corrupted transport: %s - state=%s', self, self.read_state.name, ) raise Exception( f"Corrupted transport: {self} - state={self.read_state.name}") # Set status to indicate we are waiting to read the message header self.read_state = TransportState.HEADER try: header_data = await self.read(HEADER_LEN + MAC_LEN, token) except asyncio.CancelledError: self.logger.debug( 'Transport cancelled during header read. resetting to IDLE state' ) self.read_state = TransportState.IDLE raise # Set status to indicate we are waiting to read the message body self.read_state = TransportState.BODY try: header = self._decrypt_header(header_data) except DecryptionError as err: self.logger.debug( "Bad message header from peer %s: Error: %r", self, err, ) raise MalformedMessage from err frame_size = self._get_frame_size(header) # The frame_size specified in the header does not include the padding to 16-byte boundary, # so need to do this here to ensure we read all the frame's data. read_size = roundup_16(frame_size) frame_data = await self.read(read_size + MAC_LEN, token) try: msg = self._decrypt_body(frame_data, frame_size) except DecryptionError as err: self.logger.debug( "Bad message body from peer %s: Error: %r", self, err, ) raise MalformedMessage from err # Reset status back to IDLE self.read_state = TransportState.IDLE return msg
async def read_msg(self) -> Tuple[protocol.Command, protocol.PayloadType]: header_data = await self.read(HEADER_LEN + MAC_LEN) try: header = self.decrypt_header(header_data) except DecryptionError as err: self.logger.debug( "Bad message header from peer %s: Error: %r", self, err, ) raise MalformedMessage from err frame_size = self.get_frame_size(header) # The frame_size specified in the header does not include the padding to 16-byte boundary, # so need to do this here to ensure we read all the frame's data. read_size = roundup_16(frame_size) frame_data = await self.read(read_size + MAC_LEN) try: msg = self.decrypt_body(frame_data, frame_size) except DecryptionError as err: self.logger.debug( "Bad message body from peer %s: Error: %r", self, err, ) raise MalformedMessage from err cmd = self.get_protocol_command_for(msg) # NOTE: This used to be a bottleneck but it doesn't seem to be so anymore. If we notice # too much time is being spent on this again, we need to consider running this in a # ProcessPoolExecutor(). Need to make sure we don't use all CPUs in the machine for that, # though, otherwise asyncio's event loop can't run and we can't keep up with other peers. try: decoded_msg = cast(Dict[str, Any], cmd.decode(msg)) except MalformedMessage as err: self.logger.debug( "Malformed message from peer %s: CMD:%s Error: %r", self, type(cmd).__name__, err, ) raise else: self.logger.debug2("Successfully decoded %s msg: %s", cmd, decoded_msg) self.received_msgs[cmd] += 1 return cmd, decoded_msg
def _decrypt_body(self, data: bytes, body_size: int) -> bytes: read_size = roundup_16(body_size) if len(data) < read_size + MAC_LEN: raise ValueError( f'Insufficient body length; Got {len(data)}, wanted {read_size} + {MAC_LEN}' ) frame_ciphertext = data[:read_size] frame_mac = data[read_size:read_size + MAC_LEN] self._ingress_mac.update(frame_ciphertext) fmac_seed = self._ingress_mac.digest()[:MAC_LEN] self._ingress_mac.update(sxor(self._mac_enc(fmac_seed), fmac_seed)) expected_frame_mac = self._ingress_mac.digest()[:MAC_LEN] if not hmac.compare_digest(expected_frame_mac, frame_mac): raise DecryptionError( f'Invalid frame mac: expected {expected_frame_mac}, got {frame_mac}' ) return self._aes_dec.update(frame_ciphertext)[:body_size]