deer-flow/backend/tests/test_aio_sandbox_network_proxy.py

from __future__ import annotations

import asyncio
import json
import socket
import ssl
from unittest.mock import MagicMock

import pytest

from deerflow.community.aio_sandbox import network_proxy


def test_domain_matches_exact_and_leading_wildcard_only() -> None:
    assert network_proxy.domain_matches("pypi.org", "pypi.org")
    assert not network_proxy.domain_matches("evilpypi.org", "pypi.org")
    assert network_proxy.domain_matches("files.pythonhosted.org", "*.pythonhosted.org")
    assert not network_proxy.domain_matches("pythonhosted.org", "*.pythonhosted.org")


def test_address_is_public_rejects_host_private_link_local_and_metadata() -> None:
    for address in (
        "127.0.0.1",
        "10.0.0.2",
        "172.16.0.2",
        "192.168.1.2",
        "169.254.169.254",
        "224.0.0.1",
        "::1",
        "fc00::1",
        "fec0::1",
        "fe80::1",
        "ff0e::1",
    ):
        assert not network_proxy.address_is_public(address)
    assert network_proxy.address_is_public("8.8.8.8")
    assert not network_proxy.address_is_public("198.18.1.5")
    assert network_proxy.address_is_public("198.18.1.5", allow_synthetic_dns=True)


def test_policy_denial_and_temporary_or_sandbox_grants(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    monkeypatch.setenv("DEERFLOW_NETWORK_MODE", "allowlist")
    monkeypatch.setenv("DEERFLOW_ALLOW_DOMAINS_JSON", json.dumps(["pypi.org"]))

    assert network_proxy.policy_allows("pypi.org", 443, now=100)
    assert not network_proxy.policy_allows("example.com", 443, now=100)

    temporary = network_proxy.record_denial("example.com", 443, "CONNECT")
    assert network_proxy.decide(temporary, "allow_temporary", ttl=60)
    assert network_proxy.policy_allows("example.com", 443)

    sandbox = network_proxy.record_denial("files.example.net", 443, "CONNECT")
    assert network_proxy.decide(sandbox, "allow_sandbox", ttl=60)
    assert network_proxy.policy_allows("files.example.net", 443, now=10**12)


def test_pending_events_are_consumed_once(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    request_id = network_proxy.record_denial("example.com", 443, "CONNECT")

    events = network_proxy.pending_events()

    assert events == [
        {
            "request_id": request_id,
            "host": "example.com",
            "port": 443,
            "method": "CONNECT",
            "created_at": events[0]["created_at"],
        }
    ]
    assert network_proxy.pending_events() == []


def test_pending_events_surface_only_one_destination_per_approval(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    first = network_proxy.record_denial("one.example", 443, "CONNECT")
    network_proxy.record_denial("two.example", 443, "CONNECT")

    assert [event["request_id"] for event in network_proxy.pending_events()] == [first]
    # The sibling is superseded so a retry can create a fresh approvable event.
    assert network_proxy.pending_events() == []
    fresh = network_proxy.record_denial("two.example", 443, "CONNECT")
    assert [event["request_id"] for event in network_proxy.pending_events()] == [fresh]


def test_pending_events_supersede_every_sibling_without_a_batch_limit(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    request_ids = [network_proxy.record_denial(f"host-{index}.example", 443, "CONNECT") for index in range(17)]

    assert [event["request_id"] for event in network_proxy.pending_events()] == [request_ids[0]]
    assert network_proxy.pending_events() == []


def test_deny_pending_events_atomically_denies_every_unsurfaced_event(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    for index in range(17):
        network_proxy.record_denial(f"host-{index}.example", 443, "CONNECT")

    assert network_proxy.deny_pending_events() == 17
    assert network_proxy.pending_events() == []

    with network_proxy._connect_db() as db:
        assert db.execute("SELECT COUNT(*) FROM events WHERE decision = 'deny'").fetchone() == (17,)


def test_deny_pending_events_preserves_an_already_surfaced_user_decision(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    surfaced = network_proxy.record_denial("interactive.example", 443, "CONNECT")
    assert [event["request_id"] for event in network_proxy.pending_events()] == [surfaced]
    unsurfaced = network_proxy.record_denial("scheduled.example", 443, "CONNECT")

    assert network_proxy.deny_pending_events() == 1
    assert network_proxy.decide(surfaced, "allow_temporary", ttl=60)
    assert network_proxy.decide(unsurfaced, "deny", ttl=60)


def test_pending_events_claims_old_unsurfaced_denial_on_retry(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    now = [100.0]
    monkeypatch.setattr(network_proxy.time, "time", lambda: now[0])
    request_id = network_proxy.record_denial("late.example", 443, "CONNECT")

    now[0] = 105.0
    assert network_proxy.record_denial("late.example", 443, "CONNECT") == request_id
    assert [event["request_id"] for event in network_proxy.pending_events()] == [request_id]


@pytest.mark.anyio
async def test_resolve_public_fails_closed_when_dns_contains_private_answer(monkeypatch) -> None:
    loop = __import__("asyncio").get_running_loop()

    async def fake_getaddrinfo(*_args, **_kwargs):
        return [
            (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 443)),
            (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.5", 443)),
        ]

    monkeypatch.setattr(loop, "getaddrinfo", fake_getaddrinfo)
    assert await network_proxy.resolve_public("example.com", 443) is None


@pytest.mark.anyio
async def test_resolve_public_returns_every_validated_answer_and_open_retries(monkeypatch) -> None:
    loop = asyncio.get_running_loop()
    answers = [
        (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 443)),
        (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("1.1.1.1", 443)),
    ]

    async def fake_getaddrinfo(*_args, **_kwargs):
        return answers

    attempts: list[str] = []
    connected = (MagicMock(), MagicMock())

    async def fake_open_connection(host: str, _port: int, *, family: int):
        attempts.append(host)
        assert family == socket.AF_INET
        if host == "8.8.8.8":
            raise OSError("first address unavailable")
        return connected

    monkeypatch.setattr(loop, "getaddrinfo", fake_getaddrinfo)
    monkeypatch.setattr(network_proxy.asyncio, "open_connection", fake_open_connection)

    resolved = await network_proxy.resolve_public("example.com", 443)

    assert resolved == (
        (socket.AF_INET, ("8.8.8.8", 443)),
        (socket.AF_INET, ("1.1.1.1", 443)),
    )
    assert await network_proxy._open_public(resolved, 443) is connected
    assert attempts == ["8.8.8.8", "1.1.1.1"]


@pytest.mark.anyio
@pytest.mark.parametrize("mode", ["isolated", "allowlist"])
async def test_denied_destination_is_rejected_without_dns_resolution(tmp_path, monkeypatch, mode: str) -> None:
    monkeypatch.setattr(network_proxy, "POLICY_DB", tmp_path / "policy.sqlite3")
    monkeypatch.setenv("DEERFLOW_NETWORK_MODE", mode)
    monkeypatch.setenv("DEERFLOW_ALLOW_DOMAINS_JSON", json.dumps(["allowed.example"]))
    monkeypatch.delenv("DEERFLOW_RECORD_DENIALS", raising=False)
    resolutions: list[tuple[str, int]] = []

    async def fake_resolve_public(host: str, port: int):
        resolutions.append((host, port))
        return None

    monkeypatch.setattr(network_proxy, "resolve_public", fake_resolve_public)

    proxy = await asyncio.start_server(network_proxy.handle_proxy, "127.0.0.1", 0)
    proxy_port = proxy.sockets[0].getsockname()[1]
    try:
        reader, writer = await asyncio.open_connection("127.0.0.1", proxy_port)
        writer.write(b"CONNECT denied.example:443 HTTP/1.1\r\nHost: denied.example:443\r\n\r\n")
        await writer.drain()
        response = await asyncio.wait_for(reader.read(), timeout=2)
        writer.close()
        await writer.wait_closed()

        assert b"403 Forbidden" in response
        assert resolutions == []
    finally:
        proxy.close()
        await proxy.wait_closed()


@pytest.mark.anyio
async def test_tls_client_hello_sni_is_extracted_for_connect_enforcement() -> None:
    incoming = ssl.MemoryBIO()
    outgoing = ssl.MemoryBIO()
    context = ssl.create_default_context()
    tls = context.wrap_bio(incoming, outgoing, server_side=False, server_hostname="pypi.org")
    with pytest.raises(ssl.SSLWantReadError):
        tls.do_handshake()

    reader = __import__("asyncio").StreamReader()
    wire = outgoing.read()
    reader.feed_data(wire)
    reader.feed_eof()

    parsed = await network_proxy._read_tls_client_hello(reader)

    assert parsed == ("pypi.org", wire)


def test_http_request_framing_rejects_ambiguous_or_duplicate_lengths() -> None:
    with pytest.raises(ValueError, match="cannot be combined"):
        fields = network_proxy._parse_http_header_fields(["Content-Length: 4", "Transfer-Encoding: chunked"])
        network_proxy._http_request_body_framing(fields)
    with pytest.raises(ValueError, match="one non-negative"):
        fields = network_proxy._parse_http_header_fields(["Content-Length: 4", "Content-Length: 4"])
        network_proxy._http_request_body_framing(fields)


@pytest.mark.anyio
async def test_chunked_request_body_rejects_non_hex_size() -> None:
    reader = asyncio.StreamReader()
    reader.feed_data(b"+1\r\na\r\n0\r\n\r\n")
    reader.feed_eof()
    writer = MagicMock()

    with pytest.raises(ValueError, match="Invalid chunk size"):
        await network_proxy._copy_chunked_request_body(reader, writer)


@pytest.mark.anyio
async def test_http_proxy_relays_exactly_one_request_per_connection(monkeypatch) -> None:
    received: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()

    async def upstream_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        header = await reader.readuntil(b"\r\n\r\n")
        body = await reader.readexactly(4)
        received.set_result(header + body)
        writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
        await writer.drain()
        writer.close()
        await writer.wait_closed()

    upstream = await asyncio.start_server(upstream_handler, "127.0.0.1", 0)
    upstream_port = upstream.sockets[0].getsockname()[1]

    async def fake_resolve_public(_host: str, _port: int):
        return (socket.AF_INET, ("127.0.0.1", upstream_port))

    async def fake_open_public(_resolved, _port: int):
        return await asyncio.open_connection("127.0.0.1", upstream_port)

    monkeypatch.setattr(network_proxy, "resolve_public", fake_resolve_public)
    monkeypatch.setattr(network_proxy, "_open_public", fake_open_public)
    monkeypatch.setattr(network_proxy, "policy_allows", lambda host, port: (host, port) == ("allowed.example", 80))

    proxy = await asyncio.start_server(network_proxy.handle_proxy, "127.0.0.1", 0)
    proxy_port = proxy.sockets[0].getsockname()[1]
    try:
        reader, writer = await asyncio.open_connection("127.0.0.1", proxy_port)
        writer.write(b"POST http://allowed.example/first HTTP/1.1\r\nHost: allowed.example\r\nContent-Length: 4\r\nConnection: keep-alive\r\n\r\ndataGET http://denied.example/second HTTP/1.1\r\nHost: denied.example\r\n\r\n")
        await writer.drain()
        response = await asyncio.wait_for(reader.read(), timeout=2)
        writer.close()
        await writer.wait_closed()

        upstream_request = await asyncio.wait_for(received, timeout=2)
        assert b"POST /first HTTP/1.1" in upstream_request
        assert b"Connection: close" in upstream_request
        assert b"Connection: keep-alive" not in upstream_request
        assert upstream_request.endswith(b"data")
        assert b"denied.example" not in upstream_request
        assert b"200 OK" in response
    finally:
        proxy.close()
        upstream.close()
        await proxy.wait_closed()
        await upstream.wait_closed()


@pytest.mark.anyio
async def test_sandbox_api_relay_requires_per_sandbox_token(monkeypatch) -> None:
    received: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()

    async def upstream_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        header = await reader.readuntil(b"\r\n\r\n")
        received.set_result(header)
        writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
        await writer.drain()
        writer.close()
        await writer.wait_closed()

    upstream = await asyncio.start_server(upstream_handler, "127.0.0.1", 0)
    upstream_port = upstream.sockets[0].getsockname()[1]
    monkeypatch.setenv(network_proxy.RELAY_TOKEN_ENV, "test-relay-token")
    monkeypatch.setenv("DEERFLOW_SANDBOX_TARGET", f"127.0.0.1:{upstream_port}")
    relay = await asyncio.start_server(network_proxy.handle_relay, "127.0.0.1", 0)
    relay_port = relay.sockets[0].getsockname()[1]
    try:
        denied_reader, denied_writer = await asyncio.open_connection("127.0.0.1", relay_port)
        denied_writer.write(b"GET /v1/sandbox HTTP/1.1\r\nHost: sandbox\r\n\r\n")
        await denied_writer.drain()
        denied_response = await asyncio.wait_for(denied_reader.read(), timeout=2)
        denied_writer.close()
        await denied_writer.wait_closed()

        assert b"403 Forbidden" in denied_response
        assert not received.done()

        reader, writer = await asyncio.open_connection("127.0.0.1", relay_port)
        writer.write(b"GET /v1/sandbox HTTP/1.1\r\nHost: sandbox\r\n" + f"{network_proxy.RELAY_AUTH_HEADER}: test-relay-token\r\n\r\n".encode())
        await writer.drain()
        response = await asyncio.wait_for(reader.read(), timeout=2)
        writer.close()
        await writer.wait_closed()

        assert b"200 OK" in response
        assert network_proxy.RELAY_AUTH_HEADER.encode() in await asyncio.wait_for(received, timeout=2)
    finally:
        relay.close()
        upstream.close()
        await relay.wait_closed()
        await upstream.wait_closed()


@pytest.mark.anyio
async def test_sandbox_api_relay_rejects_non_ascii_token(monkeypatch) -> None:
    monkeypatch.setenv(network_proxy.RELAY_TOKEN_ENV, "test-relay-token")
    relay = await asyncio.start_server(network_proxy.handle_relay, "127.0.0.1", 0)
    relay_port = relay.sockets[0].getsockname()[1]
    try:
        reader, writer = await asyncio.open_connection("127.0.0.1", relay_port)
        writer.write(b"GET /v1/sandbox HTTP/1.1\r\nHost: sandbox\r\n" + network_proxy.RELAY_AUTH_HEADER.encode() + b": \xff\r\n\r\n")
        await writer.drain()
        response = await asyncio.wait_for(reader.read(), timeout=2)
        writer.close()
        await writer.wait_closed()

        assert b"403 Forbidden" in response
    finally:
        relay.close()
        await relay.wait_closed()


@pytest.mark.anyio
@pytest.mark.parametrize(
    "malformed_header",
    [
        b"Host : denied.example",
        b"Transfer-Encoding : chunked",
        b"Bad(Header): value",
    ],
)
async def test_http_proxy_rejects_ambiguous_field_names_before_policy_check(monkeypatch, malformed_header: bytes) -> None:
    policy_checks: list[tuple[str, int]] = []

    def fake_policy_allows(host: str, port: int) -> bool:
        policy_checks.append((host, port))
        return False

    monkeypatch.setattr(network_proxy, "policy_allows", fake_policy_allows)

    proxy = await asyncio.start_server(network_proxy.handle_proxy, "127.0.0.1", 0)
    proxy_port = proxy.sockets[0].getsockname()[1]
    try:
        reader, writer = await asyncio.open_connection("127.0.0.1", proxy_port)
        writer.write(b"GET http://allowed.example/ HTTP/1.1\r\nHost: allowed.example\r\n" + malformed_header + b"\r\n\r\n")
        await writer.drain()
        response = await asyncio.wait_for(reader.read(), timeout=2)
        writer.close()
        await writer.wait_closed()

        assert b"400 Bad Request" in response
        assert policy_checks == []
    finally:
        proxy.close()
        await proxy.wait_closed()