diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index e278a796146..d4cbccfc43e 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -1,16 +1,25 @@ from __future__ import annotations +import io +import selectors +import socket import ssl import threading import time from collections.abc import Callable, Generator, Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, suppress from dataclasses import dataclass from functools import partial from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType -from typing import BinaryIO, Final +from typing import TYPE_CHECKING, BinaryIO, Final + +if TYPE_CHECKING: + from collections.abc import Buffer + +_POLL_SECONDS: Final = 0.05 +_WRITE_PIECE_BYTES: Final = 64 * 1024 @dataclass(frozen=True, slots=True) @@ -97,12 +106,108 @@ def _await_release( return gate.is_set() +class _RequestReader(io.BufferedReader): + cut = False + + def readline(self, size: int | None = -1, /) -> bytes: + line: Final = super().readline(size) + self.cut = self.cut or not line.endswith(b"\n") + return line + + +@dataclass(frozen=True, slots=True) +class _Blocked: + since: float + writing: bool + + +class _Connection(io.BufferedIOBase): + blocked: _Blocked | None = None + + def __init__(self, connection: socket.socket, closing: threading.Event) -> None: + super().__init__() + self.connection = connection + self.closing = closing + self.selector = selectors.DefaultSelector() + self.selector.register(connection, selectors.EVENT_READ) + + def readable(self) -> bool: + return True + + def writable(self) -> bool: + return True + + def readinto(self, buffer: Buffer, /) -> int: + if not self._wait_for_bytes(): + return 0 + with self._blocking(writing=False): + return self.connection.recv_into(buffer) + + def write(self, buffer: Buffer, /) -> int: + view: Final = memoryview(buffer).cast("B") + for start in range(0, view.nbytes, _WRITE_PIECE_BYTES): + with self._blocking(writing=True): + self.connection.sendall(view[start : start + _WRITE_PIECE_BYTES]) + return view.nbytes + + def close(self) -> None: + self.selector.close() + super().close() + + def stalled(self, grace: float, teardown_started: float) -> _Blocked | None: + blocked: Final = self.blocked + if blocked is None or time.monotonic() - max(blocked.since, teardown_started) < grace: + return None + return blocked + + def shut(self) -> None: + with suppress(OSError): + self.connection.shutdown(socket.SHUT_RDWR) + + def _wait_for_bytes(self) -> bool: + while not self.closing.is_set(): + if self._has_bytes(_POLL_SECONDS): + return True + return self._has_bytes(0) + + def _has_bytes(self, timeout: float) -> bool: + decrypted: Final = isinstance(self.connection, ssl.SSLSocket) and self.connection.pending() > 0 + return decrypted or bool(self.selector.select(timeout)) + + @contextmanager + def _blocking(self, writing: bool) -> Iterator[None]: + self.blocked = _Blocked(time.monotonic(), writing) + try: + yield + finally: + self.blocked = None + + +def _close_unless_stalled( + close: Callable[[], None], connections: SimpleQueue[_Connection], grace: float, errors: SimpleQueue[Exception] +) -> None: + teardown_started: Final = time.monotonic() + closer: Final = threading.Thread(target=close) + closer.start() + while closer.is_alive(): + for connection in tuple(connections.get_nowait() for _ in range(connections.qsize())): + stalled: Final = connection.stalled(grace, teardown_started) + if stalled is None: + connections.put(connection) + continue + if stalled.writing: + errors.put(TimeoutError(f"a reply write made no progress for {grace} s")) + connection.shut() + closer.join(timeout=_POLL_SECONDS) + + @contextmanager def wire_server( respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0, keep_alive: bool = False, + teardown_grace: float = 10, ) -> Generator[Wire, None, None]: """Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions.""" @@ -111,14 +216,20 @@ def wire_server( disconnected: Final[SimpleQueue[str]] = SimpleQueue() connected: Final[SimpleQueue[str]] = SimpleQueue() closed: Final[SimpleQueue[str]] = SimpleQueue() + streams: Final[SimpleQueue[_Connection]] = SimpleQueue() closing: Final = threading.Event() class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" - timeout = 5 + request: socket.socket + reader: _RequestReader def setup(self) -> None: - super().setup() + stream: Final = _Connection(self.request, closing) + self.reader = _RequestReader(stream) + self.rfile = self.reader + self.wfile = stream + streams.put(stream) connected.put(f"{self.client_address[0]}:{self.client_address[1]}") def finish(self) -> None: @@ -130,6 +241,9 @@ def wire_server( self.wfile.flush() def respond(self) -> None: + if self.reader.cut: + self.close_connection = True + return headers: Final = {name.lower(): value for name, value in self.headers.items()} try: body: Final = read_body(headers, self.rfile) @@ -221,7 +335,7 @@ def wire_server( self.socket = tls.wrap_socket(self.socket, server_side=True) with OwnedHTTPServer(("127.0.0.1", port), Handler) as server: - thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": _POLL_SECONDS}) thread.start() try: yield Wire( @@ -236,6 +350,6 @@ def wire_server( server.shutdown() thread.join(timeout=6) assert not thread.is_alive(), "Owned HTTP server survived cleanup" - server.server_close() + _close_unless_stalled(server.server_close, streams, teardown_grace, errors) failure: Final = None if errors.empty() else errors.get_nowait() assert failure is None, f"Owned HTTP peer failed: {failure!r}" diff --git a/tests/unit/integration_support/test_wire.py b/tests/unit/integration_support/test_wire.py new file mode 100644 index 00000000000..c786895720f --- /dev/null +++ b/tests/unit/integration_support/test_wire.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import datetime +import io +import socket +import ssl +import time +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from http.client import HTTPConnection, HTTPResponse, HTTPSConnection +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID + +from tests.integration._support.wire import Reply, Request, wire_server + +PAUSE_SECONDS: Final = 6 +HANDLER_READ_SECONDS: Final = 0.3 +FIRST: Final = b'{"call": "first"}' +LATE: Final = b'{"call": "late"}' +UNREAD: Final = b"x" * (32 * 1024 * 1024) + + +def _self_signed_server_context(directory: Path) -> ssl.SSLContext: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "wire peer")]) + now: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + certificate: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(1) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=36500)) + .sign(key, hashes.SHA256()) + ) + (directory / "peer.pem").write_bytes( + certificate.public_bytes(serialization.Encoding.PEM) + + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(directory / "peer.pem") + return context + + +def _trusting() -> ssl.SSLContext: + context: Final = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + return context + + +def _client(url: str) -> HTTPConnection: + address: Final = urlsplit(url) + if address.scheme == "http": + return HTTPConnection(address.hostname or "", address.port) + return HTTPSConnection(address.hostname or "", address.port, context=_trusting()) + + +def _tls_socket(url: str) -> ssl.SSLSocket: + address: Final = urlsplit(url) + return _trusting().wrap_socket(socket.create_connection((address.hostname or "", address.port or 0))) + + +def _echo(request: Request) -> Reply: + return Reply(body=request.body) + + +def _answer_slowly(request: Request) -> Reply: + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _answer_in_paused_chunks(request: Request) -> Reply: + return Reply(chunks=(request.body[:1], request.body[1:]), pause_between_chunks=HANDLER_READ_SECONDS) + + +def _hold_first(request: Request) -> Reply: + if request.target == "/first": + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _paced_pieces(client: socket.socket) -> Iterator[bytes]: + while piece := client.recv(256 * 1024): + yield piece + time.sleep(0.02) + + +def _read_slowly(client: socket.socket) -> bytes: + return b"".join(_paced_pieces(client)) + + +def _read_after(client: socket.socket, delay: float) -> bytes: + time.sleep(delay) + with client.makefile("rb") as stream: + return stream.read() + + +def _post(target: str, body: bytes) -> bytes: + return b"POST %s HTTP/1.1\r\nhost: peer\r\ncontent-length: %d\r\n\r\n%s" % (target.encode(), len(body), body) + + +def _pause_then_send(connection: HTTPConnection, body: bytes) -> HTTPResponse: + time.sleep(PAUSE_SECONDS) + connection.request("POST", "/late", body) + return connection.getresponse() + + +def _send_headers_then_pause_before_the_body(connection: HTTPConnection, body: bytes) -> HTTPResponse: + connection.putrequest("POST", "/late") + connection.putheader("content-length", str(len(body))) + connection.endheaders() + time.sleep(PAUSE_SECONDS) + connection.send(body) + return connection.getresponse() + + +@pytest.mark.parametrize( + "send_late", + (_pause_then_send, _send_headers_then_pause_before_the_body), + ids=("pause-before-the-request", "pause-before-the-body"), +) +def test_a_kept_alive_connection_answers_a_request_whose_bytes_arrive_after_a_long_pause( + send_late: Callable[[HTTPConnection, bytes], HTTPResponse], +) -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + late: Final = send_late(connection, LATE) + assert (first, late.status, late.read()) == (FIRST, 200, LATE) + assert tuple(request.body for request in wire.drain()) == (FIRST, LATE) + assert wire.connections() == 1 + assert connection.sock.recv(1) == b"" + connection.close() + + +def test_a_request_cut_off_by_teardown_never_reaches_respond() -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.sock.sendall(b"POST /cut HTTP/1.1\r\ncontent-type: application/json\r\ncontent-len") + time.sleep(HANDLER_READ_SECONDS) + assert (first, tuple(request.target for request in wire.drain())) == (FIRST, ("/first",)) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_stopped_reading_fails_teardown_instead_of_hanging() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + with pytest.raises(AssertionError, match="no progress"): + peer.close() + + +@pytest.mark.timeout(60) +def test_a_long_reply_its_client_is_still_reading_at_teardown_is_not_cut() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + reply: Final = pool.submit(_read_slowly, client) + time.sleep(1) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_starts_reading_after_teardown_starts_still_arrives_whole() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=2)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + time.sleep(2.5) + reply: Final = pool.submit(_read_after, client, 0.3) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +def test_a_complete_request_sent_as_teardown_starts_is_still_answered() -> None: + with wire_server(_echo, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.request("POST", "/complete", LATE) + late: Final = connection.getresponse() + assert (first, late.status, late.read(), tuple(request.target for request in wire.drain())) == ( + FIRST, + 200, + LATE, + ("/first", "/complete"), + ) + connection.close() + + +@pytest.mark.parametrize("secure", (False, True), ids=("plain", "tls")) +def test_a_reply_still_being_written_at_teardown_reaches_its_client_whole(secure: bool, tmp_path: Path) -> None: + slow: Final = Reply(chunks=(b"first", b"second"), pause_between_chunks=HANDLER_READ_SECONDS) + tls: Final = _self_signed_server_context(tmp_path) if secure else None + with wire_server(lambda _: slow, tls=tls, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("GET", "/slow") + response: Final = connection.getresponse() + head: Final = response.read(len(b"first")) + assert head + response.read() == b"firstsecond" + connection.close() + + +def test_a_complete_request_queued_behind_a_busy_handler_is_still_answered() -> None: + with wire_server(_hold_first, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + client: Final = socket.create_connection((address.hostname or "", address.port or 0)) + client.sendall(_post("/first", FIRST)) + while wire.received.empty(): + time.sleep(0.01) + client.sendall(_post("/second", LATE)) + with client, client.makefile("rb") as stream: + replies: Final = stream.read() + assert ( + tuple(request.target for request in wire.drain()), + replies.count(b"HTTP/1.1 200"), + replies[-len(LATE) :], + ) == ( + ("/first", "/second"), + 2, + LATE, + ) + + +@pytest.mark.parametrize("respond", (_answer_slowly, _answer_in_paused_chunks), ids=("slow-respond", "paused-chunks")) +def test_a_reply_that_takes_longer_than_the_teardown_grace_still_reaches_its_client( + respond: Callable[[Request], Reply], +) -> None: + with wire_server(respond, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS / 2) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/slow", FIRST) + while wire.received.empty(): + time.sleep(0.01) + response: Final = connection.getresponse() + assert (response.status, response.read()) == (200, FIRST) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_tls_request_larger_than_the_read_buffer_is_answered(tmp_path: Path) -> None: + large: Final = b"x" * (3 * io.DEFAULT_BUFFER_SIZE // 2) + with wire_server(_echo, tls=_self_signed_server_context(tmp_path), keep_alive=True) as wire: + client: Final = _tls_socket(wire.url) + client.sendall(_post("/large", large)) + response: Final = HTTPResponse(client) + response.begin() + assert (response.status, response.read()) == (200, large) + client.close() + + +@pytest.mark.timeout(60) +def test_a_tls_record_cut_off_by_teardown_closes_its_connection_instead_of_hanging(tmp_path: Path) -> None: + tls: Final = _self_signed_server_context(tmp_path) + with wire_server(_echo, tls=tls, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS) as wire: + client: Final = _tls_socket(wire.url) + socket.socket.sendall(client, b"\x17\x03\x03\x40\x00" + b"x" * 16) + time.sleep(HANDLER_READ_SECONDS) + assert wire.drain() == () + client.close()