Optimizations

tests/test_bridge.py

@@ -3,6 +3,7 @@ import asyncio
 import time
 import unittest
 from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
 
 from proxy._aes import Cipher, algorithms, modes
 from proxy.bridge import MsgSplitter, bridge_ws_reencrypt
@@ -40,9 +41,10 @@ def _intermediate(payload: bytes) -> bytes:
 class MsgSplitterTest(unittest.TestCase):
     def _split(self, proto_int, packets, chunk_sizes=None):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, proto_int)
+        splitter = MsgSplitter(proto_int)
         enc = _encryptor(relay_init)
-        stream = enc.update(b''.join(packets))
+        plain = b''.join(packets)
+        stream = enc.update(plain)
 
         chunks = []
         if chunk_sizes is None:
@@ -56,8 +58,10 @@ class MsgSplitterTest(unittest.TestCase):
                 chunks.append(stream[offset:])
 
         parts = []
+        offset = 0
         for chunk in chunks:
-            parts.extend(splitter.split(chunk))
+            parts.extend(splitter.split(chunk, plain[offset:offset + len(chunk)]))
+            offset += len(chunk)
         return splitter, stream, parts
 
     def test_abridged_stream_splits_into_packets(self):
@@ -94,27 +98,77 @@ class MsgSplitterTest(unittest.TestCase):
         self.assertEqual(len(parts), 8)
 
     def test_empty_chunk_yields_nothing(self):
-        splitter = MsgSplitter(_relay_init(), PROTO_INTERMEDIATE_INT)
-        self.assertEqual(splitter.split(b''), [])
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
+        self.assertEqual(splitter.split(b'', b''), [])
 
     def test_zero_length_packet_disables_splitting(self):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
         enc = _encryptor(relay_init)
-        stream = enc.update((0).to_bytes(4, 'little') + b'tail')
-        parts = splitter.split(stream)
+        plain = (0).to_bytes(4, 'little') + b'tail'
+        stream = enc.update(plain)
+        parts = splitter.split(stream, plain)
         self.assertEqual(parts, [stream])
-        self.assertEqual(splitter.split(b'raw'), [b'raw'])
+        self.assertEqual(splitter.split(b'raw', b'raw'), [b'raw'])
 
     def test_flush_returns_buffered_tail_once(self):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
         enc = _encryptor(relay_init)
-        partial = enc.update(_intermediate(b'x' * 32)[:10])
-        self.assertEqual(splitter.split(partial), [])
+        plain = _intermediate(b'x' * 32)[:10]
+        partial = enc.update(plain)
+        self.assertEqual(splitter.split(partial, plain), [])
         self.assertEqual(splitter.flush(), [partial])
         self.assertEqual(splitter.flush(), [])
 
+    def test_extended_and_quick_ack_headers_can_cross_chunk_boundaries(self):
+        for proto, frame in ((PROTO_ABRIDGED_INT, _abridged),
+                             (PROTO_INTERMEDIATE_INT, _intermediate),
+                             (PROTO_PADDED_INTERMEDIATE_INT, _intermediate)):
+            for size in (16, 512, 65536):
+                packet = bytearray(frame(b'p' * size))
+                packet[0 if proto == PROTO_ABRIDGED_INT else 3] |= 0x80
+                packets = [bytes(packet), frame(b't' * 8)]
+                for header_split in (1, 2, 3):
+                    with self.subTest(proto=proto, size=size, header_split=header_split):
+                        _, stream, parts = self._split(
+                            proto, packets, chunk_sizes=[header_split, 1, 11, size - 8])
+                        self.assertEqual(parts, [stream[:len(packet)], stream[len(packet):]])
+
+
+class WsReencryptTest(unittest.IsolatedAsyncioTestCase):
+    async def test_relay_preserves_encrypted_packets_and_incomplete_tail(self):
+        for proto, frame in ((PROTO_ABRIDGED_INT, _abridged),
+                             (PROTO_INTERMEDIATE_INT, _intermediate),
+                             (PROTO_PADDED_INTERMEDIATE_INT, _intermediate)):
+            with self.subTest(proto=proto):
+                packets = [frame(b'a' * 8), frame(b'b' * 16),
+                           frame(b'c' * 70000), frame(b'd' * 32)[:7]]
+                client_init, relay_init = _relay_init(), _relay_init()
+                reader = asyncio.StreamReader()
+                reader.feed_data(_encryptor(client_init).update(b''.join(packets)))
+                reader.feed_eof()
+                sent = []
+
+                async def send(data):
+                    sent.append(data)
+
+                async def send_batch(parts):
+                    sent.extend(parts)
+
+                async def recv():
+                    await asyncio.Future()
+
+                ws = SimpleNamespace(send=send, send_batch=send_batch, recv=recv,
+                                     close=AsyncMock())
+                writer = Mock(wait_closed=AsyncMock())
+                ctx = SimpleNamespace(clt_dec=_encryptor(client_init), tg_enc=_encryptor(relay_init))
+                await asyncio.wait_for(bridge_ws_reencrypt(
+                    reader, writer, ws, 'framing-test', ctx, splitter=MsgSplitter(proto)), 1)
+                decoder = _encryptor(relay_init)
+                self.assertEqual([decoder.update(part) for part in sent], packets)
+                ws.close.assert_awaited_once()
+
 
 class WsDiagnosticsTest(unittest.IsolatedAsyncioTestCase):
     async def test_flow_distinguishes_upstream_send_from_native_backpressure(self):

tests/test_h2_transport.py

@@ -27,6 +27,40 @@ from proxy.config import proxy_config
 from proxy.utils import PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE
 
 
+class H2ContextTest(unittest.IsolatedAsyncioTestCase):
+    async def test_lanes_reuse_verified_context_without_changing_websocket_contexts(self):
+        from proxy.h2_transport import _default_ssl_context
+        from proxy.raw_websocket import _ssl_ctx, _ssl_ctx_fronting
+        from proxy.utils import create_ssl_context
+
+        _default_ssl_context.cache_clear()
+        self.addCleanup(_default_ssl_context.cache_clear)
+        lanes = []
+        try:
+            with patch('proxy.h2_transport.create_ssl_context', wraps=create_ssl_context) as create:
+                lanes = [_HttpLane('kws4.example.org', 1), _HttpLane('kws2.example.net', 2)]
+                create.assert_called_once_with()
+            first, second = (lane.client._transport.ssl_context for lane in lanes)
+            self.assertIs(first, second)
+            self.assertIsNot(first, _ssl_ctx)
+            self.assertIsNot(first, _ssl_ctx_fronting)
+            self.assertTrue(first.check_hostname)
+            self.assertEqual(first.verify_mode, ssl.CERT_REQUIRED)
+            self.assertTrue(_ssl_ctx.check_hostname)
+            self.assertFalse(_ssl_ctx_fronting.check_hostname)
+            self.assertEqual(set(first.get_ca_certs(binary_form=True)),
+                             set(_ssl_ctx.get_ca_certs(binary_form=True)))
+
+            custom = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+            custom.check_hostname = False
+            transport = H2Transport(custom)
+            self.assertIs(transport.ssl_context, custom)
+            self.assertTrue(first.check_hostname)
+            await transport.aclose()
+        finally:
+            await asyncio.gather(*(lane.close() for lane in lanes))
+
+
 class H2WireTest(unittest.IsolatedAsyncioTestCase):
     @classmethod
     def setUpClass(cls):
@@ -163,6 +197,14 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
         await asyncio.gather(slow1, slow2)
         self.assertEqual(len(self.connections), 1)
 
+    async def test_shared_default_context_rejects_untrusted_certificate(self):
+        async with httpx.AsyncClient(transport=H2Transport(), trust_env=False) as client:
+            with self.assertRaises(httpx.ConnectError):
+                await client.get('https://' + self.lane.host + '/api')
+        task, request = await self.post()
+        self.respond(request)
+        self.assertEqual(await asyncio.wait_for(task, 1), b'r' * 40)
+
     async def test_debug_trace_follows_real_stream_without_logging_payload(self):
         channel = _HttpChannel(self.lane, 1, 'trace-test')
         body = b'private-payload-' + b'x' * 24

tests/test_raw_websocket.py

@@ -36,6 +36,15 @@ def _recv(chunks, cls=RawWebSocket):
 
 
 class XorMaskTest(unittest.TestCase):
+    def test_matches_bytewise_xor_at_fast_path_and_frame_boundaries(self):
+        for length in (0, 1, 2, 3, 4, 5, 125, 126, 127, 255, 511, 512, 513, 514, 515,
+                       16384, 65535, 65536, 65537, 1024 * 1024 + 3):
+            data = (bytes(range(256)) * (length // 256 + 1))[:length]
+            for mask in (b'\x00' * 4, b'\xff' * 4, b'\x01\x7f\x80\xff'):
+                with self.subTest(length=length, mask=mask):
+                    expected = bytes(value ^ mask[index % 4] for index, value in enumerate(data))
+                    self.assertEqual(_xor_mask(data, mask), expected)
+
     def test_roundtrip(self):
         data = bytes(range(256)) * 3
         mask = b'\x01\x02\x03\x04'
@@ -46,6 +55,21 @@ class XorMaskTest(unittest.TestCase):
 
 
 class BuildFrameTest(unittest.TestCase):
+    def test_masked_frame_wire_format_across_all_length_encodings(self):
+        mask = b'\x01\x7f\x80\xff'
+        for length in (1, 125, 126, 511, 512, 513, 65535, 65536, 65537):
+            payload = (bytes(range(256)) * (length // 256 + 1))[:length]
+            with self.subTest(length=length), patch('proxy.raw_websocket.os.urandom', return_value=mask):
+                frame = RawWebSocket._build_frame(RawWebSocket.OP_BINARY, payload, mask=True)
+                if length < 126:
+                    header = bytes([0x82, 0x80 | length])
+                elif length < 65536:
+                    header = b'\x82\xfe' + length.to_bytes(2, 'big')
+                else:
+                    header = b'\x82\xff' + length.to_bytes(8, 'big')
+                masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
+                self.assertEqual(frame, header + mask + masked)
+
     def test_short_unmasked_frame(self):
         self.assertEqual(
             RawWebSocket._build_frame(RawWebSocket.OP_BINARY, b'abc'),