Repository: pallets/werkzeug

Issue:
Chunked response support for werkzeug.serving
This is a new attempt to solve #1327, with a slightly different approach, and for a specific reason; explicit failures.

Right now, if a generator-driven response fails mid-stream, the error-handler may or may not add some error-description to the end of stream, before closing connection. A client will typically not distinguish this from a successful completed request. Unless the response is carefully inspected, the crash will go on undetected. In the worst case, for opaque response content, or when running with `passthrough_errors`, the response may be truncated in a way that cannot be detected by the client.

For HTTP-clients, such as Python stdlib, using `Transfer-Encoding: chunked` activates automatic detection of truncated response (`http.client.IncompleteRead` is thrown). Given that any HTTP/1.1-supporting client is mandated to support chunking, this should be a justifiable change in behavior.


Discussion:

Resolving patch:
diff --git a/src/werkzeug/serving.py b/src/werkzeug/serving.py
--- a/src/werkzeug/serving.py
+++ b/src/werkzeug/serving.py
@@ -245,9 +245,10 @@ def run_wsgi(self) -> None:
         headers_set: t.Optional[t.List[t.Tuple[str, str]]] = None
         status_sent: t.Optional[str] = None
         headers_sent: t.Optional[t.List[t.Tuple[str, str]]] = None
+        chunk_response: bool = False
 
         def write(data: bytes) -> None:
-            nonlocal status_sent, headers_sent
+            nonlocal status_sent, headers_sent, chunk_response
             assert status_set is not None, "write() before start_response"
             assert headers_set is not None, "write() before start_response"
             if status_sent is None:
@@ -262,16 +263,19 @@ def write(data: bytes) -> None:
                 header_keys = set()
                 for key, value in headers_sent:
                     self.send_header(key, value)
-                    key = key.lower()
-                    header_keys.add(key)
+                    header_keys.add(key.lower())
                 if not (
                     "content-length" in header_keys
                     or environ["REQUEST_METHOD"] == "HEAD"
                     or code < 200
                     or code in (204, 304)
                 ):
-                    self.close_connection = True
-                    self.send_header("Connection", "close")
+                    if self.protocol_version >= "HTTP/1.1":
+                        chunk_response = True
+                        self.send_header("Transfer-Encoding", "chunked")
+                    else:
+                        self.close_connection = True
+                        self.send_header("Connection", "close")
                 if "server" not in header_keys:
                     self.send_header("Server", self.version_string())
                 if "date" not in header_keys:
@@ -279,7 +283,17 @@ def write(data: bytes) -> None:
                 self.end_headers()
 
             assert isinstance(data, bytes), "applications must write bytes"
-            self.wfile.write(data)
+
+            if data:
+                if chunk_response:
+                    self.wfile.write(hex(len(data))[2:].encode())
+                    self.wfile.write(b"\r\n")
+
+                self.wfile.write(data)
+
+                if chunk_response:
+                    self.wfile.write(b"\r\n")
+
             self.wfile.flush()
 
         def start_response(status, headers, exc_info=None):  # type: ignore
@@ -298,11 +312,14 @@ def start_response(status, headers, exc_info=None):  # type: ignore
 
         def execute(app: "WSGIApplication") -> None:
             application_iter = app(environ, start_response)
+            nonlocal chunk_response
             try:
                 for data in application_iter:
                     write(data)
                 if not headers_sent:
                     write(b"")
+                if chunk_response:
+                    self.wfile.write(b"0\r\n\r\n")
             finally:
                 if hasattr(application_iter, "close"):
                     application_iter.close()  # type: ignore
@@ -314,6 +331,10 @@ def execute(app: "WSGIApplication") -> None:
         except Exception:
             if self.server.passthrough_errors:
                 raise
+
+            if status_sent is not None and chunk_response:
+                self.close_connection = True
+
             from .debug.tbtools import get_current_traceback
 
             traceback = get_current_traceback(ignore_system_exceptions=True)


Test patch:
diff --git a/tests/live_apps/run.py b/tests/live_apps/run.py
--- a/tests/live_apps/run.py
+++ b/tests/live_apps/run.py
@@ -4,9 +4,15 @@
 
 from werkzeug.serving import generate_adhoc_ssl_context
 from werkzeug.serving import run_simple
+from werkzeug.serving import WSGIRequestHandler
 from werkzeug.wrappers import Request
 from werkzeug.wrappers import Response
 
+
+class WSGI11RequestHandler(WSGIRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+
 name = sys.argv[1]
 mod = import_module(f"{name}_app")
 
@@ -24,6 +30,9 @@ def app(request):
 kwargs.update(json.loads(sys.argv[2]))
 ssl_context = kwargs.get("ssl_context")
 
+if kwargs.get("request_handler") == "HTTP/1.1":
+    kwargs["request_handler"] = WSGI11RequestHandler
+
 if ssl_context == "custom":
     kwargs["ssl_context"] = generate_adhoc_ssl_context()
 elif isinstance(ssl_context, list):
diff --git a/tests/live_apps/streaming_app.py b/tests/live_apps/streaming_app.py
new file mode 100644
--- /dev/null
+++ b/tests/live_apps/streaming_app.py
@@ -0,0 +1,14 @@
+from werkzeug.wrappers import Request
+from werkzeug.wrappers import Response
+
+
+@Request.application
+def app(request):
+    def gen():
+        for x in range(5):
+            yield f"{x}\n"
+
+        if request.path == "/crash":
+            raise Exception("crash requested")
+
+    return Response(gen())
diff --git a/tests/test_serving.py b/tests/test_serving.py
--- a/tests/test_serving.py
+++ b/tests/test_serving.py
@@ -159,7 +159,7 @@ def test_port_is_int():
 
 @pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
 @pytest.mark.parametrize("send_length", [False, True])
-def test_chunked_encoding(monkeypatch, dev_server, send_length):
+def test_chunked_request(monkeypatch, dev_server, send_length):
     stream, length, boundary = stream_encode_multipart(
         {
             "value": "this is text",
@@ -239,3 +239,35 @@ def test_multiline_header_folding(standard_app):
     data = json.load(r)
     r.close()
     assert data["HTTP_XYZ"] == "first\tsecond\tthird"
+
+
+@pytest.mark.parametrize("endpoint", ["", "crash"])
+def test_streaming_close_response(dev_server, endpoint):
+    """When using HTTP/1.0, chunked encoding is not supported. Fall
+    back to Connection: close, but this allows no reliable way to
+    distinguish between complete and truncated responses.
+    """
+    r = dev_server("streaming").request("/" + endpoint)
+    assert r.getheader("connection") == "close"
+    assert r.data == "".join(str(x) + "\n" for x in range(5)).encode()
+
+
+def test_streaming_chunked_response(dev_server):
+    """When using HTTP/1.1, use Transfer-Encoding: chunked for streamed
+    responses, since it can distinguish the end of the response without
+    closing the connection.
+
+    https://tools.ietf.org/html/rfc2616#section-3.6.1
+    """
+    r = dev_server("streaming", request_handler="HTTP/1.1").request("/")
+    assert r.getheader("transfer-encoding") == "chunked"
+    assert r.data == "".join(str(x) + "\n" for x in range(5)).encode()
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
+def test_streaming_chunked_truncation(dev_server):
+    """When using HTTP/1.1, chunked encoding allows the client to detect
+    content truncated by a prematurely closed connection.
+    """
+    with pytest.raises(http.client.IncompleteRead):
+        dev_server("streaming", request_handler="HTTP/1.1").request("/crash")