|
|
@@ -0,0 +1,172 @@
|
|
|
+"""Taking websocket:connect away ends the live connection it allowed.
|
|
|
+
|
|
|
+Only minting a WebSocket token checked the permission: the handshake and the
|
|
|
+scope refresh did not. A user whose group lost it kept their open dashboard
|
|
|
+(their printers' live status and every event of no printer), and their token
|
|
|
+reconnected for an hour.
|
|
|
+"""
|
|
|
+
|
|
|
+from types import SimpleNamespace
|
|
|
+from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
+
|
|
|
+import pytest
|
|
|
+from fastapi import WebSocketDisconnect
|
|
|
+from httpx import AsyncClient
|
|
|
+
|
|
|
+from backend.app.core import database
|
|
|
+from backend.app.core.printer_scope import ALL_PRINTERS
|
|
|
+from backend.app.core.websocket import ws_manager
|
|
|
+from backend.app.models.api_key import APIKey
|
|
|
+
|
|
|
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
|
|
+
|
|
|
+
|
|
|
+@pytest.fixture
|
|
|
+async def admin(async_client: AsyncClient) -> dict:
|
|
|
+ await async_client.post(
|
|
|
+ "/api/v1/auth/setup",
|
|
|
+ json={"auth_enabled": True, "admin_username": "wsadmin", "admin_password": "AdminPass1!"},
|
|
|
+ )
|
|
|
+ login = await async_client.post("/api/v1/auth/login", json={"username": "wsadmin", "password": "AdminPass1!"})
|
|
|
+ return {"Authorization": f"Bearer {login.json()['access_token']}"}
|
|
|
+
|
|
|
+
|
|
|
+async def _carol_in_group(client: AsyncClient, admin: dict) -> tuple[int, int]:
|
|
|
+ group = await client.post(
|
|
|
+ "/api/v1/groups/", headers=admin, json={"name": "Live", "permissions": ["websocket:connect", "printers:read"]}
|
|
|
+ )
|
|
|
+ assert group.status_code == 201, group.text
|
|
|
+ user = await client.post(
|
|
|
+ "/api/v1/users/",
|
|
|
+ headers=admin,
|
|
|
+ json={"username": "carol", "password": "Userpass1!", "role": "user", "group_ids": [group.json()["id"]]},
|
|
|
+ )
|
|
|
+ assert user.status_code == 201, user.text
|
|
|
+ return group.json()["id"], user.json()["id"]
|
|
|
+
|
|
|
+
|
|
|
+def _socket(username: str | None, user_id: int | None, key_id: int | None = None, prefix: str | None = None):
|
|
|
+ socket = MagicMock()
|
|
|
+ socket.state = SimpleNamespace(
|
|
|
+ bambuddy_scope_principal=(username if username else "", key_id),
|
|
|
+ bambuddy_principal_user_id=user_id,
|
|
|
+ bambuddy_api_key_prefix=prefix,
|
|
|
+ bambuddy_printer_scope=ALL_PRINTERS,
|
|
|
+ )
|
|
|
+ socket.close = AsyncMock()
|
|
|
+ return socket
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.parametrize("keeps", [False, True])
|
|
|
+async def test_a_group_losing_the_permission_closes_its_users_sockets(async_client: AsyncClient, admin, keeps):
|
|
|
+ group, carol = await _carol_in_group(async_client, admin)
|
|
|
+ socket = _socket("carol", carol)
|
|
|
+ ws_manager.active_connections.append(socket)
|
|
|
+ permissions = ["websocket:connect", "printers:read"] if keeps else ["printers:read"]
|
|
|
+ try:
|
|
|
+ with patch("backend.app.core.database.async_session", database.async_session):
|
|
|
+ response = await async_client.patch(
|
|
|
+ f"/api/v1/groups/{group}", headers=admin, json={"permissions": permissions}
|
|
|
+ )
|
|
|
+
|
|
|
+ assert response.status_code == 200, response.text
|
|
|
+ if keeps:
|
|
|
+ socket.close.assert_not_awaited()
|
|
|
+ else:
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ finally:
|
|
|
+ if socket in ws_manager.active_connections:
|
|
|
+ ws_manager.active_connections.remove(socket)
|
|
|
+
|
|
|
+
|
|
|
+async def test_a_user_removed_from_the_granting_group_loses_the_socket(async_client: AsyncClient, admin):
|
|
|
+ group, carol = await _carol_in_group(async_client, admin)
|
|
|
+ socket = _socket("carol", carol)
|
|
|
+ ws_manager.active_connections.append(socket)
|
|
|
+ try:
|
|
|
+ response = await async_client.delete(f"/api/v1/groups/{group}/users/{carol}", headers=admin)
|
|
|
+
|
|
|
+ assert response.status_code == 204, response.text
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ finally:
|
|
|
+ if socket in ws_manager.active_connections:
|
|
|
+ ws_manager.active_connections.remove(socket)
|
|
|
+
|
|
|
+
|
|
|
+async def _connect(principal: tuple):
|
|
|
+ from backend.app.api.routes import websocket as ws_route
|
|
|
+
|
|
|
+ socket = MagicMock()
|
|
|
+ socket.state = SimpleNamespace()
|
|
|
+ socket.accept = AsyncMock()
|
|
|
+ socket.close = AsyncMock()
|
|
|
+ socket.send_json = AsyncMock()
|
|
|
+ socket.receive_json = AsyncMock(side_effect=WebSocketDisconnect())
|
|
|
+ with (
|
|
|
+ patch.object(ws_route, "is_auth_enabled", AsyncMock(return_value=True)),
|
|
|
+ patch.object(ws_route, "verify_websocket_token_principal", AsyncMock(return_value=principal)),
|
|
|
+ patch.object(ws_route, "principal_printer_scope", AsyncMock(return_value=ALL_PRINTERS)),
|
|
|
+ patch.object(ws_route, "async_session", database.async_session),
|
|
|
+ ):
|
|
|
+ await ws_route.websocket_endpoint(socket, token="t")
|
|
|
+ return socket
|
|
|
+
|
|
|
+
|
|
|
+async def test_a_token_whose_user_lost_the_permission_is_refused(async_client: AsyncClient, admin):
|
|
|
+ """Minted while the user had it, still valid for an hour."""
|
|
|
+ group, _ = await _carol_in_group(async_client, admin)
|
|
|
+ await async_client.patch(f"/api/v1/groups/{group}", headers=admin, json={"permissions": ["printers:read"]})
|
|
|
+
|
|
|
+ socket = await _connect(("carol", None))
|
|
|
+
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ socket.accept.assert_not_awaited()
|
|
|
+
|
|
|
+
|
|
|
+async def test_a_token_of_a_key_without_status_access_is_refused(async_client: AsyncClient, admin, db_session):
|
|
|
+ key = APIKey(name="k", key_hash="h", key_prefix="bb_nostat", can_read_status=False)
|
|
|
+ db_session.add(key)
|
|
|
+ await db_session.commit()
|
|
|
+
|
|
|
+ socket = await _connect(("", key.id))
|
|
|
+
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ socket.accept.assert_not_awaited()
|
|
|
+
|
|
|
+
|
|
|
+async def test_a_key_losing_status_access_closes_its_sockets(async_client: AsyncClient, admin, db_session):
|
|
|
+ key = APIKey(name="k", key_hash="h", key_prefix="bb_stat")
|
|
|
+ db_session.add(key)
|
|
|
+ await db_session.commit()
|
|
|
+ socket = _socket(None, None, key_id=key.id, prefix="bb_stat")
|
|
|
+ ws_manager.active_connections.append(socket)
|
|
|
+ try:
|
|
|
+ with patch("backend.app.core.database.async_session", database.async_session):
|
|
|
+ response = await async_client.patch(
|
|
|
+ f"/api/v1/api-keys/{key.id}", headers=admin, json={"can_read_status": False}
|
|
|
+ )
|
|
|
+
|
|
|
+ assert response.status_code == 200, response.text
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ finally:
|
|
|
+ if socket in ws_manager.active_connections:
|
|
|
+ ws_manager.active_connections.remove(socket)
|
|
|
+
|
|
|
+
|
|
|
+async def test_the_permission_taken_while_a_socket_joins_closes_it(async_client: AsyncClient, admin):
|
|
|
+ """The group change's refresh ran before the socket was listed."""
|
|
|
+ from backend.app.api.routes import websocket as ws_route
|
|
|
+
|
|
|
+ await _carol_in_group(async_client, admin)
|
|
|
+ calls = []
|
|
|
+ real = ws_route.principal_may_connect
|
|
|
+
|
|
|
+ async def held_at_connect_then_gone(db, username, key_id):
|
|
|
+ calls.append(1)
|
|
|
+ return await real(db, username, key_id) if len(calls) == 1 else False
|
|
|
+
|
|
|
+ with patch.object(ws_route, "principal_may_connect", held_at_connect_then_gone):
|
|
|
+ socket = await _connect(("carol", None))
|
|
|
+
|
|
|
+ socket.close.assert_awaited_once_with(code=4401)
|
|
|
+ assert socket not in ws_manager.active_connections
|