Shutdown all active peer multiplexers when worker is stopping (#259)

* Shutdown all active peer multiplexers when worker is stopping

* Wait for them to actually disconnect

* adjust

* fix tests

* Assert that cleanup is working

---------

Co-authored-by: Pascal Vizeli <[email protected]>

snitun/server/peer_manager.py

@@ -1,4 +1,5 @@
 """Manage peer connections."""
+
 from __future__ import annotations
 
 import asyncio
@@ -8,6 +9,7 @@
 import logging
 from typing import Callable
 
+import async_timeout
 from cryptography.fernet import Fernet, InvalidToken, MultiFernet
 
 from ..exceptions import SniTunInvalidPeer
@@ -37,7 +39,7 @@ def __init__(
         self._loop = asyncio.get_event_loop()
         self._throttling = throttling
         self._event_callback = event_callback
-        self._peers = {}
+        self._peers: dict[str, Peer] = {}
 
     @property
     def connections(self) -> int:
@@ -96,7 +98,9 @@ def remove_peer(self, peer: Peer) -> None:
 
         if self._event_callback:
             self._loop.call_soon(
-                self._event_callback, peer, PeerManagerEvent.DISCONNECTED,
+                self._event_callback,
+                peer,
+                PeerManagerEvent.DISCONNECTED,
             )
 
     def peer_available(self, hostname: str) -> bool:
@@ -108,3 +112,20 @@ def peer_available(self, hostname: str) -> bool:
     def get_peer(self, hostname: str) -> Peer | None:
         """Get peer."""
         return self._peers.get(hostname)
+
+    async def close_connections(self, timeout: int = 10) -> None:
+        """Close all peer connections.
+
+        Use this function only if you do not controll the server socket.
+        """
+        peers = list(self._peers.values())
+        for peer in peers:
+            if peer.is_connected:
+                peer.multiplexer.shutdown()
+
+        if waiters := [peer.wait_disconnect() for peer in peers]:
+            try:
+                async with async_timeout.timeout(timeout):
+                    await asyncio.gather(*waiters, return_exceptions=True)
+            except asyncio.TimeoutError:
+                _LOGGER.error("Timeout while waiting for peer disconnect")

snitun/server/worker.py

@@ -1,4 +1,5 @@
 """SniTun worker for traffics."""
+
 from __future__ import annotations
 
 import asyncio
@@ -83,7 +84,10 @@ def shutdown(self) -> None:
         self.join(10)
 
     def handover_connection(
-        self, con: socket, data: bytes, sni: str | None = None,
+        self,
+        con: socket,
+        data: bytes,
+        sni: str | None = None,
     ) -> None:
         """Move new connection to worker."""
         self._new.put_nowait((con, data, sni))
@@ -110,16 +114,24 @@ def run(self) -> None:
 
             new[0].setblocking(False)
             asyncio.run_coroutine_threadsafe(
-                self._async_new_connection(*new), loop=self._loop,
+                self._async_new_connection(*new),
+                loop=self._loop,
             )
 
         # Shutdown worker
-        _LOGGER.info("Stop worker: %s", self.name)
+        _LOGGER.info("Stoping worker: %s", self.name)
+        asyncio.run_coroutine_threadsafe(
+            self._peers.close_connections(),
+            loop=self._loop,
+        ).result()
         self._loop.call_soon_threadsafe(self._loop.stop)
         running_loop.join(10)
 
     async def _async_new_connection(
-        self, con: socket, data: bytes, sni: str | None,
+        self,
+        con: socket,
+        data: bytes,
+        sni: str | None,
     ) -> None:
         """Handle incoming connection."""
         try:

tests/conftest.py

@@ -94,6 +94,7 @@ def test_server_sync(event_loop):
     shutdown = False
 
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
     sock.bind(("127.0.0.1", 8366))
     sock.listen(2)
     sock.setblocking(False)

tests/server/test_worker.py

@@ -1,4 +1,5 @@
 """Tests for the server worker."""
+
 from datetime import datetime, timedelta, timezone
 import hashlib
 import os
@@ -49,6 +50,8 @@ def test_peer_connection(test_server_sync, test_client_sync, event_loop):
 
     worker.shutdown()
 
+    assert worker.peer_size == 0
+
 
 def test_peer_connection_disconnect(test_server_sync, test_client_sync, event_loop):
     """Run a full flow of with a peer & disconnect."""
@@ -81,7 +84,10 @@ def test_peer_connection_disconnect(test_server_sync, test_client_sync, event_lo
 
 
 def test_sni_connection(
-    test_server_sync, test_client_sync, test_client_ssl_sync, event_loop
+    test_server_sync,
+    test_client_sync,
+    test_client_ssl_sync,
+    event_loop,
 ):
     """Run a full flow of with a peer."""
     worker = ServerWorker(FERNET_TOKENS)
@@ -91,7 +97,11 @@ def test_sni_connection(
     hostname = "localhost"
     alias = ["localhost.custom"]
     fernet_token = create_peer_config(
-        valid.timestamp(), hostname, aes_key, aes_iv, alias=alias
+        valid.timestamp(),
+        hostname,
+        aes_key,
+        aes_iv,
+        alias=alias,
     )
 
     worker.start()
@@ -111,4 +121,7 @@ def test_sni_connection(
     worker.handover_connection(test_server_sync[1], TLS_1_2, hostname)
     assert len(test_client_sync.recv(1048)) == 32
 
+    assert worker.peer_size == 1
     worker.shutdown()
+
+    assert worker.peer_size == 0