Optimizations

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'),