fixes #1125 due to missing aync override (#1126)

aio-libs · Jul 18, 2024 · 7154abc · 7154abc
1 parent eae9743
commit 7154abc
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 2 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1,5 +1,8 @@
 Changes
 -------
+2.13.2 (2024-07-18)
+^^^^^^^^^^^^^^^^^^^
+* fix for #1125 due to missing patch of StreamingChecksumBody
 
 2.13.1 (2024-06-24)
 ^^^^^^^^^^^^^^^^^^^

diff --git a/aiobotocore/__init__.py b/aiobotocore/__init__.py
@@ -1 +1 @@
-__version__ = '2.13.1'
+__version__ = '2.13.2'
diff --git a/aiobotocore/httpchecksum.py b/aiobotocore/httpchecksum.py
@@ -5,14 +5,14 @@
     AwsChunkedWrapper,
     FlexibleChecksumError,
     _apply_request_header_checksum,
-    _handle_streaming_response,
     base64,
     conditionally_calculate_md5,
     determine_content_length,
     logger,
 )
 
 from aiobotocore._helpers import resolve_awaitable
+from aiobotocore.response import StreamingBody
 
 
 class AioAwsChunkedWrapper(AwsChunkedWrapper):
@@ -44,6 +44,30 @@ async def __anext__(self):
         raise StopAsyncIteration()
 
 
+# unfortunately we can't inherit from botocore's StreamingChecksumBody due to
+# subclassing
+class StreamingChecksumBody(StreamingBody):
+    def __init__(self, raw_stream, content_length, checksum, expected):
+        super().__init__(raw_stream, content_length)
+        self._checksum = checksum
+        self._expected = expected
+
+    async def read(self, amt=None):
+        chunk = await super().read(amt=amt)
+        self._checksum.update(chunk)
+        if amt is None or (not chunk and amt > 0):
+            self._validate_checksum()
+        return chunk
+
+    def _validate_checksum(self):
+        if self._checksum.digest() != base64.b64decode(self._expected):
+            error_msg = (
+                f"Expected checksum {self._expected} did not match calculated "
+                f"checksum: {self._checksum.b64digest()}"
+            )
+            raise FlexibleChecksumError(error_msg=error_msg)
+
+
 async def handle_checksum_body(
     http_response, response, context, operation_model
 ):
@@ -87,6 +111,17 @@ async def handle_checksum_body(
     )
 
 
+def _handle_streaming_response(http_response, response, algorithm):
+    checksum_cls = _CHECKSUM_CLS.get(algorithm)
+    header_name = "x-amz-checksum-%s" % algorithm
+    return StreamingChecksumBody(
+        http_response.raw,
+        response["headers"].get("content-length"),
+        checksum_cls(),
+        response["headers"][header_name],
+    )
+
+
 async def _handle_bytes_response(http_response, response, algorithm):
     body = await http_response.content
     header_name = "x-amz-checksum-%s" % algorithm

diff --git a/tests/test_patches.py b/tests/test_patches.py
@@ -52,8 +52,10 @@
 from botocore.hooks import EventAliaser, HierarchicalEmitter
 from botocore.httpchecksum import (
     AwsChunkedWrapper,
+    StreamingChecksumBody,
     _apply_request_trailer_checksum,
     _handle_bytes_response,
+    _handle_streaming_response,
     apply_request_checksum,
     handle_checksum_body,
 )
@@ -647,6 +649,10 @@
     handle_checksum_body: {
         '898cee7a7a5e5a02af7e0e65dcbb8122257b85df',
     },
+    _handle_streaming_response: {'7ce971e012f9d4b04889f0af83f67281ed6a9e6e'},
+    StreamingChecksumBody: {
+        '2c6eb22268d46abae261ce386eb2deabbc3a0dcd',
+    },
     _handle_bytes_response: {'0761c4590c6addbe8c674e40fca9f7dd375a184b'},
     AwsChunkedWrapper._make_chunk: {
         '097361692f0fd6c863a17dd695739629982ef7e4'