From 0e153cdd1f3bf9680b0e5a965c70b3dafef7cad4 Mon Sep 17 00:00:00 2001 From: Aditya Jain Date: Sat, 15 Aug 2026 03:09:33 -0700 Subject: [PATCH] fix(streaming): _S3SeekableIO.__next__/__iter__ don't advance self._position, corrupting subsequent seek()/read() __next__ and __iter__ delegated straight to raw_stream.__next__()/raw_stream.__iter__(), bypassing self._position bookkeeping entirely -- unlike read()/readline()/readlines(), which all correctly advance self._position by the number of bytes actually consumed. seek() uses self._position as ground truth to compute the S3 GetObject Range header on the next raw_stream access. Consuming any data via the iterator protocol (e.g. "for line in s3_object: ...", a fully supported documented usage) left self._position stuck at its pre-iteration value. A subsequent seek()/read() would then reopen the S3 stream with a Range header computed from that stale position -- silently returning the wrong slice of the object (skipped or duplicated bytes), with no exception raised. Fix: route __next__ through the already position-tracked readline(), and have __iter__ return self, matching the standard Python file-iterator protocol. --- .../utilities/streaming/_s3_seekable_io.py | 10 +++- .../streaming/_boto3/test_s3_seekable_io.py | 59 +++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py b/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py index c4dccb4ae19..bc0dd646d61 100644 --- a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py +++ b/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py @@ -169,10 +169,16 @@ def closed(self) -> bool: return self._closed def __next__(self): - return self.raw_stream.__next__() + # Route through readline() (not raw_stream.__next__()) so self._position is kept in sync with + # bytes actually consumed. Without this, a seek() following iteration computes offsets relative + # to a stale self._position, causing subsequent reads to silently return the wrong byte range. + line = self.readline() + if not line: + raise StopIteration + return line def __iter__(self): - return self.raw_stream.__iter__() + return self def __enter__(self): return self diff --git a/tests/functional/streaming/_boto3/test_s3_seekable_io.py b/tests/functional/streaming/_boto3/test_s3_seekable_io.py index bdcbe1ca5b2..5f35e86035e 100644 --- a/tests/functional/streaming/_boto3/test_s3_seekable_io.py +++ b/tests/functional/streaming/_boto3/test_s3_seekable_io.py @@ -163,6 +163,65 @@ def test_next(s3_seekable_obj, s3_client_stub): next(s3_seekable_obj) +def test_next_advances_position(s3_seekable_obj, s3_client_stub): + payload = b"hello world\nworld hello" + streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload)) + + s3_client_stub.add_response( + "get_object", + {"Body": streaming_body}, + {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"}, + ) + + assert next(s3_seekable_obj) == b"hello world\n" + # Regression: __next__ must advance self._position the same way read()/readline() do. + # Before the fix, __next__ delegated straight to raw_stream.__next__(), leaving + # self._position stuck at 0 -- so tell() (and any seek() relative to it) would be wrong. + assert s3_seekable_obj.tell() == len(b"hello world\n") + + assert next(s3_seekable_obj) == b"world hello" + assert s3_seekable_obj.tell() == len(payload) + + with pytest.raises(StopIteration): + next(s3_seekable_obj) + + +def test_iter_returns_self(s3_seekable_obj): + # Standard Python iterator protocol: __iter__ must return the object itself, not a + # separate iterator over the underlying raw stream (which would bypass position tracking). + assert iter(s3_seekable_obj) is s3_seekable_obj + + +def test_seek_after_iteration_uses_correct_range(s3_seekable_obj, s3_client_stub): + payload = b"hello world\nworld hello" + streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload)) + + s3_client_stub.add_response( + "get_object", + {"Body": streaming_body}, + {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"}, + ) + + # Consume the stream via the iterator protocol (e.g. "for line in s3_object: ...") + consumed = list(s3_seekable_obj) + total_consumed = sum(len(line) for line in consumed) + + # Seeking forward by 5 bytes from here must compute against the TRUE current position + # (total_consumed + 5), not a stale position left over from before iteration began. + s3_seekable_obj.seek(5, io.SEEK_CUR) + + s3_client_stub.add_response( + "get_object", + {"Body": ""}, + { + "Bucket": s3_seekable_obj.bucket, + "Key": s3_seekable_obj.key, + "Range": f"bytes={total_consumed + 5}-", + }, + ) + assert s3_seekable_obj.raw_stream is not None + + def test_context_manager(s3_seekable_obj, s3_client_stub): payload = b"test" streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))