Websockets access guardian (#9)

* WebSocketAccess guardian

* WebSocketsAccessGuardian and test fot it

Author: nkiryanov

src/entrypoint.py

@@ -4,7 +4,8 @@
 import websockets
 
 from app import conf
-from handlers.websockets_handler import WebSocketsHandler
+from handlers import WebSocketsHandler
+from handlers import WebSocketsAccessGuardian
 from storage.subscription_storage import SubscriptionStorage
 
 
@@ -15,23 +16,27 @@ def create_stop_signal() -> asyncio.Future[None]:
     return stop_signal
 
 
-async def app_runner(settings: conf.Settings, websockets_handler: WebSocketsHandler) -> None:
+async def app_runner(settings: conf.Settings, websockets_handler: WebSocketsHandler, access_guardian: WebSocketsAccessGuardian) -> None:
     async with websockets.serve(
         ws_handler=websockets_handler.websockets_handler,
         host=settings.WEBSOCKETS_HOST,
         port=settings.WEBSOCKETS_PORT,
     ):
-        await asyncio.Future()
+        await access_guardian.run()
 
 
 async def main() -> None:
     settings = conf.get_app_settings()
+    stop_signal = create_stop_signal()
+
     storage = SubscriptionStorage()
     websockets_handler = WebSocketsHandler(storage=storage)
+    access_guardian = WebSocketsAccessGuardian(storage=storage, check_interval=60.0, stop_signal=stop_signal)
 
     await app_runner(
         settings=settings,
         websockets_handler=websockets_handler,
+        access_guardian=access_guardian,
     )
 
 

src/handlers/__init__.py

@@ -1,5 +1,7 @@
+from handlers.websockets_access_guardian import WebSocketsAccessGuardian
 from handlers.websockets_handler import WebSocketsHandler
 
 __all__ = [
+    "WebSocketsAccessGuardian",
     "WebSocketsHandler",
 ]

src/handlers/tests/tests_websockets_access_guardian.py

@@ -0,0 +1,83 @@
+import asyncio
+from datetime import datetime
+import json
+import pytest
+
+from app.types import DecodedValidToken
+from handlers.websockets_access_guardian import WebSocketsAccessGuardian
+from storage.storage_updaters import StorageWebSocketRegister
+from storage.storage_updaters import StorageUserSubscriber
+
+pytestmark = [
+    pytest.mark.slow,
+]
+
+
+@pytest.fixture(autouse=True)
+def mock_broadcast(mocker):
+    return mocker.patch("websockets.broadcast")
+
+
+@pytest.fixture
+def get_guardian_enough_time_to_do_its_job():
+    return lambda: asyncio.sleep(0.4)
+
+
+@pytest.fixture
+async def guardian(storage):
+    stop_signal = asyncio.Future()  # it's better to use loop.create_future() but ok for tests
+
+    yield WebSocketsAccessGuardian(storage=storage, check_interval=0.1, stop_signal=stop_signal)
+
+    stop_signal.set_result(None)
+
+
+@pytest.fixture(autouse=True)
+def guardian_as_task(guardian, event_loop):
+    return event_loop.create_task(guardian.run())
+
+
+@pytest.fixture(autouse=True)
+def ws_subscribed(ws, storage, event):
+    token_expiration_timestamp = int(datetime.fromisoformat("2023-01-01 12:23:00Z").timestamp())
+
+    valid_token = DecodedValidToken(sub="user1", exp=token_expiration_timestamp)
+
+    StorageWebSocketRegister(storage, ws, valid_token)()
+    StorageUserSubscriber(storage, ws, event)()
+
+    return ws
+
+
+def test_guardian_monitor_and_manage_access_remove_expired_websockets(guardian, storage, ws_subscribed):
+    guardian.monitor_and_manage_access()
+
+    registered_websockets = storage.get_registered_websockets()
+    assert registered_websockets == []
+
+
+async def test_remove_expired_websockets_from_storage(get_guardian_enough_time_to_do_its_job, storage, mock_broadcast, ws_subscribed, mocker):
+    await get_guardian_enough_time_to_do_its_job()
+
+    registered_websockets = storage.get_registered_websockets()
+    assert registered_websockets == []
+    assert ws_subscribed.closed is False, "Do not close connection when token expired, just remove from storage"
+    mock_broadcast.assert_called_once_with(websockets=[ws_subscribed], message=mocker.ANY)
+
+
+async def test_broadcasted_message(get_guardian_enough_time_to_do_its_job, mock_broadcast):
+    await get_guardian_enough_time_to_do_its_job()
+
+    broadcasted_message_as_json = json.loads(mock_broadcast.call_args.kwargs["message"])
+    assert len(broadcasted_message_as_json) == 2
+    assert broadcasted_message_as_json["message_type"] == "ErrorResponse"
+    assert broadcasted_message_as_json["errors"] == ["Token expired, user subscriptions disabled or removed"]
+
+
+@pytest.mark.freeze_time("2023-01-01 12:22:55Z", tick=True)  # 5 seconds before expiration
+async def test_do_not_remove_not_expired_connections(get_guardian_enough_time_to_do_its_job, storage, ws_subscribed, mock_broadcast):
+    await get_guardian_enough_time_to_do_its_job()
+
+    registered_websockets = storage.get_registered_websockets()
+    assert registered_websockets == [ws_subscribed]
+    mock_broadcast.assert_not_called()

src/handlers/websockets_access_guardian.py

@@ -0,0 +1,44 @@
+import asyncio
+from dataclasses import dataclass, field
+import logging
+
+import websockets
+
+from handlers.dto import ErrorResponseMessage
+from storage.storage_updaters import StorageWebSocketRemover
+from storage.subscription_storage import SubscriptionStorage
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class WebSocketsAccessGuardian:
+    storage: SubscriptionStorage
+    check_interval: float = 60.0  # in seconds
+    stop_signal: asyncio.Future[None] = field(default_factory=asyncio.Future)  # default feature will run forever
+
+    async def run(self) -> None:
+        while True:
+            await asyncio.sleep(self.check_interval)
+
+            self.monitor_and_manage_access()
+
+    def monitor_and_manage_access(self) -> None:
+        expired_websockets = self.storage.get_expired_websockets()
+
+        if expired_websockets:
+            logger.warning("Notify and remove from storage for expired websockets: '%s'", (", ").join([str(websocket.id) for websocket in expired_websockets]))
+
+            self.broadcast_authentication_expired(expired_websockets)
+            self.remove_expired_websockets(expired_websockets)
+
+        e = self.storage.get_expired_websockets()
+        assert e is not None
+
+    def broadcast_authentication_expired(self, expired_websockets: list[websockets.WebSocketServerProtocol]) -> None:
+        error_message = ErrorResponseMessage(errors=["Token expired, user subscriptions disabled or removed"], incoming_message=None)
+        websockets.broadcast(websockets=expired_websockets, message=error_message.model_dump_json(exclude_none=True))
+
+    def remove_expired_websockets(self, expired_websockets: list[websockets.WebSocketServerProtocol]) -> None:
+        for websocket in expired_websockets:
+            StorageWebSocketRemover(storage=self.storage, websocket=websocket)()

src/tests/functional/conftest.py

@@ -5,7 +5,8 @@
 from websockets import client
 
 from entrypoint import app_runner
-from handlers.websockets_handler import WebSocketsHandler
+from handlers import WebSocketsHandler
+from handlers import WebSocketsAccessGuardian
 
 
 @pytest.fixture
@@ -24,18 +25,31 @@ def websockets_handler(storage):
     return WebSocketsHandler(storage=storage)
 
 
+@pytest.fixture
+async def stop_signal():
+    return asyncio.Future()
+
+
+@pytest.fixture
+async def access_guardian(storage, stop_signal):
+    return WebSocketsAccessGuardian(storage=storage, check_interval=0.5, stop_signal=stop_signal)
+
+
 @pytest.fixture(autouse=True)
-async def serve_app_runner(settings, websockets_handler):
+async def serve_app_runner(settings, websockets_handler, access_guardian, stop_signal):
     serve_task = asyncio.get_running_loop().create_task(
         app_runner(
             settings=settings,
             websockets_handler=websockets_handler,
+            access_guardian=access_guardian,
         ),
     )
 
     await asyncio.sleep(0.1)  # give enough time to start the server
     assert serve_task.done() is False  # be sure server is running
-    return serve_task
+    yield serve_task
+
+    stop_signal.set_result(None)
 
 
 @pytest.fixture

src/tests/functional/tests_remove_registered_on_token_expiration.py

@@ -0,0 +1,26 @@
+import asyncio
+import pytest
+
+pytestmark = [
+    pytest.mark.slow,
+]
+
+
+@pytest.fixture
+def set_storage_connections_expired(storage):
+    def set_expired():
+        for _, websocket_meta in storage.registered_websockets.items():
+            websocket_meta.expiration_timestamp = 1000  # far in the past
+
+    return set_expired
+
+
+async def test_expired_connections_removed_from_active_connections(ws_client_authenticated, ws_client_recv_decoded, set_storage_connections_expired):
+    set_storage_connections_expired()
+    await asyncio.sleep(1.1)  # give enough time to validator to do its job
+
+    received = await ws_client_recv_decoded(ws_client_authenticated)
+
+    assert len(received) == 2
+    assert received["message_type"] == "ErrorResponse"
+    assert received["errors"] == ["Token expired, user subscriptions disabled or removed"]