diff --git a/docs/sse.md b/docs/sse.md index 5da7a626..a1a9310d 100644 --- a/docs/sse.md +++ b/docs/sse.md @@ -74,4 +74,16 @@ If the response does not have a `text/event-stream` content type, iterating the `SSEError` is a subclass of [`TransportError`](exceptions.md), so it is also caught by `except httpx2.TransportError`. +## Limiting event size + +An event is buffered until its terminating blank line arrives. To stop a stream that keeps sending data without ever completing an event from growing the buffer indefinitely, `client.sse()` caps the bytes buffered for a single event at 1 MiB by default and raises `SSEError` once an event exceeds the limit. Pass `max_event_size` to change the cap: + +```pycon +>>> with client.sse("https://example.com/sse", max_event_size=8 * 1024 * 1024) as source: +... for event in source: # raise the 1 MiB default to 8 MiB +... print(event.data) +``` + +The counter resets after each event, so the limit applies per event rather than to the stream as a whole. Set `max_event_size=None` to buffer events without any limit. + [mdn]: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events diff --git a/src/httpx2/httpx2/_client.py b/src/httpx2/httpx2/_client.py index ae73434e..ffc0de09 100644 --- a/src/httpx2/httpx2/_client.py +++ b/src/httpx2/httpx2/_client.py @@ -16,6 +16,7 @@ DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS, DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS, DEFAULT_LIMITS, + DEFAULT_MAX_EVENT_SIZE_BYTES, DEFAULT_MAX_MESSAGE_SIZE_BYTES, DEFAULT_MAX_REDIRECTS, DEFAULT_QUEUE_SIZE, @@ -871,6 +872,7 @@ def sse( follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT, timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, extensions: RequestExtensions | None = None, + max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES, ) -> Generator[EventSource]: """ Connect to a server-sent events endpoint and yield an `EventSource`. @@ -894,7 +896,7 @@ def sse( timeout=timeout, extensions=extensions, ) as response: - yield EventSource(response) + yield EventSource(response, max_event_size=max_event_size) @contextmanager def websocket( @@ -1708,6 +1710,7 @@ async def sse( follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT, timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, extensions: RequestExtensions | None = None, + max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES, ) -> AsyncGenerator[EventSource]: """ Connect to a server-sent events endpoint and yield an `EventSource`. @@ -1731,7 +1734,7 @@ async def sse( timeout=timeout, extensions=extensions, ) as response: - yield EventSource(response) + yield EventSource(response, max_event_size=max_event_size) @asynccontextmanager async def websocket( diff --git a/src/httpx2/httpx2/_config.py b/src/httpx2/httpx2/_config.py index 594f221a..ab6f25bc 100644 --- a/src/httpx2/httpx2/_config.py +++ b/src/httpx2/httpx2/_config.py @@ -248,6 +248,8 @@ def __repr__(self) -> str: DEFAULT_LIMITS = Limits(max_connections=100, max_keepalive_connections=20) DEFAULT_MAX_REDIRECTS = 20 +DEFAULT_MAX_EVENT_SIZE_BYTES = 1024 * 1024 + DEFAULT_MAX_MESSAGE_SIZE_BYTES = 65_536 DEFAULT_QUEUE_SIZE = 512 DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS = 20.0 diff --git a/src/httpx2/httpx2/_sse.py b/src/httpx2/httpx2/_sse.py index f04c4703..0a160a47 100644 --- a/src/httpx2/httpx2/_sse.py +++ b/src/httpx2/httpx2/_sse.py @@ -10,7 +10,8 @@ from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass -from ._exceptions import TransportError +from ._config import DEFAULT_MAX_EVENT_SIZE_BYTES +from ._exceptions import TransportError, request_context from ._models import Response __all__ = ["EventSource", "SSEError", "ServerSentEvent"] @@ -34,15 +35,18 @@ def json(self) -> object: class _SSEDecoder: - def __init__(self) -> None: + def __init__(self, max_event_size: int | None = None) -> None: + self._max_event_size = max_event_size self._event = "" self._data: list[str] = [] + self._event_size = 0 self._last_event_id = "" self._retry: int | None = None self._pending = False def decode(self, line: str) -> ServerSentEvent | None: if not line: + self._event_size = 0 if not self._pending: return None @@ -61,6 +65,10 @@ def decode(self, line: str) -> ServerSentEvent | None: if line.startswith(":"): return None + self._event_size += len(line.encode("utf-8")) + if self._max_event_size is not None and self._event_size > self._max_event_size: + raise SSEError(f"Server-sent event exceeded the {self._max_event_size} byte limit.") + fieldname, _, value = line.partition(":") value = value[1:] if value.startswith(" ") else value @@ -85,7 +93,8 @@ def decode(self, line: str) -> ServerSentEvent | None: class _SSELineDecoder: - def __init__(self) -> None: + def __init__(self, max_event_size: int | None = None) -> None: + self._max_event_size = max_event_size self._buffer = "" self._trailing_cr = False @@ -100,6 +109,8 @@ def decode(self, text: str) -> list[str]: text = self._buffer + text.replace("\r\n", "\n").replace("\r", "\n") lines = text.split("\n") self._buffer = lines.pop() + if self._max_event_size is not None and len(self._buffer.encode("utf-8")) > self._max_event_size: + raise SSEError(f"Server-sent event exceeded the {self._max_event_size} byte limit.") return lines def flush(self) -> list[str]: @@ -114,8 +125,9 @@ def flush(self) -> list[str]: class EventSource: - def __init__(self, response: Response) -> None: + def __init__(self, response: Response, max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES) -> None: self._response = response + self._max_event_size = max_event_size @property def response(self) -> Response: @@ -124,35 +136,34 @@ def response(self) -> Response: def _check_content_type(self) -> None: content_type, _, _ = self._response.headers.get("content-type", "").partition(";") if content_type.strip().lower() != "text/event-stream": - raise SSEError( - f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.", - request=self._response.request, - ) + raise SSEError(f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.") def __iter__(self) -> Iterator[ServerSentEvent]: - self._check_content_type() - decoder = _SSEDecoder() - lines = _SSELineDecoder() - for chunk in self._response.iter_text(): - for line in lines.decode(chunk): + with request_context(request=self._response.request): + self._check_content_type() + decoder = _SSEDecoder(self._max_event_size) + lines = _SSELineDecoder(self._max_event_size) + for chunk in self._response.iter_text(): + for line in lines.decode(chunk): + sse = decoder.decode(line) + if sse is not None: + yield sse + for line in lines.flush(): sse = decoder.decode(line) if sse is not None: yield sse - for line in lines.flush(): - sse = decoder.decode(line) - if sse is not None: - yield sse async def __aiter__(self) -> AsyncIterator[ServerSentEvent]: - self._check_content_type() - decoder = _SSEDecoder() - lines = _SSELineDecoder() - async for chunk in self._response.aiter_text(): - for line in lines.decode(chunk): + with request_context(request=self._response.request): + self._check_content_type() + decoder = _SSEDecoder(self._max_event_size) + lines = _SSELineDecoder(self._max_event_size) + async for chunk in self._response.aiter_text(): + for line in lines.decode(chunk): + sse = decoder.decode(line) + if sse is not None: + yield sse + for line in lines.flush(): sse = decoder.decode(line) if sse is not None: yield sse - for line in lines.flush(): - sse = decoder.decode(line) - if sse is not None: - yield sse diff --git a/tests/httpx2/test_sse.py b/tests/httpx2/test_sse.py index ae32fe2f..fb15e63b 100644 --- a/tests/httpx2/test_sse.py +++ b/tests/httpx2/test_sse.py @@ -275,6 +275,199 @@ def handler(request: httpx2.Request) -> httpx2.Response: assert event.data == "hi" +def test_max_event_size_rejects_unterminated_line() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * 8192 + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=4096) as source: + with pytest.raises(httpx2.SSEError, match="4096 byte limit"): + list(source) + + +def test_max_event_size_rejects_many_data_lines() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: chunk\n" * 1000 + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + with pytest.raises(httpx2.SSEError, match="100 byte limit"): + list(source) + + +def test_max_event_size_counts_empty_data_lines() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data\n" * 1000 + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + with pytest.raises(httpx2.SSEError, match="100 byte limit"): + list(source) + + +def test_max_event_size_counts_non_data_fields() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"id: " + b"A" * 8192 + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=4096) as source: + with pytest.raises(httpx2.SSEError, match="4096 byte limit"): + list(source) + + +def test_max_event_size_measures_utf8_bytes() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + "😀".encode() * 50 + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + with pytest.raises(httpx2.SSEError, match="100 byte limit"): + list(source) + + +def test_max_event_size_rejects_single_large_data_line() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * 8192 + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=4096) as source: + with pytest.raises(httpx2.SSEError, match="4096 byte limit"): + list(source) + + +def test_max_event_size_ignores_keepalive_comments() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b": keepalive\n\n" * 1000 + b"data: hi\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + (event,) = list(source) + + assert event.data == "hi" + + +def test_max_event_size_allows_event_under_limit() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=4096) as source: + (event,) = list(source) + + assert event.data == "hi" + + +def test_max_event_size_resets_between_events() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * 80 + b"\n\ndata: " + b"B" * 80 + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + events = list(source) + + assert [len(event.data) for event in events] == [80, 80] + + +def test_max_event_size_applies_by_default() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * (2 * 1024 * 1024) + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse") as source: + with pytest.raises(httpx2.SSEError, match="byte limit"): + list(source) + + +def test_max_event_size_none_disables_limit() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * (2 * 1024 * 1024) + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=None) as source: + (event,) = list(source) + + assert len(event.data) == 2 * 1024 * 1024 + + +def test_max_event_size_spans_completed_and_pending_lines() -> None: + def chunks() -> Iterator[bytes]: + yield b"data: " + b"A" * 60 + b"\n" + yield b"data: " + b"B" * 60 + + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, content=chunks(), headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + with pytest.raises(httpx2.SSEError, match="100 byte limit"): + list(source) + + +def test_max_event_size_does_not_conflate_adjacent_events() -> None: + def chunks() -> Iterator[bytes]: + yield b"data: " + b"A" * 60 + b"\n\ndata: " + b"B" * 60 + yield b"\n\n" + + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, content=chunks(), headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=100) as source: + events = list(source) + + assert [len(event.data) for event in events] == [60, 60] + + +def test_max_event_size_error_has_request() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + body = b"data: " + b"A" * 8192 + b"\n\n" + return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"}) + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with client.sse("http://testserver/sse", max_event_size=4096) as source: + with pytest.raises(httpx2.SSEError) as exc_info: + list(source) + + assert exc_info.value.request.url == "http://testserver/sse" + + +@pytest.mark.anyio +async def test_max_event_size_applies_by_default_async() -> None: + captured: list[int | None] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"}) + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client: + async with client.sse("http://testserver/sse") as source: + captured.append(source._max_event_size) + [event async for event in source] + + assert captured == [1024 * 1024] + + +@pytest.mark.anyio +async def test_max_event_size_allows_event_under_limit_async() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"}) + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client: + async with client.sse("http://testserver/sse", max_event_size=4096) as source: + events = [event async for event in source] + + assert [event.data for event in events] == ["hi"] + + def test_event_dispatched_at_eof_on_trailing_cr_sync() -> None: def chunks() -> Iterator[bytes]: yield b"data: hi\n"