Add more complete LZ4 testing

Author: whyitfor

ofrak_core/src/ofrak/core/lz4.py

@@ -1,17 +1,29 @@
+"""
+Lz4 Components.
+
+Lz4Unpacker currently supports unpacking modern LZ4 format (Lz4ModernData),
+legacy format (see Lz4LegacyData), and skippable data (Lz4SkippableData).
+
+Lz4Packer supports repacking the modern LZ4 format (Lz4ModernData), matching block/checksum
+information extracted during unpacking. Compression level can be specified via config.
+
+Lz4LegacyPacker supports repacking legacy LZ4 format (Lz4LegacyData) with compression level
+support (default/fast/high modes). Compression level can be specified via config.
+"""
 import logging
 from dataclasses import dataclass
 
-import lz4.frame  # type: ignore
 import lz4.block  # type: ignore
-
-from ofrak.component.identifier import Identifier
+import lz4.frame  # type: ignore
 from ofrak.component.packer import Packer
 from ofrak.component.unpacker import Unpacker
 from ofrak.core.binary import GenericBinary
-from ofrak.core.magic import MagicDescriptionPattern, MagicMimePattern
+from ofrak.core.magic import RawMagicPattern
+from ofrak.model.component_model import ComponentConfig
 from ofrak.resource import Resource
 from ofrak_type.range import Range
 
+
 LOGGER = logging.getLogger(__name__)
 
 # LZ4 frame magic numbers (little-endian)
@@ -33,27 +45,21 @@ class Lz4Data(GenericBinary):
 @dataclass
 class Lz4ModernData(Lz4Data):
     """
-    LZ4 modern frame format (default).
-
-    The modern LZ4 frame format includes:
-    - Frame descriptor with flags
-    - Optional content size and dictionary ID
-    - Block independence flags
-    - Optional checksums (content and block)
-    - End mark
+    LZ4 modern frame format (v1.4+).
     """
 
+    block_size: int
+    block_size_id: int
+    block_linked: bool
+    content_checksum: bool
+    block_checksum: bool
+    content_size: int
+
 
 @dataclass
 class Lz4LegacyData(Lz4Data):
     """
-    LZ4 legacy frame format.
-
-    Older LZ4 format predating the frame specification:
-    - Simpler structure
-    - No checksums or metadata
-    - Fixed 8MB max block size
-    - Deprecated but still encountered in the wild
+    LZ4 legacy frame format (v0.1-v0.9).
     """
 
 
@@ -62,118 +68,251 @@ class Lz4SkippableData(Lz4Data):
     """
     LZ4 skippable frame.
 
-    Special frame type for embedding metadata or application-specific data:
-    - Not compressed data
-    - Contains arbitrary bytes
-    - LZ4 parsers can safely skip these frames
-    - Typically used alongside regular frames
+    Special frame type for embedding metadata or application-specific data.
     """
 
 
-class Lz4Identifier(Identifier):
+class Lz4Unpacker(Unpacker[None]):
     """
-    Identify LZ4 compressed data by checking magic bytes.
+    Unpack (decompress) LZ4 modern frame format files.
 
-    Recognizes all LZ4 frame types:
-    - Modern/default frames (0x184D2204)
-    - Legacy frames (0x184C2102)
-    - Skippable frames (0x184D2A50-0x184D2A5F)
+    Supports:
+    - Modern frame format (Lz4ModernData)
+    - Skippable frames (metadata containers: Lz4SkippableData)
     """
 
-    id = b"Lz4Identifier"
-    targets = (GenericBinary,)
+    id = b"Lz4Unpacker"
+    targets = (Lz4ModernData, Lz4SkippableData)
+    children = (GenericBinary,)
 
-    async def identify(self, resource: Resource, config=None) -> None:
-        data = await resource.get_data(Range(0, 4))
+    async def unpack(self, resource: Resource, config=None):
+        """
+        Unpack LZ4 data.
+
+        :param resource: The LZ4 resource to unpack
 
-        if len(data) < 4:
-            return
+        :raises RuntimeError: if the data is not valid LZ4 format
+        """
+        resource_data = await resource.get_data()
 
-        # Check for modern frame
-        if data == LZ4_MODERN_MAGIC:
-            resource.add_tag(Lz4ModernData)
-            return
+        if resource.has_tag(Lz4ModernData):
+            # lz4.frame.get_frame_info() does not support legacy frames
+            frame_info = lz4.frame.get_frame_info(resource_data)
+            resource.add_view(
+                Lz4ModernData(
+                    block_size=frame_info["block_size"],
+                    block_size_id=frame_info["block_size_id"],
+                    block_linked=frame_info["block_linked"],
+                    content_checksum=frame_info["content_checksum"],
+                    block_checksum=frame_info["block_checksum"],
+                    content_size=frame_info["content_size"],
+                )
+            )
 
-        # Check for legacy frame
-        if data == LZ4_LEGACY_MAGIC:
-            resource.add_tag(Lz4LegacyData)
-            return
+        try:
+            decompressed_data = lz4.frame.decompress(resource_data)
+        except RuntimeError as e:
+            LOGGER.error(f"Failed to decompress LZ4 data: {e}")
+            raise
 
-        # Check for skippable frames
-        # Format: 0x5X 0x2A 0x4D 0x18 where X is 0-F
-        if data[1:4] == b"\x2a\x4d\x18" and 0x50 <= data[0] <= 0x5F:
-            resource.add_tag(Lz4SkippableData)
-            return
+        await resource.create_child(
+            tags=(GenericBinary,),
+            data=decompressed_data,
+        )
 
 
-class Lz4Unpacker(Unpacker[None]):
+class Lz4LegacyUnpacker(Unpacker[None]):
     """
-    Unpack (decompress) LZ4 files of all frame types.
+    Unpack (decompress) LZ4 legacy frame format files.
 
-    Supports:
-    - Modern frame format (most common)
-    - Legacy frame format (deprecated)
-    - Skippable frames (metadata containers)
+    Legacy format (v0.1-v0.9) uses lz4.block decompression instead of lz4.frame.
     """
 
-    id = b"Lz4Unpacker"
-    targets = (Lz4ModernData, Lz4LegacyData, Lz4SkippableData)
+    id = b"Lz4LegacyUnpacker"
+    targets = (Lz4LegacyData,)
     children = (GenericBinary,)
 
     async def unpack(self, resource: Resource, config=None):
         """
-        Unpack LZ4 data.
+        Unpack LZ4 legacy data.
 
-        :param resource: The LZ4 resource to unpack
-        :param config: Optional unpacker configuration
+        :param resource: The LZ4 legacy resource to unpack
 
-        :raises RuntimeError: if the data is not valid LZ4 format
+        :raises RuntimeError: if the data is not valid LZ4 legacy format
         """
         resource_data = await resource.get_data()
 
+        # Parse legacy header: 4 bytes magic + 4 bytes block size
+        if len(resource_data) < 8:
+            raise RuntimeError("Invalid LZ4 legacy format: file too short")
+
+        # Note: The header field is the compressed block size, not uncompressed size
+        block_size = int.from_bytes(resource_data[4:8], "little")
+        compressed_block = resource_data[8:]
+
+        # Validate block size matches actual data
+        if len(compressed_block) != block_size:
+            raise RuntimeError(
+                f"Invalid LZ4 legacy format: header says {block_size} bytes but found {len(compressed_block)}"
+            )
+
         try:
-            decompressed_data = lz4.frame.decompress(resource_data)
-        except RuntimeError as e:
-            LOGGER.error(f"Failed to decompress LZ4 data: {e}")
-            raise
+            # LZ4 legacy blocks don't store uncompressed size, so we need to provide
+            # a large enough buffer. Use a generous multiplier to handle any compression ratio.
+            max_uncompressed_size = block_size * 255  # LZ4 max compression ratio
+            decompressed_data = lz4.block.decompress(
+                compressed_block, uncompressed_size=max_uncompressed_size
+            )
+        except Exception as e:
+            LOGGER.error(f"Failed to decompress LZ4 legacy data: {e}")
+            raise RuntimeError(f"LZ4 legacy decompression failed: {e}")
 
         await resource.create_child(
             tags=(GenericBinary,),
             data=decompressed_data,
         )
 
 
-class Lz4Packer(Packer[None]):
+@dataclass
+class Lz4PackerConfig(ComponentConfig):
     """
-        Pack data into a compressed LZ4 file using modern frame format.
+    Configuration for LZ4 packer.
+
+    compression_level: Compression level to use (default: 0).
+        - Negative values: Fast acceleration (faster, less compression)
+        - 0-2: Minimum compression (default, all produce same output)
+        - 3: Minimum high-compression mode
+        - 4-16: Higher compression levels (16 is maximum)
+    """
+
+    compression_level: int = 0
 
-    Note: Only creates modern frame format. Legacy frames and skippable frames
-    cannot be repacked:
-    - Legacy format is deprecated and not supported by the Python lz4 library
-    - Skippable frames are metadata containers and don't make semantic sense to pack
 
-    If you unpack a legacy or skippable frame and repack, it will be converted
-    to modern frame format.
+class Lz4Packer(Packer[Lz4PackerConfig]):
+    """
+    Pack data into a compressed LZ4 file using modern frame format.
+
+    Implementation repacks modern frame format preserving frame metadata.
+    Compression level can be specified via config (default: 0).
     """
 
     targets = (Lz4ModernData,)
 
-    async def pack(self, resource: Resource, config=None):
+    async def pack(self, resource: Resource, config: Lz4PackerConfig = None):
         """
-        Pack data into LZ4 modern frame format.
+        Pack data into `Lz4ModernData` format.
 
         :param resource: The LZ4 resource to pack
-        :param config: Optional packer configuration
+        :param config: Optional configuration specifying compression level
         """
+        if config is None:
+            config = Lz4PackerConfig()
+
         lz4_child = await resource.get_only_child()
         child_data = await lz4_child.get_data()
 
-        lz4_compressed = lz4.frame.compress(child_data)
+        # Use stored compression settings from the view
+        lz4_view = await resource.view_as(Lz4ModernData)
+        content_checksum = lz4_view.content_checksum
+        block_checksum = lz4_view.block_checksum
+        block_size = lz4_view.block_size
+        store_size = lz4_view.content_size != 0
+
+        lz4_compressed = lz4.frame.compress(
+            child_data,
+            compression_level=config.compression_level,
+            content_checksum=content_checksum,
+            block_checksum=block_checksum,
+            block_size=block_size,
+            store_size=store_size,
+        )
 
         original_size = await resource.get_data_length()
         resource.queue_patch(Range(0, original_size), lz4_compressed)
 
 
-# Register magic patterns for automatic identification
-MagicMimePattern.register(Lz4Data, "application/x-lz4")
-MagicDescriptionPattern.register(Lz4Data, lambda s: s.lower().startswith("lz4 compressed data"))
+class Lz4LegacyPacker(Packer[Lz4PackerConfig]):
+    """
+    Pack data into compressed LZ4 legacy format.
+
+    Legacy format supports compression levels via lz4.block.compress():
+    - Negative values: Fast mode with acceleration
+    - 0: Default compression
+    - 1-12: High compression mode
+    """
+
+    targets = (Lz4LegacyData,)
+
+    async def pack(self, resource: Resource, config: Lz4PackerConfig = None):
+        """
+        Pack data into `Lz4LegacyData` format.
+
+        :param resource: The LZ4 legacy resource to pack
+        :param config: Optional configuration specifying compression level
+        """
+        if config is None:
+            config = Lz4PackerConfig()
+
+        lz4_child = await resource.get_only_child()
+        child_data = await lz4_child.get_data()
+
+        # Map compression_level to lz4.block.compress() parameters
+        # This matches the lz4 CLI behavior for legacy format:
+        # - Level < 0: fast mode with acceleration = -level
+        if config.compression_level < 0:
+            # Fast mode with acceleration
+            compressed_block = lz4.block.compress(
+                child_data,
+                mode="fast",
+                acceleration=abs(config.compression_level),
+                store_size=False,
+            )
+        # - Level 0-2: fast mode with acceleration = 0
+        elif config.compression_level < 3:
+            # Fast mode with acceleration = 0 (levels 0, 1, 2)
+            compressed_block = lz4.block.compress(
+                child_data, mode="fast", acceleration=0, store_size=False
+            )
+        # - Level >= 3: high compression mode
+        else:
+            # High compression mode (3-12)
+            compressed_block = lz4.block.compress(
+                child_data,
+                mode="high_compression",
+                compression=config.compression_level,
+                store_size=False,
+            )
+
+        # Build legacy header: magic (4 bytes) + compressed_block_size (4 bytes)
+        compressed_block_size = len(compressed_block)
+        header = LZ4_LEGACY_MAGIC + compressed_block_size.to_bytes(4, "little")
+
+        # Combine header + compressed block
+        lz4_compressed = header + compressed_block
+
+        original_size = await resource.get_data_length()
+        resource.queue_patch(Range(0, original_size), lz4_compressed)
+
+
+def match_lz4_modern_magic(data: bytes) -> bool:
+    if len(data) < 4:
+        return False
+    return data[:4] == LZ4_MODERN_MAGIC
+
+
+def match_lz4_legacy_magic(data: bytes) -> bool:
+    if len(data) < 4:
+        return False
+    return data[:4] == LZ4_LEGACY_MAGIC
+
+
+def match_lz4_skippable_magic(data: bytes) -> bool:
+    if len(data) < 4:
+        return False
+    # Format: 0x5X 0x2A 0x4D 0x18 where X is 0-F
+    return data[1:4] == b"\x2a\x4d\x18" and 0x50 <= data[0] <= 0x5F
+
+
+RawMagicPattern.register(Lz4ModernData, match_lz4_modern_magic)
+RawMagicPattern.register(Lz4LegacyData, match_lz4_legacy_magic)
+RawMagicPattern.register(Lz4SkippableData, match_lz4_skippable_magic)