LangBot/plugins/misc/URLSummary/tests/test_outbound_safety.py

"""Outbound-safety regression tests for the URLSummary listener.

These cover the two blocking certification findings:

* F1 - the URL posted in chat was an unvalidated server-side request target and
  TLS verification was switched off. The tests prove that loopback /
  link-local / private / metadata ranges, special addresses and non-HTTP(S)
  schemes are refused before a connection is made, that every redirect hop is
  re-validated, that the connection goes to the validated IP (so a later DNS
  answer cannot rebind it) and that TLS keeps the default verification
  settings.
* F2 - neither the download, the memory nor the parse work was bounded, and the
  synchronous parse ran on the shared event loop. The tests prove that a
  declared oversized body is refused without reading it, a streamed/compressed
  body is cut off at the cap without buffering it all, the markup handed to the
  parser is capped, the configured length is range-checked, and the fetch runs
  in a worker thread with a bounded number of in-flight fetches and cooperative
  cancellation.

Every assertion here fails against the pre-fix ``aiohttp`` implementation.
"""

from __future__ import annotations

import asyncio
import concurrent.futures
import gzip
import http.client
import logging
import socket
import ssl
import sys
import threading
import time
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

from langbot_plugin.api.entities.builtin.platform import message as platform_message

import pytest

PLUGIN_ROOT = Path(__file__).resolve().parents[1]
if str(PLUGIN_ROOT) not in sys.path:
    sys.path.insert(0, str(PLUGIN_ROOT))

import main as main_mod  # noqa: E402
import safe_fetch  # noqa: E402
from components.event_listener import url_detector as url_detector_module  # noqa: E402
from components.event_listener.url_detector import URLDetector  # noqa: E402
from main import URLSummary  # noqa: E402
from safe_fetch import (  # noqa: E402
    BlockedTargetError,
    FetchAbortedError,
    FetchError,
    ResponseTooLargeError,
    TruncatedResponseError,
)

PUBLIC_IP = "93.184.216.34"


# --------------------------------------------------------------------------- #
# Test HTTP server
# --------------------------------------------------------------------------- #
class _Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):  # noqa: N802 - http.server API
        owner = self.server.owner
        owner.hits.append(self.path)
        responder = owner.responses.get(self.path)
        if responder is None:
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        responder(self)

    def log_message(self, *args):  # keep the test output quiet
        pass


class _Server:
    """Local HTTP server whose per-path behaviour is set by the test."""

    def __init__(self) -> None:
        self.responses: dict[str, object] = {}
        self.hits: list[str] = []
        self._server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
        self._server.owner = self
        self.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)

    def __enter__(self) -> "_Server":
        self._thread.start()
        return self

    def __exit__(self, *exc) -> None:
        self._server.shutdown()
        self._server.server_close()
        self._thread.join(5)

    def url(self, path: str) -> str:
        return f"http://{PUBLIC_IP}:{self.port}{path}"


def _send(handler, body: bytes, *, status: int = 200, ctype: str = "text/html; charset=utf-8", headers=()):
    handler.send_response(status)
    handler.send_header("Content-Type", ctype)
    handler.send_header("Content-Length", str(len(body)))
    for name, value in headers:
        handler.send_header(name, value)
    handler.end_headers()
    handler.wfile.write(body)


def _route_connections(monkeypatch, port: int, attempted: list | None = None):
    """Send every validated connection to the local test server.

    The address the client decided to use is still recorded (``attempted``), so
    a test can prove which target passed validation.
    """
    real_create_connection = socket.create_connection

    def fake_create_connection(address, *args, **kwargs):
        if attempted is not None:
            attempted.append(address)
        return real_create_connection(("127.0.0.1", port), *args, **kwargs)

    monkeypatch.setattr(safe_fetch.socket, "create_connection", fake_create_connection)


def _page(plugin: URLSummary, server: "_Server", path: str, max_len: int = 8000):
    return asyncio.run(plugin.fetch_page(server.url(path), max_len))


_SCRIPTED_PEERS: list = []
"""Scripted peers stay open for the whole session: closing one while the client
still has unread data makes Windows reset the connection (WSAECONNABORTED)."""


def _scripted_peer(monkeypatch, raw_response: bytes, *, on_connect=None):
    """Answer every connection with ``raw_response`` from a real socketpair.

    No HTTP server is involved: the peer sends the scripted bytes and then stops
    sending.  That reproduces, on every platform, the clean EOF a shut-down
    socket shows on Linux, which is what makes the abort/truncation checks below
    platform-independent.
    """

    def fake_create_connection(address, *args, **kwargs):
        peer, client = socket.socketpair()
        peer.sendall(raw_response)
        peer.shutdown(socket.SHUT_WR)  # the peer simply stops sending
        _SCRIPTED_PEERS.append(peer)
        if on_connect is not None:
            on_connect()
        return client

    monkeypatch.setattr(safe_fetch.socket, "create_connection", fake_create_connection)


def _make_context(text: str) -> SimpleNamespace:
    context = SimpleNamespace(
        event=SimpleNamespace(message_chain=[platform_message.Plain(text=text)]),
        replies=[],
    )
    context.prevent_default = lambda: None
    context.prevent_postorder = lambda: None

    async def reply(message_chain):
        context.replies.append(message_chain)

    context.reply = reply
    return context


# --------------------------------------------------------------------------- #
# F1 - target validation
# --------------------------------------------------------------------------- #
def test_loopback_literal_is_refused_before_any_connection(monkeypatch):
    plugin = URLSummary()

    with _Server() as server:
        server.responses["/"] = lambda handler: _send(handler, b"<html><title>x</title>ok</html>")
        _route_connections(monkeypatch, server.port)

        with pytest.raises(BlockedTargetError):
            asyncio.run(plugin.fetch_page(f"http://127.0.0.1:{server.port}/", 8000))

        assert server.hits == []


def test_listener_refuses_a_loopback_url_end_to_end(monkeypatch):
    plugin = URLSummary()
    listener = URLDetector()
    listener.plugin = plugin
    plugin.get_config = lambda: {"max_content_length": 8000, "language": "zh_Hans", "model": "m"}

    with _Server() as server:
        server.responses["/"] = lambda handler: _send(handler, b"<html>ok</html>")
        context = _make_context(f"sum this http://127.0.0.1:{server.port}/ please")
        asyncio.run(listener._handle_message(context))
        assert server.hits == []

    assert context.replies == []


@pytest.mark.parametrize(
    "host",
    [
        "127.0.0.1",
        "127.1.2.3",
        "169.254.169.254",
        "10.0.0.1",
        "172.16.0.1",
        "192.168.1.1",
        "100.64.0.1",
        "0.0.0.0",
        "[::1]",
        "[fd00::1]",
        "[fe80::1]",
        "[::ffff:127.0.0.1]",
    ],
)
def test_non_public_targets_are_refused_before_connecting(host: str, monkeypatch):
    attempted: list = []

    def fail_connect(*args, **kwargs):
        attempted.append(args)
        pytest.fail("the client must not open a connection for a blocked target")

    monkeypatch.setattr(safe_fetch.socket, "create_connection", fail_connect)

    with pytest.raises(BlockedTargetError):
        safe_fetch.fetch_document(f"http://{host}/")

    assert attempted == []


@pytest.mark.parametrize("url", ["file:///etc/passwd", "ftp://example.com/x", "gopher://example.com/", "data:text/html,x"])
def test_non_http_schemes_are_refused(url: str, monkeypatch):
    monkeypatch.setattr(
        safe_fetch.socket, "create_connection", lambda *a, **k: pytest.fail("no connection expected")
    )

    with pytest.raises(BlockedTargetError):
        safe_fetch.fetch_document(url)


@pytest.mark.parametrize("url", ["http://[::1/", "http://", "https:///path"])
def test_malformed_urls_are_refused(url: str, monkeypatch):
    monkeypatch.setattr(
        safe_fetch.socket, "create_connection", lambda *a, **k: pytest.fail("no connection expected")
    )

    with pytest.raises(BlockedTargetError):
        safe_fetch.fetch_document(url)


def test_redirect_to_private_address_is_refused_and_never_connected(monkeypatch):
    plugin = URLSummary()
    attempted: list = []

    with _Server() as server:
        server.responses["/start"] = lambda handler: _send(
            handler,
            b"",
            status=302,
            headers=[("Location", "http://169.254.169.254/latest/meta-data/")],
        )
        _route_connections(monkeypatch, server.port, attempted)

        with pytest.raises(BlockedTargetError) as excinfo:
            _page(plugin, server, "/start")

        assert "169.254.169.254" in str(excinfo.value)
        assert server.hits == ["/start"]
        # Only the validated first hop was dialled; the redirect target was not.
        assert [address[0] for address in attempted] == [PUBLIC_IP]


def test_connection_uses_the_validated_ip_so_dns_cannot_rebind(monkeypatch):
    resolutions: list[str] = []

    def fake_getaddrinfo(host, port, **kwargs):
        resolutions.append(host)
        address = "93.184.216.34" if len(resolutions) == 1 else "127.0.0.1"
        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (address, port))]

    attempted: list = []

    def fake_create_connection(address, *args, **kwargs):
        attempted.append(address)
        raise OSError("unreachable in this test")

    monkeypatch.setattr(safe_fetch.socket, "getaddrinfo", fake_getaddrinfo)
    monkeypatch.setattr(safe_fetch.socket, "create_connection", fake_create_connection)

    with pytest.raises(FetchError):
        safe_fetch.fetch_document("http://rebind.example/")

    # The socket went to the address that was validated, never to the second
    # (private) answer a rebinding resolver would hand out.
    assert attempted == [("93.184.216.34", 80)]


def test_tls_uses_the_default_verification_context(monkeypatch):
    plugin = URLSummary()
    created: list = []
    hostnames: list = []
    handshakes: list = []
    deferred: list = []
    real_context = safe_fetch.ssl.create_default_context

    class _PlainTLS:
        """Minimal SSLSocket double: a plain transport plus a no-op handshake."""

        def __init__(self, sock):
            self._sock = sock

        def do_handshake(self):
            handshakes.append("handshake")

        def settimeout(self, value):
            self._sock.settimeout(value)

        def sendall(self, data):
            self._sock.sendall(data)

        def makefile(self, *args, **kwargs):
            return self._sock.makefile(*args, **kwargs)

        def shutdown(self, *args):
            self._sock.shutdown(*args)

        def close(self):
            self._sock.close()

        def fileno(self):
            return self._sock.fileno()

    class _RecordingContext:
        def __init__(self, context):
            self.context = context

        def wrap_socket(self, sock, server_hostname=None, do_handshake_on_connect=True, **kwargs):
            hostnames.append(server_hostname)
            # The handshake is deferred so the abort watch covers it too.
            deferred.append(do_handshake_on_connect is False)
            return _PlainTLS(sock)

    def fake_create_default_context(*args, **kwargs):
        context = real_context(*args, **kwargs)
        created.append(context)
        return _RecordingContext(context)

    with _Server() as server:
        server.responses["/"] = lambda handler: _send(handler, b"<html><title>T</title>secure</html>")
        _route_connections(monkeypatch, server.port)
        monkeypatch.setattr(safe_fetch.ssl, "create_default_context", fake_create_default_context)

        title, text = asyncio.run(
            plugin.fetch_page(f"https://{PUBLIC_IP}:{server.port}/", 8000)
        )

        assert title == "T" and "secure" in text

    assert created and created[0].verify_mode == ssl.CERT_REQUIRED
    assert created[0].check_hostname is True
    assert hostnames == [PUBLIC_IP]
    assert deferred == [True]
    assert handshakes == ["handshake"]


def test_public_addresses_are_recognised():
    assert safe_fetch.is_public_address("8.8.8.8")
    assert safe_fetch.is_public_address("2606:4700::1111")
    assert not safe_fetch.is_public_address("127.0.0.1")
    assert not safe_fetch.is_public_address("169.254.169.254")
    assert not safe_fetch.is_public_address("::ffff:127.0.0.1")
    assert not safe_fetch.is_public_address("example.com")


# --------------------------------------------------------------------------- #
# F2 - bounded download, memory and parse work
# --------------------------------------------------------------------------- #
def test_oversized_declared_body_is_rejected_without_reading_it(monkeypatch):
    plugin = URLSummary()

    def huge_body(handler):
        handler.send_response(200)
        handler.send_header("Content-Type", "text/html")
        handler.send_header("Content-Length", "100000000")
        handler.end_headers()
        handler.wfile.write(b"<html>")
        handler.wfile.flush()

    def forbidden(*args, **kwargs):
        pytest.fail("the body must not be read once Content-Length exceeds the cap")

    monkeypatch.setattr(safe_fetch, "_raw_chunks", forbidden)

    with _Server() as server:
        server.responses["/huge"] = huge_body
        _route_connections(monkeypatch, server.port)

        with pytest.raises(ResponseTooLargeError) as excinfo:
            _page(plugin, server, "/huge")

    assert "100000000" in str(excinfo.value)


def test_streamed_body_is_capped_without_buffering_it_all():
    pulled: list[int] = []

    def chunks():
        for index in range(100):
            pulled.append(index)
            yield b"A" * (1024 * 1024)

    with pytest.raises(ResponseTooLargeError):
        safe_fetch.collect_bounded(chunks(), 8192)

    # One chunk is already over the cap: the remaining 99 MiB are never pulled.
    assert pulled == [0]


def test_chunked_response_is_capped(monkeypatch):
    plugin = URLSummary()

    def chunked_body(handler):
        handler.send_response(200)
        handler.send_header("Content-Type", "text/html")
        handler.send_header("Transfer-Encoding", "chunked")
        handler.end_headers()
        block = b"B" * (256 * 1024)
        try:
            for _ in range(40):  # 10 MiB with no Content-Length to pre-check
                handler.wfile.write(b"%X\r\n" % len(block) + block + b"\r\n")
            handler.wfile.write(b"0\r\n\r\n")
        except OSError:
            pass

    with _Server() as server:
        server.responses["/chunked"] = chunked_body
        _route_connections(monkeypatch, server.port)

        with pytest.raises(ResponseTooLargeError):
            _page(plugin, server, "/chunked")


def test_compressed_body_is_capped_after_inflation(monkeypatch):
    plugin = URLSummary()
    bomb = gzip.compress(b"A" * (8 * 1024 * 1024))
    assert len(bomb) < safe_fetch.MAX_RESPONSE_BYTES  # the wire form passes the byte cap

    with _Server() as server:
        server.responses["/bomb"] = lambda handler: _send(
            handler, bomb, headers=[("Content-Encoding", "gzip")]
        )
        _route_connections(monkeypatch, server.port)

        with pytest.raises(ResponseTooLargeError) as excinfo:
            _page(plugin, server, "/bomb")

    assert "decompressed" in str(excinfo.value)


def test_gzip_inflation_stops_before_the_whole_stream_is_read():
    blob = gzip.compress(b"A" * (8 * 1024 * 1024))
    blocks = [blob[index : index + 4096] for index in range(0, len(blob), 4096)]
    pulled: list[bytes] = []

    def source():
        for block in blocks:
            pulled.append(block)
            yield block

    decoder = safe_fetch._BodyDecoder("gzip", 8192)
    with pytest.raises(ResponseTooLargeError):
        safe_fetch.collect_bounded(decoder.decode(source()), 8192)

    assert 0 < len(pulled) < len(blocks)


def test_parse_input_is_capped(monkeypatch):
    plugin = URLSummary()
    body = ("<html><head><title>Big</title></head><body>" + "word " * 200000 + "</body></html>").encode()

    monkeypatch.setattr(safe_fetch, "MAX_PARSE_CHARS", 4096)
    fed_sizes: list[int] = []
    real_extractor = main_mod._TextExtractor

    class _SpyExtractor(real_extractor):
        def feed(self, data):
            fed_sizes.append(len(data))
            return super().feed(data)

    with _Server() as server:
        server.responses["/big"] = lambda handler: _send(handler, body)
        _route_connections(monkeypatch, server.port)

        document = safe_fetch.fetch_document(server.url("/big"))
        assert len(document.text) > 4096  # the page really is larger than the cap

        monkeypatch.setattr(main_mod, "_TextExtractor", _SpyExtractor)
        title, text = _page(plugin, server, "/big", max_len=2000)

    assert title == "Big"
    assert len(text) == 2000
    assert fed_sizes and max(fed_sizes) <= 4096


def test_configured_content_length_is_range_checked(monkeypatch):
    assert url_detector_module._clamp_max_content_length(-1) == 1
    assert url_detector_module._clamp_max_content_length(0) == 1
    assert url_detector_module._clamp_max_content_length(8000) == 8000
    assert (
        url_detector_module._clamp_max_content_length(10**9)
        == url_detector_module.MAX_CONTENT_LENGTH
    )
    assert (
        url_detector_module._clamp_max_content_length("nonsense")
        == url_detector_module.DEFAULT_MAX_CONTENT_LENGTH
    )

    # ...and a hostile configuration value never reaches the fetch unbounded.
    plugin = URLSummary()
    plugin.get_config = lambda: {"max_content_length": 10**9, "model": "m", "language": "en_US"}
    seen: list[int] = []

    async def fake_fetch(url: str, max_len: int):
        seen.append(max_len)
        return "Title", "x" * 120

    async def fake_summarize(url, title, content, model_uuid, language):
        return "summary"

    monkeypatch.setattr(plugin, "fetch_page", fake_fetch)
    monkeypatch.setattr(plugin, "summarize", fake_summarize)

    listener = URLDetector()
    listener.plugin = plugin
    with _Server() as server:
        server.responses["/page"] = lambda handler: _send(handler, b"<html>ok</html>")
        context = _make_context(f"look {server.url('/page')}")
        asyncio.run(listener._handle_message(context))

    assert seen == [url_detector_module.MAX_CONTENT_LENGTH]


def test_fetch_deadline_stops_a_slow_peer():
    finished = threading.Event()

    def slow_body(handler):
        handler.send_response(200)
        handler.send_header("Content-Type", "text/html")
        handler.send_header("Content-Length", "1000000")
        handler.end_headers()
        try:
            for _ in range(100):
                handler.wfile.write(b"x" * 1024)
                handler.wfile.flush()
                time.sleep(0.1)
        except OSError:
            pass

    errors: list[Exception] = []
    returned: list[object] = []

    with _Server() as server:
        server.responses["/slow"] = slow_body
        with mock.patch.object(safe_fetch.socket, "create_connection", _connect_to(server.port)):

            def run():
                try:
                    returned.append(safe_fetch.fetch_document(server.url("/slow"), timeout=1.0))
                except Exception as exc:  # noqa: BLE001 - recorded for the assertion
                    errors.append(exc)
                finally:
                    finished.set()

            thread = threading.Thread(target=run, daemon=True)
            start = time.monotonic()
            thread.start()
            assert finished.wait(10), "the fetch did not honour its deadline"
            elapsed = time.monotonic() - start

    assert returned == [], "a timed-out fetch must not be returned as a document"
    assert errors and isinstance(errors[0], FetchAbortedError)
    assert elapsed < 5


def test_stop_event_aborts_an_in_flight_download():
    """A cancelled in-flight fetch ends promptly and never returns a document.

    The outcome must not depend on how the platform reports the shut-down
    socket (exception vs clean EOF): either way the caller sees an abort, never
    a partial document.
    """
    started = threading.Event()
    finished = threading.Event()
    errors: list[Exception] = []
    returned: list[object] = []
    stop_event = threading.Event()

    def slow_body(handler):
        handler.send_response(200)
        handler.send_header("Content-Type", "text/html")
        handler.send_header("Content-Length", "1000000")
        handler.end_headers()
        try:
            for _ in range(400):
                handler.wfile.write(b"x" * 1024)
                handler.wfile.flush()
                time.sleep(0.05)
        except OSError:
            pass

    with _Server() as server:
        server.responses["/slow"] = slow_body
        with mock.patch.object(safe_fetch.socket, "create_connection", _connect_to(server.port)):

            def run():
                started.set()
                try:
                    returned.append(
                        safe_fetch.fetch_document(
                            server.url("/slow"), timeout=30, stop_event=stop_event
                        )
                    )
                except Exception as exc:  # noqa: BLE001
                    errors.append(exc)
                finally:
                    finished.set()

            thread = threading.Thread(target=run, daemon=True)
            thread.start()
            assert started.wait(5)
            time.sleep(0.4)
            stop_event.set()

    assert finished.wait(10), "the stop flag did not interrupt the in-flight download"
    assert returned == [], "an aborted fetch must not be returned as a document"
    assert errors and isinstance(errors[0], FetchAbortedError)


def test_an_abort_that_reads_as_a_clean_eof_never_returns_a_body(monkeypatch):
    """The Linux symptom: a shut-down socket reads as EOF mid-body.

    The abort lands while the reader is blocked in the socket read, and that
    read comes back as a plain EOF.  The fetch must abort instead of handing
    the partial body back as a complete response.
    """
    raw = (
        b"HTTP/1.1 200 OK\r\n"
        b"Content-Type: text/html\r\n"
        b"Content-Length: 1000000\r\n"
        b"\r\n"
        b"<html>partial"
    )
    stop_event = threading.Event()
    real_read1 = http.client.HTTPResponse.read1
    reads: list[int] = []

    def read1(self, n=-1):
        chunk = real_read1(self, n)
        reads.append(len(chunk))
        if not chunk:
            # The cancel/deadline fires while this read is parked in the socket.
            stop_event.set()
        return chunk

    _scripted_peer(monkeypatch, raw)
    monkeypatch.setattr(http.client.HTTPResponse, "read1", read1)

    with pytest.raises(FetchAbortedError):
        safe_fetch.fetch_document(
            f"http://{PUBLIC_IP}/partial", timeout=30, stop_event=stop_event
        )

    assert reads and reads[0] > 0, "the body must really have started before the EOF"


def test_a_body_shorter_than_content_length_is_never_returned(monkeypatch):
    """A peer that stops early must fail loudly, not look like a page."""
    raw = (
        b"HTTP/1.1 200 OK\r\n"
        b"Content-Type: text/html\r\n"
        b"Content-Length: 1000000\r\n"
        b"\r\n"
        b"<html>partial"
    )
    _scripted_peer(monkeypatch, raw)

    with pytest.raises(TruncatedResponseError) as excinfo:
        safe_fetch.fetch_document(f"http://{PUBLIC_IP}/partial")

    assert "truncated" in str(excinfo.value)
    assert "13 of 1000000" in str(excinfo.value)


def _connect_to(port: int):
    """Replacement for ``socket.create_connection`` aimed at the test server."""
    real_create_connection = socket.create_connection

    def connect(address, *args, **kwargs):
        return real_create_connection(("127.0.0.1", port), *args, **kwargs)

    return connect


def test_cancelling_the_fetch_sets_the_worker_stop_flag(monkeypatch):
    observed: list[bool] = []
    done = threading.Event()

    def fake_guarded_fetch(url, max_len, stop_event, lease, *_rest):
        while not stop_event.is_set():
            time.sleep(0.01)
        observed.append(stop_event.is_set())
        done.set()
        return ("t", "x")

    async def scenario():
        monkeypatch.setattr(main_mod, "_guarded_fetch", fake_guarded_fetch)
        plugin = URLSummary()
        task = asyncio.create_task(plugin.fetch_page("https://slow.example/", 100))
        await asyncio.sleep(0.1)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task

    asyncio.run(scenario())
    assert asyncio.run(asyncio.to_thread(done.wait, 5)), "worker never observed the cancellation"
    assert observed == [True]


def test_concurrent_fetches_are_bounded(monkeypatch):
    monkeypatch.setattr(main_mod, "FETCH_SLOT_WAIT_SECONDS", 0.25)
    gate = threading.Event()
    slots = main_mod.MAX_CONCURRENT_FETCHES

    def fake_guarded_fetch(url, max_len, stop_event, lease, *_rest):
        gate.wait(10)
        try:
            return ("t", url)
        finally:
            lease.release()

    async def scenario():
        monkeypatch.setattr(main_mod, "_guarded_fetch", fake_guarded_fetch)
        plugin = URLSummary()
        tasks = [
            asyncio.create_task(plugin.fetch_page(f"https://example.com/{index}", 100))
            for index in range(slots + 1)
        ]
        await asyncio.sleep(0.1)
        busy = await asyncio.gather(tasks[-1], return_exceptions=True)
        gate.set()
        return busy[0], await asyncio.gather(*tasks[:-1])

    busy, finished = asyncio.run(scenario())

    assert isinstance(busy, Exception)
    assert "too many concurrent page fetches" in str(busy)
    assert finished == [("t", f"https://example.com/{index}") for index in range(slots)]


def test_two_fetches_keep_no_shared_state(monkeypatch):
    plugin = URLSummary()
    before = set(vars(plugin))

    with _Server() as server:
        server.responses["/first.html"] = lambda handler: _send(
            handler, b"<html><head><title>First</title></head><body>alpha content</body></html>"
        )
        server.responses["/second.html"] = lambda handler: _send(
            handler, b"<html><head><title>Second</title></head><body>beta content</body></html>"
        )
        _route_connections(monkeypatch, server.port)
        first = _page(plugin, server, "/first.html", 100)
        second = _page(plugin, server, "/second.html", 100)

    assert first[0] == "First" and "alpha content" in first[1]
    assert second[0] == "Second" and "beta content" in second[1]
    assert set(vars(plugin)) == before
# --------------------------------------------------------------------------- #
# F1 (third round) - the shared fetch slot has exactly one owner
# --------------------------------------------------------------------------- #
def test_fetch_cancelled_while_queued_releases_its_slot_exactly_once(monkeypatch):
    """A fetch cancelled before its worker starts must not leak the slot."""
    gate = threading.Event()
    worker_ran = threading.Event()

    def fake_guarded_fetch(url, max_len, stop_event, lease, *_rest):
        worker_ran.set()
        try:
            return ("t", url)
        finally:
            lease.release()

    def occupy(*_args):
        gate.wait(10)
        return ("t", "blocker")

    async def scenario():
        monkeypatch.setattr(main_mod, "_guarded_fetch", fake_guarded_fetch)
        loop = asyncio.get_running_loop()
        loop.set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=1))
        slots = main_mod._slots()
        free = slots._value
        plugin = URLSummary()

        # Occupy the single worker thread so the fetch below stays queued.
        blocker = loop.run_in_executor(None, occupy, "u", 1, None, slots, loop)
        await asyncio.sleep(0.1)

        task = asyncio.create_task(plugin.fetch_page("https://example.com/queued", 100))
        await asyncio.sleep(0.3)
        assert not worker_ran.is_set(), "the worker must still be queued"
        assert slots._value == free - 1, "the queued call holds its slot"

        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        # The queued worker has not ended, so its slot must not be back yet.
        assert slots._value == free - 1, "the slot was released while still queued/running"

        gate.set()
        assert await blocker == ("t", "blocker")
        for _ in range(250):
            if worker_ran.is_set() and slots._value == free:
                break
            await asyncio.sleep(0.02)
        assert worker_ran.is_set(), "the queued worker never ran"
        # Exactly one release: a second one would push the pool above its cap.
        assert slots._value == free

        # ...and the pool is still usable by the next caller.
        assert await plugin.fetch_page("https://example.com/next", 100) == (
            "t",
            "https://example.com/next",
        )

    asyncio.run(scenario())


def test_cancelled_running_fetch_keeps_its_slot_until_the_worker_ends(monkeypatch):
    """A worker that is still running must never lose its slot early."""
    started = threading.Event()
    still_running = threading.Event()

    def fake_guarded_fetch(url, max_len, stop_event, lease, *_rest):
        started.set()
        stop_event.wait(10)  # the caller's cancellation sets this
        still_running.set()
        time.sleep(0.3)  # the thread is undeniably still alive here
        try:
            return ("t", url)
        finally:
            lease.release()

    async def scenario():
        monkeypatch.setattr(main_mod, "_guarded_fetch", fake_guarded_fetch)
        slots = main_mod._slots()
        free = slots._value
        plugin = URLSummary()

        task = asyncio.create_task(plugin.fetch_page("https://slow.example/", 100))
        await asyncio.sleep(0.3)
        assert started.is_set()
        assert slots._value == free - 1

        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        await asyncio.sleep(0.1)
        assert still_running.is_set()
        assert slots._value == free - 1, "the slot was released while the worker was running"

        for _ in range(250):
            if slots._value == free:
                break
            await asyncio.sleep(0.02)
        assert slots._value == free

    asyncio.run(scenario())


def test_cancel_during_the_body_frees_the_slot_and_stays_usable(monkeypatch):
    """An aborted in-flight fetch ends, frees exactly one slot and stays usable."""
    started = threading.Event()

    def slow_body(handler):
        handler.send_response(200)
        handler.send_header("Content-Type", "text/html")
        handler.send_header("Content-Length", "1000000")
        handler.end_headers()
        try:
            for _ in range(400):  # far longer than the test: only the abort ends it
                handler.wfile.write(b"x" * 1024)
                handler.wfile.flush()
                started.set()
                time.sleep(0.05)
        except OSError:
            pass

    async def scenario(port: int):
        slots = main_mod._slots()
        free = slots._value
        plugin = URLSummary()

        task = asyncio.create_task(
            plugin.fetch_page(f"http://{PUBLIC_IP}:{port}/slow", 500)
        )
        await asyncio.to_thread(started.wait, 5)
        assert slots._value == free - 1

        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        for _ in range(300):
            if slots._value == free:
                break
            await asyncio.sleep(0.02)
        assert slots._value == free, "the aborted fetch did not release exactly one slot"

        return await plugin.fetch_page(f"http://{PUBLIC_IP}:{port}/ok", 500)

    with _Server() as server:
        server.responses["/slow"] = slow_body
        server.responses["/ok"] = lambda handler: _send(
            handler, b"<html><head><title>Ok</title></head><body>after abort</body></html>"
        )
        _route_connections(monkeypatch, server.port)
        title, text = asyncio.run(scenario(server.port))

    assert title == "Ok" and "after abort" in text


def test_slot_lease_releases_at_most_once():
    """The two owners of a slot can never give it back twice."""

    async def scenario():
        slots = main_mod._slots()
        free = slots._value
        await slots.acquire()
        assert slots._value == free - 1

        lease = main_mod._SlotLease(slots, asyncio.get_running_loop())
        lease.release()
        lease.release()  # a second owner must be a no-op
        assert slots._value == free

    asyncio.run(scenario())


def test_fetch_submit_failure_leaks_no_slot(monkeypatch):
    """A failed executor submission must give the already-acquired slot back."""

    async def scenario():
        slots = main_mod._slots()
        free = slots._value
        loop = asyncio.get_running_loop()

        def refuse_submit(*args, **kwargs):
            raise RuntimeError("cannot schedule new futures after shutdown")

        monkeypatch.setattr(loop, "run_in_executor", refuse_submit)
        plugin = URLSummary()
        with pytest.raises(RuntimeError):
            await plugin.fetch_page("https://example.com/", 100)
        assert slots._value == free

    asyncio.run(scenario())


def test_deadline_ends_a_response_whose_headers_never_arrive():
    """A peer that drips bytes must not outlive the absolute deadline."""
    finished = threading.Event()
    errors: list[Exception] = []

    def drip_headers(handler):
        # One byte at a time, always inside the socket idle timeout, never a
        # complete status line: ``getresponse()`` would block here forever.
        try:
            for _ in range(120):
                handler.wfile.write(b"x")
                handler.wfile.flush()
                time.sleep(0.05)
        except OSError:
            pass

    with _Server() as server:
        server.responses["/drip"] = drip_headers
        with mock.patch.object(safe_fetch.socket, "create_connection", _connect_to(server.port)):

            def run():
                try:
                    safe_fetch.fetch_document(server.url("/drip"), timeout=1.0)
                except Exception as exc:  # noqa: BLE001 - recorded for the assertion
                    errors.append(exc)
                finally:
                    finished.set()

            thread = threading.Thread(target=run, daemon=True)
            start = time.monotonic()
            thread.start()
            assert finished.wait(8), "the header read outlived the absolute deadline"
            elapsed = time.monotonic() - start

    assert isinstance(errors[0], FetchAbortedError)
    assert elapsed < 5


def test_stop_event_ends_a_response_whose_headers_never_arrive():
    """Cancellation must end the header phase too, not only the body."""
    started = threading.Event()
    finished = threading.Event()
    errors: list[Exception] = []
    stop_event = threading.Event()

    def drip_headers(handler):
        try:
            for _ in range(120):
                handler.wfile.write(b"x")
                handler.wfile.flush()
                time.sleep(0.05)
        except OSError:
            pass

    with _Server() as server:
        server.responses["/drip"] = drip_headers
        with mock.patch.object(safe_fetch.socket, "create_connection", _connect_to(server.port)):

            def run():
                started.set()
                try:
                    safe_fetch.fetch_document(
                        server.url("/drip"), timeout=30, stop_event=stop_event
                    )
                except Exception as exc:  # noqa: BLE001
                    errors.append(exc)
                finally:
                    finished.set()

            thread = threading.Thread(target=run, daemon=True)
            thread.start()
            assert started.wait(5)
            time.sleep(0.4)
            stop_event.set()

    assert finished.wait(8), "the stop flag did not end the header read"
    assert isinstance(errors[0], FetchAbortedError)


def test_deadline_abandons_a_resolver_that_never_answers(monkeypatch):
    """A stalled DNS lookup must not hold the fetch past the deadline."""
    unblock = threading.Event()

    def stalled_getaddrinfo(*args, **kwargs):
        unblock.wait(8)
        return []

    monkeypatch.setattr(safe_fetch.socket, "getaddrinfo", stalled_getaddrinfo)
    start = time.monotonic()
    try:
        with pytest.raises(FetchAbortedError):
            safe_fetch.fetch_document("http://stalled.example/", timeout=1.0)
        elapsed = time.monotonic() - start
    finally:
        unblock.set()
    assert elapsed < 5


# --------------------------------------------------------------------------- #
# F2 (third round) - error paths and logs never repeat the raw URL
# --------------------------------------------------------------------------- #
def _credential_url(port: int) -> str:
    return f"http://alice:[email protected]:{port}/page?token=s3cret-token#frag"


def _exception_text(exc: BaseException) -> str:
    return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))


def _assert_no_credentials(text: str) -> None:
    assert "s3cret-pw" not in text, text
    assert "s3cret-token" not in text, text
    assert "alice" not in text, text


def test_redaction_helpers_drop_userinfo_query_and_fragment():
    assert (
        safe_fetch.redact_url("http://alice:[email protected]:8443/a/b?token=xyz#frag")
        == "http://host.example:8443/a/b"
    )
    # A malformed port must not break the scrub (or leak the userinfo).
    assert safe_fetch.redact_url("http://host:bad/x?t=1") == "http://host/x"
    assert safe_fetch.redact_text("failed http://u:p@h/x?token=s") == "failed http://h/x"


def test_connection_errors_never_repeat_url_credentials_or_query(monkeypatch):
    def fail_connect(*args, **kwargs):
        raise OSError("connection refused")

    monkeypatch.setattr(safe_fetch.socket, "create_connection", fail_connect)

    with pytest.raises(FetchError) as excinfo:
        safe_fetch.fetch_document(_credential_url(1))

    text = "\n".join(
        (str(excinfo.value), repr(excinfo.value), _exception_text(excinfo.value))
    )
    _assert_no_credentials(text)
    # The useful, non-secret part of the target is still reported.
    assert "93.184.216.34" in text


def test_chained_exception_text_is_scrubbed_too(monkeypatch):
    raw = _credential_url(8080)

    def boom(self, *args, **kwargs):
        raise http.client.HTTPException(f"server rejected {raw}")

    monkeypatch.setattr(safe_fetch.http.client.HTTPConnection, "request", boom)

    with _Server() as server:
        server.responses["/page"] = lambda handler: _send(handler, b"<html>ok</html>")
        with mock.patch.object(safe_fetch.socket, "create_connection", _connect_to(server.port)):
            with pytest.raises(FetchError) as excinfo:
                safe_fetch.fetch_document(_credential_url(server.port))

    _assert_no_credentials(_exception_text(excinfo.value))


def test_listener_logs_never_contain_url_credentials_or_query(monkeypatch, caplog):
    raw = _credential_url(8080)

    async def failing_fetch_page(url, max_len):
        raise RuntimeError(f"boom while fetching {url}")

    plugin = URLSummary()
    plugin.get_config = lambda: {"max_content_length": 8000, "language": "zh_Hans", "model": "m"}
    monkeypatch.setattr(plugin, "fetch_page", failing_fetch_page)

    listener = URLDetector()
    listener.plugin = plugin
    context = _make_context(f"please summarize {raw}")

    with caplog.at_level(logging.WARNING, logger="URLSummary.detector"):
        asyncio.run(listener._handle_message(context))

    assert "Failed to summarize" in caplog.text
    _assert_no_credentials(caplog.text)


def test_listener_error_traceback_is_scrubbed(monkeypatch, caplog):
    raw = _credential_url(8080)

    def failing_config():
        raise RuntimeError(f"config failed for {raw}")

    plugin = URLSummary()
    plugin.get_config = failing_config

    listener = URLDetector()
    listener.plugin = plugin
    context = _make_context(f"please summarize {raw}")

    with caplog.at_level(logging.ERROR, logger="URLSummary.detector"):
        asyncio.run(listener._handle_message(context))

    assert "Error in _handle_message" in caplog.text
    _assert_no_credentials(caplog.text)


# --------------------------------------------------------------------------- #
# F1 (fourth round) - an abandoned lookup keeps its own bounded budget
# --------------------------------------------------------------------------- #
def _gated_resolver(gate, calls):
    """``getaddrinfo`` stand-in that blocks until ``gate`` is set."""

    def getaddrinfo(host, port, **kwargs):
        calls.append(host)
        gate.wait(10)
        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (PUBLIC_IP, port))]

    return getaddrinfo


def _dns_threads():
    return [t for t in threading.enumerate() if t.name == "safe-fetch-dns"]


def _drain_resolver_pool(timeout=5.0):
    deadline = time.monotonic() + timeout
    while safe_fetch.resolver_stats()["in_flight"] and time.monotonic() < deadline:
        time.sleep(0.02)
    return safe_fetch.resolver_stats()


def test_abandoned_waiter_does_not_free_the_resolver_budget(monkeypatch):
    """A timed-out waiter must not return budget the resolver still uses."""
    assert _drain_resolver_pool()["in_flight"] == 0
    gate = threading.Event()
    monkeypatch.setattr(safe_fetch.socket, "getaddrinfo", _gated_resolver(gate, []))
    try:
        for index in range(safe_fetch.MAX_DNS_OUTSTANDING):
            started = time.monotonic()
            with pytest.raises(FetchAbortedError):
                safe_fetch.fetch_document(f"http://slow{index}.example/", timeout=0.3)
            # The waiter returned on time even though the lookup never ended.
            assert time.monotonic() - started < 5

        stats = safe_fetch.resolver_stats()
        # Every waiter gave up, yet every resolution is still with the resolver:
        # the budget must not have been handed back early.
        assert stats["in_flight"] == safe_fetch.MAX_DNS_OUTSTANDING
        assert stats["running"] + stats["queued"] == safe_fetch.MAX_DNS_OUTSTANDING

        # While the budget is exhausted a further request is refused, not
        # queued behind the abandoned work.
        started = time.monotonic()
        with pytest.raises(safe_fetch.ResolverBusyError):
            safe_fetch.fetch_document("http://one-too-many.example/", timeout=2.0)
        assert time.monotonic() - started < 2, "the refusal was not immediate"
    finally:
        gate.set()


def test_resolver_budget_returns_only_when_the_lookup_really_ends(monkeypatch):
    """The budget comes back when the abandoned lookup really finishes."""
    assert _drain_resolver_pool()["in_flight"] == 0
    gate = threading.Event()
    monkeypatch.setattr(safe_fetch.socket, "getaddrinfo", _gated_resolver(gate, []))
    try:
        with pytest.raises(FetchAbortedError):
            safe_fetch.fetch_document("http://slow.example/", timeout=0.3)
        # The abandoned resolution still holds its budget slot.
        assert safe_fetch.resolver_stats()["in_flight"] == 1
    finally:
        gate.set()

    assert _drain_resolver_pool()["in_flight"] == 0
    # The budget is back: the next resolution runs to completion.
    addresses = safe_fetch.resolve_public_addresses(
        "next.example", 443, time.monotonic() + 5
    )
    assert addresses == [PUBLIC_IP]


def test_resolver_pool_is_bounded_under_repeated_timeouts(monkeypatch):
    """Repeated timeouts cannot grow the pool or its queue without limit."""
    assert _drain_resolver_pool()["in_flight"] == 0
    gate = threading.Event()
    monkeypatch.setattr(safe_fetch.socket, "getaddrinfo", _gated_resolver(gate, []))
    try:
        for index in range(safe_fetch.MAX_DNS_OUTSTANDING):
            with pytest.raises(FetchAbortedError):
                safe_fetch.fetch_document(f"http://flood{index}.example/", timeout=0.2)

        stats = safe_fetch.resolver_stats()
        assert stats["workers"] <= safe_fetch.MAX_DNS_WORKERS
        assert stats["in_flight"] == safe_fetch.MAX_DNS_OUTSTANDING
        assert len(_dns_threads()) <= safe_fetch.MAX_DNS_WORKERS

        threads_before = len(_dns_threads())
        for _ in range(25):
            with pytest.raises(safe_fetch.ResolverBusyError):
                safe_fetch.resolve_public_addresses(
                    "flood.example", 80, time.monotonic() + 0.5
                )

        stats = safe_fetch.resolver_stats()
        assert len(_dns_threads()) == threads_before, "a refusal created a thread"
        assert stats["workers"] <= safe_fetch.MAX_DNS_WORKERS
        assert stats["queued"] <= safe_fetch.MAX_DNS_OUTSTANDING
        assert stats["in_flight"] == safe_fetch.MAX_DNS_OUTSTANDING
    finally:
        gate.set()