diff --git a/aiohttp/client.py b/aiohttp/client.py index f91b81e5712..8a55d1b1268 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -98,6 +98,7 @@ strip_auth_from_url, ) from .http import WS_KEY, HttpVersion, WebSocketReader, WebSocketWriter +from .http2.adapter import get_version from .http_websocket import WSHandshakeError, ws_ext_gen, ws_ext_parse from .tracing import Trace, TraceConfig from .typedefs import ( @@ -234,23 +235,62 @@ class _WSConnectOptions(TypedDict, total=False): async def _connect_and_send_request(req: ClientRequest) -> ClientResponse: connector = req._session._connector assert connector is not None + key = req.connection_key try: + # only the first connection to a host blocks + # the rest of the connection requests are done + # concurrently + await connector.semaphore.acquire(key) + conn = await connector.connect(req, traces=req._traces, timeout=req._timeout) + + connector.semaphore.release(key) except asyncio.TimeoutError as exc: raise ConnectionTimeoutError(f"Connection timeout to host {req.url}") from exc + finally: + connector.semaphore.release(key) assert conn.protocol is not None - conn.protocol.set_response_params(**req._response_params) + assert conn.protocol.transport is not None + + alpn_protocol = get_version(conn.protocol) + + resp = None + started = False + + if alpn_protocol == "h2": + # release immediately to allow reuse + connector._release(conn._key, conn.protocol, should_close=False) + # the protocol corresponding to the connection + # remains (i.e., the count per host is always 1 for h2) + # This is the number of TCP connections not the number of + # streams + connector._acquired.add(conn.protocol) try: - resp = await req._send(conn) - try: - await resp.start(conn) - except BaseException: + # backwards compatibility + if alpn_protocol == "h2": + stream = await conn.protocol.create_stream() # type: ignore[attr-defined] + req.stream_id = stream.stream_id + # release again to clear the protocol from _acquired if required + connector._release(conn._key, conn.protocol, should_close=False) + resp = await req._send(conn) + resp.stream_id = stream.stream_id + else: + conn.protocol.set_response_params(**req._response_params) + resp = await req._send(conn) + await resp.start(conn) + + if alpn_protocol == "h2": + # we still have to null the protocol since we didn't close the connection + conn._protocol = None + + started = True + finally: + if resp is not None and not started: resp.close() - raise - except BaseException: - conn.close() - raise + conn.close() + if resp is None: + conn.close() return resp diff --git a/aiohttp/client_reqrep.py b/aiohttp/client_reqrep.py index 516a4cdb17f..f66776cceac 100644 --- a/aiohttp/client_reqrep.py +++ b/aiohttp/client_reqrep.py @@ -59,6 +59,7 @@ HttpVersion11, StreamWriter, ) +from .http2.adapter import Http2StreamWriter, get_version from .streams import EMPTY_PAYLOAD, StreamReader from .typedefs import DEFAULT_JSON_DECODER, JSONDecoder, RawHeaders @@ -280,6 +281,9 @@ class ClientResponse(HeadersMixin): _output_size: int = 0 _upload_complete: asyncio.Future[None] | None = None + # HTTP/2 stream id + stream_id: int | None = None + def __init__( self, method: str, @@ -525,7 +529,13 @@ async def start(self, connection: "Connection") -> "ClientResponse": # read response try: protocol = self._protocol - message, payload = await protocol.read() # type: ignore[union-attr] + # conditional branching to pass the stream id + # to the protocol + assert protocol is not None + if get_version(protocol) == "h2": + message, payload = await protocol.read_stream(self.stream_id) # type: ignore[attr-defined] + else: + message, payload = await protocol.read() except HttpProcessingError as exc: raise ClientResponseError( self.request_info, @@ -810,6 +820,9 @@ class ClientRequestBase: _skip_auto_headers: "CIMultiDict[None] | None" = None + # HTTP/2 stream id + stream_id: int | None = None + # N.B. # Adding __del__ method with self._writer closing doesn't make sense # because _writer is instance method, thus it keeps a reference to self. @@ -932,7 +945,9 @@ def _create_response( stream_writer=stream_writer, ) - def _create_writer(self, protocol: BaseProtocol) -> StreamWriter: + def _create_writer( + self, protocol: BaseProtocol + ) -> StreamWriter | Http2StreamWriter: return StreamWriter(protocol, self.loop) def _should_write(self, protocol: BaseProtocol) -> bool: @@ -1428,7 +1443,11 @@ def _create_response( stream_writer=stream_writer, ) - def _create_writer(self, protocol: BaseProtocol) -> StreamWriter: + def _create_writer( + self, protocol: BaseProtocol + ) -> StreamWriter | Http2StreamWriter: + if get_version(protocol) == "h2": + return Http2StreamWriter(protocol, self.loop, self) writer = StreamWriter( protocol, self.loop, diff --git a/aiohttp/connector.py b/aiohttp/connector.py index cdf6ca5bd57..ca221ff92d6 100644 --- a/aiohttp/connector.py +++ b/aiohttp/connector.py @@ -1,5 +1,6 @@ import asyncio import functools +import os import random import socket import sys @@ -52,6 +53,8 @@ set_exception, set_result, ) +from .http2.synchro import HostProbeSynchronizer +from .http_protocol import HttpDispatcherProtocol from .log import client_logger from .resolver import DefaultResolver @@ -138,7 +141,7 @@ async def create_connection( async def start_tls( loop: asyncio.AbstractEventLoop, transport: asyncio.Transport, - protocol: ResponseHandler, + protocol: HttpDispatcherProtocol | ResponseHandler, sslcontext: SSLContext, *, server_hostname: str | None, @@ -375,7 +378,7 @@ def __init__( ] = defaultdict(OrderedDict) self._loop = loop - self._factory = functools.partial(ResponseHandler, loop=loop) + self._factory = functools.partial(HttpDispatcherProtocol, loop=loop) # start keep-alive connection cleanup task self._cleanup_handle: asyncio.TimerHandle | None = None @@ -402,6 +405,12 @@ def __init__( self._placeholder_future.set_result(None) self._cleanup_closed() + # Semaphore for HTTP/2 connections + # avoids duplicate connections to the + # same host + # (HTTP/2 doesn't need connection pooling to send multiple requests) + self.semaphore = HostProbeSynchronizer() + def __del__(self, _warnings: Any = warnings) -> None: if self._closed: return @@ -939,7 +948,11 @@ def _make_ssl_context(verified: bool) -> SSLContext: sslcontext.verify_mode = ssl.CERT_NONE sslcontext.options |= ssl.OP_NO_COMPRESSION sslcontext.set_default_verify_paths() - sslcontext.set_alpn_protocols(("http/1.1",)) + + protocols = ["http/1.1"] + if os.getenv("AIOHTTP_ENABLE_EXPERIMENTAL_PROTOCOLS", False): + protocols += ["h2"] + sslcontext.set_alpn_protocols(tuple(protocols)) return sslcontext @@ -1499,7 +1512,8 @@ async def _start_tls_connection( tls_transport ) # Kick the state machine of the new TLS protocol - return tls_transport, tls_proto + # HACK use the correct type + return tls_transport, tls_proto # type: ignore[return-value] def _convert_hosts_to_addr_infos( self, hosts: list[ResolveResult] @@ -1591,7 +1605,6 @@ async def _create_direct_connection( bad_peer = sock.getpeername() aiohappyeyeballs.remove_addr_infos(addr_infos, bad_peer) continue - return transp, proto assert last_exc is not None raise last_exc @@ -1723,7 +1736,7 @@ async def _create_connection( raise raise UnixClientConnectorError(self.path, req.connection_key, exc) from exc - return proto + return proto # type: ignore[return-value] class NamedPipeConnector(BaseConnector): diff --git a/aiohttp/helpers.py b/aiohttp/helpers.py index e6d6b6ca1ca..c627f9da125 100644 --- a/aiohttp/helpers.py +++ b/aiohttp/helpers.py @@ -47,7 +47,6 @@ from . import hdrs from .log import client_logger -from .typedefs import PathLike # noqa if sys.version_info >= (3, 11): import asyncio as async_timeout diff --git a/aiohttp/http2/adapter.py b/aiohttp/http2/adapter.py new file mode 100644 index 00000000000..0d444c4e4ec --- /dev/null +++ b/aiohttp/http2/adapter.py @@ -0,0 +1,185 @@ +"""HTTP/1.1 to HTTP/2 adapters.""" + +import asyncio +from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, Optional, Tuple, Union + +from multidict import CIMultiDict + +from ..abc import AbstractStreamWriter +from ..base_protocol import BaseProtocol +from ..helpers import HeadersDictProxy +from ..http_parser import RawResponseMessage +from ..http_writer import HttpVersion + +if TYPE_CHECKING: + from ..client_reqrep import ClientRequest + + +def get_version(protocol: BaseProtocol) -> str: + """Helper to get the negotiated HTTP version from the protocol.""" + # backwards compatibility + if not hasattr(protocol, "transport"): + return "http/1.1" + transport = protocol.transport + assert transport is not None + return _get_version(transport) + + +def _get_version(transport: asyncio.BaseTransport) -> str: + ssl_object = transport.get_extra_info("ssl_object") + alpn_protocol = ssl_object.selected_alpn_protocol() if ssl_object else "http/1.1" + return alpn_protocol + + +def feed_headers( + headers: Iterable[tuple[str, str]], +) -> RawResponseMessage: + """Convert raw headers to the standard RawResponseMessage format""" + # Build a minimal RawResponseMessage (the fields that ClientResponse uses). + raw_headers: List[Tuple[bytes, bytes]] = [] + code = 500 + # there is no guarantee that the status code comes first + for key, value in headers: + if key == ":status": + code = int(value) + raw_headers.append((key.encode("latin-1"), value.encode("latin-1"))) + msg = RawResponseMessage( + version=HttpVersion(2, 0), # HTTP/2.0 + code=code, + # HTTP/2 has no reason phrase + reason="", + headers=HeadersDictProxy(CIMultiDict(headers)), + raw_headers=tuple(raw_headers), + should_close=False, + compression=None, # XXX propagate the configuration + upgrade=False, # not implemented + chunked=False, # n/a + ) + return msg + + +class Http2StreamWriter(AbstractStreamWriter): + def __init__( + self, + protocol: Any, + loop: asyncio.AbstractEventLoop, + req: "ClientRequest", + *, + on_chunk_sent: Optional[Any] = None, + on_headers_sent: Optional[Any] = None, + ) -> None: + del on_chunk_sent, on_headers_sent # skipped for now + + self._req = req + self._protocol = protocol + assert req.stream_id + self.stream_id: int = req.stream_id + self.loop = loop + + self._headers: Optional[List[Tuple[str, str]]] = None + self._headers_sent = False + self._eof = False + self.output_size = 0 + # constant + self.buffer_size = 0 + + @property + def protocol(self) -> Any: + return self._protocol + + @property + def transport(self) -> Any: + return self._protocol.transport + + def _send_headers(self, *, end_stream: bool) -> None: + if self._headers_sent or self._headers is None: + return + + self._protocol._connection.send_headers( + self.stream_id, + self._req.method, + self._req.url, + self._headers, + end_stream=end_stream, + ) + self._headers_sent = True + self._headers = None + + async def write_headers(self, status_line: str, headers: Mapping[str, str]) -> None: + # Buffer headers so that we can set END_STREAM on HEADERS + # when there is no request body. + self._headers = list(headers.items()) + + async def write( + self, + chunk: Union[bytes, bytearray, memoryview, "memoryview[bytes]"], + *, + drain: bool = True, + LIMIT: int = 0x10000, + ) -> None: + if self._eof: + raise RuntimeError("Cannot write after EOF") + + if not self._headers_sent: + if self._headers is None: + raise RuntimeError("Headers must be written before body") + self._send_headers(end_stream=False) + + if chunk: + await self._protocol._connection.send_data( + self.stream_id, bytes(chunk), end_stream=False + ) + self.output_size += len(chunk) + + # we would drain here but + # drain is no-op in HTTP/2 + + async def write_eof(self, chunk: bytes = b"") -> None: + if self._eof: + return + + if not self._headers_sent: + if self._headers is None: + raise RuntimeError("Headers must be written before body") + + if chunk: + self._send_headers(end_stream=False) + await self._protocol._connection.send_data( + self.stream_id, chunk, end_stream=True + ) + self.output_size += len(chunk) + else: + # Body-less request: HEADERS frame carries END_STREAM. + self._send_headers(end_stream=True) + else: + # Headers were already sent; send a final DATA frame with END_STREAM. + await self._protocol._connection.send_data( + self.stream_id, chunk, end_stream=True + ) + self.output_size += len(chunk) + + self._eof = True + + def set_eof(self) -> None: + """Called when there is no body.""" + if self._eof: + return + + if not self._headers_sent: + if self._headers is None: + raise RuntimeError("Headers must be written before EOF") + self._send_headers(end_stream=True) + + self._eof = True + + async def drain(self) -> None: + # HTTP/2 flow control is handled inside send_data(). + pass + + def enable_compression(self, *args: Any, **kwargs: Any) -> None: + # HTTP/2 compresses headers automatically; no body compression here. + pass + + def enable_chunking(self) -> None: + # HTTP/2 does not use chunked transfer encoding (or, rather, chunked transfer encoding is built-in). + pass diff --git a/aiohttp/http2/connection.py b/aiohttp/http2/connection.py new file mode 100644 index 00000000000..70c50ad9abf --- /dev/null +++ b/aiohttp/http2/connection.py @@ -0,0 +1,827 @@ +""" +Complete HTTP/2 client implementation (RFC 7540). + +This module provides: +- Frame-level binary wire protocol with debug logging. +- HPACK compression/decompression (using the `hpack` library). +- Full stream state machine (idle -> open -> half‑closed -> closed). +- Multiplexed connection handling (concurrent streams, flow control). +- Server settings tracking (MAX_CONCURRENT_STREAMS, INITIAL_WINDOW_SIZE, etc.). +- Integration point for aiohttp's TCPConnector via an asyncio.Protocol subclass. + +Dependencies: +- asyncio +- struct +- logging +- enum +- hpack (install with `pip install hpack`) +- collections.defaultdict + +Usage: + Connector replaces `ResponseHandler` with `Http2Protocol`. +""" + +import asyncio +import logging +import struct +from typing import Dict, Iterable, List, Optional, Set + +from hpack import Decoder, Encoder +from yarl import URL + +from ..base_protocol import BaseProtocol +from ..client_exceptions import ClientConnectionError, SocketTimeoutError +from ..helpers import ( + _EXC_SENTINEL, + DEFAULT_CHUNK_SIZE, + BaseTimerContext, + set_exception as _set_exc, + set_result, +) +from ..http_parser import RawResponseMessage +from ..streams import StreamReader +from .errors import ErrorCode +from .settings import ( + DEFAULT_SETTINGS, + FlagData, + FlagHeaders, + FlagPing, + FlagSettings, + FrameType, + Setting, +) +from .stream import Stream, StreamState + +# ---------------------------------------------------------------------- +# Logging – plaintext wire‑format emission for debugging +# ---------------------------------------------------------------------- +logger = logging.getLogger("aiohttp.http2.connection") +# logger.setLevel(logging.DEBUG) + +FRAME_HEADER_LENGTH = 9 # 9 octets +STREAM_ID_MASK = 0x7FFFFFFF # to avoid setting the reserved bit to 1 + + +# ---------------------------------------------------------------------- +# Connection‑level management +# ---------------------------------------------------------------------- +class Http2Connection: + """Manages a single HTTP/2 connection. + + Handles: + - Connection preface and SETTINGS handshake. + - Frame parsing and dispatch. + - HPACK encoding/decoding. + - Stream multiplexing and flow control. + - Server settings tracking. + """ + + def __init__( + self, + transport: asyncio.Transport, + loop: asyncio.AbstractEventLoop, + protocol: "Http2Protocol", + ) -> None: + self._transport = transport + self._loop = loop + self._protocol = protocol + + # HPACK + self.hpack_encoder: Encoder = Encoder() + self.hpack_decoder: Decoder = Decoder() + + # Settings + self.remote_settings: Dict[Setting, int] = DEFAULT_SETTINGS.copy() + self.local_settings: Dict[Setting, int] = DEFAULT_SETTINGS.copy() + + # Flow control + self.session_outbound_window: int = ( + 65535 # initial flow control (RFC 7540, 6.9.1) + ) + self.session_inbound_window: int = 65535 + self._flow_control_updated: asyncio.Event = asyncio.Event() + self._flow_control_updated.set() # initially writable + + # Streams + self.streams: Dict[int, Stream] = {} + self.next_stream_id: int = 1 # client streams are odd + self.max_concurrent_streams: int = DEFAULT_SETTINGS[ + Setting.MAX_CONCURRENT_STREAMS + ] + self._pending_streams: List[asyncio.Future[Stream]] = [] + self._last_peer_stream_id: int = ( + 0 # highest server‑initiated stream (even, unused for client) + ) + + # Frame buffers + self._frame_buffer: bytearray = bytearray() + + # GOAWAY state + self._goaway_received: bool = False + self._goaway_sent: bool = False + self._last_stream_id: int = 0 + self._error_code: int = 0 + + # Closed streams cleanup + self._closed_streams: Set[int] = set() + + # -------------------- Transport callbacks -------------------- + def data_received(self, data: bytes) -> None: + """Assemble frames from the byte stream and dispatch them.""" + self._frame_buffer.extend(data) + # Consume complete frames while enough bytes for the header exist + while len(self._frame_buffer) >= FRAME_HEADER_LENGTH: + # Parse 24-bit length, 8-bit type, 8-bit flags, 32-bit stream ID + length = ( + self._frame_buffer[0] << 16 + | self._frame_buffer[1] << 8 + | self._frame_buffer[2] + ) + frame_type_val = self._frame_buffer[3] + flags = self._frame_buffer[4] + stream_id = struct.unpack("!I", self._frame_buffer[5:9])[0] & STREAM_ID_MASK + + if len(self._frame_buffer) < FRAME_HEADER_LENGTH + length: + break # incomplete frame; wait for more data + + payload = bytes( + self._frame_buffer[FRAME_HEADER_LENGTH : FRAME_HEADER_LENGTH + length] + ) + del self._frame_buffer[: FRAME_HEADER_LENGTH + length] + + # invalid frames cause a value error + if 0 <= frame_type_val <= 9: + logger.debug( + "<- %s stream=%d flags=0x%02x len=%d", + FrameType(frame_type_val).name, + stream_id, + flags, + length, + ) + + try: + self._dispatch_frame(frame_type_val, flags, stream_id, payload) + except Exception as exc: + # we really don't want to swallow this exception + import traceback + + logger.error("\n".join(traceback.format_exception(exc))) + logger.error( + "Critical error when dispatching frame: exception=%s. (frame_type=%s, flags=%d, stream_id=%d, payload=%s)", + str(exc), + FrameType(frame_type_val), + flags, + stream_id, + payload, + ) + self._send_rst_stream(stream_id, ErrorCode.PROTOCOL_ERROR) + + def eof_received(self) -> bool: + logger.debug("EOF received from server") + self.close() + return False + + def connection_lost(self, exc: Optional[BaseException]) -> None: + logger.debug(f"Connection lost: {exc}") + # Cancel all pending streams (including those in the queue) + for stream in list(self.streams.values()): + if not stream.response_future.done(): + stream.response_future.set_exception(ConnectionError("Connection lost")) + for fut in self._pending_streams: + fut.set_exception(ConnectionError("Connection lost")) + self.streams.clear() + + # -------------------- Frame dispatch -------------------- + def _dispatch_frame( + self, frame_type: int, flags: int, stream_id: int, payload: bytes + ) -> None: + if frame_type == FrameType.DATA: + self._handle_data_frame(flags, stream_id, payload) + elif frame_type == FrameType.HEADERS: + self._handle_headers_frame(flags, stream_id, payload) + elif frame_type in { + FrameType.PRIORITY, + FrameType.PUSH_PROMISE, + FrameType.CONTINUATION, + }: + logger.warning("%d frame ignored (not implemented)", frame_type) + elif frame_type == FrameType.RST_STREAM: + self._handle_rst_stream_frame(flags, stream_id, payload) + elif frame_type == FrameType.SETTINGS: + self._handle_settings_frame(flags, stream_id, payload) + elif frame_type == FrameType.PING: + self._handle_ping_frame(flags, stream_id, payload) + elif frame_type == FrameType.GOAWAY: + self._handle_goaway_frame(flags, stream_id, payload) + elif frame_type == FrameType.WINDOW_UPDATE: + self._handle_window_update_frame(flags, stream_id, payload) + else: + logger.warning("Ignoring unknown frame type %d", frame_type) + + # ---------- Individual frame handlers ---------- + def _handle_data_frame(self, flags: int, stream_id: int, payload: bytes) -> None: + stream = self.streams.get(stream_id) + if stream is None: + if stream_id > self._last_peer_stream_id: + self._send_rst_stream(stream_id, ErrorCode.PROTOCOL_ERROR) + return + + pad_length = 0 + pos = 0 + if flags & FlagData.PADDED: + # use fuzzy tests to + # verify if it's an error + pad_length = payload[0] + pos = 1 + # XXX padding might be too long + # send protocol error + + # pad_length >= len(payload) + data = payload[pos : len(payload) - pad_length] + end_stream = bool(flags & FlagData.END_STREAM) + + # Update session flow control + self.session_inbound_window -= len(data) + + stream.receive_data(data, end_stream) + + def _handle_headers_frame(self, flags: int, stream_id: int, payload: bytes) -> None: + if flags & FlagHeaders.PRIORITY: + # Exclusive flag + stream dependency + weight + # exclude priority data + payload = payload[5:] + + # Decode headers with HPACK + try: + headers = self.hpack_decoder.decode(payload) + except Exception as exc: # too general? + logger.error(f"HPACK decode error: {exc}") + self._send_rst_stream(stream_id, ErrorCode.PROTOCOL_ERROR) + return + + end_stream = bool(flags & FlagHeaders.END_STREAM) + + stream = self.streams.get(stream_id) + + if stream is None: + logger.error("Unknown stream_id: %d", stream_id) + else: + stream.receive_headers(headers, end_stream) + + def _handle_rst_stream_frame( + self, flags: int, stream_id: int, payload: bytes + ) -> None: + del flags # rst doesn't use flags + + error_code = struct.unpack("!I", payload)[0] + stream = self.streams.get(stream_id) + if stream: + stream.transition(StreamState.CLOSED) + if not stream.response_future.done(): + stream.response_future.set_exception( + RuntimeError(f"Stream reset by server (code={error_code})") + ) + self._close_stream(stream) + + def _handle_settings_frame( + self, flags: int, stream_id: int, payload: bytes + ) -> None: + """Process SETTINGS frame (6.5)""" + if flags & FlagSettings.ACK: + logger.debug("Received SETTINGS ACK") + return # Our settings were acknowledged + if stream_id != 0: + logger.error( + "SETTING frame received after the first stream in violation of the protocol standard (RFC-9113, 3.4)" + ) + self._protocol_error() + return + + if len(payload) % 6 != 0: + logger.error("SETTINGS payload length not a multiple of 6") + self._protocol_error() + return + + # Parse key‑value pairs + for i in range(0, len(payload), 6): + identifier, value = struct.unpack("!H I", payload[i : i + 6]) + # the attribute is defined + # accessing __members__ avoids a costly try-catch (ValueError) + if ( + identifier not in Setting.__members__.values() + ): # ignore: type[attr-defined] + # ignoring as per the RFC + logger.warning("Unknown setting identifier %d", identifier) + continue + setting = Setting(identifier) + old_value = self.remote_settings.get(setting, value) + self.remote_settings[setting] = value + logger.info(f"Server SETTINGS: {setting.name} = {value}") + + # React to certain settings + if setting == Setting.INITIAL_WINDOW_SIZE and value != old_value: + # might become negative + # send WINDOW_UPDATE + delta = value - old_value + for s in self.streams.values(): + s.outbound_window += delta + elif setting == Setting.MAX_CONCURRENT_STREAMS: + self.max_concurrent_streams = value + self._maybe_unblock_streams() + elif setting == Setting.HEADER_TABLE_SIZE: + self.hpack_encoder.header_table_size = value + + # Acknowledge settings + self._send_settings_ack() + + def _handle_ping_frame(self, flags: int, stream_id: int, payload: bytes) -> None: + if stream_id != 0: + self._protocol_error() + return + if flags & FlagPing.ACK: + logger.debug("Received PING ACK") + else: + # Respond with ACK + logger.debug("Received PING, sending ACK") + self._send_ping(ack=True, opaque_data=payload) + + def _handle_goaway_frame(self, flags: int, stream_id: int, payload: bytes) -> None: + del flags, stream_id # interface + + self._goaway_received = True + last_stream_id, error_code = struct.unpack("!I I", payload[:8]) + extra = payload[8:] + self._last_stream_id = last_stream_id + self._error_code = error_code + logger.info( + "GOAWAY received: last_stream=%d, error=%d, extra=%s", + last_stream_id, + error_code, + extra.decode(errors="replace"), + ) + # Cancel streams with higher IDs + for sid, stream in list(self.streams.items()): + if sid > last_stream_id: + if not stream.response_future.done(): + stream.response_future.set_exception( + ConnectionError("GOAWAY received") + ) + self._close_stream(stream) + # clear pending streams? + + def _handle_window_update_frame( + self, flags: int, stream_id: int, payload: bytes + ) -> None: + increment = struct.unpack("!I", payload)[0] + if stream_id == 0: + # Session window update + self.session_outbound_window += increment + else: + stream = self.streams.get(stream_id) + if stream: + stream.outbound_window += increment + # Wake up any writer waiting for flow control + self._flow_control_updated.set() + + # -------------------- Frame sending helpers -------------------- + def _send_frame( + self, + frame_type: FrameType, + flags: int, + stream_id: int, + payload: bytes = b"", + ) -> None: + length = len(payload) & 0x00FFFFFF # 24 bits -> 3 bytes + header = struct.pack("!I", length)[ + 1: + ] + struct.pack( # drop the first (most‑significant) byte -> 3 bytes + "!B B I", frame_type, flags, stream_id + ) + + logger.debug( + f"-> FRAME type={frame_type.name:>15} flags=0x{flags:02x} " + f"stream_id={stream_id:<5} length={len(payload)}" + ) + self._transport.write(header + payload) + + def _send_settings_ack(self) -> None: + self._send_frame(FrameType.SETTINGS, FlagSettings.ACK, 0) + + def _send_ping(self, ack: bool = False, opaque_data: bytes = b"\x00" * 8) -> None: + flags = FlagPing.ACK if ack else 0 + self._send_frame(FrameType.PING, flags, 0, opaque_data) + + def _send_goaway(self, last_stream_id: int, error_code: int) -> None: + payload = struct.pack("!I I", last_stream_id, error_code) + self._send_frame(FrameType.GOAWAY, 0, 0, payload) + self._goaway_sent = True + + def _send_rst_stream(self, stream_id: int, error_code: int) -> None: + payload = struct.pack("!I", error_code) + self._send_frame(FrameType.RST_STREAM, 0, stream_id, payload) + + def _send_window_update(self, stream_id: int, increment: int) -> None: + payload = struct.pack("!I", increment) + self._send_frame(FrameType.WINDOW_UPDATE, 0, stream_id, payload) + + # -------------------- Stream lifecycle -------------------- + def _close_stream(self, stream: Stream) -> None: + self.streams.pop(stream.stream_id, None) + self._closed_streams.add(stream.stream_id) + # Release stream concurrency slot + self._maybe_unblock_streams() + + def _maybe_unblock_streams(self) -> None: + """Create streams from pending requests if concurrency allows.""" + while self._pending_streams and len(self.streams) < self.max_concurrent_streams: + fut = self._pending_streams.pop(0) + if not fut.done(): + stream = self._create_stream_internal() + fut.set_result(stream) + + def _create_stream_internal(self) -> Stream: + sid = self.next_stream_id + self.next_stream_id += 2 # next client stream + stream = Stream(sid, self, self._loop, self._protocol) + self.streams[sid] = stream + return stream + + async def create_stream(self) -> Stream: + """Return a new client stream, waiting if concurrency limit is reached.""" + if self._goaway_sent or self._goaway_received: + raise ConnectionError("Connection is shutting down") + if len(self.streams) < self.max_concurrent_streams: + return self._create_stream_internal() + # Queue the request + fut: asyncio.Future[Stream] = self._loop.create_future() + self._pending_streams.append(fut) + return await fut + + # -------------------- Request sending -------------------- + async def send_data( + self, stream_id: int, data: bytes, end_stream: bool = True + ) -> None: + """Asynchronously send DATA frames, respecting flow control windows.""" + stream = self.streams[stream_id] + max_frame_size = self.remote_settings[Setting.MAX_FRAME_SIZE] + offset = 0 + total = len(data) + + while offset < total: + # Wait until both session and stream windows have capacity + while stream.outbound_window <= 0 or self.session_outbound_window <= 0: + self._flow_control_updated.clear() + await self._flow_control_updated.wait() + + chunk_size = min( + max_frame_size, + stream.outbound_window, + self.session_outbound_window, + total - offset, + ) + flags = 0 + if offset + chunk_size >= total and end_stream: + flags |= FlagData.END_STREAM + + if stream.state == StreamState.OPEN: + stream.transition(StreamState.HALF_CLOSED_LOCAL) + elif stream.state == StreamState.HALF_CLOSED_REMOTE: + stream.transition(StreamState.CLOSED) + self._close_stream(stream) + # else: error + + self._send_frame( + FrameType.DATA, + flags, + stream.stream_id, + data[offset : offset + chunk_size], + ) + stream.outbound_window -= chunk_size + self.session_outbound_window -= chunk_size + offset += chunk_size + + # there is space available + if self.session_outbound_window and stream.outbound_window: + self._flow_control_updated.set() + + def send_headers( + self, + stream_id: int, + method: str, + url: URL, + headers: Iterable[tuple[str, str]], + end_stream: bool = False, + ) -> None: + stream = self.streams[stream_id] + path_and_query = url.path + if url.query: + path_and_query += "?" + url.raw_query_string + + # Build pseudo‑headers + assert url.scheme + assert url.host + + req_headers = [ + (":method", method), + (":path", path_and_query), + (":scheme", url.scheme), + # XXX add port + (":authority", url.host), + ] + + for name, value in headers: + lname = name.lower() + # HTTP/2 forbids connection-specific headers and the Host header + if lname in ( + "host", + "connection", + "keep-alive", + "proxy-connection", + "transfer-encoding", + "upgrade", + ): + continue + req_headers.append((lname, value)) + + hdrs = self.hpack_encoder.encode(req_headers) + + flags = FlagHeaders.END_HEADERS + + # stream transitions + stream.transition(StreamState.OPEN) + if end_stream: + flags |= FlagHeaders.END_STREAM + stream.transition(StreamState.HALF_CLOSED_LOCAL) + self._send_frame(FrameType.HEADERS, flags, stream.stream_id, hdrs) + + # -------------------- Connection lifecycle -------------------- + def initiate_connection(self) -> None: + """Send the connection preface and initial SETTINGS.""" + # Connection preface (RFC 7540, 3.5) + self._transport.write(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + + # Send initial SETTINGS (our preferences) + settings_payload = struct.pack( + "!H I", Setting.ENABLE_PUSH, 0 # disable server push + ) + self._send_frame(FrameType.SETTINGS, 0, 0, settings_payload) + + # Update local HPACK table size if needed + self.hpack_encoder.header_table_size = self.local_settings[ + Setting.HEADER_TABLE_SIZE + ] + + logger.debug("Connection preface and initial SETTINGS sent") + + def maybe_reset_window(self) -> None: + window_size = self.local_settings[Setting.INITIAL_WINDOW_SIZE] + # XXX add semaphore + # HTTP/2 is multiplexed so we might send updates from + # multiple connections and reset the window multiple times + # and if we reset the window with a 0, the server resets the connection + # creating a hard-to-catch bug + # send less updates + updated = window_size - self.session_inbound_window + half_size = window_size // 2 + if self.session_inbound_window < half_size: + self._send_window_update(0, updated) + self.session_inbound_window = window_size + + def _protocol_error(self) -> None: + self._send_goaway(0, ErrorCode.PROTOCOL_ERROR) + self._transport.close() + + # -------------------- Shutdown -------------------- + def close(self) -> None: + """Perform graceful shutdown.""" + self._transport.close() + + @property + def should_close(self) -> bool: + return self._goaway_sent or self._goaway_received + + def is_connected(self) -> bool: + return not self._transport.is_closing() + + +class Http2Protocol(BaseProtocol): + """BaseProtocol subclass bridging transport and Http2Connection.""" + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + super().__init__(loop, None) + + self._connection: Optional[Http2Connection] = None + self._closed_future: Optional[asyncio.Future[None]] = None + self._connection_lost_called = False + + self._should_close = False + self._upgraded = False + + self._read_timeout: Optional[float] = None + self._read_timeout_handle: Optional[asyncio.TimerHandle] = None + + self._auto_decompress: bool = True + + # ------------------------------------------------------------------ + # Properties expected by connector + # ------------------------------------------------------------------ + @property + def closed(self) -> Optional[asyncio.Future[None]]: + """Future that completes when the connection is closed. + + Mirrors ``ResponseHandler.closed``. The future is created lazily + to avoid creating unused futures. + """ + if self._closed_future is None and not self._connection_lost_called: + self._closed_future = self._loop.create_future() + return self._closed_future + + @property + def upgraded(self) -> bool: + return self._upgraded + + @property + def should_close(self) -> bool: + return bool( + self._should_close + or (self._connection is not None and self._connection.should_close) + ) + + # ------------------------------------------------------------------ + # Connection lifecycle + # ------------------------------------------------------------------ + def connection_made(self, transport: asyncio.BaseTransport) -> None: + self.transport = transport # type: ignore[assignment] + self._connection = Http2Connection(self.transport, self._loop, self) # type: ignore[arg-type] + self._connection.initiate_connection() + + def data_received(self, data: bytes) -> None: + if data: + self._reschedule_timeout() + + if self._connection is not None: + self._connection.data_received(data) + + def eof_received(self) -> bool: + self._drop_timeout() + if self._connection is not None: + self._connection.eof_received() + return False + + def connection_lost(self, exc: Optional[BaseException]) -> None: + self._connection_lost_called = True + self._drop_timeout() + + original_connection_error = exc + reraised_exc = exc + connection_closed_cleanly = exc is None + + # Complete the closed future if anyone is waiting on it. + if self._closed_future is not None: + if connection_closed_cleanly: + set_result(self._closed_future, None) + else: + assert original_connection_error is not None + _set_exc( + self._closed_future, + ClientConnectionError( + f"Connection lost: {original_connection_error !s}" + ), + original_connection_error, + ) + + # Let the HTTP/2 connection clean up its streams. + if self._connection is not None: + self._connection.connection_lost(exc) + + self._should_close = True + self._connection = None + self._reading_paused = False + + super().connection_lost(reraised_exc) + + # ------------------------------------------------------------------ + # Explicit close / abort + # ------------------------------------------------------------------ + def force_close(self) -> None: + self._should_close = True + + def close(self) -> None: + self._drop_timeout() + + if self._connection is not None: + self._connection.close() + # we need to null transport here + # multiplexed connections all use the same + # transport unlike in HTTP/1.1 + # this means that when connections + # are released en masse, `close` will be + # called multiple times + # `close` is not idempotent (i.e., multiple calls over the same protocol instance do + # not produce the same results and state) so multiple closes cause null pointer exceptions + # when we try to clean up an attribute of an object that's already nulled + self.transport = None + + def abort(self) -> None: + self.close() + + def is_connected(self) -> bool: + if self._connection is not None: + return self._connection.is_connected() + return False + + # ------------------------------------------------------------------ + # Timeout handling + # ------------------------------------------------------------------ + def start_timeout(self) -> None: + self._reschedule_timeout() + + def _drop_timeout(self) -> None: + if self._read_timeout_handle is not None: + self._read_timeout_handle.cancel() + self._read_timeout_handle = None + + def _reschedule_timeout(self) -> None: + timeout = self._read_timeout + if self._read_timeout_handle is not None: + self._read_timeout_handle.cancel() + + if timeout: + self._read_timeout_handle = self._loop.call_later( + timeout, self._on_read_timeout + ) + else: + self._read_timeout_handle = None + + def _on_read_timeout(self) -> None: + exc = SocketTimeoutError("Timeout on reading data from socket") + self.set_exception(exc) + + # ------------------------------------------------------------------ + # Backpressure + # ------------------------------------------------------------------ + def pause_reading(self) -> None: + # no-op + # this should never be called in HTTP/2 + # because flow-control management is built-in + # the size of the window ensures the server never sends too much data + pass + + def resume_reading(self, resume_parser: bool = True) -> None: + # even this might be unnecessary because + # just updating the window size is enough + super().resume_reading(resume_parser) + + self._reschedule_timeout() + + assert self._connection + self._connection.maybe_reset_window() + + # ------------------------------------------------------------------ + # Error injection + # ------------------------------------------------------------------ + def set_exception( + self, + exc: type[BaseException] | BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + self._should_close = True + self._drop_timeout() + + def set_response_params( + self, + *, + timer: BaseTimerContext | None = None, + skip_payload: bool = False, + read_until_eof: bool = False, + auto_decompress: bool = True, + read_timeout: float | None = None, + read_bufsize: int = DEFAULT_CHUNK_SIZE, + timeout_ceil_threshold: float = 5, + max_line_size: int = 8190, + max_field_size: int = 8190, + max_headers: int = 128, + ) -> None: + # read_bufsize should be respected + # however, the read_buffsize is controlled by the settings + # and the settings are negotiated with the server + del skip_payload, read_until_eof, read_bufsize, timeout_ceil_threshold # compat + del max_line_size, max_field_size, max_headers # HTTP/1.1 + + self._read_timeout = read_timeout + self._auto_decompress = auto_decompress + + # ------------------------------------------------------------------ + # Existing HTTP/2 API + # ------------------------------------------------------------------ + async def create_stream(self) -> Stream: + if self._connection is None: + raise ConnectionError("Connection is not active") + return await self._connection.create_stream() + + async def read_stream( + self, stream_id: int + ) -> tuple[RawResponseMessage, StreamReader]: + if self._connection is None: + raise ConnectionError("Connection is not active") + return await self._connection.streams[stream_id].response_future diff --git a/aiohttp/http2/errors.py b/aiohttp/http2/errors.py new file mode 100644 index 00000000000..854a3e16c8b --- /dev/null +++ b/aiohttp/http2/errors.py @@ -0,0 +1,14 @@ +from enum import IntEnum + + +class ProtocolError(Exception): + pass + + +# (rfc7540/rfc9113, Section 7) +class ErrorCode(IntEnum): + NO_ERROR = 0x0 + PROTOCOL_ERROR = 0x1 + INTERNAL_ERROR = 0x2 + FLOW_CONTROL_ERROR = 0x3 + CANCEL = 0x8 diff --git a/aiohttp/http2/settings.py b/aiohttp/http2/settings.py new file mode 100644 index 00000000000..e69b0c7491e --- /dev/null +++ b/aiohttp/http2/settings.py @@ -0,0 +1,61 @@ +from enum import IntEnum, IntFlag +from typing import Dict + + +# ---------------------------------------------------------------------- +# HTTP/2 Frame Definitions (RFC 7540, 4) +# ---------------------------------------------------------------------- +class FrameType(IntEnum): + DATA = 0x0 + HEADERS = 0x1 + PRIORITY = 0x2 + RST_STREAM = 0x3 + SETTINGS = 0x4 + PUSH_PROMISE = 0x5 + PING = 0x6 + GOAWAY = 0x7 + WINDOW_UPDATE = 0x8 + CONTINUATION = 0x9 + + +class FlagData(IntFlag): + END_STREAM = 0x1 + PADDED = 0x8 + + +class FlagHeaders(IntFlag): + END_STREAM = 0x1 + END_HEADERS = 0x4 + PADDED = 0x8 + PRIORITY = 0x20 + + +class FlagSettings(IntFlag): + ACK = 0x1 + + +class FlagPing(IntFlag): + ACK = 0x1 + + +# Known settings parameters +class Setting(IntEnum): + HEADER_TABLE_SIZE = 0x1 + ENABLE_PUSH = 0x2 + MAX_CONCURRENT_STREAMS = 0x3 + INITIAL_WINDOW_SIZE = 0x4 + MAX_FRAME_SIZE = 0x5 + MAX_HEADER_LIST_SIZE = 0x6 + SETTINGS_ENABLE_CONNECT_PROTOCOL = 0x8 + NO_RFC7540_PRIORITIES = 0x9 + + +# Default values (RFC 7540, 6.5.2) +DEFAULT_SETTINGS: Dict[Setting, int] = { + Setting.HEADER_TABLE_SIZE: 4096, + Setting.ENABLE_PUSH: 0, # we do not support server push + Setting.MAX_CONCURRENT_STREAMS: 2**32 - 1, + Setting.INITIAL_WINDOW_SIZE: 65535, + Setting.MAX_FRAME_SIZE: 16384, + Setting.MAX_HEADER_LIST_SIZE: 2**32 - 1, +} diff --git a/aiohttp/http2/stream.py b/aiohttp/http2/stream.py new file mode 100644 index 00000000000..5fad74be5c8 --- /dev/null +++ b/aiohttp/http2/stream.py @@ -0,0 +1,262 @@ +import asyncio +from enum import IntEnum +from typing import TYPE_CHECKING, Dict, Iterable, Optional, Set + +from hpack import HeaderTuple + +from ..helpers import DEFAULT_CHUNK_SIZE +from ..http_exceptions import ContentEncodingError +from ..http_parser import DeflateBuffer, RawResponseMessage +from ..streams import StreamReader +from .adapter import feed_headers +from .errors import ErrorCode, ProtocolError +from .settings import Setting + +if TYPE_CHECKING: + from .connection import Http2Connection, Http2Protocol + + +# ---------------------------------------------------------------------- +# Stream State Machine (RFC 7540 5.1) +# ---------------------------------------------------------------------- +class StreamState(IntEnum): + IDLE = 0 + RESERVED_LOCAL = 1 + RESERVED_REMOTE = 2 + OPEN = 3 + HALF_CLOSED_LOCAL = 4 + HALF_CLOSED_REMOTE = 5 + CLOSED = 6 + + +# Valid transitions (RFC 7540 Figure 2) +VALID_TRANSITIONS: Dict[StreamState, Set[StreamState]] = { + StreamState.IDLE: { + StreamState.OPEN, + StreamState.RESERVED_LOCAL, + StreamState.RESERVED_REMOTE, + }, + StreamState.RESERVED_LOCAL: {StreamState.HALF_CLOSED_REMOTE, StreamState.CLOSED}, + StreamState.RESERVED_REMOTE: {StreamState.HALF_CLOSED_LOCAL, StreamState.CLOSED}, + StreamState.OPEN: {StreamState.HALF_CLOSED_LOCAL, StreamState.HALF_CLOSED_REMOTE}, + StreamState.HALF_CLOSED_LOCAL: {StreamState.CLOSED}, + StreamState.HALF_CLOSED_REMOTE: {StreamState.CLOSED}, + StreamState.CLOSED: set(), +} + + +class Stream: + """A single HTTP/2 stream with streaming support and optional decompression.""" + + __slots__ = ( + "stream_id", + "state", + "conn", + "outbound_window", + "inbound_window", + "response_future", + "response_headers", + "response", + "body_reader", + "decompressor", + "_pending_data", + "_headers_received", + "_inbound_window_initial", + "_auto_decompress", + "closed_event", + ) + + def __init__( + self, + stream_id: int, + conn: "Http2Connection", + loop: asyncio.AbstractEventLoop, + protocol: "Http2Protocol", + ) -> None: + self.stream_id = stream_id + self.state = StreamState.IDLE + self.conn = conn + + # Flow-control windows + self.outbound_window: int = conn.remote_settings[Setting.INITIAL_WINDOW_SIZE] + self.inbound_window: int = conn.local_settings[Setting.INITIAL_WINDOW_SIZE] + self._inbound_window_initial: int = self.inbound_window + + self.response_future: asyncio.Future[ + tuple[RawResponseMessage, StreamReader] + ] = loop.create_future() + + self.response_headers: Optional[Iterable[tuple[str, str]]] = None + self.response: RawResponseMessage | None = None + self.decompressor: Optional[DeflateBuffer] = None + # NOTE if headers never come and the server sends a enough data to + # make us run out of memory the process might be OOM killed + self._pending_data: bytearray = bytearray() + self._headers_received = False + self._auto_decompress = protocol._auto_decompress + + self.body_reader = StreamReader( + protocol, + # low_water == window size so pause is never called + conn.local_settings[Setting.INITIAL_WINDOW_SIZE], + loop=loop, + ) + + self.closed_event: asyncio.Event = asyncio.Event() + + def transition(self, new_state: StreamState) -> None: + if ( + new_state not in VALID_TRANSITIONS[self.state] + and new_state != StreamState.CLOSED + ): + raise ProtocolError( + f"Invalid stream state transition {self.state.name} -> {new_state.name}" + ) + self.state = new_state + if new_state == StreamState.CLOSED: + self.closed_event.set() + + # ------------------------------------------------------------------ + # Data and header reception + # ------------------------------------------------------------------ + def maybe_reset_window(self) -> None: + """ + Reset stream-level window. + + Currently, this is done naively. + """ + # Fine-grained, stream-level flow-control ensures + # all streams get a fair share of the bandwidth. + # It requires synchronizing the `Stream`, the protocol, and the payload. + # Given we are using the HTTP/1.1 payload which resumes the read at a protocol level + # we can't have this kind of control without introducing a breaking change. + # Thus, flow-control is not enforced here. + # Unless the user has different consumers for the streams + # (e.g., is acting as a proxy for multiple hosts) + # this doesn't hurt throughput. + if self.inbound_window < self._inbound_window_initial // 2: + increment = self._inbound_window_initial - self.inbound_window + self.inbound_window = self._inbound_window_initial + self.conn._send_window_update(self.stream_id, increment) + + def receive_data(self, data: bytes, end_stream: bool) -> None: + """Process incoming DATA frame payload.""" + self.inbound_window -= len(data) + + # --- stream-level flow control refill --- + self.maybe_reset_window() + + if not self._headers_received: + # Buffer until we know the content-encoding. + self._pending_data.extend(data) + else: + # Feed data to the decompressor or directly to the reader. + if self.decompressor is not None: + try: + self.decompressor.feed_data(data) + except ContentEncodingError as exc: + self.body_reader.set_exception(exc) + self.conn._send_rst_stream(self.stream_id, ErrorCode.INTERNAL_ERROR) + return + else: + self.body_reader.feed_data(data) + + if end_stream: + # Flush any remaining decompressed data and signal EOF. + if self.decompressor is not None: + try: + self.decompressor.feed_eof() + except ContentEncodingError as exc: + self.body_reader.set_exception(exc) + self.conn._send_rst_stream(self.stream_id, ErrorCode.INTERNAL_ERROR) + return + else: + self.body_reader.feed_eof() + + if self.state == StreamState.OPEN: + self.transition(StreamState.HALF_CLOSED_REMOTE) + elif self.state == StreamState.HALF_CLOSED_LOCAL: + self.transition(StreamState.CLOSED) + self.conn._close_stream(self) + else: + raise ProtocolError( + f"Unexpected stream state {self.state.name} for END_STREAM" + ) + + def receive_headers( + self, + headers: Iterable[HeaderTuple], + end_stream: bool, + ) -> None: + """Process incoming HEADERS frame payload.""" + # HeaderTuple can be tuple[str, str] yet the type hint says + # it's tuple[bytes, bytes] + self.response_headers = headers # type: ignore[assignment] + self.response = feed_headers(headers) # type: ignore[arg-type] + self._headers_received = True + + if self._auto_decompress: + encoding = self.response.headers.get("content-encoding") + if encoding: + # Create a DeflateBuffer wrapping the StreamReader. + self.decompressor = DeflateBuffer( + self.body_reader, + encoding=encoding, + max_decompress_size=DEFAULT_CHUNK_SIZE, + ) + # Feed any data that arrived before headers. + if self._pending_data: + try: + self.decompressor.feed_data(bytes(self._pending_data)) + except ContentEncodingError as exc: + self.body_reader.set_exception(exc) + self.conn._send_rst_stream( + self.stream_id, ErrorCode.INTERNAL_ERROR + ) + self._pending_data.clear() + return + self._pending_data.clear() + else: + # No decompression needed; feed pending data directly. + if self._pending_data: + self.body_reader.feed_data(bytes(self._pending_data)) + self._pending_data.clear() + else: + # Autodecompress disabled: feed raw data. + if self._pending_data: + self.body_reader.feed_data(bytes(self._pending_data)) + self._pending_data.clear() + + # Deliver the reader as soon as headers are known. + self.maybe_deliver_response() + + if end_stream: + # If END_STREAM on HEADERS, signal EOF immediately. + if self.decompressor is not None: + try: + self.decompressor.feed_eof() + except ContentEncodingError as exc: + self.body_reader.set_exception(exc) + self.conn._send_rst_stream(self.stream_id, ErrorCode.INTERNAL_ERROR) + return + else: + self.body_reader.feed_eof() + + if self.state == StreamState.OPEN: + self.transition(StreamState.HALF_CLOSED_REMOTE) + elif self.state == StreamState.HALF_CLOSED_LOCAL: + self.transition(StreamState.CLOSED) + self.conn._close_stream(self) + else: + raise ProtocolError( + f"Unexpected stream state {self.state.name} for END_STREAM on headers" + ) + + def maybe_deliver_response(self) -> None: + """Resolve the response future once headers have been received.""" + if ( + self.response_headers is not None + and self.response is not None + and not self.response_future.done() + ): + self.response_future.set_result((self.response, self.body_reader)) diff --git a/aiohttp/http2/synchro.py b/aiohttp/http2/synchro.py new file mode 100644 index 00000000000..35e69541431 --- /dev/null +++ b/aiohttp/http2/synchro.py @@ -0,0 +1,122 @@ +import asyncio +from collections import defaultdict +from typing import Any, Dict, List, Set + + +class HostProbeSynchronizer: + """ + A key‑based synchronisation primitive. + + For each key (e.g. a connection host) only the first task that calls + `acquire(key)` is allowed to proceed immediately. All other tasks wait + until `release(key)` is called. When `release(key)` happens, *all* + waiting tasks are woken and may continue concurrently. After release, + the key is unlocked, so the next `acquire(key)` will again lock it and + proceed immediately. + + This is designed for the initial ALPN probe of an HTTP connection: only + one task should perform the probe per host; once the protocol is known + (h1 or h2), all waiting tasks can proceed (for h1 they may open + connections in parallel, for h2 they reuse the single connection). + """ + + def __init__(self) -> None: + # Keys that are currently locked (i.e. an acquire has succeeded and + # release has not yet been called). + self._locked: Set[Any] = set() + self._done: Set[Any] = set() + # For each locked key, a list of futures that waiting tasks are awaiting. + self._waiters: Dict[Any, List[asyncio.Future[None]]] = defaultdict(list) + + async def acquire(self, key: Any) -> None: + """ + Wait until the key is unlocked, then lock it and return. + + If the key is already locked, the task is suspended until `release` + is called for that key. When release is called, all waiting tasks are + woken and return from this method (they do not re‑acquire the lock). + """ + if key in self._done: + # Already released (no-op) + return + + if key not in self._locked: + # First to acquire: lock the key and proceed. + self._locked.add(key) + return + + # Key is locked; we must wait. + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._waiters[key].append(fut) + try: + await fut + except asyncio.CancelledError: + # Remove our future from the waiters list if we are cancelled. + waiters = self._waiters.get(key) + if waiters is not None and fut in waiters: + waiters.remove(fut) + # If the list becomes empty and the key is still locked, we may + # optionally clean up the empty list to save memory. + if waiters is not None and not waiters: + del self._waiters[key] + raise + + def release(self, key: Any) -> None: + """ + Unlock the given key and wake all tasks waiting for it. + + After this call, all tasks that had called `acquire(key)` and were + suspended will resume. The key becomes unlocked, so a subsequent + `acquire(key)` will lock it again and return immediately. + """ + if key not in self._locked: + # this happens when any request after the first calls `release` + return + + # Remove the lock. + self._locked.remove(key) + self._done.add(key) + + # Wake all waiters. + waiters = self._waiters.pop(key, []) + for fut in waiters: + if not fut.done(): + fut.set_result(None) + + def is_locked(self, key: Any) -> bool: + """Return True if the key is currently locked.""" + return key in self._locked + + def __contains__(self, key: Any) -> bool: + return self.is_locked(key) + + +async def main() -> None: + sync = HostProbeSynchronizer() + host = "example.com" + + async def worker(name: str) -> None: + print(f"{name}: before acquire") + await sync.acquire(host) + print(f"{name}: acquired") + # Simulate some work + await asyncio.sleep(1) + print(f"{name}: releasing") + sync.release(host) + + # First worker locks the host + t1 = asyncio.create_task(worker("first")) + await asyncio.sleep(0.1) # ensure first has acquired + # Two more workers wait + t2 = asyncio.create_task(worker("second")) + t3 = asyncio.create_task(worker("third")) + await asyncio.sleep(0.1) + # Now release will wake both t2 and t3 + await t1 + await t2 + await t3 + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/aiohttp/http_protocol.py b/aiohttp/http_protocol.py new file mode 100644 index 00000000000..107e8af11e6 --- /dev/null +++ b/aiohttp/http_protocol.py @@ -0,0 +1,55 @@ +import asyncio +from typing import Any, Optional + +from .client_proto import ResponseHandler +from .http2.adapter import _get_version +from .http2.connection import Http2Protocol + + +class HttpDispatcherProtocol(asyncio.Protocol): + """Protocol that switches between HTTP/1.1 and HTTP/2 based on ALPN.""" + + __slots__ = ("_loop", "_transport", "_handler") + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop = loop + self._transport: Optional[asyncio.Transport] = None + self._handler: Optional[asyncio.Protocol] = None + + # ---- Transport callbacks forwarded to the real handler ---- + def connection_made(self, transport: asyncio.BaseTransport) -> None: + if self._handler: + return self._handler.connection_made(transport) + + self._transport = transport # type: ignore[assignment] + + # Determine ALPN after TLS is established + alpn_protocol: str = _get_version(transport) + + if alpn_protocol == "h2": + self._handler = Http2Protocol(self._loop) + else: + self._handler = ResponseHandler(self._loop) + + # Hand the real transport to the handler. The handler will now own + # all incoming data and callbacks. + self._handler.connection_made(transport) + + def __getattribute__(self, name: str) -> Any: + if not name.startswith("__") and name not in { + "connection_made", + "__getattribute__", + "_handler", + "_transport", + "_loop", + }: + return getattr(self._handler, name) + return super().__getattribute__(name) + + def __setattr__(self, name: str, value: Any) -> None: + if name not in {"_handler", "_transport", "_loop"}: + return self._handler.__setattr__(name, value) + return super().__setattr__(name, value) + + def __delattr__(self, name: str) -> None: + return self._handler.__delattr__(name) diff --git a/aiohttp/http_writer.py b/aiohttp/http_writer.py index a1168cfdebb..e648b8b2462 100644 --- a/aiohttp/http_writer.py +++ b/aiohttp/http_writer.py @@ -3,13 +3,11 @@ import asyncio import re import sys -from typing import ( # noqa +from typing import ( TYPE_CHECKING, - Any, Awaitable, Callable, Iterable, - List, NamedTuple, Optional, Union, @@ -44,6 +42,7 @@ class HttpVersion(NamedTuple): HttpVersion10 = HttpVersion(1, 0) HttpVersion11 = HttpVersion(1, 1) +HttpVersion2 = HttpVersion(2, 0) _T_OnChunkSent = Optional[ diff --git a/docs/conf.py b/docs/conf.py index a3254645cfb..9cf25a6082e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -66,7 +66,7 @@ try: - import sphinxcontrib.spelling # noqa + import sphinxcontrib.spelling extensions.append("sphinxcontrib.spelling") except ImportError: diff --git a/pyproject.toml b/pyproject.toml index eb659f9b20e..f48977b0fde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "propcache >= 0.2.0", "typing_extensions >= 4.4 ; python_version < '3.13'", "yarl >= 1.17.0, < 2.0", + "hpack >= 4.2.0" ] dynamic = [ "version", diff --git a/requirements/runtime-deps.in b/requirements/runtime-deps.in index 3b57b27c57a..bfd1a1cee9a 100644 --- a/requirements/runtime-deps.in +++ b/requirements/runtime-deps.in @@ -9,6 +9,7 @@ backports.zstd; platform_python_implementation == 'CPython' and python_version < Brotli >= 1.2; platform_python_implementation == 'CPython' and sys_platform != 'android' and sys_platform != 'ios' brotlicffi >= 1.2; platform_python_implementation != 'CPython' frozenlist >= 1.1.1 +hpack >= 4.2.0 multidict >=4.5, < 7.0 propcache >= 0.2.0 typing_extensions >= 4.4 ; python_version < '3.13' diff --git a/tests/conftest.py b/tests/conftest.py index 6e18199bed6..46e442ecc2a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -38,6 +38,7 @@ from aiohttp.compression_utils import ZLibBackend, ZLibBackendProtocol, set_zlib_backend from aiohttp.helpers import TimerNoop from aiohttp.http import WS_KEY, HttpVersion11 +from aiohttp.http2.connection import Http2Connection, Http2Protocol from aiohttp.test_utils import REUSE_ADDRESS @@ -526,3 +527,42 @@ def slow_fn(*args: Any, **kwargs: Any) -> Any: executor = SlowExecutor(max_workers=10) yield executor executor.shutdown(wait=True) + + +# ---------------------------------------------------------------------- +# HTTP/2 Fixtures +# ---------------------------------------------------------------------- +@pytest.fixture +def mock_transport() -> mock.MagicMock: + """Return a mock asyncio.Transport that records writes.""" + t = mock.MagicMock(spec=asyncio.Transport) + t.is_closing.return_value = False + t.write = mock.MagicMock() + t.close = mock.MagicMock() + return t + + +@pytest.fixture +async def protocol( + mock_transport: mock.MagicMock, +) -> tuple[Http2Protocol, mock.MagicMock]: + """Create Http2Protocol and simulate connection_made.""" + event_loop = asyncio.get_running_loop() + proto = Http2Protocol(event_loop) + proto.connection_made(mock_transport) + mock_transport.write.reset_mock() + return proto, mock_transport + + +@pytest.fixture +async def connection( + mock_transport: mock.MagicMock, protocol: tuple[Http2Protocol, mock.MagicMock] +) -> tuple[Http2Connection, mock.MagicMock]: + """Set up Http2Connection with mock transport, send preface, and clear write log.""" + h2_proto, _ = protocol + event_loop = asyncio.get_running_loop() + + conn = Http2Connection(mock_transport, event_loop, h2_proto) + conn.initiate_connection() + mock_transport.write.reset_mock() # discard preface + initial SETTINGS + return conn, mock_transport diff --git a/tests/http2/__init__.py b/tests/http2/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/http2/fuzz.py b/tests/http2/fuzz.py new file mode 100644 index 00000000000..81deed65aa1 --- /dev/null +++ b/tests/http2/fuzz.py @@ -0,0 +1,414 @@ +import random +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Tuple + +from http2.utils import ( # noqa: I900 + build_data_frame, + build_goaway, + build_headers_frame, + build_ping, + build_rst_stream, + build_settings_frame, + build_window_update, +) + +from aiohttp.http2.connection import Http2Connection +from aiohttp.http2.errors import ErrorCode +from aiohttp.http2.settings import DEFAULT_SETTINGS, FrameType, Setting + + +# ---------------------------------------------------------------------- +# Server state representation +# ---------------------------------------------------------------------- +@dataclass +class ServerStreamState: + """State of a single stream from the server's perspective.""" + + stream_id: int + state: str = "idle" + recv_window: int = 65535 + send_window: int = 65535 + pending_data: int = 0 + + +@dataclass +class ServerConnectionState: + """Full HTTP/2 connection state as seen by the server.""" + + streams: Dict[int, ServerStreamState] = field(default_factory=dict) + last_stream_id: int = 0 + settings: Dict[Setting, int] = field( + default_factory=lambda: DEFAULT_SETTINGS.copy() + ) + peer_settings: Dict[Setting, int] = field( + default_factory=dict + ) # settings received from client + goaway_sent: bool = False + goaway_received: bool = False + + def get_stream(self, stream_id: int) -> ServerStreamState: + if stream_id not in self.streams: + self.streams[stream_id] = ServerStreamState(stream_id=stream_id) + return self.streams[stream_id] + + +# ---------------------------------------------------------------------- +# Fuzzer configuration +# ---------------------------------------------------------------------- +@dataclass +class FuzzerConfig: + """Probability distributions for frame generation. + + Each field is a dict mapping a choice to a weight (or a single value). + For example, frame_type_probs = {FrameType.HEADERS: 30, FrameType.DATA: 20, ...} + The fuzzer will randomly choose according to these weights. + """ + + frame_type_probs: Dict[int, int] = field( + default_factory=lambda: { + FrameType.HEADERS: 30, + FrameType.DATA: 25, + FrameType.SETTINGS: 10, + FrameType.RST_STREAM: 10, + FrameType.WINDOW_UPDATE: 10, + FrameType.PING: 10, + FrameType.GOAWAY: 5, + } + ) + # Probabilities for flags on HEADERS frames + headers_end_stream_prob: float = 0.3 + headers_end_headers_prob: float = 0.9 + headers_priority_prob: float = 0.1 + # DATA frame flags + data_end_stream_prob: float = 0.4 + data_padded_prob: float = 0.2 + # Probability that a stream ID is invalid (e.g., 0 or even) + invalid_stream_id_prob: float = 0.1 + # Probability to violate protocol deliberately (e.g., send DATA on closed stream) + violate_protocol_prob: float = 0.2 + # Maximum number of frames to send in one cycle + max_frames_per_cycle: int = 3 + # Random seed (optional) + seed: Optional[int] = None + + +# ---------------------------------------------------------------------- +# Main fuzzer class +# ---------------------------------------------------------------------- +class Http2ServerFuzzer: + """Stateful fuzzer that acts as an HTTP/2 server and sends arbitrary frames to an aiohttp client connection. + + Parameters + ---------- + connection : Http2Connection + The client under test. + transport : MagicMock + Mock transport used by the connection; its write method will be + recorded to capture client frames. + config : FuzzerConfig + Probability distributions for frame generation. + max_cycles : int + Number of fuzzing cycles to run. + verification : Callable + Function that receives ``(sent_frames, client_frames, state)`` and + returns ``True`` if no bug was found, ``False`` if a bug occurred. + It can also raise an AssertionError to indicate a bug with a message. + """ + + def __init__( + self, + connection: Http2Connection, + transport: Any, + config: FuzzerConfig = FuzzerConfig(), + max_cycles: int = 100, + verification: Optional[ + Callable[ + [List[bytes], List[Tuple[int, int, int, bytes]], ServerConnectionState], + bool, + ] + ] = None, + ): + self.connection = connection + self.transport = transport + self.config = config + self.max_cycles = max_cycles + self.verification = verification or self.default_verification + + self.state = ServerConnectionState() + self.rng = random.Random(config.seed) + + # Keep track of client frames sent via transport.write + self.client_frame_buffer: List[bytes] = [] + self._record_client_frames() + + def _record_client_frames(self) -> None: + """Patch transport.write to collect outgoing client frames.""" + original_write = self.transport.write + + def write_and_record(data: bytes) -> None: + self.client_frame_buffer.append(data) + original_write(data) + + self.transport.write = write_and_record + + def _clear_client_frames(self) -> None: + self.client_frame_buffer.clear() + + def _parse_client_frames(self, data: bytes) -> List[Tuple[int, int, int, bytes]]: + """Parse raw bytes into a list of (length, type, flags, payload).""" + frames = [] + while len(data) >= 9: + length = int.from_bytes(data[:3], "big") + ftype = data[3] + flags = data[4] + if len(data) < 9 + length: + break # incomplete frame + payload = data[9 : 9 + length] + frames.append((length, ftype, flags, payload)) + data = data[9 + length :] + return frames + + def _get_client_frames_since_last_cycle(self) -> List[Tuple[int, int, int, bytes]]: + """Extract and parse all client frames written since last clear.""" + all_raw = b"".join(self.client_frame_buffer) + self._clear_client_frames() + return self._parse_client_frames(all_raw) + + # ------------------------------------------------------------------ + # State update helpers (server perspective) + # ------------------------------------------------------------------ + def _apply_client_frames_to_state( + self, frames: List[Tuple[int, int, int, bytes]] + ) -> None: + """Update server state based on what the client sent.""" + for length, ftype, flags, payload in frames: + if ftype == FrameType.SETTINGS: + # parse settings payload (each setting is 6 bytes: id, value) + if not (flags & 0x1): # not ACK + for i in range(0, len(payload), 6): + setting_id = int.from_bytes(payload[i : i + 2], "big") + value = int.from_bytes(payload[i + 2 : i + 6], "big") + # Convert to Setting enum, ignore unknown ids + try: + setting: Setting = Setting(setting_id) + except ValueError: + continue + self.state.peer_settings[setting] = value + elif ftype == FrameType.HEADERS: + pass + elif ftype == FrameType.RST_STREAM: + stream_id = int.from_bytes(payload[:4], "big") & 0x7FFFFFFF + if stream_id in self.state.streams: + self.state.streams[stream_id].state = "closed" + elif ftype == FrameType.WINDOW_UPDATE: + stream_id = int.from_bytes(payload[:4], "big") & 0x7FFFFFFF + increment = int.from_bytes(payload[4:8], "big") + if stream_id == 0: + # connection-level window update + pass + else: + stream = self.state.get_stream(stream_id) + stream.send_window += increment + elif ftype == FrameType.GOAWAY: + self.state.goaway_received = True + + # ------------------------------------------------------------------ + # Frame generation (server -> client) + # ------------------------------------------------------------------ + def _choose_frame_type(self) -> int: + types = list(self.config.frame_type_probs.keys()) + weights = list(self.config.frame_type_probs.values()) + return self.rng.choices(types, weights=weights, k=1)[0] + + def _generate_frame(self) -> bytes: + """Generate a single frame based on current state and probabilities.""" + ftype = self._choose_frame_type() + # Determine stream ID (0 for connection-level frames, random valid/invalid otherwise) + use_invalid_stream = self.rng.random() < self.config.invalid_stream_id_prob + if ftype in {FrameType.SETTINGS, FrameType.PING, FrameType.GOAWAY}: + stream_id = 0 + else: + # Choose from existing streams or create a new one + stream_ids = [ + s for s in self.state.streams.keys() if s % 2 == 1 + ] # client-initiated odd streams + if stream_ids and self.rng.random() < 0.7: + stream_id = self.rng.choice(stream_ids) + else: + # new stream ID (odd, greater than last used) + stream_id = ( + max([0] + [s for s in self.state.streams.keys() if s % 2 == 1]) + 2 + ) + if use_invalid_stream: + # make it even or 0 to test error handling + stream_id = self.rng.choice([0, stream_id - 1]) + self.state.get_stream(stream_id) # ensure exists + + if ftype == FrameType.SETTINGS: + # Random settings (maybe invalid values) + settings_pairs = [] + for _ in range(self.rng.randint(0, 3)): + setting = self.rng.choice(list(Setting)) + value = self.rng.randint(0, 2**32 - 1) + settings_pairs.append((setting, value)) + ack = self.rng.random() < 0.2 + frame = build_settings_frame(settings_pairs, ack=ack) + + elif ftype == FrameType.HEADERS: + # Generate response headers + headers = [ + (":status", str(self.rng.choice([200, 204, 400, 500]))), + ("content-type", "text/plain"), + ( + "x-fuzz", + "".join(self.rng.choices("abcdef", k=self.rng.randint(0, 20))), + ), + ] + end_stream = self.rng.random() < self.config.headers_end_stream_prob + end_headers = self.rng.random() < self.config.headers_end_headers_prob + priority = None + if self.rng.random() < self.config.headers_priority_prob: + priority = b"\x00" * 5 # dummy priority + frame = build_headers_frame( + stream_id, + headers, + end_headers=end_headers, + end_stream=end_stream, + priority=priority, + ) + + elif ftype == FrameType.DATA: + data_len = self.rng.randint(0, self.state.settings[Setting.MAX_FRAME_SIZE]) + data = bytes(self.rng.getrandbits(8) for _ in range(data_len)) + end_stream = self.rng.random() < self.config.data_end_stream_prob + padded = self.rng.random() < self.config.data_padded_prob + frame = build_data_frame(stream_id, data, end_stream=end_stream, pad=padded) + + elif ftype == FrameType.RST_STREAM: + error_code = self.rng.choice(list(ErrorCode)).value + frame = build_rst_stream(stream_id, error_code) + + elif ftype == FrameType.WINDOW_UPDATE: + increment = self.rng.randint(1, 2**31 - 1) + frame = build_window_update(stream_id, increment) + + elif ftype == FrameType.PING: + ack = self.rng.random() < 0.5 + opaque = bytes(self.rng.getrandbits(8) for _ in range(8)) + frame = build_ping(ack=ack, opaque=opaque) + + elif ftype == FrameType.GOAWAY: + last_stream_id = self.rng.choice([0, stream_id]) + error_code = self.rng.choice(list(ErrorCode)).value + extra = b"fuzz" + frame = build_goaway(last_stream_id, error_code, extra) + self.state.goaway_sent = True + + else: + frame = b"" # unsupported for now + + # Update server state based on the frame we are about to send + self._apply_sent_frame_to_state(ftype, stream_id, frame) + return frame + + def _apply_sent_frame_to_state( + self, ftype: int, stream_id: int, frame: bytes + ) -> None: + """Update server state after sending a frame.""" + if ftype == FrameType.HEADERS: + stream = self.state.get_stream(stream_id) + # If END_STREAM flag set, stream becomes half-closed (local) + if frame[4] & 0x1: # END_STREAM flag + stream.state = "half_closed_local" + else: + stream.state = "open" + elif ftype == FrameType.DATA: + stream = self.state.get_stream(stream_id) + if frame[4] & 0x1: + stream.state = "half_closed_local" + # Reduce send window + payload_len = len(frame) - 9 + stream.send_window -= payload_len + elif ftype == FrameType.RST_STREAM: + self.state.get_stream(stream_id).state = "closed" + elif ftype == FrameType.GOAWAY: + self.state.goaway_sent = True + + # ------------------------------------------------------------------ + # Main fuzzing loop + # ------------------------------------------------------------------ + def run(self) -> List[Dict[str, Any]]: + """Run the fuzzer for max_cycles cycles. + + Returns a list of bug reports (empty if no bugs found). + """ + bugs = [] + for cycle in range(self.max_cycles): + # Step 1: Get client frames sent since last cycle + client_frames = self._get_client_frames_since_last_cycle() + self._apply_client_frames_to_state(client_frames) + + # Step 2: Generate one or more server frames + num_frames = self.rng.randint(1, self.config.max_frames_per_cycle) + sent_frames = [] + for _ in range(num_frames): + frame = self._generate_frame() + sent_frames.append(frame) + # Feed each frame to the client immediately + self.connection.data_received(frame) + # Check for immediate client reaction (e.g., exceptions) + # The client may write frames during data_received + # We'll capture them in the next iteration, but we can also + # check for exceptions here. + + # Step 3: Collect client frames produced as a reaction + reaction_frames = self._get_client_frames_since_last_cycle() + self._apply_client_frames_to_state(reaction_frames) + + # Step 4: Verify + try: + is_bug = not self.verification(sent_frames, reaction_frames, self.state) + except AssertionError as e: + is_bug = True + bug_info = str(e) + else: + bug_info = "Verification returned False" + + if is_bug: + bugs.append( + { + "cycle": cycle, + "sent_frames": sent_frames, + "client_frames": reaction_frames, + "state": self.state, + "info": bug_info, + } + ) + return bugs + + # ------------------------------------------------------------------ + # Default verification (override or replace) + # ------------------------------------------------------------------ + def default_verification( + self, + sent_frames: List[bytes], + client_frames: List[Tuple[int, int, int, bytes]], + state: ServerConnectionState, + ) -> bool: + """Basic sanity checks: + + - Client must not crash (obviously handled by test harness) + - If we send a GOAWAY, client should eventually close connection. + - If we send RST_STREAM on a stream, client should not send DATA on that stream. + - If we violate a protocol rule, client should send an error (RST_STREAM/GOAWAY). + """ + # We only check the client didn't send data to a closed stream + for _, ftype, _, payload in client_frames: + if ftype == FrameType.DATA: + stream_id = int.from_bytes(payload[:4], "big") & 0x7FFFFFFF + if ( + stream_id in state.streams + and state.streams[stream_id].state == "closed" + ): + return False + return True diff --git a/tests/http2/test_fuzz.py b/tests/http2/test_fuzz.py new file mode 100644 index 00000000000..518cb5f7afb --- /dev/null +++ b/tests/http2/test_fuzz.py @@ -0,0 +1,16 @@ +"""Fuzzy tests for HTTP/2. These aim to test a sample of the cartesian product of all possible HTTP/2 messages.""" + +from typing import Any +from unittest.mock import MagicMock + +from http2.fuzz import FuzzerConfig, Http2ServerFuzzer # noqa: I900 + + +async def test_fuzz_client( + connection: Any, mock_transport: Any, event_loop: Any +) -> None: + conn, transport = connection + config = FuzzerConfig(seed=42, max_frames_per_cycle=3) + fuzzer = Http2ServerFuzzer(conn, transport, config, max_cycles=50) + bugs = fuzzer.run() + assert not bugs, f"Bugs found: {bugs}" diff --git a/tests/http2/test_http2.py b/tests/http2/test_http2.py new file mode 100644 index 00000000000..629d37036a5 --- /dev/null +++ b/tests/http2/test_http2.py @@ -0,0 +1,1561 @@ +""" +Test suite for aiohttp.http2 + +Categories: +- integration: against a real httpbin server (skipped, requires network) +- unit / protocol: black‑box frame‑level RFC compliance +- unit / misc: race conditions, deadlocks, edge cases +""" + +import asyncio +import struct +from typing import Any, Dict, Generator, List, Optional, Tuple +from unittest.mock import MagicMock + +import pytest +from http2.utils import ( # noqa: I900 (http2 is not a dependency) + build_data_frame, + build_goaway, + build_headers_frame, + build_ping, + build_rst_stream, + build_settings_frame, + build_window_update, + frame_header, +) + +import aiohttp +from aiohttp import ClientConnectionError, SocketTimeoutError +from aiohttp.connector import TCPConnector +from aiohttp.helpers import DEFAULT_CHUNK_SIZE +from aiohttp.http2.connection import Http2Connection, Http2Protocol +from aiohttp.http2.errors import ErrorCode, ProtocolError +from aiohttp.http2.settings import ( + FlagData, + FlagHeaders, + FlagPing, + FlagSettings, + FrameType, + Setting, +) +from aiohttp.http2.stream import Stream, StreamState +from aiohttp.http_exceptions import ContentEncodingError + + +# ====================================================================== +# UNIT TESTS +# ====================================================================== +@pytest.mark.asyncio +async def test_incomplete_frame(connection: Tuple[Http2Connection, MagicMock]) -> None: + conn, _ = connection + frame = b"111111111" + conn.data_received(frame) + assert conn._frame_buffer == frame + + +# ---------------------------------------------------------------------- +# 1. Protocol compliance (black‑box, frame‑by‑frame) +# ---------------------------------------------------------------------- +class TestProtocolCompliance: + @pytest.mark.asyncio + async def test_receive_settings_updates_remote_and_acks( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + conn, transport = connection + # Send server SETTINGS (HEADER_TABLE_SIZE=8192, MAX_CONCURRENT_STREAMS=50) + frame = build_settings_frame( + [ + (Setting.HEADER_TABLE_SIZE, 8192), + (Setting.MAX_CONCURRENT_STREAMS, 50), + ] + ) + conn.data_received(frame) + assert conn.remote_settings[Setting.HEADER_TABLE_SIZE] == 8192 + assert conn.remote_settings[Setting.MAX_CONCURRENT_STREAMS] == 50 + # Must have sent an ACK + assert any( + call[0][0][3:4] == FrameType.SETTINGS.to_bytes(1, "big") + and call[0][0][4] & FlagSettings.ACK + for call in transport.write.call_args_list + ) + + @pytest.mark.asyncio + async def test_receive_data_flow_control( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + conn, transport = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN # assume request already sent + + # Send DATA with some bytes + data = b"hello" + frame = build_data_frame(stream.stream_id, data, end_stream=False) + initial_window = conn.session_inbound_window + conn.data_received(frame) + assert conn.session_inbound_window == initial_window - len(data) + # Should trigger WINDOW_UPDATE when below threshold (32768) + # Because initial window is 65535 and we just consumed 5, still above threshold + assert not any( + b"WINDOW_UPDATE" in call.args[0] for call in transport.write.call_args_list + ) + + # Send more data to drop below 32768 + big_data = b"x" * 40000 + frame2 = build_data_frame(stream.stream_id, big_data, end_stream=False) + conn.data_received(frame2) + # Now session window should have triggered an update + updates = [ + call.args[0] + for call in transport.write.call_args_list + if FrameType.WINDOW_UPDATE.to_bytes(1, "big") in call.args[0] + ] + assert len(updates) >= 1 + + @pytest.mark.asyncio + async def test_rst_stream_handling( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + + frame = build_rst_stream(stream.stream_id, error_code=0x8) # CANCEL + conn.data_received(frame) + + assert stream.state == StreamState.CLOSED + assert stream.response_future.done() + with pytest.raises(RuntimeError): + stream.response_future.result() + assert stream.stream_id not in conn.streams + + @pytest.mark.asyncio + async def test_goaway_cancels_higher_streams( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + conn, _ = connection + # create three streams (1,3,5) + s1 = await conn.create_stream() + s3 = await conn.create_stream() + s5 = await conn.create_stream() + s1.state = s3.state = s5.state = StreamState.OPEN + + # GOAWAY with last_stream_id = 3 + frame = build_goaway(last_stream_id=3, error_code=0) + conn.data_received(frame) + + # s1 (1) and s3 (3) should be unaffected, s5 (5) cancelled + assert s1.stream_id in conn.streams + assert s3.stream_id in conn.streams + assert s5.stream_id not in conn.streams + assert s5.response_future.exception() is not None + + @pytest.mark.asyncio + async def test_ping_ack( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + conn, transport = connection + frame = build_ping(ack=False, opaque=b"12345678") + conn.data_received(frame) + # Expect ACK sent back with same data + acks: List[bytes] = [] + for call in transport.write.call_args_list: + arg = call.args[0] + if FrameType.PING.to_bytes(1, "big") in arg and arg[4] & FlagPing.ACK: + acks.append(arg) + assert len(acks) == 1 + assert b"12345678" in acks[0] + + @pytest.mark.asyncio + async def test_max_concurrent_streams_blocking( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + conn, _ = connection + conn.max_concurrent_streams = 1 + s1 = await conn.create_stream() + # second stream should block + create_task = asyncio.ensure_future(conn.create_stream()) + await asyncio.sleep(0.01) + assert not create_task.done() + # close s1 to release slot + conn._close_stream(s1) + s2 = await create_task + assert s2.stream_id > s1.stream_id + assert len(conn.streams) == 1 + + @pytest.mark.asyncio + async def test_unknown_frame_ignored( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + conn, _ = connection + # send frame type 0x1a (unused) + frame = frame_header(0, 0x1A, 0, 0) + conn.data_received(frame) + # Should not raise, connection stays intact + assert not conn._goaway_sent + + @pytest.mark.asyncio + async def test_bad_hpack_triggers_protocol_error( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + conn, transport = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + # corrupted headers + payload = b"\xff\xff\xff" + frame = ( + frame_header( + len(payload), + FrameType.HEADERS, + FlagHeaders.END_HEADERS, + stream.stream_id, + ) + + payload + ) + conn.data_received(frame) + # Should have sent RST_STREAM and/or GOAWAY + rst = any( + call.args[0][3:4] == FrameType.RST_STREAM.to_bytes(1, "big") + for call in transport.write.call_args_list + ) + goaway = any( + call.args[0][3:4] == FrameType.GOAWAY.to_bytes(1, "big") + for call in transport.write.call_args_list + ) + assert rst or goaway + + @pytest.mark.asyncio + async def test_data_frame_with_padding( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Cover DATA frame with PADDED flag.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + # padded data: pad length 1, data 'x', zero padding byte + frame = build_data_frame(stream.stream_id, b"x", pad=True) + conn.data_received(frame) + # no headers yet so data is buffered + assert stream._pending_data == b"x" + + @pytest.mark.asyncio + async def test_headers_frame_with_priority( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Cover HEADERS frame with PRIORITY flag.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + # priority block: exclusive (1 byte) + dependency (4 bytes) + weight (1 byte) + priority_data = b"\x00\x00\x00\x00\x10" + headers = [(":status", "200")] + frame = build_headers_frame( + stream.stream_id, + headers, + end_headers=True, + end_stream=True, + priority=priority_data, + ) + conn.data_received(frame) + assert stream.response_future.done() + assert stream.response_headers is not None + + @pytest.mark.asyncio + async def test_rst_stream_for_unknown_stream( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """RST_STREAM on an unknown stream ID is silently ignored.""" + conn, _ = connection + frame = build_rst_stream(999, error_code=0) + conn.data_received(frame) # must not raise + + @pytest.mark.asyncio + async def test_rst_stream_when_future_done( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """RST_STREAM when response future already completed.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + stream.response_future.set_result(([], b"")) # type: ignore[arg-type] + frame = build_rst_stream(stream.stream_id, error_code=0) + conn.data_received(frame) + assert stream.state == StreamState.CLOSED + + @pytest.mark.asyncio + async def test_receive_settings_ack( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Receive SETTINGS ACK (should be a no‑op).""" + conn, _ = connection + frame = build_settings_frame(ack=True) + conn.data_received(frame) # no crash + + @pytest.mark.asyncio + async def test_settings_on_nonzero_stream( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """SETTINGS frame on stream_id != 0 triggers GOAWAY.""" + conn, transport = connection + frame = frame_header(0, FrameType.SETTINGS, 0, 5) # stream 5 + conn.data_received(frame) + assert any( + call[0][0][3:4] == FrameType.GOAWAY.to_bytes(1, "big") + for call in transport.write.call_args_list + ) + + @pytest.mark.asyncio + async def test_settings_invalid_payload_length( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """SETTINGS payload not a multiple of 6 triggers protocol error.""" + conn, transport = connection + payload = b"\x00\x01\x02" # 3 bytes + frame = frame_header(len(payload), FrameType.SETTINGS, 0, 0) + payload + conn.data_received(frame) + assert any( + call[0][0][3:4] == FrameType.GOAWAY.to_bytes(1, "big") + for call in transport.write.call_args_list + ) + + @pytest.mark.asyncio + async def test_settings_initial_window_size_update( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """INITIAL_WINDOW_SIZE setting updates stream windows.""" + conn, _ = connection + stream = await conn.create_stream() + old_window = stream.outbound_window + frame = build_settings_frame([(Setting.INITIAL_WINDOW_SIZE, 131072)]) + conn.data_received(frame) + assert stream.outbound_window == old_window + (131072 - 65535) + + @pytest.mark.asyncio + async def test_settings_header_table_size( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """HEADER_TABLE_SIZE setting is processed.""" + conn, _ = connection + frame = build_settings_frame([(Setting.HEADER_TABLE_SIZE, 4096)]) + conn.data_received(frame) + # internal effect on encoder, just confirm no crash + + @pytest.mark.asyncio + async def test_ping_on_nonzero_stream( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """PING on non‑zero stream triggers GOAWAY.""" + conn, transport = connection + frame = frame_header(8, FrameType.PING, 0, 1) + b"\x00" * 8 + conn.data_received(frame) + assert any( + call[0][0][3:4] == FrameType.GOAWAY.to_bytes(1, "big") + for call in transport.write.call_args_list + ) + + @pytest.mark.asyncio + async def test_receive_ping_ack( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Receiving a PING ACK is logged and does not cause another response.""" + conn, _ = connection + frame = build_ping(ack=True, opaque=b"12345678") + conn.data_received(frame) # no crash, no additional PING sent + + @pytest.mark.asyncio + async def test_goaway_when_future_already_done( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """GOAWAY should not fail when a stream's future is already complete.""" + conn, _ = connection + s1 = await conn.create_stream() + s3 = await conn.create_stream() + s1.state = s3.state = StreamState.OPEN + s1.response_future.set_result(([], b"")) # type: ignore[arg-type] + frame = build_goaway(last_stream_id=1, error_code=0) + conn.data_received(frame) + assert s1.stream_id in conn.streams + assert s3.stream_id not in conn.streams + + @pytest.mark.asyncio + async def test_window_update_for_unknown_stream( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """WINDOW_UPDATE on unknown stream must be ignored (no crash).""" + conn, _ = connection + frame = build_window_update(123, 100) + conn.data_received(frame) + assert conn.session_outbound_window == 65535 # unchanged + + @pytest.mark.asyncio + async def test_continuation_frame_ignored( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """CONTINUATION frame is ignored with a warning.""" + conn, _ = connection + frame = frame_header(0, FrameType.CONTINUATION, 0, 1) + conn.data_received(frame) # no crash + + @pytest.mark.asyncio + async def test_unknown_frame_type_ignored( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Unknown frame type (>9) is ignored.""" + conn, _ = connection + frame = frame_header(0, 0x1A, 0, 0) + conn.data_received(frame) # no crash + + @pytest.mark.asyncio + async def test_send_data_end_stream_last_chunk( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """send_data with end_stream=True sets END_STREAM on the last chunk.""" + conn, transport = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + conn.session_outbound_window = 100 + stream.outbound_window = 100 + await conn.send_data(stream.stream_id, b"x" * 10, end_stream=True) + data_frames = [ + call[0][0] + for call in transport.write.call_args_list + if FrameType.DATA.to_bytes(1, "big") in call[0][0] + ] + assert len(data_frames) == 1 + assert data_frames[0][4] & FlagData.END_STREAM + + @pytest.mark.asyncio + async def test_create_stream_after_goaway( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """create_stream raises ConnectionError when GOAWAY has been sent.""" + conn, _ = connection + conn._goaway_sent = True + with pytest.raises(ConnectionError): + await conn.create_stream() + + @pytest.mark.asyncio + async def test_maybe_unblock_streams_done_future( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """_maybe_unblock_streams skips futures that are already done.""" + conn, _ = connection + fut = conn._loop.create_future() + fut.set_result(None) + conn._pending_streams.append(fut) + conn._maybe_unblock_streams() + assert len(conn.streams) == 0 # no new stream created + + @pytest.mark.asyncio + async def test_connection_lost_done_futures( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """connection_lost must handle streams whose futures are already done.""" + conn, _ = connection + stream = await conn.create_stream() + stream.response_future.set_result((stream.stream_id, [], b"")) # type: ignore[arg-type] + conn.connection_lost(None) # no exception + + +# ---------------------------------------------------------------------- +# 2. Miscellaneous tests (race conditions, deadlocks, edge cases) +# ---------------------------------------------------------------------- +class TestMiscellaneous: + @pytest.mark.asyncio + async def test_concurrent_send_data_does_not_deadlock( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """Multiple tasks sending data on the same stream should not deadlock.""" + conn, transport = connection + stream = await conn.create_stream() + # Set window large enough + conn.session_outbound_window = 1_000_000 + stream.outbound_window = 1_000_000 + + async def send_chunk() -> None: + await conn.send_data(stream.stream_id, b"x" * 100, end_stream=False) + + tasks = [asyncio.create_task(send_chunk()) for _ in range(5)] + await asyncio.gather(*tasks, return_exceptions=True) + # All writes should complete eventually + assert transport.write.call_count >= 5 + + @pytest.mark.asyncio + async def test_window_update_wakes_all_waiters( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """When window is zero, multiple blocked tasks resume on WINDOW_UPDATE.""" + conn, transport = connection + stream = await conn.create_stream() + conn.session_outbound_window = 0 + stream.outbound_window = 0 + + async def blocked_send() -> None: + await conn.send_data(stream.stream_id, b"hello", end_stream=False) + + task1 = asyncio.create_task(blocked_send()) + task2 = asyncio.create_task(blocked_send()) + await asyncio.sleep(0.01) # both waiting + + # Simulate WINDOW_UPDATE that opens 10 bytes + frame = build_window_update(0, 10) # session + conn.data_received(frame) + frame2 = build_window_update(stream.stream_id, 10) + conn.data_received(frame2) + await asyncio.sleep(0.01) + assert task1.done() + assert task2.done() + + @pytest.mark.asyncio + async def test_stream_cancelled_before_response( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Pending create_stream futures are cancelled on connection loss.""" + conn, _ = connection + conn.max_concurrent_streams = 1 + await conn.create_stream() + # Queue a second stream + fut = asyncio.ensure_future(conn.create_stream()) + await asyncio.sleep(0.01) + assert not fut.done() + # Simulate connection loss + conn.connection_lost(ConnectionError("test")) + await asyncio.sleep(0.01) + assert fut.done() + with pytest.raises(ConnectionError): + fut.result() + + @pytest.mark.asyncio + async def test_close_stream_on_rst_without_headers( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """RST_STREAM before headers deliver must resolve future with error.""" + conn, _ = connection + stream = await conn.create_stream() + frame = build_rst_stream(stream.stream_id, error_code=0x8) + conn.data_received(frame) + assert stream.response_future.done() + with pytest.raises(RuntimeError): + stream.response_future.result() + + @pytest.mark.asyncio + async def test_data_received_without_connection( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """Http2Protocol.data_received is a no‑op before connection_made.""" + proto, _ = protocol + proto._connection = None + proto.data_received(b"anything") # must not raise + + +class TestStreamStateMachine: + @pytest.mark.asyncio + async def test_invalid_transition_raises( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """Invalid state transition raises ProtocolError.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + with pytest.raises(ProtocolError): + stream.transition(StreamState.RESERVED_LOCAL) + + @pytest.mark.asyncio + async def test_receive_data_end_stream_half_closed_local( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """DATA END_STREAM when HALF_CLOSED_LOCAL -> CLOSED and stream removed.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.HALF_CLOSED_LOCAL + stream.response_headers = [(":status", "200")] + stream.receive_data(b"body", end_stream=True) + assert stream.state == StreamState.CLOSED + assert stream.stream_id not in conn.streams + + @pytest.mark.asyncio + async def test_receive_data_end_stream_invalid_state( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """DATA END_STREAM in CLOSED state raises ProtocolError.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.CLOSED + with pytest.raises(ProtocolError): + stream.receive_data(b"x", end_stream=True) + + @pytest.mark.asyncio + async def test_receive_headers_end_stream_half_closed_local( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """HEADERS END_STREAM when HALF_CLOSED_LOCAL -> CLOSED and stream removed.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.HALF_CLOSED_LOCAL + stream.receive_headers([(":status", "200")], end_stream=True) # type: ignore[list-item] + assert stream.state == StreamState.CLOSED + assert stream.stream_id not in conn.streams + + @pytest.mark.asyncio + async def test_receive_headers_end_stream_invalid_state( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """HEADERS END_STREAM in CLOSED state raises ProtocolError.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.CLOSED + with pytest.raises(ProtocolError): + stream.receive_headers([(":status", "200")], end_stream=True) # type: ignore[list-item] + + @pytest.mark.asyncio + async def test_data_stream_before_headers( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """DATA before headers does NOT set future until headers arrive.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + # double end stream is invalid + stream.receive_data(b"body", end_stream=False) + assert not stream.response_future.done() + stream.receive_headers([(":status", "200")], end_stream=True) # type: ignore[list-item] + assert stream.response_future.done() + headers, body = stream.response_future.result() + assert await body.read() == b"body" + + @pytest.mark.asyncio + async def test_future_already_done_data_end_stream( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """receive_data with END_STREAM does not double‑set an already done future.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + stream.response_headers = [(":status", "200")] + stream.response_future.set_result((stream.response_headers, b"")) # type: ignore[arg-type] + stream.receive_data(b"more", end_stream=True) # no exception + + @pytest.mark.asyncio + async def test_future_already_done_headers_end_stream( + self, connection: Tuple[Http2Connection, MagicMock] + ) -> None: + """receive_headers with END_STREAM does not double‑set an already done future.""" + conn, _ = connection + stream = await conn.create_stream() + stream.state = StreamState.OPEN + stream.response_future.set_result((stream.stream_id, [], b"")) # type: ignore[arg-type] + stream.receive_headers([(":status", "200")], end_stream=True) # type: ignore[list-item] + + +# ---------------------------------------------------------------------- +# Mock transport – records writes and lies about ALPN +# ---------------------------------------------------------------------- +class MockH2Transport(asyncio.Transport): + def __init__(self, extra_info: Optional[Dict[str, Any]] = None) -> None: + super().__init__() + self.written = bytearray() + self._closing = False + self._extra = extra_info or {} + self._protocol: Optional[Http2Protocol] = None + + def write(self, data: bytes | bytearray | memoryview) -> None: + self.written.extend(data) + + def close(self) -> None: + self._closing = True + + def is_closing(self) -> bool: + return self._closing + + def get_extra_info(self, name: str, default: Any = None) -> Any: + if name == "ssl_object": + return self._extra.get("ssl_object", MagicMock()) + return self._extra.get(name, default) + + +# ---------------------------------------------------------------------- +# Custom connector – always returns Http2Protocol for h2 connections +# ---------------------------------------------------------------------- +class H2TestConnector(TCPConnector): + def _get_protocol(self, loop: asyncio.AbstractEventLoop) -> type: + # Return the class; aiohttp will instantiate it + return Http2Protocol + + async def close(self, *, abort_ssl: bool = False) -> None: + self._closed = True + return None + + +# ---------------------------------------------------------------------- +# Fixture: session + mock transport + captured protocol +# ---------------------------------------------------------------------- +@pytest.fixture +async def h2_client() -> Any: # returns a generator of (session, transport, protocol) + """Create a ClientSession that uses our Http2Protocol over a mock transport.""" + # Mock SSL object that tells aiohttp we’ve negotiated h2 + mock_ssl = MagicMock() + mock_ssl.selected_alpn_protocol.return_value = "h2" + transport = MockH2Transport(extra_info={"ssl_object": mock_ssl}) + + protocol_instance: Optional[Http2Protocol] = None + + async def fake_create_connection( + protocol_factory: Any, *args: Any, **kwargs: Any + ) -> Tuple[MockH2Transport, Http2Protocol]: + nonlocal protocol_instance + protocol_instance = protocol_factory() # Http2Protocol() + protocol_instance.connection_made(transport) + transport._protocol = protocol_instance + return transport, protocol_instance + + connector = H2TestConnector() + connector._wrap_create_connection = fake_create_connection # type: ignore[assignment] + async with aiohttp.ClientSession(connector=connector) as session: + yield session, transport, protocol_instance + + +# if it's a real URL it hangs (missing mock?) +URL = "https://127.3.3.3" + + +class TestIncomingResponses: + @pytest.mark.asyncio + async def test_get_200_response(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) # request sent + + # Feed a minimal 200 response + hframe = build_headers_frame(1, [(":status", "200")], end_stream=True) + proto = transport._protocol + proto.data_received(hframe) + + resp = await task + assert resp.status == 200 + assert await resp.read() == b"" + + @pytest.mark.asyncio + async def test_response_with_body(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + # Send HEADERS (no END_STREAM) then DATA with body + hframe = build_headers_frame(1, [(":status", "200")], end_stream=False) + dframe = build_data_frame(1, b"Hello, h2!", end_stream=True) + proto = transport._protocol + proto.data_received(hframe) + proto.data_received(dframe) + + resp = await task + assert resp.status == 200 + assert await resp.text() == "Hello, h2!" + + @pytest.mark.asyncio + async def test_json_response(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + headers = [(":status", "200"), ("content-type", "application/json")] + body = b'{"key":"value"}' + proto = transport._protocol + proto.data_received(build_headers_frame(1, headers, end_stream=False)) + proto.data_received(build_data_frame(1, body, end_stream=True)) + + resp = await task + assert await resp.json() == {"key": "value"} + + @pytest.mark.asyncio + async def test_response_cookies(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + headers = [(":status", "200"), ("set-cookie", "session=abc123; Path=/")] + proto = transport._protocol + proto.data_received(build_headers_frame(1, headers, end_stream=True)) + + resp = await task + assert "session" in resp.cookies + assert resp.cookies["session"].value == "abc123" + + @pytest.mark.asyncio + async def test_concurrent_requests_mux(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + t1 = asyncio.create_task(session.get(URL)) + t2 = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + # Stream 1 gets response, stream 3 gets response + proto = transport._protocol + proto.data_received( + build_headers_frame(1, [(":status", "200")], end_stream=True) + ) + proto.data_received( + build_headers_frame(3, [(":status", "201")], end_stream=True) + ) + + r1, r2 = await asyncio.gather(t1, t2) + assert r1.status == 200 + assert r2.status == 201 + + @pytest.mark.asyncio + async def test_redirect_headers(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + headers = [(":status", "302"), ("location", "/new")] + proto = transport._protocol + proto.data_received(build_headers_frame(1, headers, end_stream=True)) + + await asyncio.sleep(0.01) + + headers = [(":status", "200")] + proto.data_received(build_headers_frame(3, headers, end_stream=True)) + + resp = await task + first = resp._history[0] + assert first.status == 302 + assert first.headers.get("location") == "/new" + + assert resp.status == 200 + + @pytest.mark.asyncio + async def test_error_response_raises(self, h2_client: Any) -> None: # type: ignore[misc] + session, transport, _ = h2_client + task = asyncio.create_task(session.get(URL)) + await asyncio.sleep(0.01) + + proto = transport._protocol + proto.data_received( + build_headers_frame(1, [(":status", "404")], end_stream=True) + ) + resp = await task + with pytest.raises(aiohttp.ClientResponseError): + resp.raise_for_status() + + +# ---------------------------------------------------------------------- +# Additional tests for full coverage of missing paths +# ---------------------------------------------------------------------- + + +class TestConnectionEdgeCases: + @pytest.mark.asyncio + async def test_eof_received_calls_close( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """EOF triggers connection close.""" + conn, transport = connection + result = conn.eof_received() + assert result is False + transport.close.assert_called_once() + + @pytest.mark.asyncio + async def test_protocol_eof_received( + self, protocol: Tuple[Http2Protocol, MagicMock], mock_transport: MagicMock + ) -> None: + """Http2Protocol.eof_received delegates to connection and returns False.""" + proto, transport = protocol + # connection is established + assert proto._connection is not None + result = proto.eof_received() + assert result is False + # transport.close should be called via conn.eof_received -> conn.close() + transport.close.assert_called_once() + + @pytest.mark.asyncio + async def test_data_frame_unknown_stream_above_last_peer( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """DATA frame for unknown stream_id > _last_peer_stream_id sends RST_STREAM.""" + conn, transport = connection + # last_peer_stream_id is initially 0 + unknown_id = 5 # > 0 + frame = build_data_frame(unknown_id, b"x", end_stream=False) + conn.data_received(frame) + # Should have sent RST_STREAM with PROTOCOL_ERROR (1) + rst_frames = [ + call.args[0] + for call in transport.write.call_args_list + if FrameType.RST_STREAM.to_bytes(1, "big") in call.args[0] + ] + assert len(rst_frames) == 1 + payload = rst_frames[0][9:] # after 9-byte header + error_code = struct.unpack("!I", payload)[0] + assert error_code == 1 # PROTOCOL_ERROR + + @pytest.mark.asyncio + async def test_data_frame_unknown_stream_not_above_last_peer( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """DATA frame for unknown stream_id <= _last_peer_stream_id is ignored.""" + conn, transport = connection + # artificially raise last_peer_stream_id + conn._last_peer_stream_id = 10 + frame = build_data_frame(5, b"x", end_stream=False) + conn.data_received(frame) + # No RST_STREAM sent + assert not any( + FrameType.RST_STREAM.to_bytes(1, "big") in call.args[0] + for call in transport.write.call_args_list + ) + + @pytest.mark.asyncio + async def test_send_data_end_stream_when_half_closed_remote( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """When stream is HALF_CLOSED_REMOTE, final DATA with END_STREAM closes it.""" + conn, transport = connection + stream = await conn.create_stream() + # Simulate remote half-close + stream.state = StreamState.HALF_CLOSED_REMOTE + conn.session_outbound_window = 1000 + stream.outbound_window = 1000 + await conn.send_data(stream.stream_id, b"done", end_stream=True) + # Stream should be CLOSED and removed + assert stream.state == StreamState.CLOSED + assert stream.stream_id not in conn.streams + + @pytest.mark.asyncio + async def test_connection_lost_cancels_open_streams( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """connection_lost sets ConnectionError on all unfinished streams.""" + conn, _ = connection + s1 = await conn.create_stream() + s2 = await conn.create_stream() + s1.state = s2.state = StreamState.OPEN + # Both futures not done + conn.connection_lost(ConnectionError("test")) + assert s1.response_future.exception() is not None + assert s2.response_future.exception() is not None + # Also check pending streams cleared + assert len(conn._pending_streams) == 0 + + @pytest.mark.asyncio + async def test_connection_lost_clears_pending_futures( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """Pending create_stream futures are cancelled on connection loss.""" + conn, _ = connection + conn.max_concurrent_streams = 1 + await conn.create_stream() # fills one slot + fut = asyncio.ensure_future(conn.create_stream()) + await asyncio.sleep(0.01) + assert not fut.done() + conn.connection_lost(ConnectionError("test")) + await asyncio.sleep(0.01) + assert fut.done() + with pytest.raises(ConnectionError): + fut.result() + + @pytest.mark.asyncio + async def test_close_transport( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """close() calls transport.close().""" + conn, transport = connection + conn.close() + transport.close.assert_called_once() + + @pytest.mark.asyncio + async def test_should_close_reflects_goaway( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """should_close returns True if GOAWAY sent or received.""" + conn, _ = connection + assert not conn.should_close + conn._goaway_received = True + assert conn.should_close + + @pytest.mark.asyncio + async def test_is_connected( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """is_connected mirrors transport.is_closing().""" + conn, transport = connection + assert conn.is_connected() is True + transport.is_closing.return_value = True + assert conn.is_connected() is False + + @pytest.mark.asyncio + async def test_protocol_close_and_abort( + self, protocol: Tuple[Http2Protocol, MagicMock], mock_transport: MagicMock + ) -> None: + """Http2Protocol.close() and abort() propagate to transport.""" + proto, transport = protocol + proto.close() + transport.close.assert_called_once() + # abort is same as close + proto.abort() + assert transport.close.call_count == 2 + + @pytest.mark.asyncio + async def test_protocol_should_close_no_connection( + self, protocol: Tuple[Http2Protocol, MagicMock], mock_transport: MagicMock + ) -> None: + """should_close returns False when connection is None.""" + proto, _ = protocol + proto._connection = None + assert not proto.should_close + + @pytest.mark.asyncio + async def test_protocol_is_connected_no_connection( + self, protocol: Tuple[Http2Protocol, MagicMock], mock_transport: MagicMock + ) -> None: + """is_connected returns False when connection is None.""" + proto, _ = protocol + proto._connection = None + assert not proto.is_connected() + + @pytest.mark.asyncio + async def test_initiate_connection_sends_preface_and_settings( + self, mock_transport: MagicMock, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """initiate_connection writes HTTP/2 preface and initial SETTINGS.""" + h2_proto, _ = protocol + loop = asyncio.get_running_loop() + conn = Http2Connection(mock_transport, loop, h2_proto) + conn.initiate_connection() + written = b"".join(call.args[0] for call in mock_transport.write.call_args_list) + assert b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n" in written + assert FrameType.SETTINGS.to_bytes(1, "big") in written + + @pytest.mark.asyncio + async def test_send_data_flow_control_releases_waiters( + self, connection: Tuple[Http2Connection, MagicMock], mock_transport: MagicMock + ) -> None: + """After a send, if window space remains, _flow_control_updated is set.""" + conn, transport = connection + stream = await conn.create_stream() + conn.session_outbound_window = 100 + stream.outbound_window = 100 + # set event cleared before wait + conn._flow_control_updated.clear() + await conn.send_data(stream.stream_id, b"x" * 10, end_stream=False) + # After send, there is still window space, so event should be set + assert conn._flow_control_updated.is_set() + + +# ---------------------------------------------------------------------- +# Fixture: create a Stream with mocked connection and protocol +# ---------------------------------------------------------------------- +@pytest.fixture +def stream_setup() -> Generator[Any, None, None]: + """Create a Stream with mocked dependencies, return stream and mocks.""" + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + # Mock connection + conn = MagicMock() + conn.remote_settings = {Setting.INITIAL_WINDOW_SIZE: 65535} + conn.local_settings = {Setting.INITIAL_WINDOW_SIZE: 65535} + conn._send_window_update = MagicMock() + conn._send_rst_stream = MagicMock() + conn._close_stream = MagicMock() + + # Mock protocol + protocol = MagicMock() + protocol._auto_decompress = True + + stream = Stream(stream_id=1, conn=conn, loop=loop, protocol=protocol) + stream.body_reader = MagicMock() + # Reset mocks after initialization (constructor may have used them) + conn._send_window_update.reset_mock() + conn._send_rst_stream.reset_mock() + conn._close_stream.reset_mock() + + yield stream, conn, protocol, loop + loop.close() + + +# ---------------------------------------------------------------------- +# Helper to create headers with optional content-encoding +# ---------------------------------------------------------------------- +def make_headers( + encoding: Optional[str] = None, end_stream: bool = False +) -> List[Tuple[str, str]]: + headers: List[Tuple[str, str]] = [(":status", "200")] + if encoding: + headers.append(("content-encoding", encoding)) + return headers + + +# ====================================================================== +# Tests for receive_headers +# ====================================================================== +class TestReceiveHeaders: + async def test_no_end_stream_no_encoding_no_pending( + self, stream_setup: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers() + + stream.receive_headers(headers, end_stream=False) + + assert stream._headers_received is True + assert stream.response_headers == headers + assert stream.response is not None + assert stream.decompressor is None + assert stream._pending_data == b"" + assert stream.response_future.done() + resp, reader = stream.response_future.result() + assert resp is stream.response + assert reader is stream.body_reader + + async def test_no_end_stream_encoding_no_pending( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers(encoding="gzip") + + # Patch DeflateBuffer to avoid real decompression + mock_deflate_cls = MagicMock() + mock_deflate = mock_deflate_cls.return_value + monkeypatch.setattr("aiohttp.http2.stream.DeflateBuffer", mock_deflate_cls) + + stream.receive_headers(headers, end_stream=False) + + assert stream.decompressor is mock_deflate + mock_deflate_cls.assert_called_once_with( + stream.body_reader, encoding="gzip", max_decompress_size=DEFAULT_CHUNK_SIZE + ) + mock_deflate.feed_data.assert_not_called() + assert stream._pending_data == b"" + + async def test_no_end_stream_encoding_with_pending( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + pending = b"compressed-bytes" + stream._pending_data = bytearray(pending) + headers = make_headers(encoding="gzip") + + mock_deflate_cls = MagicMock() + mock_deflate = mock_deflate_cls.return_value + monkeypatch.setattr("aiohttp.http2.stream.DeflateBuffer", mock_deflate_cls) + + stream.receive_headers(headers, end_stream=False) + + mock_deflate.feed_data.assert_called_once_with(pending) + assert stream._pending_data == b"" + + async def test_no_end_stream_no_encoding_with_pending( + self, stream_setup: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + pending = b"raw-bytes" + stream._pending_data = bytearray(pending) + headers = make_headers() + + stream.receive_headers(headers, end_stream=False) + + assert stream.decompressor is None + assert stream._pending_data == b"" + stream.body_reader.feed_data.assert_called_once_with(pending) + + async def test_no_end_stream_auto_decompress_false(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + protocol._auto_decompress = False + pending = b"raw-bytes" + stream._pending_data = bytearray(pending) + headers = make_headers() + + stream.receive_headers(headers, end_stream=False) + + assert stream.decompressor is None + assert stream._pending_data == b"" + stream.body_reader.feed_data.assert_called_once_with(pending) + + async def test_end_stream_no_encoding_open_state(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers() + stream.state = StreamState.OPEN + + stream.receive_headers(headers, end_stream=True) + + assert stream.state == StreamState.HALF_CLOSED_REMOTE + stream.body_reader.feed_eof.assert_called_once() + + async def test_end_stream_encoding_open_state( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers(encoding="gzip") + stream.state = StreamState.OPEN + + mock_deflate_cls = MagicMock() + mock_deflate = mock_deflate_cls.return_value + monkeypatch.setattr("aiohttp.http2.stream.DeflateBuffer", mock_deflate_cls) + + stream.receive_headers(headers, end_stream=True) + + mock_deflate.feed_eof.assert_called_once() + assert stream.state == StreamState.HALF_CLOSED_REMOTE + + async def test_end_stream_half_closed_local_closes(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers() + stream.state = StreamState.HALF_CLOSED_LOCAL + + stream.receive_headers(headers, end_stream=True) + + assert stream.state == StreamState.CLOSED + conn._close_stream.assert_called_once_with(stream) + + async def test_end_stream_unexpected_state_raises(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers() + stream.state = StreamState.IDLE # invalid for END_STREAM handling + + with pytest.raises(ProtocolError): + stream.receive_headers(headers, end_stream=True) + + async def test_error_during_decompression_pending( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + stream._pending_data = bytearray(b"bad-data") + headers = make_headers(encoding="gzip") + error = ContentEncodingError("bad encoding") + + mock_deflate_cls = MagicMock() + mock_deflate = mock_deflate_cls.return_value + mock_deflate.feed_data.side_effect = error + monkeypatch.setattr("aiohttp.http2.stream.DeflateBuffer", mock_deflate_cls) + + stream.receive_headers(headers, end_stream=False) + + stream.body_reader.set_exception.assert_called_once_with(error) + conn._send_rst_stream.assert_called_once_with( + stream.stream_id, ErrorCode.INTERNAL_ERROR + ) + assert stream._pending_data == b"" + + async def test_error_during_decompression_eof( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + headers = make_headers(encoding="gzip") + stream.state = StreamState.OPEN + error = ContentEncodingError("bad eof") + + mock_deflate_cls = MagicMock() + mock_deflate = mock_deflate_cls.return_value + mock_deflate.feed_eof.side_effect = error + monkeypatch.setattr("aiohttp.http2.stream.DeflateBuffer", mock_deflate_cls) + + stream.receive_headers(headers, end_stream=True) + + stream.body_reader.set_exception.assert_called_once_with(error) + conn._send_rst_stream.assert_called_once_with( + stream.stream_id, ErrorCode.INTERNAL_ERROR + ) + + +# ====================================================================== +# Tests for receive_data +# ====================================================================== +class TestReceiveData: + async def test_no_headers_received_buffers_data(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + data = b"early-data" + initial_window = stream.inbound_window + + stream.receive_data(data, end_stream=False) + + assert stream._pending_data == bytearray(data) + assert stream.inbound_window == initial_window - len(data) + stream.body_reader.feed_data.assert_not_called() + + async def test_no_headers_received_end_stream_feeds_eof( + self, stream_setup: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + data = b"early-data" + stream.state = StreamState.OPEN + + stream.receive_data(data, end_stream=True) + + assert stream._pending_data == bytearray(data) + stream.body_reader.feed_eof.assert_called_once() + assert stream.state == StreamState.HALF_CLOSED_REMOTE + + async def test_headers_received_no_decompressor(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + stream._headers_received = True + data = b"plain-data" + + stream.receive_data(data, end_stream=False) + + stream.body_reader.feed_data.assert_called_once_with(data) + assert stream._pending_data == b"" + + async def test_headers_received_with_decompressor( + self, stream_setup: Any, monkeypatch: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + stream._headers_received = True + data = b"compressed" + mock_deflate = MagicMock() + stream.decompressor = mock_deflate + + stream.receive_data(data, end_stream=False) + + mock_deflate.feed_data.assert_called_once_with(data) + stream.body_reader.feed_data.assert_not_called() + + async def test_decompressor_error_resets_stream(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + stream._headers_received = True + data = b"bad" + mock_deflate = MagicMock() + mock_deflate.feed_data.side_effect = ContentEncodingError("bad") + stream.decompressor = mock_deflate + + stream.receive_data(data, end_stream=False) + + stream.body_reader.set_exception.assert_called_once() + conn._send_rst_stream.assert_called_once_with( + stream.stream_id, ErrorCode.INTERNAL_ERROR + ) + + async def test_end_stream_with_decompressor(self, stream_setup: Any) -> None: + stream, conn, protocol, _ = stream_setup + stream._headers_received = True + stream.state = StreamState.OPEN + data = b"compressed" + mock_deflate = MagicMock() + stream.decompressor = mock_deflate + + stream.receive_data(data, end_stream=True) + + mock_deflate.feed_eof.assert_called_once() + assert stream.state == StreamState.HALF_CLOSED_REMOTE + + async def test_end_stream_decompressor_error_resets_stream( + self, stream_setup: Any + ) -> None: + stream, conn, protocol, _ = stream_setup + stream._headers_received = True + stream.state = StreamState.OPEN + data = b"bad" + mock_deflate = MagicMock() + mock_deflate.feed_eof.side_effect = ContentEncodingError("bad eof") + stream.decompressor = mock_deflate + + stream.receive_data(data, end_stream=True) + + +class TestProtocolMethods: + @pytest.mark.asyncio + async def test_closed_property_before_connection_made( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """closed returns a future that is created lazily.""" + proto, _ = protocol + # Before connection_made, closed should create a future + closed_fut = proto.closed + assert closed_fut is not None + assert isinstance(closed_fut, asyncio.Future) + assert not closed_fut.done() + + @pytest.mark.asyncio + async def test_closed_property_after_clean_connection_lost( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """After connection_lost(None), closed future is resolved with None.""" + proto, _ = protocol + # Access closed to create the future + closed_fut = proto.closed + assert closed_fut is not None + assert not closed_fut.done() + proto.connection_lost(None) + assert closed_fut.done() + assert closed_fut.result() is None + + @pytest.mark.asyncio + async def test_closed_property_after_error_connection_lost( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """After connection_lost with an exception, closed future gets ClientConnectionError.""" + proto, _ = protocol + closed_fut = proto.closed + assert closed_fut is not None + exc = ConnectionError("test error") + proto.connection_lost(exc) + assert closed_fut.done() + with pytest.raises(ClientConnectionError) as exc_info: + closed_fut.result() + assert "test error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_closed_property_after_connection_lost_twice( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """Once connection_lost has been called, closed returns None.""" + proto, _ = protocol + # First connection_lost + proto.connection_lost(None) + # Now _connection_lost_called is True, so closed should be None + assert proto.closed is None + + @pytest.mark.asyncio + async def test_read_timeout_setter( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto.read_timeout = 5.0 # type: ignore[attr-defined] + assert proto.read_timeout == 5.0 # type: ignore[attr-defined] + proto.read_timeout = None # type: ignore[attr-defined] + assert proto.read_timeout is None # type: ignore[attr-defined] + + @pytest.mark.asyncio + async def test_connection_lost_cleanup( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """connection_lost performs necessary cleanup.""" + proto, _ = protocol + # Mock the connection's connection_lost + proto._connection = MagicMock() + conn = proto._connection # save reference since we clean it up during close + # Access closed to create future + closed_fut = proto.closed + assert closed_fut is not None + proto.connection_lost(None) + # Check internal state + assert proto._connection_lost_called is True + assert proto._should_close is True + assert proto._connection is None + assert proto._reading_paused is False # type: ignore[unreachable] + # closed future resolved + assert closed_fut.done() + assert closed_fut.result() is None + # Connection's connection_lost called + conn.connection_lost.assert_called_once_with(None) + + @pytest.mark.asyncio + async def test_connection_lost_with_exception_cleanup( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + """connection_lost with exception sets error on future.""" + proto, _ = protocol + proto._connection = MagicMock() + closed_fut = proto.closed + assert closed_fut is not None + exc = TimeoutError("timeout") + proto.connection_lost(exc) + assert proto._connection_lost_called is True + assert proto._should_close is True + assert proto._connection is None + # connection can be None so it should be (and is) reachable + assert closed_fut.done() # type: ignore[unreachable] + with pytest.raises(ClientConnectionError) as exc_info: + closed_fut.result() + assert "timeout" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_force_close_sets_should_close( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + assert not proto._should_close + proto.force_close() + assert proto._should_close is True + + @pytest.mark.asyncio + async def test_set_exception_causes_cleanup( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + # Set a read_timeout_handle to be dropped + handle = MagicMock() + proto._read_timeout_handle = handle + proto.set_exception(ConnectionError("boom")) + assert proto._should_close is True + handle.cancel.assert_called_once() + assert proto._read_timeout_handle is None + + @pytest.mark.asyncio + async def test_resume_reading_reschedules_timeout( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + # Simulate paused state + proto._reading_paused = True + proto._read_timeout = 5.0 + # Mock _reschedule_timeout + proto._reschedule_timeout = MagicMock() # type: ignore[method-assign] + proto.resume_reading() + assert proto._reading_paused is False # after super().resume_reading + proto._reschedule_timeout.assert_called_once() + + @pytest.mark.asyncio + async def test_start_timeout_with_read_timeout( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto._read_timeout = 3.0 + # Mock loop.call_later + proto._loop = MagicMock() + proto._loop.call_later.return_value = MagicMock() + proto.start_timeout() + proto._loop.call_later.assert_called_once_with(3.0, proto._on_read_timeout) + assert proto._read_timeout_handle is not None + + @pytest.mark.asyncio + async def test_start_timeout_without_read_timeout( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto._read_timeout = None + proto._read_timeout_handle = MagicMock() + proto.start_timeout() + assert proto._read_timeout_handle is None + + @pytest.mark.asyncio + async def test_on_read_timeout_sets_exception( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto.set_exception = MagicMock() # type: ignore[method-assign] + proto._on_read_timeout() + # Check that set_exception called with SocketTimeoutError + proto.set_exception.assert_called_once() + exc_arg = proto.set_exception.call_args[0][0] + assert isinstance(exc_arg, SocketTimeoutError) + + @pytest.mark.asyncio + async def test_set_response_params_sets_timeout_and_decompress( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto.set_response_params( + read_timeout=7.5, + auto_decompress=False, + # other parameters are ignored + skip_payload=True, + read_until_eof=True, + ) + assert proto._read_timeout == 7.5 + assert proto._auto_decompress is False + + @pytest.mark.asyncio + async def test_create_stream_when_no_connection( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto._connection = None + with pytest.raises(ConnectionError): + await proto.create_stream() + + @pytest.mark.asyncio + async def test_read_stream_when_no_connection( + self, protocol: Tuple[Http2Protocol, MagicMock] + ) -> None: + proto, _ = protocol + proto._connection = None + with pytest.raises(ConnectionError): + await proto.read_stream(1) diff --git a/tests/http2/test_http2_adapter.py b/tests/http2/test_http2_adapter.py new file mode 100644 index 00000000000..e22a5a2008d --- /dev/null +++ b/tests/http2/test_http2_adapter.py @@ -0,0 +1,238 @@ +import asyncio +from typing import Any, Generator, List, Tuple +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from aiohttp.http2.adapter import Http2StreamWriter + + +@pytest.fixture +def writer() -> Generator[Any, None, None]: + """Create an Http2StreamWriter with mocked dependencies.""" + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + protocol = MagicMock() + protocol.transport = MagicMock() + protocol._connection = MagicMock() + protocol._connection.send_headers = MagicMock() + protocol._connection.send_data = AsyncMock() + + req = MagicMock() + req.stream_id = 123 + req.method = "POST" + req.url = "https://example.com/path" + + writer = Http2StreamWriter(protocol, loop, req) + + yield writer, protocol, req, loop + loop.close() + + +# ---------------------------------------------------------------------- +# Helper to check send_headers call arguments +# ---------------------------------------------------------------------- +def assert_send_headers_called( + protocol: Any, + stream_id: int, + method: str, + url: Any, + headers: List[Tuple[str, str]], + end_stream: bool, +) -> None: + protocol._connection.send_headers.assert_called_once_with( + stream_id, method, url, headers, end_stream=end_stream + ) + + +# ====================================================================== +# write_headers +# ====================================================================== +class TestWriteHeaders: + def test_buffers_headers(self, writer: Any) -> None: + w, protocol, req, loop = writer + headers = {"content-type": "application/json", "x-test": "value"} + asyncio.run(w.write_headers("HTTP/2.0 200", headers)) + + assert w._headers == [("content-type", "application/json"), ("x-test", "value")] + assert w._headers_sent is False + protocol._connection.send_headers.assert_not_called() + + +# ====================================================================== +# write +# ====================================================================== +class TestWrite: + def test_write_after_eof_raises(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._eof = True + with pytest.raises(RuntimeError, match="Cannot write after EOF"): + asyncio.run(w.write(b"data")) + + def test_write_without_headers_raises(self, writer: Any) -> None: + w, protocol, req, loop = writer + # _headers is None and _headers_sent is False + with pytest.raises(RuntimeError, match="Headers must be written before body"): + asyncio.run(w.write(b"data")) + + def test_write_sends_headers_and_data(self, writer: Any) -> None: + w, protocol, req, loop = writer + headers = [("content-type", "text/plain")] + w._headers = headers + data = b"hello" + + asyncio.run(w.write(data)) + + # Headers sent first + assert_send_headers_called( + protocol, 123, "POST", req.url, headers, end_stream=False + ) + # Data sent + protocol._connection.send_data.assert_called_once_with( + 123, data, end_stream=False + ) + assert w.output_size == len(data) + assert w._headers_sent is True + assert w._headers is None + + def test_write_empty_chunk_after_headers(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._headers = [("x", "y")] + w._headers_sent = True # simulate already sent + w._headers = None + + asyncio.run(w.write(b"")) + + protocol._connection.send_data.assert_not_called() + assert w.output_size == 0 + + def test_write_after_headers_sent(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._headers_sent = True + data = b"more" + + asyncio.run(w.write(data)) + + protocol._connection.send_headers.assert_not_called() + protocol._connection.send_data.assert_called_once_with( + 123, data, end_stream=False + ) + assert w.output_size == len(data) + + +# ====================================================================== +# write_eof +# ====================================================================== +class TestWriteEof: + def test_write_eof_twice_noop(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._eof = True + asyncio.run(w.write_eof(b"")) + protocol._connection.send_headers.assert_not_called() + protocol._connection.send_data.assert_not_called() + + def test_write_eof_without_headers_raises(self, writer: Any) -> None: + w, protocol, req, loop = writer + with pytest.raises(RuntimeError, match="Headers must be written before body"): + asyncio.run(w.write_eof(b"")) + + def test_write_eof_headers_not_sent_with_chunk(self, writer: Any) -> None: + w, protocol, req, loop = writer + headers = [("content-type", "text/plain")] + w._headers = headers + chunk = b"final-data" + + asyncio.run(w.write_eof(chunk)) + + # Headers sent without END_STREAM + assert_send_headers_called( + protocol, 123, "POST", req.url, headers, end_stream=False + ) + # Data with END_STREAM + protocol._connection.send_data.assert_called_once_with( + 123, chunk, end_stream=True + ) + assert w.output_size == len(chunk) + assert w._eof is True + + def test_write_eof_headers_not_sent_no_chunk(self, writer: Any) -> None: + w, protocol, req, loop = writer + headers = [("content-type", "text/plain")] + w._headers = headers + + asyncio.run(w.write_eof(b"")) + + # Headers with END_STREAM + assert_send_headers_called( + protocol, 123, "POST", req.url, headers, end_stream=True + ) + protocol._connection.send_data.assert_not_called() + assert w.output_size == 0 + assert w._eof is True + + def test_write_eof_headers_already_sent(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._headers_sent = True + chunk = b"final" + + asyncio.run(w.write_eof(chunk)) + + protocol._connection.send_headers.assert_not_called() + protocol._connection.send_data.assert_called_once_with( + 123, chunk, end_stream=True + ) + assert w.output_size == len(chunk) + assert w._eof is True + + def test_write_eof_headers_already_sent_empty_chunk(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._headers_sent = True + + asyncio.run(w.write_eof(b"")) + + protocol._connection.send_headers.assert_not_called() + protocol._connection.send_data.assert_called_once_with( + 123, b"", end_stream=True + ) + assert w.output_size == 0 + assert w._eof is True + + +# ====================================================================== +# set_eof +# ====================================================================== +class TestSetEof: + def test_set_eof_twice_noop(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._eof = True + w.set_eof() + protocol._connection.send_headers.assert_not_called() + + def test_set_eof_without_headers_raises(self, writer: Any) -> None: + w, protocol, req, loop = writer + with pytest.raises(RuntimeError, match="Headers must be written before EOF"): + w.set_eof() + + def test_set_eof_headers_not_sent(self, writer: Any) -> None: + w, protocol, req, loop = writer + headers = [("content-type", "text/plain")] + w._headers = headers + + w.set_eof() + + assert_send_headers_called( + protocol, 123, "POST", req.url, headers, end_stream=True + ) + assert w._eof is True + assert w._headers is None + assert w._headers_sent is True + + def test_set_eof_headers_already_sent(self, writer: Any) -> None: + w, protocol, req, loop = writer + w._headers_sent = True + + w.set_eof() + + protocol._connection.send_headers.assert_not_called() + assert w._eof is True diff --git a/tests/http2/test_synchro.py b/tests/http2/test_synchro.py new file mode 100644 index 00000000000..fd267bc8c0d --- /dev/null +++ b/tests/http2/test_synchro.py @@ -0,0 +1,244 @@ +import asyncio + +import pytest + +from aiohttp.http2.synchro import HostProbeSynchronizer + +SLEEP_TIME = 0.0001 + + +@pytest.fixture +def sync() -> HostProbeSynchronizer: + return HostProbeSynchronizer() + + +@pytest.mark.asyncio +async def test_first_acquire_immediate_and_second_blocks( + sync: HostProbeSynchronizer, +) -> None: + key = "host" + acquired_first = asyncio.Event() + release_now = asyncio.Event() + + async def first_worker() -> None: + await sync.acquire(key) + acquired_first.set() + await release_now.wait() + sync.release(key) + + async def second_worker() -> bool: + await sync.acquire(key) + # If we get here, the second acquire succeeded + return True + + # Start first worker; it should acquire immediately. + t1 = asyncio.create_task(first_worker()) + await asyncio.wait_for(acquired_first.wait(), timeout=1.0) + assert sync.is_locked(key) + + # Start second worker; it should block. + t2 = asyncio.create_task(second_worker()) + await asyncio.sleep(SLEEP_TIME) + assert not t2.done() + + # Release and let second worker proceed. + release_now.set() + result = await asyncio.wait_for(t2, timeout=1.0) + assert result is True + assert not sync.is_locked(key) + + # Clean up first worker (already finished release, but ensure it completes) + await asyncio.wait_for(t1, timeout=1.0) + + +@pytest.mark.asyncio +async def test_release_wakes_all_waiters(sync: HostProbeSynchronizer) -> None: + key = "host" + num_waiters = 3 + acquired_events = [asyncio.Event() for _ in range(num_waiters)] + release_now = asyncio.Event() + + async def worker(i: int) -> None: + await sync.acquire(key) + acquired_events[i].set() + await release_now.wait() + sync.release(key) + + # First worker acquires, others wait. + t_first = asyncio.create_task(worker(0)) + await asyncio.wait_for(acquired_events[0].wait(), timeout=1.0) + tasks = [t_first] + for i in range(1, num_waiters): + t = asyncio.create_task(worker(i)) + tasks.append(t) + await asyncio.sleep(SLEEP_TIME) + + # Ensure none of the waiters acquired yet. + for i in range(1, num_waiters): + assert not acquired_events[i].is_set() + + # Release, all waiters should wake and set their events. + release_now.set() + for i in range(1, num_waiters): + await asyncio.wait_for(acquired_events[i].wait(), timeout=1.0) + + # All tasks should complete. + for t in tasks: + await asyncio.wait_for(t, timeout=1.0) + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_cancellation_while_waiting(sync: HostProbeSynchronizer) -> None: + key = "host" + acquired = asyncio.Event() + release_now = asyncio.Event() + + async def first_worker() -> None: + await sync.acquire(key) + acquired.set() + await release_now.wait() + sync.release(key) + + t1 = asyncio.create_task(first_worker()) + await asyncio.wait_for(acquired.wait(), timeout=1.0) + + # Start a waiter that will be cancelled. + async def waiter() -> None: + await sync.acquire(key) + + t2 = asyncio.create_task(waiter()) + await asyncio.sleep(SLEEP_TIME) + assert not t2.done() + + # Cancel the waiter. + t2.cancel() + with pytest.raises(asyncio.CancelledError): + await t2 + + # Ensure the key is still locked by first worker. + assert sync.is_locked(key) + + # Release and verify no errors, key unlocked. + release_now.set() + await asyncio.wait_for(t1, timeout=1.0) + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_reacquire_after_release(sync: HostProbeSynchronizer) -> None: + key = "host" + await sync.acquire(key) + assert sync.is_locked(key) + sync.release(key) + assert not sync.is_locked(key) + + # Acquire again – should not block. + await sync.acquire(key) + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_is_locked_reflects_state(sync: HostProbeSynchronizer) -> None: + key = "host" + assert not sync.is_locked(key) + await sync.acquire(key) + assert sync.is_locked(key) + sync.release(key) + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_different_keys_are_independent(sync: HostProbeSynchronizer) -> None: + key1 = "host1" + key2 = "host2" + + await sync.acquire(key1) + assert sync.is_locked(key1) + assert not sync.is_locked(key2) + + # Acquiring key2 should not block because it's a different key. + await sync.acquire(key2) + assert sync.is_locked(key2) + + # Release key1; key2 remains locked. + sync.release(key1) + assert not sync.is_locked(key1) + assert sync.is_locked(key2) + + sync.release(key2) + assert not sync.is_locked(key2) + + +@pytest.mark.asyncio +async def test_release_without_acquire_does_not_raise( + sync: HostProbeSynchronizer, +) -> None: + key = "never_acquired" + # Should not raise. + sync.release(key) + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_release_after_cancelled_waiter(sync: HostProbeSynchronizer) -> None: + key = "host" + acquired = asyncio.Event() + release_now = asyncio.Event() + + async def first_worker() -> None: + await sync.acquire(key) + acquired.set() + await release_now.wait() + sync.release(key) + + t1 = asyncio.create_task(first_worker()) + await asyncio.wait_for(acquired.wait(), timeout=1.0) + + # Start a waiter, let it block, then cancel it. + async def waiter() -> None: + await sync.acquire(key) + + t2 = asyncio.create_task(waiter()) + await asyncio.sleep(SLEEP_TIME) + assert not t2.done() + t2.cancel() + with pytest.raises(asyncio.CancelledError): + await t2 + + # Release and ensure no exception. + release_now.set() + await asyncio.wait_for(t1, timeout=1.0) + assert not sync.is_locked(key) + + # A new acquire should work. + await sync.acquire(key) + # This is not the first request + # therefore it does not lock + assert not sync.is_locked(key) + + +@pytest.mark.asyncio +async def test_concurrent_acquire_after_release(sync: HostProbeSynchronizer) -> None: + key = "host" + # First acquire and release immediately. + await sync.acquire(key) + sync.release(key) + + # Now start two tasks that both try to acquire. + # Both should run without locking + t1 = asyncio.create_task(sync.acquire(key)) + t2 = asyncio.create_task(sync.acquire(key)) + + # Allow both to run; one will lock, the other will wait. + await asyncio.sleep(SLEEP_TIME) + assert not sync.is_locked(key) + # All tasks are done + done = [t for t in (t1, t2) if t.done()] + assert len(done) == 2 + + # Release, the second task should complete. + sync.release(key) + for t in (t1, t2): + await asyncio.wait_for(t, timeout=1.0) + assert not sync.is_locked(key) diff --git a/tests/http2/utils.py b/tests/http2/utils.py new file mode 100644 index 00000000000..e94f701aaca --- /dev/null +++ b/tests/http2/utils.py @@ -0,0 +1,103 @@ +import struct +from typing import Any, Optional + +from hpack import Encoder + +from aiohttp.http2.settings import ( + FlagData, + FlagHeaders, + FlagPing, + FlagSettings, + FrameType, + Setting, +) + + +# ---------------------------------------------------------------------- +# Helper: minimal URL mock +# ---------------------------------------------------------------------- +def url_mock(path: str = "/") -> Any: + """Create a simple URL-like object expected by the implementation.""" + return type( + "URL", + (), + {"scheme": "https", "host": "example.com", "path": path, "query": None}, + ) + + +# ---------------------------------------------------------------------- +# Frame construction helpers +# ---------------------------------------------------------------------- +def frame_header(length: int, ftype: int, flags: int, stream_id: int) -> bytes: + # 24-bit length (3 bytes) + type + flags + stream_id + return struct.pack("!I", length)[1:] + struct.pack( + "!B B I", ftype, flags, stream_id + ) + + +def build_settings_frame( + settings_pairs: Optional[list[tuple[Setting, int]]] = None, ack: bool = False +) -> bytes: + payload = b"" + if not ack and settings_pairs: + for setting_id, value in settings_pairs: + payload += struct.pack("!H I", setting_id, value) + flags = FlagSettings.ACK if ack else 0 + return frame_header(len(payload), FrameType.SETTINGS, flags, 0) + payload + + +def build_headers_frame( + stream_id: int, + headers: list[tuple[str, str]], + end_headers: bool = True, + end_stream: bool = False, + priority: Optional[bytes] = None, +) -> bytes: + encoder = Encoder() + header_block = encoder.encode(headers) + flags = 0 + if end_headers: + flags |= FlagHeaders.END_HEADERS + if end_stream: + flags |= FlagHeaders.END_STREAM + if priority is not None: + flags |= FlagHeaders.PRIORITY + header_block = priority + header_block + return ( + frame_header(len(header_block), FrameType.HEADERS, flags, stream_id) + + header_block + ) + + +def build_data_frame( + stream_id: int, data: bytes, end_stream: bool = False, pad: bool = False +) -> bytes: + payload = data + flags = 0 + if end_stream: + flags |= FlagData.END_STREAM + if pad: + pad_len = 1 # minimal padding for test + payload = bytes([pad_len]) + data + b"\x00" * pad_len + flags |= FlagData.PADDED + return frame_header(len(payload), FrameType.DATA, flags, stream_id) + payload + + +def build_rst_stream(stream_id: int, error_code: int = 0) -> bytes: + payload = struct.pack("!I", error_code) + return frame_header(4, FrameType.RST_STREAM, 0, stream_id) + payload + + +def build_goaway(last_stream_id: int, error_code: int, extra: bytes = b"") -> bytes: + payload = struct.pack("!I I", last_stream_id, error_code) + extra + return frame_header(len(payload), FrameType.GOAWAY, 0, 0) + payload + + +def build_window_update(stream_id: int, increment: int) -> bytes: + payload = struct.pack("!I", increment) + return frame_header(4, FrameType.WINDOW_UPDATE, 0, stream_id) + payload + + +def build_ping(ack: bool = False, opaque: bytes = b"\x00" * 8) -> bytes: + flags = FlagPing.ACK if ack else 0 + return frame_header(8, FrameType.PING, flags, 0) + opaque