codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2import unittest3import struct4from unittest.mock import Mock, patch, MagicMock5 6from websocket._abnf import ABNF7from websocket._core import WebSocket8from websocket._exceptions import WebSocketProtocolException, WebSocketPayloadException9from websocket._ssl_compat import SSLError10 11"""12test_large_payloads.py13websocket - WebSocket client library for Python14 15Copyright 2025 engn33r16 17Licensed under the Apache License, Version 2.0 (the "License");18you may not use this file except in compliance with the License.19You may obtain a copy of the License at20 21    http://www.apache.org/licenses/LICENSE-2.022 23Unless required by applicable law or agreed to in writing, software24distributed under the License is distributed on an "AS IS" BASIS,25WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.26See the License for the specific language governing permissions and27limitations under the License.28"""29 30class LargePayloadTest(unittest.TestCase):31    def test_frame_length_encoding_boundaries(self):32        """Test WebSocket frame length encoding at various boundaries"""33 34        # Test length encoding boundaries as per RFC 645535        test_cases = [36            (125, "Single byte length"),  # Max for 7-bit length37            (126, "Two byte length start"),  # Start of 16-bit length38            (127, "Two byte length"),39            (65535, "Two byte length max"),  # Max for 16-bit length40            (65536, "Eight byte length start"),  # Start of 64-bit length41            (16384, "16KB boundary"),  # The problematic size42            (16385, "Just over 16KB"),43            (32768, "32KB"),44            (131072, "128KB"),45        ]46 47        for length, description in test_cases:48            with self.subTest(length=length, description=description):49                # Create payload of specified length50                payload = b"A" * length51 52                # Create frame53                frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)54 55                # Verify frame can be formatted without error56                formatted = frame.format()57 58                # Verify the frame header is correctly structured59                self.assertIsInstance(formatted, bytes)60                self.assertTrue(len(formatted) >= length)  # Header + payload61 62                # Verify payload length is preserved63                self.assertEqual(len(frame.data), length)64 65    def test_recv_large_payload_chunked(self):66        """Test receiving large payloads in chunks (simulating the 16KB recv issue)"""67 68        # Create a large payload that would trigger chunked reading69        large_payload = b"B" * 32768  # 32KB70 71        # Mock recv function that returns data in 16KB chunks72        chunks = []73        chunk_size = 1638474        for i in range(0, len(large_payload), chunk_size):75            chunks.append(large_payload[i : i + chunk_size])76 77        call_count = 078 79        def mock_recv(bufsize):80            nonlocal call_count81            if call_count >= len(chunks):82                return b""83            result = chunks[call_count]84            call_count += 185            return result86 87        # Test the frame buffer's recv_strict method88        from websocket._abnf import frame_buffer89 90        fb = frame_buffer(mock_recv, skip_utf8_validation=True)91 92        # This should handle large payloads by chunking93        result = fb.recv_strict(len(large_payload))94 95        self.assertEqual(result, large_payload)96        # Verify multiple recv calls were made97        self.assertGreater(call_count, 1)98 99    def test_ssl_large_payload_simulation(self):100        """Simulate SSL BAD_LENGTH error scenario"""101 102        # This test demonstrates that the 16KB limit in frame buffer protects against SSL issues103        payload_size = 16385104 105        recv_calls = []106 107        def mock_recv_with_ssl_limit(bufsize):108            recv_calls.append(bufsize)109            # This simulates the SSL issue: BAD_LENGTH when trying to recv > 16KB110            if bufsize > 16384:111                raise SSLError("[SSL: BAD_LENGTH] unknown error")112            return b"C" * min(bufsize, 16384)113 114        from websocket._abnf import frame_buffer115 116        fb = frame_buffer(mock_recv_with_ssl_limit, skip_utf8_validation=True)117 118        # The frame buffer handles this correctly by chunking recv calls119        result = fb.recv_strict(payload_size)120 121        # Verify it worked and chunked the calls properly122        self.assertEqual(len(result), payload_size)123        # Verify no single recv call was > 16KB124        self.assertTrue(all(call <= 16384 for call in recv_calls))125        # Verify multiple calls were made126        self.assertGreater(len(recv_calls), 1)127 128    def test_frame_format_large_payloads(self):129        """Test frame formatting with various large payload sizes"""130 131        # Test sizes around potential problem areas132        test_sizes = [16383, 16384, 16385, 32768, 65535, 65536]133 134        for size in test_sizes:135            with self.subTest(size=size):136                payload = b"D" * size137                frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)138 139                # Should not raise any exceptions140                formatted = frame.format()141 142                # Verify structure143                self.assertIsInstance(formatted, bytes)144                self.assertEqual(len(frame.data), size)145 146                # Verify length encoding is correct based on size147                # Note: frames from create_frame() include masking by default (4 extra bytes)148                mask_size = 4  # WebSocket frames are masked by default149                if size < ABNF.LENGTH_7:  # < 126150                    # Length should be encoded in single byte151                    expected_header_size = (152                        2 + mask_size153                    )  # 1 byte opcode + 1 byte length + 4 byte mask154                elif size < ABNF.LENGTH_16:  # < 65536155                    # Length should be encoded in 2 bytes156                    expected_header_size = (157                        4 + mask_size158                    )  # 1 byte opcode + 1 byte marker + 2 bytes length + 4 byte mask159                else:160                    # Length should be encoded in 8 bytes161                    expected_header_size = (162                        10 + mask_size163                    )  # 1 byte opcode + 1 byte marker + 8 bytes length + 4 byte mask164 165                self.assertEqual(len(formatted), expected_header_size + size)166 167    def test_send_large_payload_chunking(self):168        """Test that large payloads are sent in chunks to avoid SSL issues"""169 170        mock_sock = Mock()171 172        # Track how data is sent173        sent_chunks = []174 175        def mock_send(data):176            sent_chunks.append(len(data))177            return len(data)178 179        mock_sock.send = mock_send180        mock_sock.gettimeout.return_value = 30.0181 182        # Create WebSocket with mocked socket183        ws = WebSocket()184        ws.sock = mock_sock185        ws.connected = True186 187        # Create large payload188        large_payload = b"E" * 32768  # 32KB189 190        # Send the payload191        with patch("websocket._core.send") as mock_send_func:192            mock_send_func.side_effect = lambda sock, data: len(data)193 194            # This should work without SSL errors195            result = ws.send_binary(large_payload)196 197            # Verify payload was accepted198            self.assertGreater(result, 0)199 200    def test_utf8_validation_large_text(self):201        """Test UTF-8 validation with large text payloads"""202 203        # Create large valid UTF-8 text204        large_text = "Hello 世界! " * 2000  # About 26KB with Unicode205 206        # Test frame creation207        frame = ABNF.create_frame(large_text, ABNF.OPCODE_TEXT)208 209        # Should not raise validation errors210        formatted = frame.format()211        self.assertIsInstance(formatted, bytes)212 213        # Test with close frame that has invalid UTF-8 (this is what validate() actually checks)214        invalid_utf8_close_data = struct.pack("!H", 1000) + b"\xff\xfe invalid utf8"215 216        # Create close frame with invalid UTF-8 data217        frame = ABNF(1, 0, 0, 0, ABNF.OPCODE_CLOSE, 1, invalid_utf8_close_data)218 219        # Validation should catch the invalid UTF-8 in close frame reason220        with self.assertRaises(WebSocketProtocolException):221            frame.validate(skip_utf8_validation=False)222 223    def test_frame_buffer_edge_cases(self):224        """Test frame buffer with edge cases that could trigger bugs"""225 226        # Test scenario: exactly 16KB payload split across recv calls227        payload_16k = b"F" * 16384228 229        # Simulate receiving in smaller chunks230        chunks = [payload_16k[i : i + 4096] for i in range(0, len(payload_16k), 4096)]231 232        call_count = 0233 234        def mock_recv(bufsize):235            nonlocal call_count236            if call_count >= len(chunks):237                return b""238            result = chunks[call_count]239            call_count += 1240            return result241 242        from websocket._abnf import frame_buffer243 244        fb = frame_buffer(mock_recv, skip_utf8_validation=True)245        result = fb.recv_strict(16384)246 247        self.assertEqual(result, payload_16k)248        # Verify multiple recv calls were made249        self.assertEqual(call_count, 4)  # 16KB / 4KB = 4 chunks250 251    def test_max_frame_size_limits(self):252        """Test behavior at WebSocket maximum frame size limits"""253 254        # Test just under the maximum theoretical frame size255        # (This is a very large test, so we'll use a smaller representative size)256 257        # Test with a reasonably large payload that represents the issue258        large_size = 1024 * 1024  # 1MB259        payload = b"G" * large_size260 261        # This should work without issues262        frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)263 264        # Verify the frame can be formatted265        formatted = frame.format()266        self.assertIsInstance(formatted, bytes)267 268        # Verify payload is preserved269        self.assertEqual(len(frame.data), large_size)270 271 272if __name__ == "__main__":273    unittest.main()274