learningcircuit--local-deep-research/tests/web/test_socket_subscribe_ownership.py

"""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