Rewind file bodies when internally retrying a request (aio-libs#13330)

Author: aiolibsbot

CHANGES/13329.bugfix.rst

@@ -0,0 +1,2 @@
+Fixed internally retried requests sending a truncated body when the request
+data was a file object -- by :user:`aiolibsbot`.

CHANGES/13330.bugfix.rst

@@ -0,0 +1 @@
+13329.bugfix.rst

aiohttp/client.py

@@ -733,10 +733,18 @@ async def _request(
                     ):
                         raise
                     except (ClientOSError, ServerDisconnectedError):
-                        if retry_persistent_connection:
-                            retry_persistent_connection = False
-                            continue
-                        raise
+                        if not retry_persistent_connection:
+                            raise
+                        retry_persistent_connection = False
+                        if data is not None:
+                            # Rebuilding from `data` would resend only the unread
+                            # remainder of a file object; reuse the payload, which
+                            # rewinds itself once the cancelled writer has settled.
+                            await req._close()
+                            if req._body.consumed:
+                                raise
+                            data = req._body
+                        continue
                     except ClientError:
                         raise
                     except OSError as exc:

aiohttp/multipart.py

@@ -663,6 +663,10 @@ async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> byt
 
     async def write(self, writer: AbstractStreamWriter) -> None:
         field = self._value
+        # Reading the part drains the underlying stream irreversibly, so mark the
+        # payload consumed up front: even an interrupted write leaves nothing that
+        # a retry or redirect could replay.
+        self._consumed = True
         while chunk := await field.read_chunk(size=DEFAULT_CHUNK_SIZE):
             async for d in field.decode_iter(chunk):
                 await writer.write(d)
@@ -944,6 +948,11 @@ def __exit__(
     ) -> None:
         pass
 
+    @property
+    def consumed(self) -> bool:
+        """Whether the writer or any of its parts can no longer be replayed."""
+        return self._consumed or any(part.consumed for part, _, _ in self._parts)
+
     def __iter__(self) -> Iterator[_Part]:
         return iter(self._parts)
 

aiohttp/payload.py

@@ -1049,6 +1049,10 @@ async def write_with_length(
 
         # Stream from the iterator
         remaining_bytes = content_length
+        # Nothing is cached, so advancing the iterator is irreversible: mark the
+        # payload consumed up front so an interrupted write cannot be replayed
+        # from a partially drained iterator.
+        self._consumed = True
 
         try:
             while True:
@@ -1066,7 +1070,6 @@ async def write_with_length(
         except StopAsyncIteration:
             # Iterator is exhausted
             self._iter = None
-            self._consumed = True  # Mark as consumed when streamed without caching
 
     def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str:
         """Decode the payload content as a string if cached chunks are available."""

tests/test_client_functional.py

@@ -5895,6 +5895,74 @@ async def handler(request: web.Request) -> web.Response:
         await asyncio.to_thread(f.close)
 
 
+async def test_file_upload_retry_persistent_connection(
+    aiohttp_client: AiohttpClient, tmp_path: pathlib.Path
+) -> None:
+    """A retried request must resend the whole file, not the unread remainder."""
+    received_bodies: list[bytes] = []
+    num_requests = 0
+
+    async def handler(request: web.Request) -> web.Response:
+        nonlocal num_requests
+        num_requests += 1
+        if num_requests == 1:
+            assert request.transport is not None
+            request.transport.close()
+            return web.Response()
+
+        received_bodies.append(await request.read())
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_put("/upload", handler)
+
+    client = await aiohttp_client(app)
+    client.session._retry_connection = True
+
+    test_file = tmp_path / "test_retry_upload.txt"
+    content = b"This is test file content for a retried upload."
+    await asyncio.to_thread(test_file.write_bytes, content)
+
+    f = await asyncio.to_thread(open, test_file, "rb")
+    try:
+        async with client.put("/upload", data=f) as resp:
+            assert resp.status == 200
+    finally:
+        await asyncio.to_thread(f.close)
+
+    assert num_requests == 2
+    assert received_bodies == [content]
+
+
+async def test_upload_retry_persistent_connection_unseekable_body(
+    aiohttp_client: AiohttpClient,
+) -> None:
+    """An unreplayable body must not be silently resent truncated on retry."""
+    num_requests = 0
+
+    async def handler(request: web.Request) -> web.Response:
+        nonlocal num_requests
+        num_requests += 1
+        assert request.transport is not None
+        request.transport.close()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_put("/upload", handler)
+
+    client = await aiohttp_client(app)
+    client.session._retry_connection = True
+
+    async def gen() -> AsyncIterator[bytes]:
+        yield b"chunk1"
+        yield b"chunk2"
+
+    with pytest.raises((aiohttp.ServerDisconnectedError, aiohttp.ClientOSError)):
+        await client.put("/upload", data=gen())
+
+    assert num_requests == 1
+
+
 async def test_stream_reader_total_raw_bytes(aiohttp_client: AiohttpClient) -> None:
     """Test whether StreamReader.total_raw_bytes returns the number of bytes downloaded"""
     source_data = b"@dKal^pH>1h|YW1:c2J$" * 4096

tests/test_multipart.py

@@ -4,6 +4,7 @@
 import json
 import pathlib
 import sys
+from collections.abc import AsyncIterator
 from types import TracebackType
 from unittest import mock
 
@@ -1835,3 +1836,59 @@ async def test_multipart_writer_close_with_exceptions() -> None:
     await writer.close()
     assert part1.close.call_count == 1
     assert part2.close.call_count == 1
+
+
+async def test_multipart_writer_consumed_follows_parts() -> None:
+    """A writer holding an unreplayable part must report itself as consumed."""
+
+    async def gen() -> AsyncIterator[bytes]:
+        yield b"chunk1"
+        yield b"chunk2"
+
+    writer = aiohttp.MultipartWriter()
+    writer.append(b"replayable")
+    assert writer.consumed is False
+
+    part = writer.append(gen())
+    assert writer.consumed is False
+
+    stream = mock.Mock()
+    stream.write = mock.AsyncMock()
+    await part.write_with_length(stream, None)
+
+    assert part.consumed is True
+    assert writer.consumed is True
+
+
+async def test_body_part_reader_payload_consumed_after_write() -> None:
+    """A drained body part reader must report itself as consumed."""
+    with Stream(b"Hello, world!\r\n--:--") as stream:
+        body_part = aiohttp.BodyPartReader(
+            BOUNDARY, HeadersDictProxy(CIMultiDict()), stream
+        )
+        payload = BodyPartReaderPayload(body_part)
+        assert payload.consumed is False
+
+        writer = mock.Mock()
+        writer.write = mock.AsyncMock()
+        await payload.write(writer)
+
+        assert payload.consumed is True
+
+
+async def test_multipart_writer_consumed_follows_body_part_reader() -> None:
+    """A writer holding a drained body part reader must report itself consumed."""
+    with Stream(b"Hello, world!\r\n--:--") as stream:
+        body_part = aiohttp.BodyPartReader(
+            BOUNDARY, HeadersDictProxy(CIMultiDict()), stream
+        )
+        writer = aiohttp.MultipartWriter()
+        part = writer.append(body_part)
+        assert writer.consumed is False
+
+        out = mock.Mock()
+        out.write = mock.AsyncMock()
+        await part.write_with_length(out, None)
+
+        assert part.consumed is True
+        assert writer.consumed is True

tests/test_payload.py

@@ -964,6 +964,33 @@ async def gen() -> AsyncIterator[bytes]:
     assert writer2.get_written_bytes() == b""
 
 
+async def test_async_iterable_payload_consumed_on_interrupted_write() -> None:
+    """An interrupted write must still mark an uncached payload as consumed."""
+
+    async def gen() -> AsyncIterator[bytes]:
+        yield b"chunk1"
+        yield b"chunk2"
+
+    class FailingWriter(MockStreamWriter):
+        async def write(
+            self,
+            chunk: Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"],
+        ) -> None:
+            if self.written:
+                raise ConnectionResetError("connection lost")
+            await super().write(chunk)
+
+    p = payload.AsyncIterablePayload(gen())
+    writer = FailingWriter()
+
+    with pytest.raises(ConnectionResetError):
+        await p.write_with_length(writer, None)
+
+    # The iterator was partially drained, so the payload cannot be replayed.
+    assert writer.get_written_bytes() == b"chunk1"
+    assert p.consumed is True
+
+
 async def test_bytes_io_payload_close_does_not_close_io() -> None:
     """Test that BytesIOPayload close() does not close the underlying BytesIO."""
     bytes_io = io.BytesIO(b"data")