fix Transfer-Encoding behavior in input for SimpleRequestsClient (#39)

Author: bentsku

rolo/client.py

@@ -128,6 +128,10 @@ def request(self, request: Request, server: str | None = None) -> Response:
         if not request.headers.get("accept-encoding"):
             headers["accept-encoding"] = urllib3.util.SKIP_HEADER
 
+        # we need to remove `Transfer-Encoding` because it is a hop-by-hop header, and most of the `Request` objects
+        # passed are coming from a webserver, which already decoded the request
+        headers.pop("Transfer-Encoding", None)
+
         response = self.session.request(
             method=request.method,
             # use raw base url to preserve path url encoding

tests/test_client.py

@@ -58,3 +58,57 @@ def multi_values_handler(_request: WerkzeugRequest) -> Response:
         with SimpleRequestsClient() as client:
             response = client.request(request, url)
             assert response.headers.getlist("Set-Cookie") == ["value1", "value2"]
+
+    def test_chunked_encoded_request(self, httpserver: HTTPServer):
+        # because SimpleRequestsClient mostly forwards an existing incoming request, and it uses `restore_payload`
+        # this means that any `Transfer-Encoding` headers must be stripped before sending
+
+        def transfer_encoded_handler(_request: WerkzeugRequest) -> Response:
+            return Response(response=_request.data)
+
+        httpserver.expect_request("/").respond_with_handler(transfer_encoded_handler)
+
+        url = httpserver.url_for("/")
+        body = b"hello world"
+        request = Request(
+            path="/",
+            method="POST",
+            body=body,
+            headers={
+                "Transfer-Encoding": "chunked",
+                "Content-Length": str(len(body)),
+            },
+        )
+
+        with SimpleRequestsClient() as client:
+            response = client.request(request, url)
+
+        assert response.data == body
+
+    def test_gzip_encoded_request(self, httpserver: HTTPServer):
+        def transfer_encoded_handler(_request: WerkzeugRequest) -> Response:
+            return Response(response=_request.data)
+
+        httpserver.expect_request("/").respond_with_handler(transfer_encoded_handler)
+
+        url = httpserver.url_for("/")
+        raw_body = b"hello world"
+        request = Request(
+            path="/",
+            method="POST",
+            # we need to use the raw body here, because in real world use case, the webserver would have read and
+            # decoded the payload
+            body=raw_body,
+            headers={
+                "Transfer-Encoding": "gzip",
+                "Content-Length": str(len(raw_body)),
+            },
+        )
+
+        with SimpleRequestsClient() as client:
+            response = client.request(request, url)
+
+        assert response.data == raw_body
+        # we're making sure we're not passing the `Transfer-Encoding` gzip along, as we're read the body and sent it
+        # decoded over the wire
+        assert "Transfer-Encoding" not in httpserver.log[0][0].headers

tests/test_proxy.py

@@ -183,7 +183,9 @@ def _handler(request: WerkzeugRequest):
         }
 
     @pytest.mark.parametrize("chunked", [True, False])
-    def test_proxy_handler_transfer_encoding(self, router_server, httpserver: HTTPServer, chunked):
+    def test_proxy_handler_transfer_encoding_output(
+        self, router_server, httpserver: HTTPServer, chunked
+    ):
         router, proxy = router_server
         backend = httpserver
         body = "enough-for-content-length"
@@ -328,6 +330,28 @@ def _handler(_request: Request) -> Response:
 
         assert response.data == body
 
+    def test_proxy_handler_transfer_encoding_input(self, router_server, httpserver: HTTPServer):
+        router, proxy = router_server
+        backend = httpserver
+
+        def _handler(_request: WerkzeugRequest):
+            # the handler/httpserver fixture will automatically decode the request. if the request was malformed, it
+            # would fail
+            return Response(_request.data)
+
+        backend.expect_request("").respond_with_handler(_handler)
+
+        router.add("/", ProxyHandler(backend.url_for("/")))
+
+        def gen():
+            yield b"fizz\n"
+            yield b"buzz\n"
+            yield b"done"
+
+        response = requests.post(proxy.url, data=gen())
+
+        assert response.text == "fizz\nbuzz\ndone"
+
 
 @pytest.mark.parametrize("consume_data", [True, False])
 def test_forward_files_and_form_data_proxy_consumes_data(