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

"""Tests for the PEP 578 sys.audit hook that gates socket.connect.

The hook is the secondary line of defense for the egress policy. These
tests pin the contract at three levels:

1. Hook plumbing (install idempotency, context get/set/clear, the
   ``active_egress_context`` manager).
2. Per-scope behaviour on real ``socket.connect`` (not via requests/
   urllib — the point of the hook is to catch raw socket use, so the
   test exercises the raw socket path it is meant to catch).
3. Thread isolation: setting a context in one thread does NOT leak to
   another.

We use 127.0.0.1 / 8.8.8.8 / 192.168.x.x as classification targets. We
never need the connect to succeed — the hook fires BEFORE the kernel
ever sees the SYN, so a refused or timed-out connect is fine; we only
check whether the hook raised ``PolicyDeniedError`` or not.
"""

from __future__ import annotations

import socket
import threading

import pytest

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


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _ctx(scope: EgressScope) -> EgressContext:
    """Build a minimal EgressContext for the given scope. No DNS cache
    or denial counter state — each test gets a fresh ctx so denial
    quotas across tests don't accumulate.
    """
    return EgressContext(
        scope=scope,
        primary_engine="wikipedia",
        require_local_llm=False,
        require_local_embeddings=False,
    )


@pytest.fixture(autouse=True)
def _clear_context_between_tests():
    """The hook reads from a process-wide thread-local. Even though
    tests run on a single thread, leaking a context across tests would
    produce confusing failures (one test's PRIVATE_ONLY context masking
    the next test's setup). Clear it before and after every test.
    """
    clear_active_context()
    yield
    clear_active_context()


def _attempt_connect(host: str, port: int = 80, timeout: float = 0.5):
    """Open a raw AF_INET socket and call connect. Returns the exception
    raised (or None if the connect somehow succeeded — unlikely against
    a port we don't expect to be open, but we return None either way).

    The kernel-level outcome (refused, timed-out, ECONNRESET) is NOT what
    we test — we only care about ``PolicyDeniedError`` raised by the
    audit hook BEFORE the SYN is sent.
    """
    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()


# ---------------------------------------------------------------------------
# 1. Hook plumbing
# ---------------------------------------------------------------------------


def test_install_audit_hook_is_idempotent():
    """Calling install_audit_hook multiple times must not stack hooks
    (PEP 578 hooks are append-only and CANNOT be removed; stacking
    would mean every connect fires the policy check N times).
    """
    install_audit_hook()
    install_audit_hook()
    install_audit_hook()
    assert is_audit_hook_installed() is True


def test_set_and_get_active_context():
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    assert get_active_context() is None
    set_active_context(ctx)
    assert get_active_context() is ctx
    clear_active_context()
    assert get_active_context() is None


def test_set_active_context_none_is_clear():
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    set_active_context(ctx)
    set_active_context(None)
    assert get_active_context() is None


def test_clear_active_context_safe_without_set():
    """Clearing a context that was never set must not raise (workers in
    finally blocks call this unconditionally).
    """
    clear_active_context()
    clear_active_context()


def test_active_egress_context_manager_clears_on_exit():
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    with active_egress_context(ctx):
        assert get_active_context() is ctx
    assert get_active_context() is None


def test_active_egress_context_clears_on_exception():
    """Cleanup must happen even when the block raises — otherwise a
    crashed worker leaks its context to whatever thread-pool task runs
    next.
    """
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    with pytest.raises(RuntimeError):
        with active_egress_context(ctx):
            assert get_active_context() is ctx
            raise RuntimeError("boom")
    assert get_active_context() is None


# ---------------------------------------------------------------------------
# 2. Per-scope behaviour on real socket.connect
# ---------------------------------------------------------------------------


def test_no_active_context_lets_everything_through():
    """The hook is install-once-and-keep-forever per PEP 578, so the
    no-context behaviour is the only thing standing between unrelated
    code (test runners, import-time helpers, scripts) and a global
    deny-everything regression.
    """
    assert get_active_context() is None
    err = _attempt_connect("8.8.8.8", 53)
    # Socket-level outcome can vary (timeout vs refused vs ECONNRESET)
    # but it MUST NOT be PolicyDeniedError.
    assert not isinstance(err, PolicyDeniedError), (
        f"hook fired without an active context: {err!r}"
    )


def test_private_only_blocks_public_ip():
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    with active_egress_context(ctx):
        err = _attempt_connect("8.8.8.8", 53)
    assert isinstance(err, PolicyDeniedError), (
        f"expected PolicyDeniedError, got {err!r}"
    )
    assert err.decision.reason == "scope_mismatch_private_only"
    assert err.target == "8.8.8.8"


def test_private_only_allows_private_ip():
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    with active_egress_context(ctx):
        err = _attempt_connect("127.0.0.1", 1)  # very likely closed
    # Whatever socket-level error happens, it must not be a policy denial.
    assert not isinstance(err, PolicyDeniedError), (
        f"hook wrongly blocked localhost under PRIVATE_ONLY: {err!r}"
    )


def test_public_only_is_a_noop_at_the_audit_hook():
    """PUBLIC_ONLY governs which SEARCH ENGINES the user picks (no
    local engines), not which hosts the process is allowed to reach at
    the socket level. Local Ollama, local embeddings, the user's own
    sqlite settings DB all use private IPs — blocking them here would
    be a false positive against the user's actual intent. The
    search-engine PEP (evaluate_engine) still gates engine selection.
    """
    ctx = _ctx(EgressScope.PUBLIC_ONLY)
    with active_egress_context(ctx):
        for host in ("8.8.8.8", "192.168.42.1", "127.0.0.1"):
            err = _attempt_connect(host, 1)
            assert not isinstance(err, PolicyDeniedError), (
                f"PUBLIC_ONLY wrongly blocked {host} at audit hook: {err!r}"
            )


def test_both_scope_is_a_noop_at_the_audit_hook():
    """BOTH is the permissive default — every classified host is fine.
    The audit hook has nothing to enforce.
    """
    ctx = _ctx(EgressScope.BOTH)
    with active_egress_context(ctx):
        for host in ("8.8.8.8", "192.168.42.1", "127.0.0.1"):
            err = _attempt_connect(host, 1)
            assert not isinstance(err, PolicyDeniedError), (
                f"BOTH scope wrongly blocked {host}: {err!r}"
            )


def test_strict_scope_blocks_public_host():
    """STRICT permits private hosts (the user's local engine) but
    refuses public ones — same contract as evaluate_url.
    """
    ctx = _ctx(EgressScope.STRICT)
    with active_egress_context(ctx):
        err = _attempt_connect("8.8.8.8", 53)
    assert isinstance(err, PolicyDeniedError)
    assert err.decision.reason == "strict_public_host"


def test_strict_scope_allows_private_host():
    ctx = _ctx(EgressScope.STRICT)
    with active_egress_context(ctx):
        err = _attempt_connect("127.0.0.1", 1)
    assert not isinstance(err, PolicyDeniedError), (
        f"STRICT wrongly blocked localhost: {err!r}"
    )


def test_ipv6_loopback_allowed_under_private_only():
    """IPv6 socket family must take the same code path. ``::1`` is the
    IPv6 loopback — has to be classified as private.
    """
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    with active_egress_context(ctx):
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.settimeout(0.5)
        try:
            sock.connect(("::1", 1))
            err = None
        except PolicyDeniedError as exc:
            err = exc
        except OSError as exc:
            err = exc
        finally:
            sock.close()
    assert not isinstance(err, PolicyDeniedError), (
        f"IPv6 loopback wrongly blocked under PRIVATE_ONLY: {err!r}"
    )


def test_af_unix_not_gated():
    """AF_UNIX is off-network. The hook must not classify or block —
    sqlite, syslog, dbus, etc. would all break otherwise.
    """
    import tempfile
    from pathlib import Path

    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    # Pick a path that does not exist so connect errors with ENOENT.
    # The point is to assert no PolicyDeniedError, not a successful
    # connection.
    tmpdir = Path(tempfile.mkdtemp())
    sock_path = str(tmpdir / "nonexistent.sock")
    try:
        with active_egress_context(ctx):
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            sock.settimeout(0.5)
            try:
                sock.connect(sock_path)
                err = None
            except PolicyDeniedError as exc:
                err = exc
            except OSError as exc:
                err = exc
            finally:
                sock.close()
        assert not isinstance(err, PolicyDeniedError), (
            f"AF_UNIX wrongly gated: {err!r}"
        )
    finally:
        tmpdir.rmdir()


# ---------------------------------------------------------------------------
# 3. Thread isolation
# ---------------------------------------------------------------------------


def test_thread_local_context_does_not_leak_across_threads():
    """Setting a PRIVATE_ONLY context in one thread must not block a
    sibling thread's public-IP connect — the thread-local design is
    what makes concurrent multi-research-per-process safe.
    """
    ctx = _ctx(EgressScope.PRIVATE_ONLY)
    set_active_context(ctx)

    sibling_result = []

    def sibling_work():
        # Sibling has no active context → connect should not raise
        # PolicyDeniedError. We capture whatever exception it does
        # raise and assert it is not a policy denial.
        err = _attempt_connect("8.8.8.8", 53)
        sibling_result.append(err)

    t = threading.Thread(target=sibling_work)
    t.start()
    t.join()

    assert sibling_result, "sibling thread did not run"
    assert not isinstance(sibling_result[0], PolicyDeniedError), (
        f"main-thread context leaked into sibling: {sibling_result[0]!r}"
    )

    # Main thread still sees the context.
    assert get_active_context() is ctx


def test_concurrent_threads_can_have_different_scopes():
    """Two threads with two different scopes must each see their own
    contract — the policy is per-thread, not per-process.
    """
    results = {}

    def worker(scope, host, label):
        ctx = _ctx(scope)
        with active_egress_context(ctx):
            results[label] = _attempt_connect(host, 1)

    # PRIVATE_ONLY thread reaches a public IP → must be blocked.
    # STRICT thread reaches a public IP → must also be blocked. Using
    # two different enforcing scopes makes sure each thread is reading
    # ITS context, not someone else's.
    t_private = threading.Thread(
        target=worker,
        args=(EgressScope.PRIVATE_ONLY, "8.8.8.8", "private"),
    )
    t_strict = threading.Thread(
        target=worker,
        args=(EgressScope.STRICT, "1.1.1.1", "strict"),
    )
    t_private.start()
    t_strict.start()
    t_private.join()
    t_strict.join()

    assert isinstance(results["private"], PolicyDeniedError), (
        f"PRIVATE_ONLY thread did not block public IP: {results['private']!r}"
    )
    assert isinstance(results["strict"], PolicyDeniedError), (
        f"STRICT thread did not block public IP: {results['strict']!r}"
    )
    # Both threads should have cleared their context on exit.
    assert get_active_context() is None


# ---------------------------------------------------------------------------
# Re-entrancy guard
# ---------------------------------------------------------------------------


class _FakeSock:
    """Minimal stand-in carrying just the .family the hook inspects."""

    def __init__(self, family=socket.AF_INET):
        self.family = family


def test_reentrancy_guard_makes_hook_a_noop_while_active():
    """The hook sets _hook_reentry.active around its own evaluate path so
    its internal DNS/socket work can't recursively re-trigger itself. When
    the flag is set, the hook must short-circuit even for a context+address
    that would otherwise be denied.

    Regression: the re-entrancy guard had no direct coverage. Without it
    the hook would recurse (or self-deny) on its own lookups.
    """
    from local_deep_research.security.egress.audit_hook import (
        _audit_hook,
        _hook_reentry,
    )

    # PRIVATE_ONLY would normally block a public IP at the hook.
    set_active_context(_ctx(EgressScope.PRIVATE_ONLY))
    args = (_FakeSock(), ("8.8.8.8", 80))

    # Guard active => no-op, even though 8.8.8.8 is public under PRIVATE_ONLY.
    _hook_reentry.active = True
    try:
        # Must NOT raise.
        _audit_hook("socket.connect", args)
    finally:
        _hook_reentry.active = False

    # Sanity: with the guard cleared, the SAME call is denied — proving the
    # guard (not some other no-op condition) is what suppressed the block.
    with pytest.raises(PolicyDeniedError):
        _audit_hook("socket.connect", args)


# ---------------------------------------------------------------------------
# Review round 2 (audit-hook correctness) — regression tests
# ---------------------------------------------------------------------------

from local_deep_research.security.egress.audit_hook import (  # noqa: E402
    _audit_hook,
    _extract_host,
)


def test_bytes_host_is_classified_not_bypassed():
    """R2 #11: CPython fires socket.connect with a bytes host; the hook must
    decode and classify it, not pass it through. A bytes metadata literal
    under PRIVATE_ONLY must be blocked."""
    assert _extract_host((b"169.254.169.254", 80)) == "169.254.169.254"
    assert _extract_host((b"10.0.0.5", 80)) == "10.0.0.5"

    set_active_context(_ctx(EgressScope.PRIVATE_ONLY))
    # bytes public host under PRIVATE_ONLY must be denied (was a bypass).
    with pytest.raises(PolicyDeniedError):
        _audit_hook("socket.connect", (_FakeSock(), (b"8.8.8.8", 80)))


def test_extract_host_undecodable_bytes_passes_through():
    """Non-ASCII bytes host (not a real IP/host) returns None — preserves the
    'failing to extract is never a reason to block' contract."""
    assert _extract_host((b"\xff\xfe", 80)) is None


def test_set_active_context_rejects_adaptive_scope():
    """R2 #1: an unresolved ADAPTIVE context must not be storable — the hook
    only gates PRIVATE_ONLY/STRICT, so ADAPTIVE would silently disarm it.
    Fail fast instead."""
    with pytest.raises(ValueError):
        set_active_context(_ctx(EgressScope.ADAPTIVE))
    # Nothing was stored.
    assert get_active_context() is None


def test_active_egress_context_nests_and_restores_parent():
    """R2 #13: a nested active_egress_context must restore the parent on exit,
    not wipe it (which would leave the parent's remaining work unprotected)."""
    parent = _ctx(EgressScope.PRIVATE_ONLY)
    child = _ctx(EgressScope.STRICT)
    with active_egress_context(parent):
        assert get_active_context() is parent
        with active_egress_context(child):
            assert get_active_context() is child
        # Parent must be restored, not cleared.
        assert get_active_context() is parent
    assert get_active_context() is None