|
| 1 | +"""Response-body cap enforcement: validate, read-capped (sync + async).""" |
| 2 | + |
| 3 | +import typing |
| 4 | +from collections.abc import Mapping |
| 5 | +from http import HTTPStatus |
| 6 | + |
| 7 | +import httpx2 |
| 8 | + |
| 9 | +from httpware.errors import ResponseTooLargeError |
| 10 | + |
| 11 | + |
| 12 | +_MAX_RESPONSE_BODY_BYTES_INVALID = "max_response_body_bytes must be >= 1" |
| 13 | + |
| 14 | + |
| 15 | +def _validate_max_response_body_bytes(cap: int | None) -> None: |
| 16 | + """Reject a non-None cap below 1. None means unbounded (the default).""" |
| 17 | + if cap is not None and cap < 1: |
| 18 | + raise ValueError(_MAX_RESPONSE_BODY_BYTES_INVALID) |
| 19 | + |
| 20 | + |
| 21 | +def _parse_content_length(raw: str | None) -> int | None: |
| 22 | + """Return a non-negative int Content-Length, or None for missing/garbage. Never raises.""" |
| 23 | + if raw is None: |
| 24 | + return None |
| 25 | + try: |
| 26 | + value = int(raw) |
| 27 | + except ValueError: |
| 28 | + return None |
| 29 | + return value if value >= 0 else None |
| 30 | + |
| 31 | + |
| 32 | +class _CapExceeded(Exception): # noqa: N818 — internal control-flow signal, not a user-facing error |
| 33 | + """Internal signal: decoded bytes crossed the cap mid-read. Carries bytes read so far.""" |
| 34 | + |
| 35 | + def __init__(self, *, read: int) -> None: |
| 36 | + self.read = read |
| 37 | + super().__init__(f"decoded body exceeded cap after {read} bytes") |
| 38 | + |
| 39 | + |
| 40 | +def _accumulate_capped(chunks: typing.Iterable[bytes], cap: int) -> bytes: |
| 41 | + """Concatenate `chunks`, raising `_CapExceeded` the moment the running total exceeds `cap`. |
| 42 | +
|
| 43 | + Counts decoded bytes (the in-memory footprint). Grown in a single bytearray |
| 44 | + so there is no transient list-plus-join double allocation. |
| 45 | + """ |
| 46 | + buf = bytearray() |
| 47 | + for chunk in chunks: |
| 48 | + buf += chunk |
| 49 | + if len(buf) > cap: |
| 50 | + raise _CapExceeded(read=len(buf)) |
| 51 | + return bytes(buf) |
| 52 | + |
| 53 | + |
| 54 | +def _safe_extensions(extensions: Mapping[str, typing.Any]) -> dict[str, typing.Any]: |
| 55 | + """Copy response extensions, dropping the now-stale `network_stream`. |
| 56 | +
|
| 57 | + The rebuilt buffered Response never touches its network stream, so carrying a |
| 58 | + consumed/closed one wholesale is sloppy. `http_version`/`reason_phrase` and |
| 59 | + any other keys are preserved. |
| 60 | + """ |
| 61 | + return {key: value for key, value in extensions.items() if key != "network_stream"} |
| 62 | + |
| 63 | + |
| 64 | +# Headers describing the wire encoding of the body. The accumulator yields the |
| 65 | +# DECODED body, so these no longer apply; httpx2 recomputes content-length from |
| 66 | +# the buffered content. Carrying content-encoding forward makes httpx2 try to |
| 67 | +# re-decode already-decoded bytes and raise. |
| 68 | +_WIRE_BODY_HEADERS = ("content-encoding", "content-length", "transfer-encoding") |
| 69 | +_BODILESS_STATUS = frozenset({HTTPStatus.NO_CONTENT, HTTPStatus.NOT_MODIFIED}) # 204, 304 |
| 70 | + |
| 71 | + |
| 72 | +def _buffered_headers(headers: httpx2.Headers) -> httpx2.Headers: |
| 73 | + """Copy `headers`, stripping wire-encoding headers stale after decoding+buffering.""" |
| 74 | + out = httpx2.Headers(headers) |
| 75 | + for name in _WIRE_BODY_HEADERS: |
| 76 | + if name in out: |
| 77 | + del out[name] |
| 78 | + return out |
| 79 | + |
| 80 | + |
| 81 | +def _response_has_body(method: str, status_code: int) -> bool: |
| 82 | + """Whether a response carries a message body (RFC 9110 §6.4.1). |
| 83 | +
|
| 84 | + HEAD responses and 204/304 never have a body regardless of a declared |
| 85 | + Content-Length, so they must never trip the cap. |
| 86 | + """ |
| 87 | + return method.upper() != "HEAD" and status_code not in _BODILESS_STATUS |
| 88 | + |
| 89 | + |
| 90 | +def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) -> httpx2.Response: |
| 91 | + """Buffer a streaming sync `response` under `cap` decoded bytes; return a buffered Response. |
| 92 | +
|
| 93 | + Raises `ResponseTooLargeError` (reason="declared") if the declared |
| 94 | + Content-Length already exceeds `cap` — before any byte is read — and |
| 95 | + (reason="streamed") if the decoded body crosses `cap` mid-read. Does not |
| 96 | + close `response`; the caller owns the stream lifecycle. |
| 97 | + """ |
| 98 | + if not _response_has_body(request.method, response.status_code): |
| 99 | + response.read() # empty body; preserve the original response (and its headers) |
| 100 | + return response |
| 101 | + content_length = _parse_content_length(response.headers.get("content-length")) |
| 102 | + if content_length is not None and content_length > cap: |
| 103 | + raise ResponseTooLargeError( |
| 104 | + status_code=response.status_code, limit=cap, content_length=content_length, reason="declared" |
| 105 | + ) |
| 106 | + try: |
| 107 | + content = _accumulate_capped(response.iter_bytes(), cap) |
| 108 | + except _CapExceeded: |
| 109 | + raise ResponseTooLargeError( |
| 110 | + status_code=response.status_code, limit=cap, content_length=content_length, reason="streamed" |
| 111 | + ) from None |
| 112 | + return httpx2.Response( |
| 113 | + status_code=response.status_code, |
| 114 | + headers=_buffered_headers(response.headers), |
| 115 | + content=content, |
| 116 | + request=request, |
| 117 | + extensions=_safe_extensions(response.extensions), |
| 118 | + history=response.history, |
| 119 | + ) |
| 120 | + |
| 121 | + |
| 122 | +async def _read_capped_async(response: httpx2.Response, cap: int, request: httpx2.Request) -> httpx2.Response: |
| 123 | + """Async mirror of `_read_capped` (counts decoded bytes from `aiter_bytes`).""" |
| 124 | + if not _response_has_body(request.method, response.status_code): |
| 125 | + await response.aread() # empty body; preserve the original response (and its headers) |
| 126 | + return response |
| 127 | + content_length = _parse_content_length(response.headers.get("content-length")) |
| 128 | + if content_length is not None and content_length > cap: |
| 129 | + raise ResponseTooLargeError( |
| 130 | + status_code=response.status_code, limit=cap, content_length=content_length, reason="declared" |
| 131 | + ) |
| 132 | + buf = bytearray() |
| 133 | + async for chunk in response.aiter_bytes(): |
| 134 | + buf += chunk |
| 135 | + if len(buf) > cap: |
| 136 | + raise ResponseTooLargeError( |
| 137 | + status_code=response.status_code, limit=cap, content_length=content_length, reason="streamed" |
| 138 | + ) |
| 139 | + return httpx2.Response( |
| 140 | + status_code=response.status_code, |
| 141 | + headers=_buffered_headers(response.headers), |
| 142 | + content=bytes(buf), |
| 143 | + request=request, |
| 144 | + extensions=_safe_extensions(response.extensions), |
| 145 | + history=response.history, |
| 146 | + ) |
0 commit comments