[pull] master from aio-libs:master

CHANGES/13509.bugfix.rst

@@ -0,0 +1 @@
+Fixed some edge case handling in multipart parts using base 64 encoding -- by :user:`Dreamsorcerer`.

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.

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)
+        )

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"