diff --git a/CHANGES/13509.bugfix.rst b/CHANGES/13509.bugfix.rst new file mode 100644 index 00000000000..8e74540415d --- /dev/null +++ b/CHANGES/13509.bugfix.rst @@ -0,0 +1 @@ +Fixed some edge case handling in multipart parts using base 64 encoding -- by :user:`Dreamsorcerer`. diff --git a/aiohttp/multipart.py b/aiohttp/multipart.py index bb217643426..e965756dad1 100644 --- a/aiohttp/multipart.py +++ b/aiohttp/multipart.py @@ -63,6 +63,13 @@ ) +# The base64 alphabet plus the padding character. +_BASE64_CHARS = frozenset( + b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" +) +_NON_BASE64_BYTES = bytes(b for b in range(256) if b not in _BASE64_CHARS) + + if TYPE_CHECKING: from .client_reqrep import ClientResponse @@ -304,6 +311,7 @@ def __init__( raise ValueError(f"invalid Content-Length: {length!r}") self._length = int(length) if length is not None else None self._read_bytes = 0 + self._b64_carry = b"" self._unread: deque[bytes] = deque() self._prev_chunk: bytes | None = None self._content_eof = 0 @@ -358,43 +366,59 @@ async def read_chunk(self, size: int = chunk_size) -> bytes: """ if self._at_eof: return b"" + carry = self._b64_carry + want = size - len(carry) + if carry: + self._b64_carry = b"" + want = max(want, self._boundary_len) if self._length: - chunk = await self._read_chunk_from_length(size) + fresh = await self._read_chunk_from_length(want) else: - chunk = await self._read_chunk_from_stream(size) + fresh = await self._read_chunk_from_stream(want) + chunk = carry + fresh + self._read_bytes += len(fresh) - # For the case of base64 data, we must read a fragment of size with a - # remainder of 0 by dividing by 4 for string without symbols \n or \r + # base64 decodes in quartets and every chunk is decoded on its own, so + # a chunk should not end mid-quartet. encoding = self.headers.get(CONTENT_TRANSFER_ENCODING) if encoding and encoding.lower() == "base64": - stripped_chunk = b"".join(chunk.split()) - remainder = len(stripped_chunk) % 4 - - while remainder != 0 and not self.at_eof(): - over_chunk_size = 4 - remainder - over_chunk = b"" - - if self._prev_chunk: - over_chunk = self._prev_chunk[:over_chunk_size] - self._prev_chunk = self._prev_chunk[len(over_chunk) :] - - if len(over_chunk) != over_chunk_size: - over_chunk += await self._content.read(4 - len(over_chunk)) + chunk = self._align_base64_chunk(chunk, len(carry) + want) - if not over_chunk: - self._at_eof = True - - stripped_chunk += b"".join(over_chunk.split()) - chunk += over_chunk - remainder = len(stripped_chunk) % 4 - - self._read_bytes += len(chunk) if self._read_bytes == self._length: self._at_eof = True if self._at_eof and await self._content.readline() != b"\r\n": raise ValueError("Reader did not read all the data or it is malformed") return chunk + def _align_base64_chunk(self, chunk: bytes, size: int) -> bytes: + at_end = self._at_eof or ( + self._length is not None and self._read_bytes >= self._length + ) + if not at_end and len(chunk) > size: + self._b64_carry = chunk[size:] + chunk = chunk[:size] + + remainder = len(chunk.translate(None, _NON_BASE64_BYTES)) % 4 + if not remainder or at_end: + return chunk + + # Walk back over the trailing partial quartet and carry it into the + # next chunk. + cut = len(chunk) + left = remainder + while left: + cut -= 1 + if chunk[cut] in _BASE64_CHARS: + left -= 1 + if not cut: + # No whole quartet to hand back, and carrying the lot would make + # no progress: the caller asked for this many bytes, and a part + # that holds no quartet within them holds none to give. + return chunk + + self._b64_carry = chunk[cut:] + self._b64_carry + return chunk[:cut] + async def _read_chunk_from_length(self, size: int) -> bytes: # Reads body part content chunk of the specified size. # The body part must has Content-Length header with proper value. diff --git a/tests/test_benchmarks_multipart.py b/tests/test_benchmarks_multipart.py new file mode 100644 index 00000000000..4596f00ae7f --- /dev/null +++ b/tests/test_benchmarks_multipart.py @@ -0,0 +1,72 @@ +"""codspeed benchmarks for multipart body part reading.""" + +import asyncio +import base64 +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +from multidict import CIMultiDict + +from aiohttp.hdrs import CONTENT_TRANSFER_ENCODING +from aiohttp.helpers import DEFAULT_CHUNK_SIZE, HeadersDictProxy +from aiohttp.multipart import BodyPartReader +from aiohttp.streams import StreamReader + +if TYPE_CHECKING: + from pytest_codspeed import BenchmarkFixture +else: + pytest_codspeed = pytest.importorskip("pytest_codspeed") + BenchmarkFixture = pytest_codspeed.BenchmarkFixture + +BOUNDARY = b"--:" +BASE64_HEADERS: CIMultiDict[str] = CIMultiDict({CONTENT_TRANSFER_ENCODING: "base64"}) + + +def _part(body: bytes, loop: asyncio.AbstractEventLoop) -> BodyPartReader: + stream = StreamReader( + mock.Mock(_reading_paused=False), DEFAULT_CHUNK_SIZE, loop=loop + ) + stream.feed_data(body) + stream.feed_eof() + return BodyPartReader( + BOUNDARY, + HeadersDictProxy(BASE64_HEADERS), + stream, + client_max_size=10 * 1024**2, + ) + + +def test_read_base64_part( + event_loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture +) -> None: + """Read a line-wrapped base64 part to completion. + + Every 8 KiB chunk lands mid-quartet, so this covers the common cost of + the base64 realignment in ``read_chunk`` on well-formed input. + """ + body = base64.encodebytes(b"x" * (256 * 1024)).replace(b"\n", b"\r\n") + body += b"\r\n--:--" + + @benchmark + def _run() -> None: + event_loop.run_until_complete(_part(body, event_loop).read()) + + +def test_read_chunk_base64_realignment( + event_loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture +) -> None: + """Complete a base64 quartet across the longest run it will tolerate. + + ``read_chunk`` extends a chunk until its significant-character count is a + multiple of 4. Insignificant bytes never advance that count, so this walks + close to the whole allowance in one call -- the worst legitimate case, and + the shape that used to cost quadratic time. + """ + body = b"A" + b" " * (12 * 1024) + b"BBB\r\n--:--" + + @benchmark + def _run() -> None: + event_loop.run_until_complete( + _part(body, event_loop).read_chunk(BodyPartReader.chunk_size) + ) diff --git a/tests/test_multipart.py b/tests/test_multipart.py index fe7903d48c9..d5c3c00592d 100644 --- a/tests/test_multipart.py +++ b/tests/test_multipart.py @@ -1,4 +1,5 @@ import asyncio +import base64 import gzip import io import json @@ -424,7 +425,7 @@ async def test_decode_with_content_transfer_encoding_base64(self) -> None: obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) result = b"" while not obj.at_eof(): - chunk = await obj.read_chunk(size=6) + chunk = await obj.read_chunk(size=8) result += obj.decode(chunk) assert b"Time to Relax!" == result @@ -434,11 +435,127 @@ async def test_decode_iter_with_content_transfer_encoding_base64(self) -> None: obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) result = b"" while not obj.at_eof(): - chunk = await obj.read_chunk(size=6) + chunk = await obj.read_chunk(size=8) async for decoded_chunk in obj.decode_iter(chunk): result += decoded_chunk assert b"Time to Relax!" == result + async def test_read_chunk_base64_content_length_no_overread(self) -> None: + # A base64 part whose Content-Length is not a multiple of 4 must not + # let the base64 realignment loop read past the declared length into + # the boundary and the parts that follow (regression for a negative + # `_read_chunk_from_length` chunk_size / StreamReader.read(-1) bypass + # of client_max_size). + b64 = b"VGltZSB0byBSZWxheCE" # 19 chars, 19 % 4 == 3 (unpadded) + secret = b"secret-from-next-part" + rest = b"\r\n--:\r\nContent-Length: %d\r\n\r\n%s\r\n--:--\r\n" % ( + len(secret), + secret, + ) + h = HeadersDictProxy( + CIMultiDict( + {"CONTENT-LENGTH": str(len(b64)), CONTENT_TRANSFER_ENCODING: "base64"} + ) + ) + with Stream(b64 + rest) as stream: + obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) + chunk = await obj.read_chunk(8192) + assert chunk == b64 + assert obj.at_eof() + assert secret not in chunk + # The following part is left intact on the stream. + assert secret in await stream.read() + + async def test_read_chunk_base64_bounded_by_requested_size(self) -> None: + # Whole quartets are handed back and the trailing partial quartet is + # carried into the next read, so size acts as a cap rather than a + # starting point. The carry is subtracted from the next read instead + # of being pushed back, so nothing accumulates over a long part: the + # lookahead would otherwise gain a couple of bytes on every call. + payload = b"x" * 8192 + h = HeadersDictProxy(CIMultiDict({CONTENT_TRANSFER_ENCODING: "base64"})) + body = base64.encodebytes(payload).replace(b"\n", b"\r\n") + with Stream(body + b"\r\n--:--") as stream: + obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) + decoded = b"" + chunks = [] + carried = 0 + while not obj.at_eof(): + chunks.append(await obj.read_chunk(64)) + decoded += obj.decode(chunks[-1]) + carried += bool(obj._b64_carry) + assert len(obj._prev_chunk or b"") <= 64 + assert len(obj._b64_carry) < 64 + assert decoded == payload + assert all(len(c) <= 64 for c in chunks[:-1]) + # The bound above is only meaningful if the carry is exercised, and + # over enough calls for a per-call drift to become visible. + assert len(chunks) > 100 + assert carried > len(chunks) // 2 + + async def test_read_chunk_base64_padding_run_does_not_amplify(self) -> None: + # A part padded with a long run of insignificant bytes used to make a + # single call swallow the whole run. Such a run holds no whole quartet, + # so this covers the branch that hands the chunk straight back: it must + # still cap the chunk at the size asked for, and must hand back + # something every time so callers looping on a truthy chunk make + # progress. Nothing is carried here, so this says nothing about the + # deferral path -- that is bounded by + # test_read_chunk_base64_bounded_by_requested_size. + h = HeadersDictProxy(CIMultiDict({CONTENT_TRANSFER_ENCODING: "base64"})) + body = b"A" + b" " * (64 * 1024) + b"BBB\r\n--:--" + with Stream(body) as stream: + obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) + while not obj.at_eof(): + chunk = await obj.read_chunk(8192) + assert chunk + assert len(chunk) <= 8192 + + async def test_read_chunk_base64_length_delimited_carries_quartet(self) -> None: + # A length-delimited part has no lookahead (_prev_chunk stays None), so + # this covers carrying a partial quartet on that path. The carry is + # subtracted from the next read, which the Content-Length accounting + # has to agree with, or the part overruns its declared length and the + # trailing CRLF check rejects it. Sizes are chosen to land mid-quartet. + payload = b"z" * 300 + b64 = base64.b64encode(payload) + h = HeadersDictProxy( + CIMultiDict( + {"CONTENT-LENGTH": str(len(b64)), CONTENT_TRANSFER_ENCODING: "base64"} + ) + ) + for size in (6, 7, 9, 13): + with Stream(b64 + b"\r\n--:--") as stream: + obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) + assert obj._prev_chunk is None + decoded = b"" + carried = 0 + while not obj.at_eof(): + decoded += obj.decode(await obj.read_chunk(size)) + carried += bool(obj._b64_carry) + assert decoded == payload, size + # Guard against this quietly becoming a no-op: the point is + # that the carry is exercised, not merely that reading works. + assert carried > 10, (size, carried) + + @pytest.mark.parametrize("size", (5, 6, 7, 8)) + async def test_read_chunk_base64_small_size_carry_drains(self, size: int) -> None: + # _read_chunk_from_stream refuses to read fewer than boundary_len + # bytes, so for a size this small the carry cannot always be + # subtracted from the next read. The chunk then has to be allowed over + # size: capping it back would read more than it hands back on every + # call, and the carry would climb without ever draining. + payload = b"q" * 600 + h = HeadersDictProxy(CIMultiDict({CONTENT_TRANSFER_ENCODING: "base64"})) + body = base64.encodebytes(payload).replace(b"\n", b"\r\n") + with Stream(body + b"\r\n--:--") as stream: + obj = aiohttp.BodyPartReader(BOUNDARY, h, stream) + decoded = b"" + while not obj.at_eof(): + decoded += obj.decode(await obj.read_chunk(size)) + assert len(obj._b64_carry) <= size + obj._boundary_len + assert decoded == payload + async def test_decode_with_content_encoding_deflate(self) -> None: h = HeadersDictProxy(CIMultiDict({CONTENT_ENCODING: "deflate"})) data = b"\x0b\xc9\xccMU(\xc9W\x08J\xcdI\xacP\x04\x00"