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)