From b5ee15267ae505938c988b07fffff2fefffb41c2 Mon Sep 17 00:00:00 2001
From: Paco Yang <yangbaigao@gmail.com>
Date: Fri, 28 Jan 2022 16:45:48 +0800
Subject: [PATCH] fix: move all data handle to protocol & ensure connection is
 closed (#1332)

---
 uvicorn/_handlers/http.py                | 26 +++++++++---------------
 uvicorn/protocols/http/h11_impl.py       |  4 ++++
 uvicorn/protocols/http/httptools_impl.py |  2 ++
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/uvicorn/_handlers/http.py b/uvicorn/_handlers/http.py
index e810f6626..996127e34 100644
--- a/uvicorn/_handlers/http.py
+++ b/uvicorn/_handlers/http.py
@@ -7,6 +7,9 @@
     from uvicorn.server import ServerState
 
 
+MAX_RECV = 2 ** 16
+
+
 async def handle_http(
     reader: asyncio.StreamReader,
     writer: asyncio.StreamWriter,
@@ -33,13 +36,14 @@ async def handle_http(
     # Use a future to coordinate between the protocol and this handler task.
     # https://docs.python.org/3/library/asyncio-protocol.html#connecting-existing-sockets
     loop = asyncio.get_event_loop()
-    connection_lost = loop.create_future()
+    reader_read = loop.create_task(reader.read(MAX_RECV))
+
 
     # Switch the protocol from the stream reader to our own HTTP protocol class.
     protocol = config.http_protocol_class(  # type: ignore[call-arg, operator]
         config=config,
         server_state=server_state,
-        on_connection_lost=lambda: connection_lost.set_result(True),
+        on_connection_lost=reader_read.cancel,
     )
     transport = writer.transport
     transport.set_protocol(protocol)
@@ -55,7 +59,7 @@ async def handle_http(
 
     @task.add_done_callback
     def retrieve_exception(task: asyncio.Task) -> None:
-        exc = task.exception()
+        exc = task.exception() if not task.cancelled() else None
 
         if exc is None:
             return
@@ -73,19 +77,8 @@ def retrieve_exception(task: asyncio.Task) -> None:
 
     # Kick off the HTTP protocol.
     protocol.connection_made(transport)
-
-    # Pass any data already in the read buffer.
-    # The assumption here is that we haven't read any data off the stream reader
-    # yet: all data that the client might have already sent since the connection has
-    # been established is in the `_buffer`.
-    data = reader._buffer  # type: ignore
-    if data:
-        protocol.data_received(data)
-
-    # Let the transport run in the background. When closed, this future will complete
-    # and we'll exit here.
-    await connection_lost
-
+    data = await reader_read
+    protocol.data_received(data)
 
 def _get_current_task() -> asyncio.Task:
     try:
@@ -97,3 +90,4 @@ def _get_current_task() -> asyncio.Task:
     task = current_task()
     assert task is not None
     return task
+
diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py
index 04b997ffe..c518616c5 100644
--- a/uvicorn/protocols/http/h11_impl.py
+++ b/uvicorn/protocols/http/h11_impl.py
@@ -222,6 +222,10 @@ def handle_events(self):
                     continue
                 self.cycle.more_body = False
                 self.cycle.message_event.set()
+            elif event_type is h11.ConnectionClosed:
+                break
+        if self.conn.our_state is h11.MUST_CLOSE and not self.transport.is_closing():
+            self.transport.close()
 
     def handle_upgrade(self, event):
         upgrade_value = None
diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py
index 33e21dc3b..0415ac49e 100644
--- a/uvicorn/protocols/http/httptools_impl.py
+++ b/uvicorn/protocols/http/httptools_impl.py
@@ -136,6 +136,8 @@ def data_received(self, data):
             self.transport.close()
         except httptools.HttpParserUpgrade:
             self.handle_upgrade()
+        if data == b"" and not self.transport.is_closing():
+            self.transport.close()
 
     def handle_upgrade(self):
         upgrade_value = None