def enqueue_disconnect(self, socket_connection: SocketConnection, should_retry: Optional[bool] = None): """ Queues up a disconnect for the event loop to close the socket and destroy the connection object for. This should always be called with a value provided for `should_retry`, unless the connection object is unknown (e.g. from the event loop or SocketConnection classes). """ fileno = socket_connection.fileno() logger.trace("Enqueuing disconnect from {}", fileno) if should_retry is None: conn = self.connection_pool.get_by_fileno(fileno) if conn is None: logger.debug( "Unexpectedly tried to enqueue a disconnect without a connection object on fileno: {}", fileno) should_retry = False else: conn.log_debug( "Connection close triggered by socket layer or event loop." ) conn.mark_for_close() should_retry = conn.from_me socket_connection.mark_for_close() self.disconnect_queue.append(DisconnectRequest(fileno, should_retry))
def setUp(self) -> None: self.socket_instance = MagicMock(spec=socket.socket) self.socket_instance.fileno = MagicMock(return_value=1) self.node = AbstractNode(helpers.get_common_opts(1234)) self.sut = SocketConnection(self.socket_instance, self.node, False) self.connection = helpers.create_connection(MockConnection, self.node, from_me=True) self.connection.socket_connection = self.sut
def _register_socket(self, new_socket, address, is_server=False, initialized=True, from_me=False): socket_connection = SocketConnection(new_socket, self._node, is_server) if initialized: socket_connection.set_state(SocketConnectionState.INITIALIZED) self._socket_connections[new_socket.fileno()] = socket_connection if not is_server: self._node.on_connection_added(socket_connection, address[0], address[1], from_me) if initialized: self._node.on_connection_initialized(new_socket.fileno())
def __init__(self, sock: SocketConnection, address, node, from_me=False): self.socket_connection = sock self.fileno = sock.fileno() # (IP, Port) at time of socket creation. We may get a new application level port in # the version message if the connection is not from me. self.peer_ip, self.peer_port = address self.peer_id = None self.my_ip = node.opts.external_ip self.my_port = node.opts.external_port self.from_me = from_me # Whether or not I initiated the connection self.outputbuf = OutputBuffer() self.inputbuf = InputBuffer() self.node = node self.is_persistent = False 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.message_handlers = None self.network_num = node.opts.blockchain_network_num self.enqueued_messages = []
class SocketConnectionTest(AbstractTestCase): def setUp(self) -> None: self.socket_instance = MagicMock(spec=socket.socket) self.socket_instance.fileno = MagicMock(return_value=1) self.node = AbstractNode(helpers.get_common_opts(1234)) self.sut = SocketConnection(self.socket_instance, self.node, False) self.connection = helpers.create_connection(MockConnection, self.node, from_me=True) self.connection.socket_connection = self.sut def test_close_from_socket(self): self.socket_instance.recv_into = MagicMock(return_value=0) self.sut.receive() self.assertIn((1, True), self.node.disconnect_queue) self.assertTrue(self.sut.state & SocketConnectionState.MARK_FOR_CLOSE) self.node.on_connection_closed(1, True) self.assertEqual(0, len(self.node.connection_pool))
def on_connection_added(self, socket_connection: SocketConnection, ip: str, port: int, from_me: bool): """ Notifies the node that a connection is coming in. """ fileno = socket_connection.fileno() # If we're already connected to the remote peer, log the event and request disconnect. if self.connection_exists(ip, port): logger.debug( "Duplicate connection attempted to: {0}:{1}. Dropping.", ip, port) # Schedule dropping the added connection and keep the old one. self.enqueue_disconnect(socket_connection, False) else: self._initialize_connection(socket_connection, ip, port, from_me)
def _initialize_connection(self, socket_connection: SocketConnection, ip: str, port: int, from_me: bool): conn_obj = self.build_connection(socket_connection, ip, port, from_me) if conn_obj is not None: logger.info("Connecting to: {}...", conn_obj) self.alarm_queue.register_alarm(constants.CONNECTION_TIMEOUT, self._connection_timeout, conn_obj) self.connection_pool.add(socket_connection.fileno(), ip, port, conn_obj) if conn_obj.CONNECTION_TYPE == ConnectionType.SDN: self.sdn_connection = conn_obj else: logger.warning( "Could not determine expected connection type for {}:{}. Disconnecting...", ip, port) self.enqueue_disconnect(socket_connection, from_me)