Fix missing signed empty end frame on event stream close (#691)

* Send signed empty end frame on event stream close

* Update signer to use sign_empty method

* Add sign_empty method to EventSigner protocol

* Update unit test

* Address PR feedback

* Update changelog entries to feature type

Author: alexgromero

packages/aws-sdk-signers/.changes/next-release/aws-sdk-signers-enhancement-d97ea923259e448da2f865cd60bb7768.json

@@ -1,4 +1,4 @@
 {
-  "type": "enhancement",
+  "type": "feature",
   "description": "Added `AsyncEventSigner.sign_empty` for SigV4-signed Amazon Event Stream terminator frames. This supports services that require a final signed empty message before the HTTP body stream closes."
 }

packages/smithy-aws-event-stream/.changes/next-release/smithy-aws-event-stream-bugfix-e80cb499b6b04551a1792cdaa2211752.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added a signed empty end frame to `AWSEventPublisher.close` to signal event stream completion. This supports services that require a final signed empty message before the HTTP body stream closes."
+}

packages/smithy-aws-event-stream/src/smithy_aws_event_stream/aio/__init__.py

@@ -17,7 +17,7 @@
 
 from .._private.deserializers import EventDeserializer as _EventDeserializer
 from .._private.serializers import EventSerializer as _EventSerializer
-from ..events import Event
+from ..events import Event, EventMessage
 from ..exceptions import EventError
 
 logger = logging.getLogger(__name__)
@@ -85,9 +85,26 @@ async def close(self) -> None:
             return
         self._closed = True
 
-        if (close := getattr(self._writer, "close", None)) is not None:
-            if asyncio.iscoroutine(result := close()):
-                await result
+        try:
+            # Send a signed empty frame to signal stream completion.
+            if self._signing_config is not None:
+                end_frame = EventMessage()
+                identity = await self._signing_config.identity_resolver.get_identity(
+                    properties=self._signing_config.identity_properties
+                )
+                end_frame = await self._signing_config.signer.sign_empty(
+                    event=end_frame,
+                    identity=identity,
+                    properties=self._signing_config.signing_properties,
+                )
+                logger.debug(
+                    "Sending signed empty message to terminate the event stream."
+                )
+                await self._writer.write(end_frame.encode())
+        finally:
+            if (close := getattr(self._writer, "close", None)) is not None:
+                if asyncio.iscoroutine(result := close()):
+                    await result
 
     @property
     def closed(self) -> bool:

packages/smithy-aws-event-stream/tests/unit/_private/test_serializers.py

@@ -1,11 +1,14 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+import asyncio
+from io import BytesIO
 from typing import Any
+from unittest.mock import AsyncMock
 
 import pytest
 from smithy_aws_event_stream._private.serializers import EventSerializer
 from smithy_aws_event_stream.aio import AWSEventPublisher
-from smithy_aws_event_stream.events import EventMessage
+from smithy_aws_event_stream.events import Event, EventMessage
 from smithy_core.aio.types import AsyncBytesProvider
 from smithy_core.serializers import SerializeableShape
 from smithy_json import JSONCodec
@@ -77,3 +80,33 @@ async def test_send_to_closed_writer():
         await publisher.send(EVENT_STREAM_SERDE_CASES[0][0])
 
     assert publisher.closed
+
+
+async def test_close_sends_empty_end_frame_when_signing():
+    writer = AsyncBytesProvider()
+    end_frame = EventMessage()
+    signing_config = AsyncMock()
+    signing_config.signer.sign_empty.return_value = end_frame
+
+    publisher: AWSEventPublisher[Any] = AWSEventPublisher(
+        payload_codec=JSONCodec(), async_writer=writer, signing_config=signing_config
+    )
+
+    # Read from the writer concurrently with close(), since close() flushes
+    # and blocks until all chunks are consumed.
+    reader = asyncio.create_task(_read(writer))
+    await publisher.close()
+    written = await reader
+
+    signing_config.signer.sign_empty.assert_awaited_once()
+    assert publisher.closed
+    assert writer.closed
+
+    decoded = Event.decode(BytesIO(written))
+    assert decoded is not None
+    assert decoded.message.payload == b""
+    assert decoded.message.headers == {}
+
+
+async def _read(source: AsyncBytesProvider) -> bytes:
+    return b"".join([chunk async for chunk in source])

packages/smithy-core/.changes/next-release/smithy-core-enhancement-857985c5f3a848ddbab3179678a57aa8.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added `EventSigner.sign_empty` to the protocol for signing empty events."
+}

packages/smithy-core/src/smithy_core/aio/interfaces/auth.py

@@ -31,6 +31,17 @@ async def sign(self, *, event: Any, identity: I, properties: SP) -> Any:
         """Get a signed version of the event.
 
         :param event: The event to be signed.
+        :param identity: The identity to use to sign the event.
+        :param properties: Additional properties used to sign the event.
+        """
+        ...
+
+    async def sign_empty(self, *, event: Any, identity: I, properties: SP) -> Any:
+        """Get a signed version of an empty event.
+
+        :param event: The empty event to be signed.
+        :param identity: The identity to use to sign the event.
+        :param properties: Additional properties used to sign the event.
         """
         ...