"""Unit test for the WebSocket subscribe research-ownership guard. The `__handle_subscribe` handler must reject subscription attempts when the authenticated user does not own the research being subscribed to. Without this guard, any logged-in user could subscribe to a guessed or leaked research UUID and receive its progress events — a cross-user information disclosure. We exercise the `_user_owns_research` static helper directly so we do not need to stand up a Flask app or the singleton wrapper. """ from unittest.mock import patch, MagicMock from src.local_deep_research.web.services.socket_service import SocketIOService def _patched_session(row): """Return a context-manager mock that yields a DB whose first() = row.""" mock_db = MagicMock() mock_db.query.return_value.filter.return_value.first.return_value = row class _Ctx: def __enter__(self_inner): return mock_db def __exit__(self_inner, *args): return False return _Ctx() def test_owns_research_true_when_row_exists(): with patch( "src.local_deep_research.database.session_context.get_user_db_session", return_value=_patched_session(("abc",)), ): assert SocketIOService._user_owns_research("alice", "abc") is True def test_owns_research_false_when_row_missing(): """Cross-user disclosure regression: bob must not see alice's research.""" with patch( "src.local_deep_research.database.session_context.get_user_db_session", return_value=_patched_session(None), ): assert ( SocketIOService._user_owns_research("bob", "alices-research") is False ) def test_owns_research_denies_on_exception(): """A DB-open / query failure must deny, not silently allow.""" def raising(*_args, **_kwargs): raise RuntimeError("DB unavailable") with patch( "src.local_deep_research.database.session_context.get_user_db_session", raising, ): assert SocketIOService._user_owns_research("alice", "anything") is False def _patched_session_seq(rows): """Context-manager mock whose successive .first() calls yield `rows`.""" mock_db = MagicMock() mock_db.query.return_value.filter.return_value.first.side_effect = list( rows ) class _Ctx: def __enter__(self_inner): return mock_db def __exit__(self_inner, *args): return False return _Ctx() def test_owns_benchmark_run_true_for_integer_id(): """Benchmark pages subscribe with an integer BenchmarkRun.id, so the ownership gate must recognize the user's own benchmark runs — otherwise benchmark live progress is dropped. Regression: the gate previously only checked ResearchHistory (UUID ids), so an integer benchmark id never matched and the subscribe was rejected. First query (ResearchHistory) misses; second (BenchmarkRun) hits. """ with patch( "src.local_deep_research.database.session_context.get_user_db_session", return_value=_patched_session_seq([None, (5,)]), ): assert SocketIOService._user_owns_research("alice", "5") is True def test_owns_benchmark_run_false_when_not_owned(): """An integer id matching no ResearchHistory and no BenchmarkRun row in the user's own DB is rejected (no cross-user widening).""" with patch( "src.local_deep_research.database.session_context.get_user_db_session", return_value=_patched_session_seq([None, None]), ): assert SocketIOService._user_owns_research("bob", "5") is False def test_unsubscribe_rejected_when_user_does_not_own_research(): """The unsubscribe handler is symmetric with subscribe: non-owners must NOT be able to mutate the per-research subscription set. Without this guard the impact is bounded (no data exfiltration; the subscribe path is already locked down), but a malicious client could spam unsubscribes for guessed research_ids and flood the log / create lock contention. Closing the asymmetry keeps the authz boundary consistent. """ mock_app = MagicMock() mock_app.config = {"SECRET_KEY": "test-secret"} SocketIOService._instance = None with ( patch( "src.local_deep_research.web.services.socket_service.SocketIO" ) as mock_socketio_class, patch( "src.local_deep_research.web.services.socket_service.session", {"username": "alice"}, ), # Ownership check returns False → non-owner trying to unsubscribe. patch.object( SocketIOService, "_user_owns_research", return_value=False ), ): mock_socketio_class.return_value = MagicMock() service = SocketIOService(app=mock_app) # Seed an existing subscription that the attacker should NOT be # able to evict. legit_sid = "legit-owner-sid" service._SocketIOService__socket_subscriptions["target-research"] = { legit_sid } attacker_request = MagicMock() attacker_request.sid = "attacker-sid" # Drive the unsubscribe handler directly. service._SocketIOService__handle_unsubscribe( {"research_id": "target-research"}, attacker_request ) # The legit subscriber's sid must still be there, and the # attacker's sid must not have been added or removed (it was # never a subscriber to begin with). subs = service._SocketIOService__socket_subscriptions assert "target-research" in subs assert legit_sid in subs["target-research"] SocketIOService._instance = None def test_unsubscribe_allowed_when_user_owns_research(): """Owners must still be able to unsubscribe from their own research.""" mock_app = MagicMock() mock_app.config = {"SECRET_KEY": "test-secret"} SocketIOService._instance = None with ( patch( "src.local_deep_research.web.services.socket_service.SocketIO" ) as mock_socketio_class, patch( "src.local_deep_research.web.services.socket_service.session", {"username": "alice"}, ), patch.object(SocketIOService, "_user_owns_research", return_value=True), ): mock_socketio_class.return_value = MagicMock() service = SocketIOService(app=mock_app) owner_sid = "alice-sid" service._SocketIOService__socket_subscriptions["my-research"] = { owner_sid } owner_request = MagicMock() owner_request.sid = owner_sid service._SocketIOService__handle_unsubscribe( {"research_id": "my-research"}, owner_request ) # The owner's sid was removed; because it was the only subscriber, # the research_id key should also have been pruned. assert ( "my-research" not in service._SocketIOService__socket_subscriptions ) SocketIOService._instance = None