learningcircuit--local-deep-research/tests/security/test_egress_threadpool_propagation.py

"""Integration tests: egress audit-hook backstop propagation into worker threads.

These are NOT PDP-unit tests. They drive the REAL runtime wiring that re-arms
the PEP-578 socket audit hook inside worker threads, because the hook reads its
active context from a ``threading.local`` that stdlib ``ThreadPoolExecutor``
workers do NOT inherit. Two real call-site mechanisms are exercised:

1. ``utilities.thread_context.preserve_research_context`` — the decorator the
   global search pool wraps tasks with. It captures the submitter thread's
   egress context at decoration time and re-arms/clears it per task.
2. The hand-rolled capture-then-rearm pattern used by ``news_strategy`` and
   ``parallel_search_engine`` (``get_active_context()`` on submit, then
   ``active_egress_context()`` / ``set_active_context()`` + finally-clear in the
   worker).

The security property under test: under an armed PRIVATE_ONLY context, a raw
``socket.connect`` to a public IP must raise ``PolicyDeniedError`` *inside the
worker* — and must NOT when the context was never propagated (proving the
re-arm is what supplies the protection, not ambient state). The hook fires
BEFORE the SYN, so no test depends on a real outbound connection completing.

Each test uses an allow+deny pair so it fails if the propagation were reverted.
Armed context is always cleared in a finally so the thread-local does not leak
to other tests.
"""

from __future__ import annotations

import concurrent.futures
import socket
import threading

import pytest

from local_deep_research.security import (
    active_egress_context,
    clear_active_context,
    get_active_context,
    install_audit_hook,
    set_active_context,
)
from local_deep_research.security.egress.policy import (
    EgressContext,
    EgressScope,
    PolicyDeniedError,
)
from local_deep_research.utilities.thread_context import (
    preserve_research_context,
)


def _ctx(scope: EgressScope) -> EgressContext:
    """Minimal EgressContext for a given scope (fresh per test, no shared
    denial-quota or DNS-cache state)."""
    return EgressContext(
        scope=scope,
        primary_engine="wikipedia",
        require_local_llm=False,
        require_local_embeddings=False,
    )


def _attempt_connect(host: str, port: int = 53, timeout: float = 0.3):
    """Raw AF_INET connect; return the exception raised (or None on success).

    Under an armed enforcing context the audit hook raises PolicyDeniedError
    BEFORE any network I/O, so the kernel-level outcome for the allow-cases
    (refused/timeout) is irrelevant — we only classify whether the policy
    fired.
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    try:
        sock.connect((host, port))
        return None
    except PolicyDeniedError as exc:
        return exc
    except OSError as exc:
        return exc
    finally:
        sock.close()


@pytest.fixture(autouse=True)
def _hook_and_clean_context():
    """Hook must be installed (it is at security import time, but be explicit
    and idempotent) and the thread-local must be clean around each test so a
    leaked PRIVATE_ONLY context can't mask another test's setup."""
    install_audit_hook()
    clear_active_context()
    yield
    clear_active_context()


_PUBLIC_HOST = "8.8.8.8"
_PRIVATE_HOST = "127.0.0.1"


# ---------------------------------------------------------------------------
# preserve_research_context (the global-pool decorator)
# ---------------------------------------------------------------------------


def test_preserve_research_context_rearms_hook_in_worker_and_blocks_public():
    """A task wrapped with preserve_research_context, submitted to a stdlib
    ThreadPoolExecutor under an armed PRIVATE_ONLY context, must (a) run on a
    DIFFERENT thread, (b) see the SAME armed context inside the worker, and
    (c) have a public socket.connect blocked by the hook in that worker.
    """
    armed = _ctx(EgressScope.PRIVATE_ONLY)
    main_tid = threading.get_ident()
    set_active_context(armed)
    try:
        # Decoration happens here, WHILE armed — this is where the egress
        # context is captured off the submitter thread.
        def _worker():
            return {
                "tid": threading.get_ident(),
                "ctx": get_active_context(),
                "err": _attempt_connect(_PUBLIC_HOST),
            }

        wrapped = preserve_research_context(_worker)

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            result = ex.submit(wrapped).result()
    finally:
        clear_active_context()

    assert result["tid"] != main_tid, "task did not run on a worker thread"
    assert result["ctx"] is armed, (
        "wrapper did not re-arm the submitter's egress context in the worker"
    )
    assert isinstance(result["err"], PolicyDeniedError), (
        f"hook did not fire in re-armed worker: {result['err']!r}"
    )
    assert result["err"].decision.reason == "scope_mismatch_private_only"


def test_unwrapped_worker_has_no_context_so_hook_is_inactive():
    """Deny-pair to the test above: the SAME public connect, the SAME armed
    main thread, but the task is NOT wrapped. Because threading.local is not
    inherited by pool workers, the worker has no active context and the hook
    is inactive there — demonstrating exactly why the wrapper is required.
    """
    armed = _ctx(EgressScope.PRIVATE_ONLY)
    set_active_context(armed)
    try:

        def _worker():
            return {
                "ctx": get_active_context(),
                "err": _attempt_connect(_PUBLIC_HOST),
            }

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            result = ex.submit(_worker).result()  # unwrapped on purpose
    finally:
        clear_active_context()

    assert result["ctx"] is None, (
        "unexpected ambient context in fresh worker — test is not isolating "
        "the propagation mechanism"
    )
    assert not isinstance(result["err"], PolicyDeniedError), (
        "hook fired in a worker that never received the context — would mean "
        "the context leaked rather than being explicitly propagated"
    )


def test_preserve_research_context_clears_context_so_reused_worker_is_clean():
    """The wrapper must clear the egress context in its finally so a reused
    pool worker does not inherit PRIVATE_ONLY into the next (possibly
    unrelated) task. Run a wrapped armed task, then a bare probe on the SAME
    single-worker executor and assert the probe lands on the same thread with
    no context and an un-gated public connect.
    """
    armed = _ctx(EgressScope.PRIVATE_ONLY)
    set_active_context(armed)
    try:

        def _armed_task():
            return {
                "tid": threading.get_ident(),
                "err": _attempt_connect(_PUBLIC_HOST),
            }

        wrapped = preserve_research_context(_armed_task)

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            first = ex.submit(wrapped).result()

            # Bare probe (NOT wrapped) reused on the same idle worker.
            def _probe():
                return {
                    "tid": threading.get_ident(),
                    "ctx": get_active_context(),
                    "err": _attempt_connect(_PUBLIC_HOST),
                }

            second = ex.submit(_probe).result()
    finally:
        clear_active_context()

    # The armed task really was protected (otherwise the "clear" claim is moot).
    assert isinstance(first["err"], PolicyDeniedError)
    # Same worker thread reused.
    assert second["tid"] == first["tid"], (
        "probe ran on a different worker — cannot prove no per-worker leak"
    )
    assert second["ctx"] is None, "wrapper leaked egress context to next task"
    assert not isinstance(second["err"], PolicyDeniedError), (
        "PRIVATE_ONLY leaked into the reused worker's next task"
    )


# ---------------------------------------------------------------------------
# Capture-then-rearm pattern (news_strategy / parallel_search_engine)
# ---------------------------------------------------------------------------


def _run_news_style(submit_armed: bool, host: str):
    """Replicate news_strategy: capture get_active_context() on the submitter
    thread, then re-arm it in the worker via active_egress_context(). Returns
    the connect outcome observed inside the worker.
    """
    if submit_armed:
        set_active_context(_ctx(EgressScope.PRIVATE_ONLY))
    try:
        # Capture on the submitter thread (mirrors news_strategy line 262).
        captured = get_active_context()

        def _analyze():
            with active_egress_context(captured):
                return _attempt_connect(host)

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            return ex.submit(_analyze).result()
    finally:
        clear_active_context()


def test_news_strategy_capture_then_rearm_blocks_public_under_private_only():
    """news_strategy's capture+active_egress_context re-arm must block a public
    connect in the worker when armed PRIVATE_ONLY, and must NOT block when the
    submitter had no context (captured is None → no-op)."""
    blocked = _run_news_style(submit_armed=True, host=_PUBLIC_HOST)
    assert isinstance(blocked, PolicyDeniedError), (
        f"re-armed worker failed to block public host: {blocked!r}"
    )
    assert blocked.decision.reason == "scope_mismatch_private_only"

    # Deny-pair: nothing captured → re-arm is a no-op → not blocked.
    allowed = _run_news_style(submit_armed=False, host=_PUBLIC_HOST)
    assert not isinstance(allowed, PolicyDeniedError), (
        f"hook fired with no captured context (active_egress_context(None) "
        f"should be a no-op): {allowed!r}"
    )


def test_news_strategy_rearm_still_allows_private_host_under_private_only():
    """Re-arming PRIVATE_ONLY in the worker must not over-block: a private
    (loopback) target stays allowed, same contract as on the main thread."""
    err = _run_news_style(submit_armed=True, host=_PRIVATE_HOST)
    assert not isinstance(err, PolicyDeniedError), (
        f"re-armed PRIVATE_ONLY wrongly blocked a private host: {err!r}"
    )


def test_parallel_engine_capture_set_then_finally_clear_no_worker_leak():
    """Replicate parallel_search_engine._run_with_context: capture the egress
    context on submit, set_active_context() in the worker, and clear it in a
    finally. Verify (a) the worker blocks a public connect while armed, and
    (b) the finally-clear leaves the reused worker clean for the next task.
    """
    set_active_context(_ctx(EgressScope.PRIVATE_ONLY))
    try:
        submitter_egress_ctx = get_active_context()

        def _run_with_context(host):
            # Mirrors the engine's wrapper: set on entry, clear in finally.
            if submitter_egress_ctx is not None:
                set_active_context(submitter_egress_ctx)
            try:
                return {
                    "tid": threading.get_ident(),
                    "ctx": get_active_context(),
                    "err": _attempt_connect(host),
                }
            finally:
                if submitter_egress_ctx is not None:
                    clear_active_context()

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            armed_result = ex.submit(_run_with_context, _PUBLIC_HOST).result()

            # Bare probe reused on the same worker — must see no leak.
            def _probe():
                return {
                    "tid": threading.get_ident(),
                    "ctx": get_active_context(),
                    "err": _attempt_connect(_PUBLIC_HOST),
                }

            probe_result = ex.submit(_probe).result()
    finally:
        clear_active_context()

    assert armed_result["ctx"] is submitter_egress_ctx
    assert isinstance(armed_result["err"], PolicyDeniedError), (
        f"engine-style worker failed to block public host: "
        f"{armed_result['err']!r}"
    )
    assert probe_result["tid"] == armed_result["tid"], (
        "probe ran on a different worker — cannot prove the finally-clear"
    )
    assert probe_result["ctx"] is None, (
        "engine wrapper's finally-clear failed: context leaked to next task"
    )
    assert not isinstance(probe_result["err"], PolicyDeniedError), (
        "PRIVATE_ONLY leaked into reused worker via missing finally-clear"
    )