[TlsInterception] Fix broken ChunkedResponsePlugin for Python < 3.10 (#986)

* Add TLS interception integration tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix CI

* unused

* Well 3.9 just worked locally

* Dispatching empty byte to client results in OSError for Python < 3.10

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Uncomment old integration tests

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: abhinavsingh

Makefile

@@ -12,9 +12,10 @@ HTTPS_CERT_FILE_PATH := https-cert.pem
 HTTPS_CSR_FILE_PATH := https-csr.pem
 HTTPS_SIGNED_CERT_FILE_PATH := https-signed-cert.pem
 
-CA_KEY_FILE_PATH := ca-key.pem
-CA_CERT_FILE_PATH := ca-cert.pem
-CA_SIGNING_KEY_FILE_PATH := ca-signing-key.pem
+CA_CERT_SUFFIX :=
+CA_KEY_FILE_PATH := ca-key$(CA_CERT_SUFFIX).pem
+CA_CERT_FILE_PATH := ca-cert$(CA_CERT_SUFFIX).pem
+CA_SIGNING_KEY_FILE_PATH := ca-signing-key$(CA_CERT_SUFFIX).pem
 
 # Dummy invalid hardcoded value
 PROXYPY_PKG_PATH := dist/proxy.py.whl

proxy/common/flag.py

@@ -242,12 +242,14 @@ def initialize(
             ),
         )
         args.disable_headers = disabled_headers if disabled_headers is not None else DEFAULT_DISABLE_HEADERS
+
         args.certfile = cast(
             Optional[str], opts.get(
                 'cert_file', args.cert_file,
             ),
         )
         args.keyfile = cast(Optional[str], opts.get('key_file', args.key_file))
+
         args.ca_key_file = cast(
             Optional[str], opts.get(
                 'ca_key_file', args.ca_key_file,
@@ -272,6 +274,7 @@ def initialize(
                 args.ca_file,
             ),
         )
+
         args.hostname = cast(
             IpAddress,
             opts.get('hostname', ipaddress.ip_address(args.hostname)),

proxy/core/connection/client.py

@@ -46,7 +46,6 @@ def wrap(self, keyfile: str, certfile: str) -> None:
         self._conn = ssl.wrap_socket(
             self.connection,
             server_side=True,
-            # ca_certs=self.flags.ca_cert_file,
             certfile=certfile,
             keyfile=keyfile,
             ssl_version=ssl.PROTOCOL_TLS,

proxy/http/proxy/plugin.py

@@ -131,11 +131,14 @@ def handle_client_request(
     # No longer abstract since 2.4.0
     #
     # @abstractmethod
-    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+    def handle_upstream_chunk(self, chunk: memoryview) -> Optional[memoryview]:
         """Handler called right after receiving raw response from upstream server.
 
         For HTTPS connections, chunk will be encrypted unless
-        TLS interception is also enabled."""
+        TLS interception is also enabled.
+
+        Return None if you don't want to sent this chunk to the client.
+        """
         return chunk  # pragma: no cover
 
     # No longer abstract since 2.4.0

proxy/http/proxy/server.py

@@ -280,24 +280,29 @@ async def read_from_descriptors(self, r: Readables) -> bool:
 
             for plugin in self.plugins.values():
                 raw = plugin.handle_upstream_chunk(raw)
+                if raw is None:
+                    break
 
             # parse incoming response packet
             # only for non-https requests and when
             # tls interception is enabled
-            if not self.request.is_https_tunnel \
-                    or self.tls_interception_enabled:
-                if self.response.is_complete:
-                    self.handle_pipeline_response(raw)
+            if raw is not None:
+                if (
+                    not self.request.is_https_tunnel
+                    or self.tls_interception_enabled
+                ):
+                    if self.response.is_complete:
+                        self.handle_pipeline_response(raw)
+                    else:
+                        # TODO(abhinavsingh): Remove .tobytes after parser is
+                        # memoryview compliant
+                        chunk = raw.tobytes()
+                        self.response.parse(chunk)
+                        self.emit_response_events(len(chunk))
                 else:
-                    # TODO(abhinavsingh): Remove .tobytes after parser is
-                    # memoryview compliant
-                    chunk = raw.tobytes()
-                    self.response.parse(chunk)
-                    self.emit_response_events(len(chunk))
-            else:
-                self.response.total_size += len(raw)
-            # queue raw data for client
-            self.client.queue(raw)
+                    self.response.total_size += len(raw)
+                # queue raw data for client
+                self.client.queue(raw)
         return False
 
     def on_client_connection_close(self) -> None:

proxy/plugin/cache/base.py

@@ -54,7 +54,7 @@ def handle_client_request(
         assert self.store
         return self.store.cache_request(request)
 
-    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+    def handle_upstream_chunk(self, chunk: memoryview) -> Optional[memoryview]:
         assert self.store
         return self.store.cache_response_chunk(chunk)
 

proxy/plugin/man_in_the_middle.py

@@ -8,12 +8,14 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
+from typing import Optional
+
 from ..http.responses import okResponse
 from ..http.proxy import HttpProxyBasePlugin
 
 
 class ManInTheMiddlePlugin(HttpProxyBasePlugin):
     """Modifies upstream server responses."""
 
-    def handle_upstream_chunk(self, _chunk: memoryview) -> memoryview:
+    def handle_upstream_chunk(self, _chunk: memoryview) -> Optional[memoryview]:
         return okResponse(content=b'Hello from man in the middle')

proxy/plugin/modify_chunk_response.py

@@ -8,7 +8,7 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-from typing import Any
+from typing import Any, Optional
 
 from ..http.parser import HttpParser, httpParserTypes
 from ..http.proxy import HttpProxyBasePlugin
@@ -29,7 +29,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Create a new http protocol parser for response payloads
         self.response = HttpParser(httpParserTypes.RESPONSE_PARSER)
 
-    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+    def handle_upstream_chunk(self, chunk: memoryview) -> Optional[memoryview]:
         # Parse the response.
         # Note that these chunks also include headers
         self.response.parse(chunk.tobytes())
@@ -40,4 +40,4 @@ def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
             if self.response.body_expected:
                 self.response.body = b'\n'.join(self.DEFAULT_CHUNKS) + b'\n'
             self.client.queue(memoryview(self.response.build_response()))
-        return memoryview(b'')
+        return None

proxy/plugin/proxy_pool.py

@@ -182,7 +182,7 @@ def handle_client_data(self, raw: memoryview) -> Optional[memoryview]:
         self.upstream.queue(raw)
         return raw
 
-    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+    def handle_upstream_chunk(self, chunk: memoryview) -> Optional[memoryview]:
         """Will never be called since we didn't establish an upstream connection."""
         if not self.upstream:
             return chunk

tests/integration/test_integration.py

@@ -10,10 +10,10 @@
 
     Test the simplest proxy use scenario for smoke.
 """
-import sys
 import time
 import pytest
 import tempfile
+import subprocess
 
 from pathlib import Path
 from typing import Any, Generator
@@ -22,11 +22,13 @@
 from proxy.common.constants import IS_WINDOWS
 
 
-TLS_INTERCEPTION_FLAGS = ' '.join((
-    '--ca-cert-file', 'ca-cert.pem',
-    '--ca-key-file', 'ca-key.pem',
-    '--ca-signing-key', 'ca-signing-key.pem',
-))
+def _tls_interception_flags(ca_cert_suffix: str = '') -> str:
+    return ' '.join((
+        '--ca-cert-file', 'ca-cert%s.pem' % ca_cert_suffix,
+        '--ca-key-file', 'ca-key%s.pem' % ca_cert_suffix,
+        '--ca-signing-key', 'ca-signing-key%s.pem' % ca_cert_suffix,
+    ))
+
 
 PROXY_PY_FLAGS_INTEGRATION = (
     ('--threadless'),
@@ -35,45 +37,55 @@
 )
 
 PROXY_PY_FLAGS_TLS_INTERCEPTION = (
-    ('--threadless ' + TLS_INTERCEPTION_FLAGS),
-    ('--threadless --local-executor 0 ' + TLS_INTERCEPTION_FLAGS),
-    ('--threaded ' + TLS_INTERCEPTION_FLAGS),
+    ('--threadless ' + _tls_interception_flags()),
+    ('--threadless --local-executor 0 ' + _tls_interception_flags()),
+    ('--threaded ' + _tls_interception_flags()),
 )
 
 PROXY_PY_FLAGS_MODIFY_CHUNK_RESPONSE_PLUGIN = (
     (
         '--threadless --plugin proxy.plugin.ModifyChunkResponsePlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-chunk')
     ),
     (
         '--threadless --local-executor 0 --plugin proxy.plugin.ModifyChunkResponsePlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-chunk')
     ),
     (
         '--threaded --plugin proxy.plugin.ModifyChunkResponsePlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-chunk')
     ),
 )
 
 PROXY_PY_FLAGS_MODIFY_POST_DATA_PLUGIN = (
     (
         '--threadless --plugin proxy.plugin.ModifyPostDataPlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-post')
     ),
     (
         '--threadless --local-executor 0 --plugin proxy.plugin.ModifyPostDataPlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-post')
     ),
     (
         '--threaded --plugin proxy.plugin.ModifyPostDataPlugin ' +
-        TLS_INTERCEPTION_FLAGS
+        _tls_interception_flags('-post')
     ),
 )
 
 
 @pytest.fixture(scope='session', autouse=True)  # type: ignore[misc]
-def _gen_ca_certificates() -> None:
-    check_output(['make', 'ca-certificates'])
+def _gen_ca_certificates(request: Any) -> None:
+    check_output([
+        'make', 'ca-certificates',
+    ])
+    check_output([
+        'make', 'ca-certificates',
+        '-e', 'CA_CERT_SUFFIX=-chunk',
+    ])
+    check_output([
+        'make', 'ca-certificates',
+        '-e', 'CA_CERT_SUFFIX=-post',
+    ])
 
 
 # FIXME: Ignore is necessary for as long as pytest hasn't figured out
@@ -104,7 +116,7 @@ def proxy_py_subprocess(request: Any) -> Generator[int, None, None]:
         '--ca-cert-dir', str(ca_cert_dir),
         '--log-level', 'd',
     ) + tuple(request.param.split())
-    proxy_proc = Popen(proxy_cmd)
+    proxy_proc = Popen(proxy_cmd, stderr=subprocess.STDOUT)
     # Needed because port file might not be available immediately
     while not port_file.exists():
         time.sleep(1)
@@ -163,10 +175,6 @@ def test_integration_with_interception_flags(proxy_py_subprocess: int) -> None:
     IS_WINDOWS,
     reason='OSError: [WinError 193] %1 is not a valid Win32 application',
 )  # type: ignore[misc]
-@pytest.mark.skipif(
-    sys.version_info < (3, 10),
-    reason='For version < 3.10, GHA integration run into OSError when flushing to clients',
-)  # type: ignore[misc]
 def test_modify_chunk_response_integration(proxy_py_subprocess: int) -> None:
     """An acceptance test for :py:class:`~proxy.plugin.ModifyChunkResponsePlugin`
     interception using ``curl`` through proxy.py."""