def get_connection_protocol_version(self,
input_buffer: InputBuffer) -> int:
if (input_buffer.length <
(constants.STARTING_SEQUENCE_BYTES_LEN +
constants.BX_HDR_COMMON_OFF + constants.VERSION_NUM_LEN)):
return self.CURRENT_PROTOCOL_VERSION
header_buf = input_buffer.peek_message(VersionMessage.HEADER_LENGTH)
if header_buf[:constants.
STARTING_SEQUENCE_BYTES_LEN] == constants.STARTING_SEQUENCE_BYTES:
command, payload_len = VersionMessage.unpack(header_buf)
header_len = VersionMessage.HEADER_LENGTH
else:
command = bytearray(header_buf[:constants.MSG_TYPE_LEN])
payload_len = constants.MSG_TYPE_LEN
header_len = constants.BX_HDR_COMMON_OFF
if command != self.version_message_command:
if constants.HTTP_MESSAGE in command:
raise NonVersionMessageError(
msg=
"Instead of a version hello message, we received an HTTP request: {} with payload length: {} "
"Ignoring and closing connection. ".format(
header_buf, payload_len),
is_known=True)
if command in constants.BITCOIN_MESSAGES:
raise NonVersionMessageError(
msg="Received some kind of bitcoin peering message: {}. "
"Ignoring and closing connection.".format(command),
is_known=True)
logger.debug(
"Received message of type {} instead of hello message. Use current version of protocol {}.",
command, self.CURRENT_PROTOCOL_VERSION)
return self.CURRENT_PROTOCOL_VERSION
if payload_len < self.VERSION_MESSAGE_MAIN_LENGTH:
return 1
version_buf = input_buffer.get_slice(header_len,
header_len + VERSION_NUM_LEN)
version, = struct.unpack_from("<L", version_buf, 0)
return version
def get_broadcast_message_preview(
self, input_buffer: InputBuffer) -> BroadcastMessagePreview:
"""
Peeks the hash and network number from hashed messages.
Currently, only Broadcast messages are supported here.
:param input_buffer
:return: is full header, message hash, network number, source id, payload length
"""
# -1 for control flag length
broadcast_header_length = self.base_message_type.HEADER_LENGTH + AbstractBroadcastMessage.PAYLOAD_LENGTH - \
constants.CONTROL_FLAGS_LEN
is_full_header = input_buffer.length >= broadcast_header_length
if not is_full_header:
return BroadcastMessagePreview(False, None, None, None, None)
else:
_is_full_message, _command, payload_length = self.get_message_header_preview_from_input_buffer(
input_buffer)
broadcast_header = input_buffer.peek_message(
broadcast_header_length)
offset = self.base_message_type.HEADER_LENGTH
message_hash = broadcast_header[offset:offset +
crypto.SHA256_HASH_LEN]
offset += crypto.SHA256_HASH_LEN
network_num, = struct.unpack_from(
"<L",
broadcast_header[offset:offset + constants.NETWORK_NUM_LEN])
offset += constants.NETWORK_NUM_LEN
source_id = uuid_pack.from_bytes(
struct.unpack_from(
"<16s",
broadcast_header[offset:offset +
constants.NODE_ID_SIZE_IN_BYTES])[0])
return BroadcastMessagePreview(is_full_header,
Sha256Hash(message_hash),
network_num, source_id,
payload_length)
def get_hashed_message_preview_from_input_buffer(self, input_buffer: InputBuffer) \
-> Tuple[bool, Optional[Sha256Hash], Optional[int], Optional[int]]:
"""
Peeks the hash and network number from hashed messages.
Currently, only Broadcast messages are supported here.
:param input_buffer
:return: is full header, message hash, network number, payload length
"""
hash_header_length = self.base_message_type.HEADER_LENGTH + crypto.SHA256_HASH_LEN + constants.NETWORK_NUM_LEN
_is_full_message, _command, payload_length = self.get_message_header_preview_from_input_buffer(input_buffer)
is_full_header = input_buffer.length >= hash_header_length
if payload_length is None or not is_full_header:
return False, None, None, None
else:
hash_header = input_buffer.peek_message(hash_header_length)
offset = self.base_message_type.HEADER_LENGTH
message_hash = hash_header[offset:offset + crypto.SHA256_HASH_LEN]
offset += crypto.SHA256_HASH_LEN
network_num, = struct.unpack_from("<L", hash_header[offset:offset + constants.NETWORK_NUM_LEN])
return is_full_header, Sha256Hash(message_hash), network_num, payload_length
class TestInputBuffer(unittest.TestCase):
def setUp(self):
self.in_buf = InputBuffer()
self.length_to_add = 20
self.data1 = bytearray([i for i in range(1, 21)])
self.data2 = bytearray([i for i in range(21, 41)])
self.data3 = bytearray([i for i in range(41, 61)])
def test_endswith(self):
self.assertFalse(self.in_buf.endswith(5))
self.in_buf.add_bytes(bytearray([1, 2, 3, 4, 5, 6]))
self.assertTrue(self.in_buf.endswith(bytearray([5, 6])))
self.assertFalse(self.in_buf.endswith(bytearray([4, 6])))
with self.assertRaises(ValueError):
self.in_buf.endswith([5, 6])
self.assertEqual(
self.in_buf.endswith(bytearray([0, 1, 2, 3, 4, 5, 6])), False)
self.assertEqual(True, self.in_buf.endswith(bytearray(0)))
def test_add_bytes(self):
self.in_buf.add_bytes(self.data1)
self.in_buf.add_bytes(self.data2)
self.assertEqual(2 * self.length_to_add, self.in_buf.length)
self.assertEqual(deque([self.data1, self.data2]),
self.in_buf.input_list)
def test_remove_bytes(self):
with self.assertRaises(AssertionError):
self.in_buf.remove_bytes(5)
self.in_buf.remove_bytes(0)
self.assertEqual(0, self.in_buf.length, 0)
self.make_input_buffer()
self.assertEqual(bytearray([i for i in range(1, 6)]),
self.in_buf.remove_bytes(5))
self.assertEqual(55, self.in_buf.length, 55)
self.assertEqual(bytearray([i for i in range(6, 26)]),
self.in_buf.remove_bytes(20))
self.assertEqual(35, self.in_buf.length)
self.assertEqual(bytearray([i for i in range(26, 61)]),
self.in_buf.remove_bytes(35))
self.assertEqual(0, self.in_buf.length)
def test_peek_message(self):
self.make_input_buffer()
# Edge Case: peek_message returns all bytes when it peeks a number greater than the message length.
self.assertEqual(bytearray([i for i in range(1, 61)]),
self.in_buf.peek_message(70))
self.assertIn(bytearray([i for i in range(1, 5)]),
self.in_buf.peek_message(5))
self.assertIn(bytearray([i for i in range(1, 31)]),
self.in_buf.peek_message(30))
self.assertIn(bytearray([i for i in range(1, 61)]),
self.in_buf.peek_message(60))
def test_get_slice(self):
self.make_input_buffer()
with self.assertRaises(ValueError):
self.in_buf.get_slice(5, None)
with self.assertRaises(ValueError):
self.in_buf.get_slice(None, 6)
with self.assertRaises(ValueError):
self.in_buf.get_slice(100, 60)
self.assertEqual(bytearray([i for i in range(1, 6)]),
self.in_buf.get_slice(0, 5))
self.assertEqual(self.in_buf[:5], self.in_buf.get_slice(0, 5))
self.assertEqual(bytearray([i for i in range(11, 32)]),
self.in_buf.get_slice(10, 31))
self.assertEqual(self.in_buf[10:31], self.in_buf.get_slice(10, 31))
self.assertEqual(bytearray([i for i in range(35, 61)]),
self.in_buf.get_slice(34, 60))
self.assertEqual(self.in_buf[34:], self.in_buf.get_slice(34, 60))
def make_input_buffer(self):
self.in_buf.add_bytes(self.data1)
self.in_buf.add_bytes(self.data2)
self.in_buf.add_bytes(self.data3)
class AbstractConnection(Generic[Node]):
__metaclass__ = ABCMeta
CONNECTION_TYPE: ClassVar[ConnectionType] = ConnectionType.NONE
node: Node
def __init__(self, socket_connection, address, node: Node, from_me=False):
if not isinstance(socket_connection, SocketConnection):
raise ValueError(
"SocketConnection type is expected for socket_connection arg but was {0}."
.format(type(socket_connection)))
self.socket_connection = socket_connection
self.fileno = socket_connection.fileno()
# (IP, Port) at time of socket creation.
# If the version/hello message contains a different port (i.e. connection is not from me), this will
# be updated to the one in the message.
self.peer_ip, self.peer_port = address
self.peer_id: Optional[str] = None
self.external_ip = node.opts.external_ip
self.external_port = node.opts.external_port
self.from_me = from_me # Whether or not I initiated the connection
if node.opts.track_detailed_sent_messages:
self.message_tracker = MessageTracker(self)
self.outputbuf = OutputBuffer()
self.inputbuf = InputBuffer()
self.node = node
self.state = ConnectionState.CONNECTING
# Number of bad messages I've received in a row.
self.num_bad_messages = 0
self.peer_desc = "%s %d" % (self.peer_ip, self.peer_port)
self.can_send_pings = False
self.hello_messages = []
self.header_size = 0
self.message_factory = None
self.message_handlers = None
self.log_throughput = True
self.ping_message = None
self.pong_message = None
self.ack_message = None
# Default network number to network number of current node. But it can change after hello message is received
self.network_num = node.network_num
self.message_validator = DefaultMessageValidator()
self._debug_message_tracker = defaultdict(int)
self._last_debug_message_log_time = time.time()
self.ping_interval_s: int = constants.PING_INTERVAL_S
self.peer_model: Optional[OutboundPeerModel] = None
self.log_debug("Connection initialized.")
def __repr__(self):
if logger.isEnabledFor(LogLevel.DEBUG):
details = f"fileno: {self.fileno}, address: {self.peer_desc}, network_num: {self.network_num}"
else:
details = f"fileno: {self.fileno}, address: {self.peer_desc}"
return f"{self.CONNECTION_TYPE}({details})"
def _log_message(self, level: LogLevel, message, *args, **kwargs):
logger.log(level, f"[{self}] {message}", *args, **kwargs)
def log_trace(self, message, *args, **kwargs):
self._log_message(LogLevel.TRACE, message, *args, **kwargs)
def log_debug(self, message, *args, **kwargs):
self._log_message(LogLevel.DEBUG, message, *args, **kwargs)
def log_info(self, message, *args, **kwargs):
self._log_message(LogLevel.INFO, message, *args, **kwargs)
def log_warning(self, message, *args, **kwargs):
self._log_message(LogLevel.WARNING, message, *args, **kwargs)
def log_error(self, message, *args, **kwargs):
self._log_message(LogLevel.ERROR, message, *args, **kwargs)
def is_active(self):
"""
Indicates whether the connection is established and not marked for close.
"""
return self.state & ConnectionState.ESTABLISHED == ConnectionState.ESTABLISHED and \
not self.state & ConnectionState.MARK_FOR_CLOSE
def is_sendable(self):
"""
Indicates whether the connection should send bytes on broadcast.
"""
return self.is_active()
def on_connection_established(self):
self.state |= ConnectionState.ESTABLISHED
self.log_info("Connection established.")
def add_received_bytes(self, bytes_received: int):
"""
Adds bytes received from socket connection to input buffer
:param bytes_received: new bytes received from socket connection
"""
assert not self.state & ConnectionState.MARK_FOR_CLOSE
self.inputbuf.add_bytes(bytes_received)
def get_bytes_to_send(self):
assert not self.state & ConnectionState.MARK_FOR_CLOSE
return self.outputbuf.get_buffer()
def advance_sent_bytes(self, bytes_sent):
self.advance_bytes_on_buffer(self.outputbuf, bytes_sent)
if self.message_tracker:
self.message_tracker.advance_bytes(bytes_sent)
def enqueue_msg(self, msg: AbstractMessage, prepend: bool = False):
"""
Enqueues the contents of a Message instance, msg, to our outputbuf and attempts to send it if the underlying
socket has room in the send buffer.
:param msg: message
:param prepend: if the message should be bumped to the front of the outputbuf
"""
self._log_message(msg.log_level(), "Enqueued message: {}", msg)
if self.message_tracker:
full_message = msg
else:
full_message = None
self.enqueue_msg_bytes(msg.rawbytes(), prepend, full_message)
def enqueue_msg_bytes(self,
msg_bytes: Union[bytearray, memoryview],
prepend: bool = False,
full_message: Optional[AbstractMessage] = None):
"""
Enqueues the raw bytes of a message, msg_bytes, to our outputbuf and attempts to send it if the
underlying socket has room in the send buffer.
:param msg_bytes: message bytes
:param prepend: if the message should be bumped to the front of the outputbuf
:param full_message: full message for detailed logging
"""
if self.state & ConnectionState.MARK_FOR_CLOSE:
return
size = len(msg_bytes)
self.log_trace("Enqueued {} bytes.", size)
if prepend:
self.outputbuf.prepend_msgbytes(msg_bytes)
if self.message_tracker:
self.message_tracker.prepend_message(len(msg_bytes),
full_message)
else:
self.outputbuf.enqueue_msgbytes(msg_bytes)
if self.message_tracker:
self.message_tracker.append_message(len(msg_bytes),
full_message)
# TODO: temporary fix for some situations where, see https://bloxroute.atlassian.net/browse/BX-1153
self.socket_connection.can_send = True
self.socket_connection.send()
def pre_process_msg(self):
is_full_msg, msg_type, payload_len = self.message_factory.get_message_header_preview_from_input_buffer(
self.inputbuf)
return is_full_msg, msg_type, payload_len
def process_msg_type(self, message_type, is_full_msg, payload_len):
"""
Processes messages that require changes to the regular message handling flow
(pop off single message, process it, continue on with the stream)
:param message_type: message type
:param is_full_msg: flag indicating if full message is available on input buffer
:param payload_len: length of payload
:return:
"""
pass
def process_message(self):
"""
Processes the next bytes on the socket's inputbuffer.
Returns 0 in order to avoid being rescheduled if this was an alarm.
"""
start_time = time.time()
messages_processed = defaultdict(int)
while True:
input_buffer_len_before = self.inputbuf.length
is_full_msg = False
payload_len = 0
msg = None
msg_type = None
try:
# abort message processing if connection has been closed
if self.state & ConnectionState.MARK_FOR_CLOSE:
return
is_full_msg, msg_type, payload_len = self.pre_process_msg()
self.message_validator.validate(is_full_msg, msg_type,
self.header_size, payload_len,
self.inputbuf)
self.process_msg_type(msg_type, is_full_msg, payload_len)
if not is_full_msg:
break
msg = self.pop_next_message(payload_len)
# If there was some error in parsing this message, then continue the loop.
if msg is None:
if self._report_bad_message():
return
continue
# Full messages must be one of the handshake messages if the connection isn't established yet.
if not (self.state & ConnectionState.ESTABLISHED == ConnectionState.ESTABLISHED) \
and msg_type not in self.hello_messages:
self.log_warning(
"Received unexpected message ({}) before handshake completed. Closing.",
msg_type)
self.mark_for_close()
return
if self.log_throughput:
hooks.add_throughput_event(Direction.INBOUND, msg_type,
len(msg.rawbytes()),
self.peer_desc)
if not logger.isEnabledFor(
msg.log_level()) and logger.isEnabledFor(
LogLevel.INFO):
self._debug_message_tracker[msg_type] += 1
elif len(self._debug_message_tracker) > 0:
self.log_debug(
"Processed the following messages types: {} over {:.2f} seconds.",
self._debug_message_tracker,
time.time() - self._last_debug_message_log_time)
self._debug_message_tracker.clear()
self._last_debug_message_log_time = time.time()
self._log_message(msg.log_level(), "Processing message: {}",
msg)
if msg_type in self.message_handlers:
msg_handler = self.message_handlers[msg_type]
msg_handler(msg)
messages_processed[msg_type] += 1
# TODO: Investigate possible solutions to recover from PayloadLenError errors
except PayloadLenError as e:
self.log_error("Could not parse message. Error: {}", e.msg)
self.mark_for_close()
return
except MemoryError as e:
self.log_error(
"Out of memory error occurred during message processing. Error: {}. ",
e,
exc_info=True)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before,
payload_len))
raise
except UnauthorizedMessageError as e:
self.log_error("Unauthorized message {} from {}.",
e.msg.MESSAGE_TYPE, self.peer_desc)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before,
payload_len))
# give connection a chance to restore its state and get ready to process next message
self.clean_up_current_msg(
payload_len,
input_buffer_len_before == self.inputbuf.length)
if self._report_bad_message():
return
except MessageValidationError as e:
self.log_warning(
"Message validation failed for {} message: {}.", msg_type,
e.msg)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before,
payload_len))
if is_full_msg:
self.clean_up_current_msg(
payload_len,
input_buffer_len_before == self.inputbuf.length)
else:
self.log_error(
"Unable to recover after message that failed validation. Closing connection."
)
self.mark_for_close()
return
if self._report_bad_message():
return
# TODO: Throw custom exception for any errors that come from input that has not been validated and only catch that subclass of exceptions
except Exception as e:
# Attempt to recover connection by removing bad full message
if is_full_msg:
self.log_error(
"Message processing error; trying to recover. Error: {}.",
e,
exc_info=True)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before,
payload_len))
# give connection a chance to restore its state and get ready to process next message
self.clean_up_current_msg(
payload_len,
input_buffer_len_before == self.inputbuf.length)
# Connection is unable to recover from message processing error if incomplete message is received
else:
self.log_error(
"Message processing error; unable to recover. Error: {}.",
e,
exc_info=True)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before,
payload_len))
self.mark_for_close()
return
if self._report_bad_message():
return
else:
self.num_bad_messages = 0
time_elapsed = time.time() - start_time
self.log_trace("Processed {} messages in {:.2f} seconds",
messages_processed, time_elapsed)
def pop_next_message(self, payload_len):
"""
Pop the next message off of the buffer given the message length.
Preserve invariant of self.inputbuf always containing the start of a valid message.
:param payload_len: length of payload
:return: message object
"""
msg_len = self.message_factory.base_message_type.HEADER_LENGTH + payload_len
msg_contents = self.inputbuf.remove_bytes(msg_len)
return self.message_factory.create_message_from_buffer(msg_contents)
def advance_bytes_on_buffer(self, buf, bytes_written):
hooks.add_throughput_event(Direction.OUTBOUND, None, bytes_written,
self.peer_desc)
try:
buf.advance_buffer(bytes_written)
except ValueError as e:
raise RuntimeError(
"Connection: {}, Failed to advance buffer".format(self)) from e
def send_ping(self):
"""
Send a ping (and reschedule if called from alarm queue)
"""
if self.can_send_pings and not self.state & ConnectionState.MARK_FOR_CLOSE:
self.enqueue_msg(self.ping_message)
return self.ping_interval_s
return constants.CANCEL_ALARMS
def msg_hello(self, msg):
self.state |= ConnectionState.HELLO_RECVD
if msg.node_id() is None:
self.log_debug("Received hello message without peer id.")
self.peer_id = msg.node_id()
self.node.connection_pool.index_conn_node_id(self.peer_id, self)
if len(self.node.connection_pool.get_by_node_id(self.peer_id)) > 1:
if self.from_me:
self.log_info(
"Received duplicate connection from: {}. Closing.",
self.peer_id)
self.mark_for_close()
return
self.enqueue_msg(self.ack_message)
if self.is_active():
self.on_connection_established()
def msg_ack(self, _msg):
"""
Handle an Ack Message
"""
self.state |= ConnectionState.HELLO_ACKD
if self.is_active():
self.on_connection_established()
def msg_ping(self, msg):
self.enqueue_msg(self.pong_message)
def msg_pong(self, _msg):
pass
def mark_for_close(self):
"""
Marks a connection for close. Prefer using this method to close a connection over
AbstractConnection#destroy_conn, as this allows a cleaner showdown and finish processing messages.
"""
self.state |= ConnectionState.MARK_FOR_CLOSE
self.log_debug("Marking connection for close.")
def close(self):
"""
Cleans up connection state after socket has been terminated.
Do not call this directly from connection event handlers.
"""
assert self.state & ConnectionState.MARK_FOR_CLOSE
def clean_up_current_msg(self, payload_len: int,
msg_is_in_input_buffer: bool) -> None:
"""
Removes current message from the input buffer and resets connection to a state ready to process next message.
Called during the handling of message processing exceptions.
:param payload_len: length of the payload of the currently processing message
:param msg_is_in_input_buffer: flag indicating if message bytes are still in the input buffer
:return:
"""
if msg_is_in_input_buffer:
self.inputbuf.remove_bytes(self.header_size + payload_len)
def on_input_received(self) -> bool:
"""handles an input event from the event loop
:return: True if the connection is receivable, otherwise False
"""
return True
def log_connection_mem_stats(self) -> None:
"""
logs the connection's memory stats
"""
class_name = self.__class__.__name__
hooks.add_obj_mem_stats(
class_name,
self.network_num,
self.inputbuf,
"input_buffer",
memory_utils.ObjectSize("input_buffer",
memory_utils.get_special_size(
self.inputbuf).size,
is_actual_size=True),
object_item_count=len(self.inputbuf.input_list),
object_type=memory_utils.ObjectType.BASE,
size_type=memory_utils.SizeType.TRUE)
hooks.add_obj_mem_stats(
class_name,
self.network_num,
self.outputbuf,
"output_buffer",
memory_utils.ObjectSize("output_buffer",
memory_utils.get_special_size(
self.outputbuf).size,
is_actual_size=True),
object_item_count=len(self.outputbuf.output_msgs),
object_type=memory_utils.ObjectType.BASE,
size_type=memory_utils.SizeType.TRUE)
def update_model(self, model: OutboundPeerModel):
self.log_trace("Updated connection model: {}", model)
self.peer_model = model
def _report_bad_message(self):
"""
Increments counter for bad messages. Returns True if connection should be closed.
:return: if connection should be closed
"""
if self.num_bad_messages == constants.MAX_BAD_MESSAGES:
self.log_warning("Received too many bad messages. Closing.")
self.mark_for_close()
return True
else:
self.num_bad_messages += 1
return False
def _get_last_msg_bytes(self, msg, input_buffer_len_before, payload_len):
if msg is not None:
return convert.bytes_to_hex(
msg.rawbytes()[:constants.MAX_LOGGED_BYTES_LEN])
# bytes still available on input buffer
if input_buffer_len_before == self.inputbuf.length and payload_len is not None:
return convert.bytes_to_hex(
self.inputbuf.peek_message(
min(self.header_size + payload_len,
constants.MAX_LOGGED_BYTES_LEN)))
return "<not available>"
class AbstractConnection(Generic[Node]):
__metaclass__ = ABCMeta
CONNECTION_TYPE: ClassVar[ConnectionType] = ConnectionType.NONE
node: Node
message_factory: AbstractMessageFactory
format_connection_desc: str
connection_repr: str
# performance critical attribute, has been pulled out of connection state
established: bool
def __init__(self, socket_connection: AbstractSocketConnectionProtocol, node: Node) -> None:
self.socket_connection = socket_connection
self.file_no = socket_connection.file_no
# (IP, Port) at time of socket creation.
# If the version/hello message contains a different port (i.e. connection is not from me), this will
# be updated to the one in the message.
self.endpoint = self.socket_connection.endpoint
self.peer_ip, self.peer_port = self.endpoint
self.peer_id: Optional[str] = None
self.external_ip = node.opts.external_ip
self.external_port = node.opts.external_port
self.direction = self.socket_connection.direction
self.from_me = self.direction == NetworkDirection.OUTBOUND
self.outputbuf = OutputBuffer()
self.inputbuf = InputBuffer()
self.node = node
self.state = ConnectionState.CONNECTING
self.established = False
# Number of bad messages I've received in a row.
self.num_bad_messages = 0
self.peer_desc = repr(self.endpoint)
self.can_send_pings = False
self.pong_timeout_enabled = False
self.hello_messages = []
self.header_size = 0
self.message_factory = self.connection_message_factory()
self.message_handlers = {}
self.log_throughput = True
self.pong_message = None
self.ack_message = None
self.ping_alarm_id: Optional[AlarmId] = None
self.ping_interval_s = constants.PING_INTERVAL_S
self.pong_timeout_alarm_id: Optional[AlarmId] = None
# Default network number to network number of current node. But it can change after hello message is received
self.network_num = node.network_num
self.message_validator = DefaultMessageValidator()
self._debug_message_tracker = defaultdict(int)
self._last_debug_message_log_time = time.time()
self.processing_message_index = 0
self.peer_model: Optional[OutboundPeerModel] = None
self._is_authenticated = False
self.account_id: Optional[str] = None
self.tier_name: Optional[str] = None
self._close_waiter: Optional[Future] = None
self.format_connection()
self.log_debug("Connection initialized.")
def __repr__(self):
if logger.isEnabledFor(LogLevel.DEBUG):
details = f"file_no: {self.file_no}, address: {self.peer_desc}, network_num: {self.network_num}"
else:
details = f"file_no: {self.file_no}, address: {self.peer_desc}"
return f"{self.CONNECTION_TYPE} ({details})"
@abstractmethod
def ping_message(self) -> AbstractMessage:
"""
Define ping message characteristics for pinging on connection.
This function may have side-effects; only call this if the ping message will
be used.
"""
@abstractmethod
def connection_message_factory(self) -> AbstractMessageFactory:
pass
def log_trace(self, message, *args, **kwargs):
if logger.isEnabledFor(LogLevel.TRACE):
self._log_message(LogLevel.TRACE, message, *args, **kwargs)
def log_debug(self, message, *args, **kwargs):
if logger.isEnabledFor(LogLevel.DEBUG):
self._log_message(LogLevel.DEBUG, message, *args, **kwargs)
def log_info(self, message, *args, **kwargs):
self._log_message(LogLevel.INFO, message, *args, **kwargs)
def log_warning(self, message, *args, **kwargs):
self._log_message(LogLevel.WARNING, message, *args, **kwargs)
def log_error(self, message, *args, **kwargs):
self._log_message(LogLevel.ERROR, message, *args, **kwargs)
def log(self, level: LogLevel, message, *args, **kwargs):
self._log_message(level, message, *args, **kwargs)
def is_active(self) -> bool:
"""
Indicates whether the connection is established and ready for normal messages.
This function is very frequently called. Avoid doing any sort of complex
operations, inline function calls, and avoid flags.
"""
return self.established and self.socket_connection is not None and self.socket_connection.alive
def is_alive(self) -> bool:
"""
Indicates whether the connection's socket is alive.
This function is very frequently called. Avoid doing any sort of complex
operations, inline function calls, and avoid flags.
"""
if self.socket_connection is None:
return False
return self.socket_connection.alive
def on_connection_established(self):
if not self.is_active():
self.state |= ConnectionState.HELLO_RECVD
self.state |= ConnectionState.HELLO_ACKD
self.state |= ConnectionState.ESTABLISHED
self.established = True
self.log_info("Connection established.")
# Reset num_retries when a connection established in order to support resetting the Fibonnaci logic
# to determine next retry
self.node.num_retries_by_ip[(self.peer_ip, self.peer_port)] = 0
for peer_model in self.node.outbound_peers:
if (
(
peer_model.ip == self.peer_ip
and peer_model.port == self.peer_port
)
or peer_model.node_id == self.peer_id
):
self.peer_model = peer_model
def add_received_bytes(self, bytes_received: Union[bytearray, bytes]):
"""
Adds bytes received from socket connection to input buffer
:param bytes_received: new bytes received from socket connection
"""
assert self.is_alive()
self.inputbuf.add_bytes(bytes_received)
def get_bytes_to_send(self):
assert self.is_alive()
return self.outputbuf.get_buffer()
def advance_sent_bytes(self, bytes_sent):
self.advance_bytes_on_buffer(self.outputbuf, bytes_sent)
def enqueue_msg(self, msg: AbstractMessage, prepend: bool = False):
"""
Enqueues the contents of a Message instance, msg, to our outputbuf and attempts to send it if the underlying
socket has room in the send buffer.
:param msg: message
:param prepend: if the message should be bumped to the front of the outputbuf
"""
self._log_message(msg.log_level(), "Enqueued message: {}", msg)
self.enqueue_msg_bytes(msg.rawbytes(), prepend)
def enqueue_msg_bytes(
self,
msg_bytes: Union[bytearray, memoryview],
prepend: bool = False,
):
"""
Enqueues the raw bytes of a message, msg_bytes, to our outputbuf and attempts to send it if the
underlying socket has room in the send buffer.
This function is very frequently called. Avoid doing any sort of complex
operations, inline function calls, and avoid flags.
:param msg_bytes: message bytes
:param prepend: if the message should be bumped to the front of the outputbuf
"""
if not self.socket_connection.alive:
return
self.log_trace("Enqueued {} bytes.", len(msg_bytes))
if prepend:
self.outputbuf.prepend_msgbytes(msg_bytes)
else:
self.outputbuf.enqueue_msgbytes(msg_bytes)
self.socket_connection.send()
def pre_process_msg(self) -> ConnectionMessagePreview:
is_full_msg, msg_type, payload_len = self.message_factory.get_message_header_preview_from_input_buffer(
self.inputbuf
)
return ConnectionMessagePreview(is_full_msg, True, msg_type, payload_len)
def process_msg_type(self, message_type, is_full_msg, payload_len):
"""
Processes messages that require changes to the regular message handling flow
(pop off single message, process it, continue on with the stream)
:param message_type: message type
:param is_full_msg: flag indicating if full message is available on input buffer
:param payload_len: length of payload
:return:
"""
def process_message(self):
"""
Processes the next bytes on the socket's inputbuffer.
Returns 0 in order to avoid being rescheduled if this was an alarm.
"""
# pylint: disable=too-many-return-statements, too-many-branches, too-many-statements
logger.trace("START PROCESSING from {}", self)
start_time = time.time()
messages_processed = defaultdict(int)
total_bytes_processed = 0
self.processing_message_index = 0
while True:
input_buffer_len_before = self.inputbuf.length
is_full_msg = False
payload_len = None
msg = None
msg_type = None
try:
# abort message processing if connection has been closed
if not self.socket_connection.alive:
return
is_full_msg, should_process, msg_type, payload_len = self.pre_process_msg()
if not should_process and is_full_msg:
self.pop_next_bytes(payload_len)
continue
self.message_validator.validate(
is_full_msg,
msg_type,
self.header_size,
payload_len,
self.inputbuf
)
self.process_msg_type(msg_type, is_full_msg, payload_len)
if not is_full_msg:
break
msg = self.pop_next_message(payload_len)
total_bytes_processed += len(msg.rawbytes())
# If there was some error in parsing this message, then continue the loop.
if msg is None:
if self._report_bad_message():
return
continue
# Full messages must be one of the handshake messages if the connection isn't established yet.
if (
not self.established and msg_type not in self.hello_messages
):
self.log_warning(log_messages.UNEXPECTED_MESSAGE, msg_type)
self.mark_for_close()
return
if self.log_throughput:
hooks.add_throughput_event(
NetworkDirection.INBOUND,
msg_type,
len(msg.rawbytes()),
self.peer_desc,
self.peer_id
)
if not logger.isEnabledFor(msg.log_level()) and logger.isEnabledFor(LogLevel.INFO):
self._debug_message_tracker[msg_type] += 1
elif len(self._debug_message_tracker) > 0:
self.log_debug(
"Processed the following messages types: {} over {:.2f} seconds.",
self._debug_message_tracker,
time.time() - self._last_debug_message_log_time
)
self._debug_message_tracker.clear()
self._last_debug_message_log_time = time.time()
self._log_message(msg.log_level(), "Processing message: {}", msg)
if msg_type in self.message_handlers:
msg_handler = self.message_handlers[msg_type]
handler_start = time.time()
msg_handler(msg)
performance_utils.log_operation_duration(
msg_handling_logger,
"Single message handler",
handler_start,
constants.MSG_HANDLERS_CYCLE_DURATION_WARN_THRESHOLD_S,
connection=self,
handler=msg_handler,
message=msg
)
messages_processed[msg_type] += 1
# TODO: Investigate possible solutions to recover from PayloadLenError errors
except PayloadLenError as e:
self.log_error(log_messages.COULD_NOT_PARSE_MESSAGE, e.msg)
self.mark_for_close()
return
except MemoryError as e:
self.log_error(log_messages.OUT_OF_MEMORY, e, exc_info=True)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len)
)
raise
except UnauthorizedMessageError as e:
self.log_error(log_messages.UNAUTHORIZED_MESSAGE, e.msg.MESSAGE_TYPE, self.peer_desc)
self.log_debug(
"Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len)
)
# give connection a chance to restore its state and get ready to process next message
self.clean_up_current_msg(payload_len, input_buffer_len_before == self.inputbuf.length)
if self._report_bad_message():
return
except MessageValidationError as e:
if self.node.NODE_TYPE not in NodeType.GATEWAY_TYPE:
if isinstance(e, ControlFlagValidationError):
if e.is_cancelled_cut_through:
self.log_debug(
"Message validation failed for {} message: {}. Probably cut-through cancellation",
msg_type, e.msg)
else:
self.log_warning(log_messages.MESSAGE_VALIDATION_FAILED, msg_type, e.msg)
else:
self.log_warning(log_messages.MESSAGE_VALIDATION_FAILED, msg_type, e.msg)
else:
self.log_debug("Message validation failed for {} message: {}.", msg_type, e.msg)
self.log_debug("Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len))
if is_full_msg:
self.clean_up_current_msg(payload_len, input_buffer_len_before == self.inputbuf.length)
else:
self.log_error(log_messages.UNABLE_TO_RECOVER_PARTIAL_MESSAGE)
self.mark_for_close()
return
if self._report_bad_message():
return
except NonVersionMessageError as e:
if e.is_known:
self.log_debug("Received invalid handshake request on {}:{}, {}", self.peer_ip, self.peer_port,
e.msg)
else:
self.log_warning(log_messages.INVALID_HANDSHAKE, self.peer_ip, self.peer_port, e.msg)
self.log_debug("Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len))
self.mark_for_close()
return
# TODO: Throw custom exception for any errors that come from input that has not been
# validated and only catch that subclass of exceptions
# pylint: disable=broad-except
except Exception as e:
# Attempt to recover connection by removing bad full message
if is_full_msg:
self.log_error(log_messages.TRYING_TO_RECOVER_MESSAGE, e, exc_info=True)
self.log_debug("Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len))
# give connection a chance to restore its state and get ready to process next message
self.clean_up_current_msg(payload_len, input_buffer_len_before == self.inputbuf.length)
# Connection is unable to recover from message processing error if incomplete message is received
else:
self.log_error(log_messages.UNABLE_TO_RECOVER_FULL_MESSAGE, e, exc_info=True)
self.log_debug("Failed message bytes: {}",
self._get_last_msg_bytes(msg, input_buffer_len_before, payload_len))
self.mark_for_close()
return
if self._report_bad_message():
return
else:
self.num_bad_messages = 0
self.processing_message_index += 1
performance_utils.log_operation_duration(msg_handling_logger,
"Message handlers",
start_time,
constants.MSG_HANDLERS_DURATION_WARN_THRESHOLD_S,
connection=self, count=messages_processed)
duration_ms = (time.time() - start_time) * 1000
logger.trace("DONE PROCESSING from {}. Bytes processed: {}. Messages processed: {}. Duration: {}",
self, total_bytes_processed, messages_processed, stats_format.duration(duration_ms))
def pop_next_message(self, payload_len: int) -> AbstractMessage:
"""
Pop the next full message off of the buffer given the message length.
Preserves invariant of self.inputbuf always containing the start of a
valid message. The caller of this function is responsible for ensuring
there is a valid message on the buffer.
:param payload_len: length of payload
:return: message object
"""
return self.message_factory.create_message_from_buffer(
self.pop_next_bytes(payload_len)
)
def pop_next_bytes(self, payload_len: int) -> Union[memoryview, bytearray, bytes]:
msg_len = self.message_factory.base_message_type.HEADER_LENGTH + payload_len
return self.inputbuf.remove_bytes(msg_len)
def advance_bytes_on_buffer(self, buf, bytes_written):
hooks.add_throughput_event(NetworkDirection.OUTBOUND, None, bytes_written, self.peer_desc, self.peer_id)
try:
buf.advance_buffer(bytes_written)
except ValueError as e:
raise RuntimeError("Connection: {}, Failed to advance buffer".format(self)) from e
def schedule_pings(self) -> None:
"""
Schedules ping on the connection. Multiple calls of this method will
override the existing alarm.
"""
self._unschedule_pings()
if self.can_send_pings:
self.ping_alarm_id = self.node.alarm_queue.register_alarm(
self.ping_interval_s, self.send_ping
)
def send_ping(self) -> float:
"""
Send a ping (and reschedule if called from alarm queue)
"""
if self.can_send_pings and self.is_alive():
self.enqueue_msg(self.ping_message())
if self.pong_timeout_enabled:
self.schedule_pong_timeout()
return self.ping_interval_s
return constants.CANCEL_ALARMS
def msg_hello(self, msg):
self.state |= ConnectionState.HELLO_RECVD
if msg.node_id() is None:
self.log_debug("Received hello message without peer id.")
if self.peer_id is None:
self.peer_id = msg.node_id()
# This should only be necessary for pre1.6 connections.
if self.peer_id in self.node.connection_pool.by_node_id and \
self.socket_connection.endpoint.ip_address[0:3] != "172":
existing_connection = self.node.connection_pool.get_by_node_id(
self.peer_id
)
if existing_connection != self:
if self.from_me:
self.mark_for_close()
if existing_connection.from_me:
existing_connection.mark_for_close()
self.log_warning(log_messages.DUPLICATE_CONNECTION, self.peer_id, existing_connection)
self.enqueue_msg(self.ack_message)
if ConnectionState.INITIALIZED | ConnectionState.HELLO_ACKD in self.state:
self.on_connection_established()
def msg_ack(self, _msg):
"""
Handle an Ack Message
"""
self.state |= ConnectionState.HELLO_ACKD
if ConnectionState.INITIALIZED | ConnectionState.HELLO_RECVD in self.state:
self.on_connection_established()
def msg_ping(self, _msg):
self.enqueue_msg(self.pong_message)
def msg_pong(self, _msg):
self.cancel_pong_timeout()
def mark_for_close(self, should_retry: Optional[bool] = None):
"""
Marks a connection for close, so AbstractNode can dispose of this class.
Use this where possible for a clean shutdown.
"""
loop = asyncio.get_event_loop()
self._close_waiter = loop.create_future()
if should_retry is None:
should_retry = self.from_me
self.log_debug("Marking connection for close, should_retry: {}.", should_retry)
self.socket_connection.mark_for_close(should_retry)
def dispose(self):
"""
Performs any need operations after connection object has been discarded by the AbstractNode.
"""
if self._close_waiter is not None:
self._close_waiter.set_result(True)
self._unschedule_pings()
self.node = None
self.socket_connection = None
def clean_up_current_msg(self, payload_len: int, msg_is_in_input_buffer: bool) -> None:
"""
Removes current message from the input buffer and resets connection to a state ready to process next message.
Called during the handling of message processing exceptions.
:param payload_len: length of the payload of the currently processing message
:param msg_is_in_input_buffer: flag indicating if message bytes are still in the input buffer
:return:
"""
if msg_is_in_input_buffer:
self.inputbuf.remove_bytes(self.header_size + payload_len)
def on_input_received(self) -> bool:
"""handles an input event from the event loop
:return: True if the connection is receivable, otherwise False
"""
return True
def log_connection_mem_stats(self) -> None:
"""
logs the connection's memory stats
"""
class_name = self.__class__.__name__
hooks.add_obj_mem_stats(
class_name,
self.network_num,
self.inputbuf,
"input_buffer",
memory_utils.ObjectSize("input_buffer", 0,
is_actual_size=False),
object_item_count=len(self.inputbuf.input_list),
object_type=memory_utils.ObjectType.BASE,
size_type=memory_utils.SizeType.TRUE
)
hooks.add_obj_mem_stats(
class_name,
self.network_num,
self.outputbuf,
"output_buffer",
memory_utils.ObjectSize("output_buffer", 0,
is_actual_size=False),
object_item_count=len(self.outputbuf.output_msgs),
object_type=memory_utils.ObjectType.BASE,
size_type=memory_utils.SizeType.TRUE
)
def update_model(self, model: OutboundPeerModel):
self.log_trace("Updated connection model: {}", model)
self.peer_model = model
def schedule_pong_timeout(self) -> None:
if self.pong_timeout_alarm_id is None:
self.log_trace(
"Schedule pong reply timeout for ping message in {} seconds",
constants.PING_PONG_REPLY_TIMEOUT_S
)
self.pong_timeout_alarm_id = self.node.alarm_queue.register_alarm(
constants.PING_PONG_REPLY_TIMEOUT_S, self._pong_msg_timeout
)
def cancel_pong_timeout(self):
if self.pong_timeout_alarm_id is not None:
self.node.alarm_queue.unregister_alarm(self.pong_timeout_alarm_id)
self.pong_timeout_alarm_id = None
def on_connection_authenticated(self, peer_info: AuthenticatedPeerInfo) -> None:
self.peer_id = peer_info.peer_id
if self.CONNECTION_TYPE != peer_info.connection_type:
self.node.connection_pool.update_connection_type(self, peer_info.connection_type)
self.account_id = peer_info.account_id
self._is_authenticated = True
async def wait_closed(self):
if self._close_waiter is not None:
await self._close_waiter
self._close_waiter = None
else:
await asyncio.sleep(0)
if self.is_alive():
raise ConnectionStateError("Connection is still alive after closed", self)
def set_account_id(self, account_id: Optional[str], tier_name: str):
if self._is_authenticated and account_id != self.account_id:
raise ConnectionAuthenticationError(
f"Invalid account id {account_id} is different than connection account id: {self.account_id}")
if not self._is_authenticated:
self.account_id = account_id
self.tier_name = tier_name
def get_backlog_size(self) -> int:
output_buffer_backlog = self.outputbuf.length
socket_buffer_backlog = self.socket_connection.get_write_buffer_size()
self.log_trace("Output backlog: {}, socket backlog: {}", output_buffer_backlog, socket_buffer_backlog)
return output_buffer_backlog + socket_buffer_backlog
def format_connection(self) -> None:
self.format_connection_desc = "{} - {}".format(
self.peer_desc, self.CONNECTION_TYPE.format_short()
)
self.connection_repr = repr(self)
def _unschedule_pings(self) -> None:
existing_alarm = self.ping_alarm_id
if existing_alarm:
self.node.alarm_queue.unregister_alarm(existing_alarm)
self.ping_alarm_id = None
def _pong_msg_timeout(self) -> None:
if self.is_alive():
self.log_info(
"Connection appears to be broken. Peer did not reply to PING message within allocated time. "
"Closing connection."
)
self.mark_for_close()
self.pong_timeout_alarm_id = None
def _report_bad_message(self):
"""
Increments counter for bad messages. Returns True if connection should be closed.
:return: if connection should be closed
"""
if self.num_bad_messages == constants.MAX_BAD_MESSAGES:
self.log_warning(log_messages.TOO_MANY_BAD_MESSAGES)
self.mark_for_close()
return True
else:
self.num_bad_messages += 1
return False
def _get_last_msg_bytes(self, msg, input_buffer_len_before, payload_len):
if msg is not None:
return convert.bytes_to_hex(msg.rawbytes()[:constants.MAX_LOGGED_BYTES_LEN])
# bytes still available on input buffer
if input_buffer_len_before == self.inputbuf.length and payload_len is not None:
return convert.bytes_to_hex(
self.inputbuf.peek_message(min(self.header_size + payload_len, constants.MAX_LOGGED_BYTES_LEN)))
return "<not available>"
def _log_message(self, level: LogLevel, message, *args, **kwargs):
logger.log(
level, message, HAS_PREFIX, self.connection_repr, *args, **kwargs
)