fix termination of generators when client connection is closed (localstack#8582)

dfangl · web-flow · commit 0088cd7c8e27 · 2023-06-27T21:39:08.000+02:00
diff --git a/localstack/aws/handlers/logging.py b/localstack/aws/handlers/logging.py
@@ -82,7 +82,9 @@ def _prepare_logger(self, logger: logging.Logger, formatter: Type):
     def _log(self, context: RequestContext, response: Response):
         aws_logger = self.aws_logger
         http_logger = self.http_logger
-        is_internal_call = is_internal_call_context(context.request.headers)
+        is_internal_call = (
+            is_internal_call_context(context.request.headers) or context.is_internal_call
+        )
         if is_internal_call:
             aws_logger = self.internal_aws_logger
             http_logger = self.internal_http_logger
@@ -143,7 +145,7 @@ def _log(self, context: RequestContext, response: Response):
                     "request_headers": dict(context.request.headers),
                     # response
                     "output_type": "Response",
-                    "output": response.data,
+                    "output": "StreamingBody(unknown)" if response.is_streamed else response.data,
                     "response_headers": dict(response.headers),
                 },
             )
diff --git a/localstack/http/asgi.py b/localstack/http/asgi.py
@@ -253,6 +253,8 @@ def write_sync(self, data: bytes) -> None:
     async def write(self, data: bytes) -> None:
         if not self.started:
             raise ValueError("not started the response yet")
+        if getattr(self.send.__self__, "closed", None):
+            raise BrokenPipeError("Connection closed")
         await self.send({"type": "http.response.body", "body": data, "more_body": True})
         self.sent += len(data)
         if self.sent >= self.content_length:
@@ -374,7 +376,15 @@ async def handle_http(
                 else:
                     for packet in iterable:
                         await response.write(packet)
+        except ConnectionError as e:
+            client_info = "unknown"
+            if client := scope.get("client"):
+                address, port = client
+                client_info = f"{address}:{port}"
+            LOG.debug("Error while writing responses: %s (client_info: %s)", e, client_info)
         finally:
+            if iterable and hasattr(iterable, "aclose"):
+                await iterable.aclose()
             await response.close()
 
     async def handle_lifespan(
diff --git a/tests/unit/http_/test_asgi.py b/tests/unit/http_/test_asgi.py
@@ -1,5 +1,6 @@
 import json
 import logging
+import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from queue import Queue
@@ -96,6 +97,43 @@ def _gen():
     assert next(it) == b"baz"
 
 
+def test_chunked_transfer_encoding_client_timeout(serve_asgi_adapter):
+    # this test makes sure that creating a response with a generator automatically creates a
+    # transfer-encoding=chunked response
+
+    generator_exited = threading.Event()
+    continue_request = threading.Event()
+
+    @Request.application
+    def app(_request: Request) -> Response:
+        def _gen():
+            try:
+                yield "foo"
+                yield "bar\n"
+                continue_request.wait()
+                # only three are needed, let's send some more to make sure
+                for _ in range(10):
+                    yield "baz\n"
+            except GeneratorExit:
+                generator_exited.set()
+
+        return Response(_gen(), 200)
+
+    server = serve_asgi_adapter(app)
+
+    with requests.get(server.url, stream=True) as response:
+        assert response.headers["Transfer-Encoding"] == "chunked"
+
+        it = response.iter_lines()
+
+        assert next(it) == b"foobar"
+
+    # request is now closed, continue the response generator
+    continue_request.set()
+    # this flag is only set when generator is exited
+    assert generator_exited.wait(timeout=10)
+
+
 def test_chunked_transfer_encoding_request(serve_asgi_adapter):
     request_list: List[Request] = []
 
