From d960cd7fbca3f10b7533547a26fdec55ff3b2c4e Mon Sep 17 00:00:00 2001 From: xxx <18797809093@163.com> Date: Sun, 30 Nov 2025 14:06:32 +0800 Subject: [PATCH] test --- debug_import.py | 10 ++ msgpack/__init__.py | 1 + msgpack/chunked.py | 248 +++++++++++++++++++++++++++++++++++++++++ msgpack/exceptions.py | 8 ++ msgpack/test_import.py | 10 ++ test/test_chunked.py | 126 +++++++++++++++++++++ test_chunked_import.py | 14 +++ test_chunked_simple.py | 34 ++++++ test_python.py | 9 ++ 9 files changed, 460 insertions(+) create mode 100644 debug_import.py create mode 100644 msgpack/chunked.py create mode 100644 msgpack/test_import.py create mode 100644 test/test_chunked.py create mode 100644 test_chunked_import.py create mode 100644 test_chunked_simple.py create mode 100644 test_python.py diff --git a/debug_import.py b/debug_import.py new file mode 100644 index 00000000..b5aeec4e --- /dev/null +++ b/debug_import.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python + +import traceback + +try: + import msgpack.chunked + print('Import successful!') +except Exception as e: + print(f'Import failed: {e}') + traceback.print_exc() \ No newline at end of file diff --git a/msgpack/__init__.py b/msgpack/__init__.py index f3266b70..0180628f 100644 --- a/msgpack/__init__.py +++ b/msgpack/__init__.py @@ -3,6 +3,7 @@ from .exceptions import * # noqa: F403 from .ext import ExtType, Timestamp +from . import chunked version = (1, 1, 2) __version__ = "1.1.2" diff --git a/msgpack/chunked.py b/msgpack/chunked.py new file mode 100644 index 00000000..a2fb93e4 --- /dev/null +++ b/msgpack/chunked.py @@ -0,0 +1,248 @@ +"""Chunked MessagePack streaming with verification and resume support""" + +import struct +import hashlib +from .fallback import Unpacker, Packer +from .exceptions import ChunkVerifyError, BreakpointCorruptError, FormatError + + +class ChunkedUnpacker(Unpacker): + """Streaming unpacker with chunk verification and resume support""" + + def __init__( + self, + file_like=None, + *, + read_size=0, + use_list=True, + raw=False, + timestamp=0, + strict_map_key=True, + object_hook=None, + object_pairs_hook=None, + list_hook=None, + unicode_errors=None, + max_buffer_size=100 * 1024 * 1024, + ext_hook=None, + max_str_len=-1, + max_bin_len=-1, + max_array_len=-1, + max_map_len=-1, + max_ext_len=-1, + chunk_size=4 * 1024 * 1024, + verify_algorithm='crc32', + breakpoint_callback=None, + ): + super().__init__( + file_like, + read_size=read_size, + use_list=use_list, + raw=raw, + timestamp=timestamp, + strict_map_key=strict_map_key, + object_hook=object_hook, + object_pairs_hook=object_pairs_hook, + list_hook=list_hook, + unicode_errors=unicode_errors, + max_buffer_size=max_buffer_size, + ext_hook=ext_hook, + max_str_len=max_str_len, + max_bin_len=max_bin_len, + max_array_len=max_array_len, + max_map_len=max_map_len, + max_ext_len=max_ext_len, + ) + + # Validate chunk size (1MB to 64MB) + if not (1 * 1024 * 1024 <= chunk_size <= 64 * 1024 * 1024): + raise ValueError("chunk_size must be between 1MB and 64MB") + self.chunk_size = chunk_size + + # Validate verify algorithm + if verify_algorithm not in ['crc32', 'md5', 'sha1']: + raise ValueError("verify_algorithm must be one of 'crc32', 'md5', 'sha1'") + self.verify_algorithm = verify_algorithm + + # Calculate header length based on verify algorithm + if self.verify_algorithm == 'crc32': + self.header_length = 8 + elif self.verify_algorithm == 'md5': + self.header_length = 20 # 4 (chunk size) + 16 (md5) + elif self.verify_algorithm == 'sha1': + self.header_length = 24 # 4 (chunk size) + 20 (sha1) + + self.breakpoint_callback = breakpoint_callback + + # Chunk tracking + self.current_chunk = bytearray() + self.current_chunk_number = 0 + self.breakpoint = None + + def _update_breakpoint(self): + """Update breakpoint information""" + breakpoint_info = { + 'chunk_number': self.current_chunk_number, + 'stream_offset': self._stream_offset, + 'buffer_offset': self._buff_i, + 'checkpoint': self._buf_checkpoint, + } + if self.breakpoint_callback: + self.breakpoint_callback(breakpoint_info) + self.breakpoint = breakpoint_info + + def _verify_chunk(self, chunk, checksum): + """Verify chunk integrity""" + if self.verify_algorithm == 'crc32': + import zlib + calculated = zlib.crc32(chunk) & 0xffffffff + if calculated != checksum: + raise ChunkVerifyError(f"CRC32 check failed: expected 0x{checksum:08x}, got 0x{calculated:08x}") + elif self.verify_algorithm == 'md5': + calculated = hashlib.md5(chunk).digest() + if calculated != checksum: + raise ChunkVerifyError("MD5 check failed") + elif self.verify_algorithm == 'sha1': + calculated = hashlib.sha1(chunk).digest() + if calculated != checksum: + raise ChunkVerifyError("SHA1 check failed") + + def feed(self, next_bytes): + """Feed data to the unpacker, handling chunked verification""" + if not self._feeding: + raise RuntimeError("feed() cannot be used with file_like") + + # If we have a breakpoint, skip ahead + if self.breakpoint: + # For simplicity, we'll assume the breakpoint is at chunk boundary + # In a real implementation, we'd handle partial chunks + self.current_chunk_number = self.breakpoint['chunk_number'] + self._stream_offset = self.breakpoint['stream_offset'] + self._buff_i = self.breakpoint['buffer_offset'] + self._buf_checkpoint = self.breakpoint['checkpoint'] + self.breakpoint = None + + # Combine current_chunk with new data + combined_data = self.current_chunk + next_bytes + view = memoryview(combined_data) + chunk_length = len(view) + offset = 0 + + while offset < chunk_length: + # Check if we have enough data for header + if offset + self.header_length > chunk_length: + # Not enough data for header + self.current_chunk = view[offset:] + break + + # Read header + header = view[offset:offset+self.header_length] + chunk_size = struct.unpack_from('!I', header)[0] + + # Validate chunk size + if chunk_size < 0 or chunk_size > self.chunk_size: + raise FormatError(f"Invalid chunk size: {chunk_size}") + + # Check if we have enough data for full chunk + if offset + self.header_length + chunk_size > chunk_length: + # Not enough data for full chunk + self.current_chunk = view[offset:] + break + + # Read chunk data + chunk_data = view[offset+self.header_length:offset+self.header_length+chunk_size] + + # Read checksum from header + if self.verify_algorithm == 'crc32': + checksum = struct.unpack_from('!I', header, 4)[0] + elif self.verify_algorithm == 'md5': + checksum = header[4:20] + elif self.verify_algorithm == 'sha1': + checksum = header[4:24] + + # Verify chunk + self._verify_chunk(chunk_data, checksum) + + # Feed chunk to parent unpacker + super().feed(chunk_data) + + # Update breakpoint + self._update_breakpoint() + + # Move to next chunk + offset += self.header_length + chunk_size + self.current_chunk_number += 1 + + # If offset reached chunk_length, clear current_chunk + if offset == chunk_length: + self.current_chunk = bytearray() + + +class ChunkedPacker(Packer): + """Packer that generates chunked data with verification""" + + def __init__( + self, + chunk_size=4 * 1024 * 1024, + verify_algorithm='crc32', + *args, + **kwargs + ): + super().__init__(*args, **kwargs) + + # Validate chunk size + if not (1 * 1024 * 1024 <= chunk_size <= 64 * 1024 * 1024): + raise ValueError("chunk_size must be between 1MB and 64MB") + self.chunk_size = chunk_size + + # Validate verify algorithm + if verify_algorithm not in ['crc32', 'md5', 'sha1']: + raise ValueError("verify_algorithm must be one of 'crc32', 'md5', 'sha1'") + self.verify_algorithm = verify_algorithm + + # Buffer for data + self.buffer = bytearray() + + def _finalize_chunk(self): + """Finalize current chunk and return it""" + if not self.buffer: + return b'' + + chunk_data = bytes(self.buffer) + + # Calculate checksum + if self.verify_algorithm == 'crc32': + import zlib + checksum = zlib.crc32(chunk_data) & 0xffffffff + # Header: chunk size (4 bytes), checksum (4 bytes) + header = struct.pack('!II', len(chunk_data), checksum) + elif self.verify_algorithm == 'md5': + checksum = hashlib.md5(chunk_data).digest() + # Header: chunk size (4 bytes), checksum (16 bytes) + header = struct.pack('!I', len(chunk_data)) + checksum + elif self.verify_algorithm == 'sha1': + checksum = hashlib.sha1(chunk_data).digest() + # Header: chunk size (4 bytes), checksum (20 bytes) + header = struct.pack('!I', len(chunk_data)) + checksum + + self.buffer.clear() + return header + chunk_data + + def pack(self, obj): + """Pack an object into chunked data""" + # Use super().pack() to get raw data + raw_data = super().pack(obj) + + # Split into chunks + for i in range(0, len(raw_data), self.chunk_size): + chunk = raw_data[i:i+self.chunk_size] + self.buffer.extend(chunk) + if len(self.buffer) >= self.chunk_size: + yield self._finalize_chunk() + + # Finalize any remaining data + if self.buffer: + yield self._finalize_chunk() + + def packb(self, obj): + """Pack an object into a single chunk (if fits) or multiple chunks""" + return b''.join(self.pack(obj)) \ No newline at end of file diff --git a/msgpack/exceptions.py b/msgpack/exceptions.py index d6d2615c..7bf6f628 100644 --- a/msgpack/exceptions.py +++ b/msgpack/exceptions.py @@ -42,6 +42,14 @@ def __str__(self): return "unpack(b) received extra data." +class ChunkVerifyError(UnpackException): + """Chunk verification failed""" + + +class BreakpointCorruptError(UnpackException): + """Breakpoint data is corrupt""" + + # Deprecated. Use Exception instead to catch all exception during packing. PackException = Exception PackValueError = ValueError diff --git a/msgpack/test_import.py b/msgpack/test_import.py new file mode 100644 index 00000000..00a30e25 --- /dev/null +++ b/msgpack/test_import.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python + +import traceback + +try: + import chunked + print('Import successful!') +except Exception as e: + print(f'Import failed: {e}') + traceback.print_exc() \ No newline at end of file diff --git a/test/test_chunked.py b/test/test_chunked.py new file mode 100644 index 00000000..d95dd97e --- /dev/null +++ b/test/test_chunked.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python + +import zlib +import pytest +import hashlib +from io import BytesIO +from msgpack.chunked import ChunkedUnpacker, ChunkedPacker +from msgpack.exceptions import ChunkVerifyError + + +def test_chunked_pack_unpack(): + """Test basic chunked packing and unpacking""" + test_data = { + 'key1': 'value1', + 'key2': [1, 2, 3, 4], + 'key3': {'a': 1, 'b': 2}, + 'key4': b'hello world', + 'key5': None, + 'key6': True, + 'key7': False, + } + + # Test with different chunk sizes + for chunk_size in [1024, 4096, 8192]: + for verify_algorithm in ['crc32', 'md5', 'sha1']: + # Pack with chunked packer + packer = ChunkedPacker(chunk_size=chunk_size, verify_algorithm=verify_algorithm) + chunks = list(packer.pack(test_data)) + + # Unpack with chunked unpacker + unpacker = ChunkedUnpacker(verify_algorithm=verify_algorithm) + for chunk in chunks: + unpacker.feed(chunk) + + result = list(unpacker) + assert len(result) == 1 + assert result[0] == test_data + + +def test_chunk_verification(): + """Test chunk verification fails when data is corrupted""" + test_data = {'key': 'value' * 1000} + + # Pack data + packer = ChunkedPacker(chunk_size=1024, verify_algorithm='crc32') + chunks = list(packer.pack(test_data)) + + # Corrupt the first chunk + corrupted_chunk = chunks[0][:4] + b'\x00\x00\x00\x00' + chunks[0][8:] + + # Try to unpack corrupted data + unpacker = ChunkedUnpacker(verify_algorithm='crc32') + unpacker.feed(corrupted_chunk) + + with pytest.raises(ChunkVerifyError): + list(unpacker) + + +def test_breakpoint_resume(): + """Test breakpoint resume functionality""" + test_data = [{'key': f'value{i}'} for i in range(100)] + + # Pack data + packer = ChunkedPacker(chunk_size=1024, verify_algorithm='crc32') + chunks = [] + for item in test_data: + chunks.extend(packer.pack(item)) + + # Simulate partial unpacking + breakpoint_info = None + def save_breakpoint(info): + nonlocal breakpoint_info + breakpoint_info = info + + unpacker = ChunkedUnpacker(breakpoint_callback=save_breakpoint, verify_algorithm='crc32') + + # Feed only first 3 chunks + for i in range(3): + unpacker.feed(chunks[i]) + + # Get partial result + partial_result = list(unpacker) + assert len(partial_result) > 0 + + # Save breakpoint + assert breakpoint_info is not None + + # Create new unpacker with breakpoint + new_unpacker = ChunkedUnpacker(verify_algorithm='crc32') + new_unpacker.breakpoint = breakpoint_info + + # Feed remaining chunks + for i in range(3, len(chunks)): + new_unpacker.feed(chunks[i]) + + # Get complete result + complete_result = partial_result + list(new_unpacker) + assert len(complete_result) == len(test_data) + assert all(a == b for a, b in zip(complete_result, test_data)) + + +def test_invalid_chunk_size(): + """Test invalid chunk size raises ValueError""" + with pytest.raises(ValueError, match="chunk_size must be between 1MB and 64MB"): + ChunkedPacker(chunk_size=512 * 1024) + + with pytest.raises(ValueError, match="chunk_size must be between 1MB and 64MB"): + ChunkedPacker(chunk_size=65 * 1024 * 1024) + + +def test_invalid_verify_algorithm(): + """Test invalid verify algorithm raises ValueError""" + with pytest.raises(ValueError, match="verify_algorithm must be one of"): + ChunkedPacker(verify_algorithm='invalid') + + with pytest.raises(ValueError, match="verify_algorithm must be one of"): + ChunkedUnpacker(verify_algorithm='invalid') + + +if __name__ == '__main__': + test_chunked_pack_unpack() + test_chunk_verification() + test_breakpoint_resume() + test_invalid_chunk_size() + test_invalid_verify_algorithm() + print("All tests passed!") \ No newline at end of file diff --git a/test_chunked_import.py b/test_chunked_import.py new file mode 100644 index 00000000..875d26a8 --- /dev/null +++ b/test_chunked_import.py @@ -0,0 +1,14 @@ +#!/usr/bin/env python + +import traceback +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +try: + import msgpack.chunked + print('Successfully imported msgpack.chunked!') +except Exception as e: + print(f'Error importing msgpack.chunked: {type(e).__name__}: {e}') + traceback.print_exc() \ No newline at end of file diff --git a/test_chunked_simple.py b/test_chunked_simple.py new file mode 100644 index 00000000..b28060b5 --- /dev/null +++ b/test_chunked_simple.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python + +import zlib +import hashlib +from io import BytesIO +from msgpack.chunked import ChunkedUnpacker, ChunkedPacker +from msgpack.exceptions import ChunkVerifyError + +print("Testing ChunkedPacker and ChunkedUnpacker...") + +# Test basic functionality +test_data = { + 'key1': 'value1', + 'key2': [1, 2, 3, 4], + 'key3': {'a': 1, 'b': 2}, +} + +# Create chunked packer and unpacker +packer = ChunkedPacker(chunk_size=1024) +unpacker = ChunkedUnpacker() + +# Pack data +chunks = list(packer.pack(test_data)) +print(f"Packed into {len(chunks)} chunks") + +# Unpack data +for chunk in chunks: + unpacker.feed(chunk) + +result = list(unpacker) +print(f"Unpacked result: {result}") +print(f"Result matches original: {result == [test_data]}") + +print("\nAll tests completed!") \ No newline at end of file diff --git a/test_python.py b/test_python.py new file mode 100644 index 00000000..430a003a --- /dev/null +++ b/test_python.py @@ -0,0 +1,9 @@ +#!/usr/bin/env python + +print('Python is working!') +import sys +print(f'Python version: {sys.version}') +import struct +print('struct module imported successfully') +import hashlib +print('hashlib module imported successfully')