From afee2fda5d011219973890c68c23016fa8f002a2 Mon Sep 17 00:00:00 2001 From: Ben Simon Hartung <42031100+bentsku@users.noreply.github.com> Date: Wed, 23 Oct 2024 14:41:43 +0200 Subject: [PATCH] remove Transfer-Encoding from Proxy responses (#24) --- rolo/client.py | 10 ++++- tests/test_proxy.py | 100 +++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 103 insertions(+), 7 deletions(-) diff --git a/rolo/client.py b/rolo/client.py index 3e18170..f693d3f 100644 --- a/rolo/client.py +++ b/rolo/client.py @@ -124,8 +124,16 @@ def request(self, request: Request, server: str | None = None) -> Response: return final_response response_headers = Headers(dict(response.headers)) - if "chunked" in response_headers.get("Transfer-Encoding", ""): + if "chunked" in (transfer_encoding := response_headers.get("Transfer-Encoding", "")): response_headers.pop("Content-Length", None) + # We should not set `Transfer-Encoding` in a Response, because it is the responsibility of the webserver + # to do so, if there are no Content-Length. However, gzip behavior is more related to the actual content of + # the response, so we keep that one. + transfer_encoding_values = [v.strip() for v in transfer_encoding.split(",")] + transfer_encoding_no_chunked = [ + v for v in transfer_encoding_values if v.lower() != "chunked" + ] + response_headers.setlist("Transfer-Encoding", transfer_encoding_no_chunked) final_response = Response( response=(chunk for chunk in response.raw.stream(1024, decode_content=False)), diff --git a/tests/test_proxy.py b/tests/test_proxy.py index 2ab8bed..34b3870 100644 --- a/tests/test_proxy.py +++ b/tests/test_proxy.py @@ -1,3 +1,4 @@ +import gzip import json import pytest @@ -181,6 +182,36 @@ def _handler(request: WerkzeugRequest): "query": "%E4%B8%8A%E6%B5%B7%2B%E4%B8%AD%E5%9C%8B=%E4%B8%8A%E6%B5%B7%2B%E4%B8%AD%E5%9C%8B", } + @pytest.mark.parametrize("chunked", [True, False]) + def test_proxy_handler_transfer_encoding(self, router_server, httpserver: HTTPServer, chunked): + router, proxy = router_server + backend = httpserver + body = "enough-for-content-length" + + def _handler(_: WerkzeugRequest): + # if the response is chunked, return a generator instead, which will return `Transfer-Encoding: chunked` + if chunked: + _body = (c for c in body) + else: + _body = body + + return Response(_body, status=200) + + backend.expect_request("").respond_with_handler(_handler) + + router.add("/", ProxyHandler(backend.url_for("/"))) + + response = requests.get(proxy.url) + + if chunked: + assert response.headers["Transfer-Encoding"] == "chunked" + assert "Content-Length" not in response.headers + else: + assert response.headers["Content-Length"] == str(len(body)) + assert "Transfer-Encoding" not in response.headers + + assert response.text == body + class TestProxy: def test_proxy_with_custom_client(self, httpserver: HTTPServer): @@ -217,12 +248,14 @@ def test_proxy_for_transfer_encoding_chunked( ): body = "enough-for-content-length" - def _handler(_: WerkzeugRequest): - headers = ( - {"Content-Length": len(body)} if not chunked else {"Transfer-Encoding": "chunked"} - ) + def _handler(_request: Request) -> Response: + # if the response is chunked, return a generator instead, which will return `Transfer-Encoding: chunked` + if chunked: + _body = (c for c in body) + else: + _body = body - return Response(body, headers=headers) + return Response(_body, status=200) httpserver.expect_request("").respond_with_handler(_handler) @@ -233,12 +266,67 @@ def _handler(_: WerkzeugRequest): response = proxy.request(request) if chunked: - assert response.headers["Transfer-Encoding"] == "chunked" + # the proxy should not return a Transfer-Encoding, as this is something the webserver should set + assert "Transfer-Encoding" not in response.headers assert "Content-Length" not in response.headers else: assert response.headers["Content-Length"] == str(len(body)) assert "Transfer-Encoding" not in response.headers + assert response.data.decode() == body + + @pytest.mark.parametrize( + "chunked,gzipped", + [ + (False, False), + (False, True), + (True, False), + (True, True), + ], + ) + def test_proxy_for_transfer_encoding_chunked_and_gzip( + self, + httpserver: HTTPServer, + chunked, + gzipped, + ): + body = b"enough-for-content-length" + if gzipped: + body = gzip.compress(body, mtime=0) + + def _handler(_request: Request) -> Response: + # if the response is chunked, return a generator instead, which will return `Transfer-Encoding: chunked` + headers = {} + _body = body + if gzipped: + headers["Transfer-Encoding"] = "gzip" + + if chunked: + _body = (chr(c).encode("latin-1") for c in body) + + return Response(_body, status=200, headers=headers) + + httpserver.expect_request("/proxy").respond_with_handler(_handler) + + proxy = Proxy(httpserver.url_for("/").lstrip("/")) + + request = Request(path="/proxy", method="GET", headers={"Host": "127.0.0.1:80"}) + + response = proxy.request(request) + + if gzipped: + assert response.headers["Transfer-Encoding"] == "gzip" + + if chunked: + assert "Content-Length" not in response.headers + assert "chunked" not in response.headers.get("Transfer-Encoding", "") + else: + assert response.headers["Content-Length"] == str(len(body)) + if not gzipped: + assert "Transfer-Encoding" not in response.headers + + assert response.data == body + @pytest.mark.parametrize("consume_data", [True, False]) def test_forward_files_and_form_data_proxy_consumes_data(