Fix chunked forwarding response

 Please enter the commit message for your changes. Lines starting

Author: julien-duponchelle

examples/persistent_ca_usage.py

@@ -50,6 +50,10 @@ def get_or_create_ca():
 class LoggingForwardProxyHandler(HTTPSForwardProxyHandler):
     """Example forward proxy handler with logging."""
 
+    def __init__(self):
+        super().__init__()
+        self.response_size = 0
+
     async def on_client_connected(self):
         print(f"Client connected: {self.request}")
         await super().on_client_connected()
@@ -63,6 +67,15 @@ async def on_error(self, error: Exception):
         print(f"❌ Proxy error: {type(error).__name__}: {error}")
         await super().on_error(error)
 
+    async def on_response_chunk(self, chunk: bytes) -> bytes:
+        """Process each response chunk - log size and analyze content."""
+        chunk_size = len(chunk)
+        self.response_size += chunk_size
+        print(
+            f"  Received chunk: {chunk_size} bytes, total so far: {self.response_size} bytes"
+        )
+        return chunk
+
 
 async def main():
     """Run a proxy with persistent CA."""

src/asyncio_https_proxy/chunked_encoding.py

@@ -0,0 +1,146 @@
+"""
+HTTP Chunked Transfer Encoding Module
+
+This module provides utilities for handling HTTP/1.1 chunked transfer encoding,
+which allows HTTP messages to be sent in chunks when the total content length
+is not known in advance.
+
+## How Chunked Encoding Works
+
+In HTTP chunked transfer encoding, the message body is sent as a series of chunks:
+
+1. **Chunk Format**: Each chunk consists of:
+   - Chunk size in hexadecimal followed by CRLF
+   - Chunk data
+   - Trailing CRLF
+
+2. **Example**:
+   ```
+   5\r\n        <- chunk size (5 bytes in hex)
+   Hello\r\n    <- chunk data + trailing CRLF
+   6\r\n        <- next chunk size (6 bytes)
+   World!\r\n   <- chunk data + trailing CRLF
+   0\r\n        <- terminating chunk (size 0)
+   \r\n         <- final CRLF (no trailers)
+   ```
+
+3. **Termination**: The last chunk has size 0, followed by optional trailers
+   and a final CRLF.
+"""
+
+import asyncio
+from typing import AsyncIterator, Awaitable, Callable, Optional
+
+
+class ChunkedReader:
+    """Handles reading HTTP chunked transfer encoding."""
+
+    def __init__(self, reader: asyncio.StreamReader):
+        self.reader = reader
+
+    async def read_chunks(self) -> AsyncIterator[bytes]:
+        """
+        Read chunked response data.
+
+        Yields:
+            Chunks of data as bytes.
+        """
+        while True:
+            # Read chunk size line
+            chunk_size_line = await self.reader.readline()
+            if not chunk_size_line:
+                break
+
+            # Parse chunk size (hex digits before optional ';' and extensions)
+            chunk_size_str = chunk_size_line.decode("ascii").strip()
+            try:
+                # Split on ';' to handle chunk extensions, take only the size part
+                chunk_size = int(chunk_size_str.split(";")[0], 16)
+            except (ValueError, UnicodeDecodeError):
+                break
+
+            if chunk_size == 0:
+                # Terminating chunk - read trailers until empty line
+                while True:
+                    trailer = await self.reader.readline()
+                    if not trailer or trailer == b"\r\n":
+                        break
+                break
+
+            # Read exactly chunk_size bytes of data
+            chunk_data = b""
+            remaining = chunk_size
+            while remaining > 0:
+                data = await self.reader.read(remaining)
+                if not data:
+                    # Connection closed unexpectedly
+                    break
+                chunk_data += data
+                remaining -= len(data)
+
+            if len(chunk_data) != chunk_size:
+                # Incomplete read - connection was closed
+                break
+
+            # Read the trailing CRLF after chunk data
+            trailing_crlf = b""
+            while len(trailing_crlf) < 2:
+                data = await self.reader.read(2 - len(trailing_crlf))
+                if not data:
+                    break
+                trailing_crlf += data
+
+            if trailing_crlf != b"\r\n":
+                # Invalid chunk format
+                break
+
+            yield chunk_data
+
+
+class ChunkedWriter:
+    """Handles writing HTTP chunked transfer encoding."""
+
+    def __init__(self, write_func: Callable[[bytes], None]):
+        self.write_func = write_func
+
+    def write_chunk(self, data: bytes) -> None:
+        """
+        Write data as a chunked transfer encoding chunk.
+
+        Args:
+            data: The chunk data to write
+        """
+        if data:
+            chunk_size_hex = hex(len(data))[2:].encode()  # Remove '0x' prefix
+            self.write_func(chunk_size_hex + b"\r\n")
+            self.write_func(data)
+            self.write_func(b"\r\n")
+
+    def write_final_chunk(self) -> None:
+        self.write_func(b"0\r\n\r\n")
+
+
+async def forward_chunked_response(
+    reader: asyncio.StreamReader,
+    write_func: Callable[[bytes], None],
+    chunk_hook: Callable[[bytes], Awaitable[Optional[bytes]]],
+) -> None:
+    """
+    Forward a chunked response from reader to writer.
+
+    Args:
+        reader: The stream reader to read from
+        write_func: Function to write response data
+        chunk_hook: Async function to process each chunk
+    """
+    chunked_reader = ChunkedReader(reader)
+    chunked_writer = ChunkedWriter(write_func)
+
+    async for chunk_data in chunked_reader.read_chunks():
+        # Process chunk data through hook
+        processed_chunk = await chunk_hook(chunk_data)
+
+        if processed_chunk is not None:
+            chunked_writer.write_chunk(processed_chunk)
+
+    chunked_writer.write_final_chunk()

src/asyncio_https_proxy/https_forward_proxy_handler.py

@@ -1,7 +1,7 @@
 import asyncio
 import ssl
-from typing import AsyncIterator
 
+from .chunked_encoding import forward_chunked_response
 from .http_response import HTTPResponse
 from .https_proxy_handler import HTTPSProxyHandler
 
@@ -203,7 +203,11 @@ async def _forward_response_body(self):
                     transfer_encoding = self.response.headers.first("Transfer-Encoding")
 
                 if transfer_encoding and "chunked" in transfer_encoding.lower():
-                    await self._forward_chunked_response()
+                    await forward_chunked_response(
+                        self.upstream_reader,
+                        self.write_response,
+                        self.on_response_chunk,
+                    )
                 else:
                     # Read until EOF
                     while True:
@@ -222,84 +226,3 @@ async def _forward_response_body(self):
             if not self._response_complete:
                 self._response_complete = True
                 await self.on_response_complete()
-
-    async def _forward_chunked_response(self):
-        """Forward chunked response data."""
-
-        if self.upstream_reader is None:
-            return
-
-        while True:
-            # Read chunk size line
-            chunk_size_line = await self.upstream_reader.readline()
-            self.write_response(chunk_size_line)
-
-            # Parse chunk size
-            try:
-                chunk_size = int(chunk_size_line.decode().split(";")[0], 16)
-            except (ValueError, UnicodeDecodeError):
-                break
-
-            if chunk_size == 0:
-                # Read trailers until empty line
-                while True:
-                    trailer = await self.upstream_reader.readline()
-                    self.write_response(trailer)
-                    if trailer == b"\r\n":
-                        break
-                break
-
-            # Read chunk data + CRLF
-            chunk_data_with_crlf = await self.upstream_reader.read(chunk_size + 2)
-            if len(chunk_data_with_crlf) >= 2:
-                # Separate actual data from CRLF
-                chunk_data = chunk_data_with_crlf[:-2]  # Remove CRLF
-                crlf = chunk_data_with_crlf[-2:]  # Extract CRLF
-
-                # Process chunk data through hook
-                processed_chunk = await self.on_response_chunk(chunk_data)
-                if processed_chunk is not None:
-                    # Write the processed data + CRLF
-                    self.write_response(processed_chunk + crlf)
-                else:
-                    # If chunk was filtered out, we still need to write CRLF for protocol compliance
-                    self.write_response(crlf)
-            else:
-                # Handle edge case where we didn't get enough data
-                self.write_response(chunk_data_with_crlf)
-
-    async def read_response_body(self) -> AsyncIterator[bytes]:
-        """
-        Read the response body from the upstream server.
-
-        This method can be called from on_response_received() to read and potentially
-        modify the response body before it's forwarded to the client.
-
-        Yields:
-            Chunks of the response body as bytes.
-        """
-        if not self.upstream_reader or not self.response:
-            return
-
-        content_length = None
-        if self.response.headers:
-            content_length_header = self.response.headers.first("Content-Length")
-            if content_length_header:
-                content_length = int(content_length_header)
-
-        if content_length is not None:
-            remaining = content_length
-            while remaining > 0:
-                chunk_size = min(remaining, MAX_CHUNK_SIZE)
-                chunk = await self.upstream_reader.read(chunk_size)
-                if not chunk:
-                    break
-                yield chunk
-                remaining -= len(chunk)
-        else:
-            # Read until EOF (for non-chunked responses without Content-Length)
-            while True:
-                chunk = await self.upstream_reader.read(MAX_CHUNK_SIZE)
-                if not chunk:
-                    break
-                yield chunk

tests/test_https_forward_proxy_handler.py

@@ -1,4 +1,5 @@
 import ssl
+import unittest.mock
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -181,40 +182,25 @@ async def test_forward_chunked_response(self):
         handler.upstream_reader.readline.side_effect = [
             b"5\r\n",  # chunk size
             b"0\r\n",  # end chunk
-            b"\r\n",  # final CRLF
+            b"\r\n",  # final CRLF after zero chunk (no trailers)
         ]
         handler.upstream_reader.read.side_effect = [
-            b"Hello\r\n"  # chunk data + CRLF
+            b"Hello",  # chunk data (5 bytes)
+            b"\r\n",  # trailing CRLF after chunk data
         ]
 
         with patch.object(handler, "write_response") as mock_write:
             with patch.object(handler, "flush_response", new_callable=AsyncMock):
                 await handler._forward_response_body()
 
-                # Verify chunked data was forwarded
-                mock_write.assert_any_call(b"5\r\n")
-                mock_write.assert_any_call(b"Hello\r\n")
-                mock_write.assert_any_call(b"0\r\n")
-
-    @pytest.mark.asyncio
-    async def test_read_response_body(self):
-        """Test reading response body for modification."""
-        handler = HTTPSForwardProxyHandler()
-
-        # Mock response with Content-Length
-        handler.response = HTTPResponse()
-        handler.response.headers = HTTPHeader(b"Content-Length: 11\r\n\r\n")
-
-        # Mock upstream reader
-        handler.upstream_reader = AsyncMock()
-        handler.upstream_reader.read.side_effect = [b"Hello", b" World", b""]
-
-        # Collect chunks
-        chunks = []
-        async for chunk in handler.read_response_body():
-            chunks.append(chunk)
-
-        assert chunks == [b"Hello", b" World"]
+                # Verify proper chunked format was sent: size\r\ndata\r\n0\r\n\r\n
+                expected_calls = [
+                    unittest.mock.call(b"5\r\n"),  # chunk size
+                    unittest.mock.call(b"Hello"),  # chunk data
+                    unittest.mock.call(b"\r\n"),  # chunk trailing CRLF
+                    unittest.mock.call(b"0\r\n\r\n"),  # final zero chunk
+                ]
+                mock_write.assert_has_calls(expected_calls)
 
     @pytest.mark.asyncio
     async def test_error_handling_connection_failed(self):
@@ -377,21 +363,20 @@ async def test_on_response_complete_called_after_forwarding(self):
     @pytest.mark.asyncio
     async def test_chunked_response_only_calls_hook_for_data_chunks(self):
         """Test that chunked responses only call hook for actual data, not protocol overhead."""
+        from asyncio_https_proxy.chunked_encoding import forward_chunked_response
+
         handler = HTTPSForwardProxyHandler()
         handler.upstream_reader = AsyncMock()
 
-        # Mock response with chunked encoding
-        handler.response = HTTPResponse()
-        handler.response.headers = HTTPHeader(b"Transfer-Encoding: chunked\r\n\r\n")
-
-        # Mock chunked data: "5\r\nHello\r\n0\r\n\r\n"
+        # Mock chunked data properly
         handler.upstream_reader.readline.side_effect = [
             b"5\r\n",  # chunk size
             b"0\r\n",  # end chunk
-            b"\r\n",  # final CRLF
+            b"\r\n",  # final CRLF after zero chunk (no trailers)
         ]
         handler.upstream_reader.read.side_effect = [
-            b"Hello\r\n"  # chunk data + CRLF
+            b"Hello",  # chunk data (5 bytes)
+            b"\r\n",  # trailing CRLF after chunk data
         ]
 
         # Track chunk calls
@@ -405,7 +390,11 @@ async def mock_on_response_chunk(chunk):
         handler.on_response_chunk = mock_on_response_chunk
 
         with patch.object(handler, "write_response"):
-            await handler._forward_chunked_response()
+            await forward_chunked_response(
+                handler.upstream_reader,
+                handler.write_response,
+                handler.on_response_chunk,
+            )
 
             # Verify only actual data was passed to hook, not protocol overhead
             assert chunks_received == [b"Hello"]