Files
RustProxy/scripts/performance/transport_clients.py
T
2026-10-17 11:57:16 +02:00

192 lines
8.3 KiB
Python
Executable File

"""Strict H2/TLS and H3/QUIC clients for local RustProxy acceptance tests.
The generated fixture certificate is explicitly trusted. Neither transport
silently falls back to HTTP/1.1 or disables certificate verification.
"""
from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
import ssl
import httpx
from aioquic.asyncio import QuicConnectionProtocol, connect
from aioquic.h3.connection import H3Connection, MessageError
from aioquic.h3.events import DataReceived, HeadersReceived
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.events import ConnectionTerminated, ProtocolNegotiated, StreamReset
@dataclass
class Response:
status: int
headers: dict[str, str]
body: bytes
protocol: str
alpn: str | None
stream_id: int | None = None
class H2Client:
def __init__(self, client: httpx.AsyncClient, port: int):
self.client, self.port = client, port
async def request(self, method, path, headers=None, body=None):
async def chunks():
for offset in range(0, len(body), 16384):
yield body[offset:offset + 16384]
await asyncio.sleep(0)
response = await self.client.request(
method, f"https://127.0.0.1:{self.port}{path}",
headers=headers or {}, content=chunks() if body is not None else None,
)
network = response.extensions.get("network_stream")
tls = network.get_extra_info("ssl_object") if network else None
alpn = tls.selected_alpn_protocol() if tls else None
if response.http_version != "HTTP/2" or alpn != "h2":
raise AssertionError(f"H2 required, got {response.http_version}, ALPN={alpn!r}")
return Response(response.status_code, dict(response.headers), response.content,
"HTTP/2", alpn, response.extensions.get("stream_id"))
@dataclass
class Pending:
future: asyncio.Future
status: int = 0
headers: dict[str, str] = field(default_factory=dict)
body: bytearray = field(default_factory=bytearray)
class HeadAwareH3Connection(H3Connection):
"""Supply request-method context missing from aioquic 1.7.0's length check.
RFC 9110 section 8.6 permits HEAD's Content-Length to describe a GET body.
Still reject any actual HEAD DATA; keep aioquic's exact length check for all
other requests. This changes only the test client, never the target server.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.head_streams = set()
def _check_content_length(self, stream):
if stream.stream_id in self.head_streams:
if stream.content_length:
raise MessageError("HEAD response contained DATA")
return
super()._check_content_length(stream)
def validate_head_adapter():
"""Keep adapter checks separate from the real-server acceptance totals."""
from types import SimpleNamespace
connection = HeadAwareH3Connection.__new__(HeadAwareH3Connection)
connection.head_streams = {0}
for stream_id, declared, received, valid in [(0, 1040, 0, True), (0, 1040, 1, False),
(4, 1040, 1040, True), (4, 1040, 0, False)]:
stream = SimpleNamespace(stream_id=stream_id, expected_content_length=declared,
content_length=received)
try:
connection._check_content_length(stream)
accepted = True
except MessageError:
accepted = False
assert accepted == valid, "H3 fixture HEAD adapter changed ordinary framing validation"
return {"cases": 4, "passed": 4, "kind": "test-client self-test, not RustProxy coverage"}
class H3Client(QuicConnectionProtocol):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.http = HeadAwareH3Connection(self._quic)
self.pending: dict[int, Pending] = {}
self.alpn = None
self.authority = None
def quic_event_received(self, event):
if isinstance(event, ProtocolNegotiated):
self.alpn = event.alpn_protocol
if isinstance(event, ConnectionTerminated):
for item in self.pending.values():
if not item.future.done():
item.future.set_exception(ConnectionError(
f"QUIC closed code={event.error_code}: {event.reason_phrase}"))
if isinstance(event, StreamReset):
item = self.pending.get(event.stream_id)
if item is not None and not item.future.done():
item.future.set_exception(ConnectionError(
f"H3 stream reset code={event.error_code}"))
for http_event in self.http.handle_event(event):
if not isinstance(http_event, (HeadersReceived, DataReceived)):
continue
item = self.pending.get(http_event.stream_id)
if item is None or item.future.done():
continue
if isinstance(http_event, HeadersReceived):
for name, value in http_event.headers:
if name == b":status":
item.status = int(value)
elif not name.startswith(b":"):
item.headers[name.decode("ascii").lower()] = value.decode("latin-1")
else:
item.body.extend(http_event.data)
if len(item.body) > 4 * 1024 * 1024:
item.future.set_exception(ValueError("fixture response exceeded 4 MiB bound"))
continue
if http_event.stream_ended:
item.future.set_result(Response(item.status, item.headers, bytes(item.body),
"HTTP/3", self.alpn, http_event.stream_id))
async def request(self, method, path, headers=None, body=None):
if self.alpn != "h3":
raise AssertionError(f"H3 required, got ALPN={self.alpn!r}")
stream_id = self._quic.get_next_available_stream_id()
item = Pending(asyncio.get_running_loop().create_future())
self.pending[stream_id] = item
if method == "HEAD":
self.http.head_streams.add(stream_id)
fields = [(b":method", method.encode("ascii")), (b":scheme", b"https"),
(b":authority", self.authority.encode("ascii")),
(b":path", path.encode("ascii"))]
fields += [(name.lower().encode("ascii"), value.encode("latin-1"))
for name, value in (headers or {}).items()]
try:
self.http.send_headers(stream_id, fields, end_stream=body is None)
self.transmit()
if body is not None:
if not body:
self.http.send_data(stream_id, b"", end_stream=True)
self.transmit()
for offset in range(0, len(body), 16384):
end = min(offset + 16384, len(body))
self.http.send_data(stream_id, body[offset:end], end_stream=end == len(body))
self.transmit()
await asyncio.sleep(0)
return await asyncio.wait_for(item.future, timeout=10)
finally:
self.pending.pop(stream_id, None)
self.http.head_streams.discard(stream_id)
@asynccontextmanager
async def client(protocol: str, port: int, certificate: Path):
if protocol == "h2":
context = ssl.create_default_context(cafile=str(certificate))
async with httpx.AsyncClient(http2=True, verify=context, trust_env=False,
timeout=10, limits=httpx.Limits(max_connections=1)) as session:
yield H2Client(session, port)
elif protocol == "h3":
config = QuicConfiguration(is_client=True, alpn_protocols=["h3"],
server_name="localhost", idle_timeout=10)
config.verify_mode = ssl.CERT_REQUIRED
config.load_verify_locations(cafile=str(certificate))
async with connect("127.0.0.1", port, configuration=config,
create_protocol=H3Client) as session:
session.authority = f"127.0.0.1:{port}"
yield session
else:
raise ValueError(f"unsupported fixture protocol {protocol!r}")