def test_normal_frame(self):
cipher1, cipher2 = self.setup_ciphers()
dummy_msg_type = 1
dummy_payload = helpers.generate_bytearray(123)
dummy_protocol = 0
frames = frame_utils.get_frames(
dummy_msg_type, memoryview(dummy_payload), dummy_protocol,
eth_common_constants.DEFAULT_FRAME_SIZE)
input_buffer = InputBuffer()
for frame in frames:
encrypted_frame = cipher1.encrypt_frame(frame)
encrypted_frame_bytes = bytearray(
rlp_utils.str_to_bytes(encrypted_frame))
input_buffer.add_bytes(encrypted_frame_bytes)
# adding some dummy bytes but less than header size len
dummy_bytes_len = int(eth_common_constants.FRAME_HDR_TOTAL_LEN / 2)
input_buffer.add_bytes(helpers.generate_bytearray(dummy_bytes_len))
framed_input_buffer = FramedInputBuffer(cipher2)
is_full, msg_type = framed_input_buffer.peek_message(input_buffer)
self.assertTrue(is_full)
self.assertEqual(msg_type, dummy_msg_type)
message, full_msg_type = framed_input_buffer.get_full_message()
self.assertEqual(message, dummy_payload)
self.assertEqual(full_msg_type, dummy_msg_type)
self.assertEqual(input_buffer.length, dummy_bytes_len)
def __init__(self, rlpx_cipher):
if not isinstance(rlpx_cipher, RLPxCipher):
raise TypeError(
"Argument rlpx_cipher is expected to be of type RLPxCipher but was {}"
.format(type(rlpx_cipher)))
super(EthProtocolMessageFactory, self).__init__()
self.message_type_mapping = self._MESSAGE_TYPE_MAPPING
self._framed_input_buffer = FramedInputBuffer(rlpx_cipher)
self._expected_msg_type = None
def test_chunked_frames(self):
cipher1, cipher2 = self.setup_ciphers()
dummy_msg_type = 10
expected_frames_count = 10
dummy_payload = helpers.generate_bytearray(self.TEST_FRAME_SIZE *
(expected_frames_count - 1))
dummy_protocol = 0
frames = frame_utils.get_frames(dummy_msg_type,
memoryview(dummy_payload),
dummy_protocol, self.TEST_FRAME_SIZE)
self.assertEqual(len(frames), expected_frames_count)
frames_bytes = bytearray(0)
for frame in frames:
encrypted_frame = cipher1.encrypt_frame(frame)
encrypted_frame_bytes = bytearray(
rlp_utils.str_to_bytes(encrypted_frame))
frames_bytes += encrypted_frame_bytes
# adding some dummy bytes but less than header size len
dummy_bytes_len = int(eth_common_constants.FRAME_HDR_TOTAL_LEN / 2)
frames_bytes += helpers.generate_bytearray(dummy_bytes_len)
input_buffer = InputBuffer()
framed_input_buffer = FramedInputBuffer(cipher2)
read_start = 0
read_size = eth_common_constants.FRAME_HDR_TOTAL_LEN
is_full = False
msg_type = None
while not is_full and read_start < len(frames_bytes):
input_buffer.add_bytes(frames_bytes[read_start:read_start +
read_size])
is_full, msg_type = framed_input_buffer.peek_message(input_buffer)
read_start += read_size
self.assertTrue(is_full)
self.assertEqual(msg_type, dummy_msg_type)
class EthProtocolMessageFactory(AbstractMessageFactory):
_MESSAGE_TYPE_MAPPING = {
EthProtocolMessageType.HELLO: HelloEthProtocolMessage,
EthProtocolMessageType.DISCONNECT: DisconnectEthProtocolMessage,
EthProtocolMessageType.PING: PingEthProtocolMessage,
EthProtocolMessageType.PONG: PongEthProtocolMessage,
EthProtocolMessageType.STATUS: StatusEthProtocolMessage,
EthProtocolMessageType.TRANSACTIONS: TransactionsEthProtocolMessage,
EthProtocolMessageType.NEW_BLOCK_HASHES:
NewBlockHashesEthProtocolMessage,
EthProtocolMessageType.GET_BLOCK_HEADERS:
GetBlockHeadersEthProtocolMessage,
EthProtocolMessageType.BLOCK_HEADERS: BlockHeadersEthProtocolMessage,
EthProtocolMessageType.GET_BLOCK_BODIES:
GetBlockBodiesEthProtocolMessage,
EthProtocolMessageType.BLOCK_BODIES: BlockBodiesEthProtocolMessage,
EthProtocolMessageType.NEW_BLOCK: NewBlockEthProtocolMessage,
EthProtocolMessageType.GET_NODE_DATA: GetNodeDataEthProtocolMessage,
EthProtocolMessageType.NODE_DATA: NodeDataEthProtocolMessage,
EthProtocolMessageType.GET_RECEIPTS: GetReceiptsEthProtocolMessage,
EthProtocolMessageType.RECEIPTS: ReceiptsEthProtocolMessage,
}
def __init__(self, rlpx_cipher):
if not isinstance(rlpx_cipher, RLPxCipher):
raise TypeError(
"Argument rlpx_cipher is expected to be of type RLPxCipher but was {}"
.format(type(rlpx_cipher)))
super(EthProtocolMessageFactory, self).__init__()
self.message_type_mapping = self._MESSAGE_TYPE_MAPPING
self._framed_input_buffer = FramedInputBuffer(rlpx_cipher)
self._expected_msg_type = None
def get_base_message_type(self) -> Type[AbstractMessage]:
return EthProtocolMessage
def set_expected_msg_type(self, msg_type):
if not msg_type in [
EthProtocolMessageType.AUTH, EthProtocolMessageType.AUTH_ACK
]:
raise ValueError("msg_type can be AUTH or AUTH_ACK")
self._expected_msg_type = msg_type
def reset_expected_msg_type(self):
self._expected_msg_type = None
def get_message_header_preview_from_input_buffer(
self, input_buffer: InputBuffer) -> MessagePreview:
"""
Peeks at a message, determining if its full.
Returns (is_full_message, command, payload_length)
"""
if self._expected_msg_type == EthProtocolMessageType.AUTH:
return MessagePreview(True, EthProtocolMessageType.AUTH,
input_buffer.length)
elif self._expected_msg_type == EthProtocolMessageType.AUTH_ACK and \
input_buffer.length >= eth_common_constants.ENC_AUTH_ACK_MSG_LEN:
return MessagePreview(True, EthProtocolMessageType.AUTH_ACK,
eth_common_constants.ENC_AUTH_ACK_MSG_LEN)
elif self._expected_msg_type is None:
decryption_start_time = time.time()
is_full_msg, command = self._framed_input_buffer.peek_message(
input_buffer)
eth_gateway_stats_service.log_decrypted_message(
time.time() - decryption_start_time)
if is_full_msg:
return MessagePreview(True, command, 0)
return MessagePreview(False, None, None)
def create_message_from_buffer(self, buf):
"""
Parses a full message from a buffer based on its command into one of the loaded message types.
"""
if self._expected_msg_type is not None:
return self.base_message_type.initialize_class(
RawEthProtocolMessage, buf, None)
if len(buf) != 0:
raise ParseError(
"All bytes are expected to be already read by self._framed_input_buffer"
)
message_bytes, command = self._framed_input_buffer.get_full_message()
return self.create_message(command, message_bytes)