DexDemo/frontend/venv/lib/python3.12/site-packages/websocket/tests/test_large_payloads.py

# -*- coding: utf-8 -*-
import unittest
import struct
from unittest.mock import Mock, patch, MagicMock

from websocket._abnf import ABNF
from websocket._core import WebSocket
from websocket._exceptions import WebSocketProtocolException, WebSocketPayloadException
from websocket._ssl_compat import SSLError

"""
test_large_payloads.py
websocket - WebSocket client library for Python

Copyright 2025 engn33r

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

class LargePayloadTest(unittest.TestCase):
    def test_frame_length_encoding_boundaries(self):
        """Test WebSocket frame length encoding at various boundaries"""

        # Test length encoding boundaries as per RFC 6455
        test_cases = [
            (125, "Single byte length"),  # Max for 7-bit length
            (126, "Two byte length start"),  # Start of 16-bit length
            (127, "Two byte length"),
            (65535, "Two byte length max"),  # Max for 16-bit length
            (65536, "Eight byte length start"),  # Start of 64-bit length
            (16384, "16KB boundary"),  # The problematic size
            (16385, "Just over 16KB"),
            (32768, "32KB"),
            (131072, "128KB"),
        ]

        for length, description in test_cases:
            with self.subTest(length=length, description=description):
                # Create payload of specified length
                payload = b"A" * length

                # Create frame
                frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)

                # Verify frame can be formatted without error
                formatted = frame.format()

                # Verify the frame header is correctly structured
                self.assertIsInstance(formatted, bytes)
                self.assertTrue(len(formatted) >= length)  # Header + payload

                # Verify payload length is preserved
                self.assertEqual(len(frame.data), length)

    def test_recv_large_payload_chunked(self):
        """Test receiving large payloads in chunks (simulating the 16KB recv issue)"""

        # Create a large payload that would trigger chunked reading
        large_payload = b"B" * 32768  # 32KB

        # Mock recv function that returns data in 16KB chunks
        chunks = []
        chunk_size = 16384
        for i in range(0, len(large_payload), chunk_size):
            chunks.append(large_payload[i : i + chunk_size])

        call_count = 0

        def mock_recv(bufsize):
            nonlocal call_count
            if call_count >= len(chunks):
                return b""
            result = chunks[call_count]
            call_count += 1
            return result

        # Test the frame buffer's recv_strict method
        from websocket._abnf import frame_buffer

        fb = frame_buffer(mock_recv, skip_utf8_validation=True)

        # This should handle large payloads by chunking
        result = fb.recv_strict(len(large_payload))

        self.assertEqual(result, large_payload)
        # Verify multiple recv calls were made
        self.assertGreater(call_count, 1)

    def test_ssl_large_payload_simulation(self):
        """Simulate SSL BAD_LENGTH error scenario"""

        # This test demonstrates that the 16KB limit in frame buffer protects against SSL issues
        payload_size = 16385

        recv_calls = []

        def mock_recv_with_ssl_limit(bufsize):
            recv_calls.append(bufsize)
            # This simulates the SSL issue: BAD_LENGTH when trying to recv > 16KB
            if bufsize > 16384:
                raise SSLError("[SSL: BAD_LENGTH] unknown error")
            return b"C" * min(bufsize, 16384)

        from websocket._abnf import frame_buffer

        fb = frame_buffer(mock_recv_with_ssl_limit, skip_utf8_validation=True)

        # The frame buffer handles this correctly by chunking recv calls
        result = fb.recv_strict(payload_size)

        # Verify it worked and chunked the calls properly
        self.assertEqual(len(result), payload_size)
        # Verify no single recv call was > 16KB
        self.assertTrue(all(call <= 16384 for call in recv_calls))
        # Verify multiple calls were made
        self.assertGreater(len(recv_calls), 1)

    def test_frame_format_large_payloads(self):
        """Test frame formatting with various large payload sizes"""

        # Test sizes around potential problem areas
        test_sizes = [16383, 16384, 16385, 32768, 65535, 65536]

        for size in test_sizes:
            with self.subTest(size=size):
                payload = b"D" * size
                frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)

                # Should not raise any exceptions
                formatted = frame.format()

                # Verify structure
                self.assertIsInstance(formatted, bytes)
                self.assertEqual(len(frame.data), size)

                # Verify length encoding is correct based on size
                # Note: frames from create_frame() include masking by default (4 extra bytes)
                mask_size = 4  # WebSocket frames are masked by default
                if size < ABNF.LENGTH_7:  # < 126
                    # Length should be encoded in single byte
                    expected_header_size = (
                        2 + mask_size
                    )  # 1 byte opcode + 1 byte length + 4 byte mask
                elif size < ABNF.LENGTH_16:  # < 65536
                    # Length should be encoded in 2 bytes
                    expected_header_size = (
                        4 + mask_size
                    )  # 1 byte opcode + 1 byte marker + 2 bytes length + 4 byte mask
                else:
                    # Length should be encoded in 8 bytes
                    expected_header_size = (
                        10 + mask_size
                    )  # 1 byte opcode + 1 byte marker + 8 bytes length + 4 byte mask

                self.assertEqual(len(formatted), expected_header_size + size)

    def test_send_large_payload_chunking(self):
        """Test that large payloads are sent in chunks to avoid SSL issues"""

        mock_sock = Mock()

        # Track how data is sent
        sent_chunks = []

        def mock_send(data):
            sent_chunks.append(len(data))
            return len(data)

        mock_sock.send = mock_send
        mock_sock.gettimeout.return_value = 30.0

        # Create WebSocket with mocked socket
        ws = WebSocket()
        ws.sock = mock_sock
        ws.connected = True

        # Create large payload
        large_payload = b"E" * 32768  # 32KB

        # Send the payload
        with patch("websocket._core.send") as mock_send_func:
            mock_send_func.side_effect = lambda sock, data: len(data)

            # This should work without SSL errors
            result = ws.send_binary(large_payload)

            # Verify payload was accepted
            self.assertGreater(result, 0)

    def test_utf8_validation_large_text(self):
        """Test UTF-8 validation with large text payloads"""

        # Create large valid UTF-8 text
        large_text = "Hello 世界! " * 2000  # About 26KB with Unicode

        # Test frame creation
        frame = ABNF.create_frame(large_text, ABNF.OPCODE_TEXT)

        # Should not raise validation errors
        formatted = frame.format()
        self.assertIsInstance(formatted, bytes)

        # Test with close frame that has invalid UTF-8 (this is what validate() actually checks)
        invalid_utf8_close_data = struct.pack("!H", 1000) + b"\xff\xfe invalid utf8"

        # Create close frame with invalid UTF-8 data
        frame = ABNF(1, 0, 0, 0, ABNF.OPCODE_CLOSE, 1, invalid_utf8_close_data)

        # Validation should catch the invalid UTF-8 in close frame reason
        with self.assertRaises(WebSocketProtocolException):
            frame.validate(skip_utf8_validation=False)

    def test_frame_buffer_edge_cases(self):
        """Test frame buffer with edge cases that could trigger bugs"""

        # Test scenario: exactly 16KB payload split across recv calls
        payload_16k = b"F" * 16384

        # Simulate receiving in smaller chunks
        chunks = [payload_16k[i : i + 4096] for i in range(0, len(payload_16k), 4096)]

        call_count = 0

        def mock_recv(bufsize):
            nonlocal call_count
            if call_count >= len(chunks):
                return b""
            result = chunks[call_count]
            call_count += 1
            return result

        from websocket._abnf import frame_buffer

        fb = frame_buffer(mock_recv, skip_utf8_validation=True)
        result = fb.recv_strict(16384)

        self.assertEqual(result, payload_16k)
        # Verify multiple recv calls were made
        self.assertEqual(call_count, 4)  # 16KB / 4KB = 4 chunks

    def test_max_frame_size_limits(self):
        """Test behavior at WebSocket maximum frame size limits"""

        # Test just under the maximum theoretical frame size
        # (This is a very large test, so we'll use a smaller representative size)

        # Test with a reasonably large payload that represents the issue
        large_size = 1024 * 1024  # 1MB
        payload = b"G" * large_size

        # This should work without issues
        frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)

        # Verify the frame can be formatted
        formatted = frame.format()
        self.assertIsInstance(formatted, bytes)

        # Verify payload is preserved
        self.assertEqual(len(frame.data), large_size)


if __name__ == "__main__":
    unittest.main()