From e637b26b3bc6268dd589fa1439fecf99e49a565b Mon Sep 17 00:00:00 2001
From: Maximilian Knespel <mxmlnkn@users.noreply.github.com>
Date: Sun, 21 Apr 2024 20:15:14 +0200
Subject: [PATCH] Fix lz4 and lzma compressed images (#27)

* Add simple automated test

* Add missing LZMA support

* Fix broken LZ4 support

Trying to access an LZ4-compressed squashfs file would result in:

    LZ4F_getFrameInfo failed with code: ERROR_frameType_unknown
---
 .github/workflows/tests.yml           | 51 +++++++++++++++++++++++++++
 PySquashfsImage/compressor.py         | 37 +++++++++++++++++--
 PySquashfsImage/tests/test_library.py | 38 ++++++++++++++++++++
 3 files changed, 123 insertions(+), 3 deletions(-)
 create mode 100644 .github/workflows/tests.yml
 create mode 100644 PySquashfsImage/tests/test_library.py

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 0000000..7dcad82
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,51 @@
+name: Tests
+
+on:
+  push:
+    branches: '**'
+    tags-ignore: '**'
+  pull_request:
+
+jobs:
+  Tests:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ['3.7', '3.12']
+
+    defaults:
+      run:
+        shell: bash
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - uses: msys2/setup-msys2@v2
+      if: startsWith( matrix.os, 'windows' )
+
+    - name: Install Dependencies (Linux)
+      if: startsWith( matrix.os, 'ubuntu' )
+      run: |
+        sudo apt-get -y install liblzo2-dev
+
+    - name: Install pip Dependencies
+      run: |
+        python3 -m pip install --upgrade pip
+        python3 -m pip install --upgrade wheel setuptools twine pytest python-lzo lz4 zstandard
+
+    - name: Test Installation From Tarball
+      run: |
+        python3 setup.py clean check build sdist bdist_egg bdist_wheel
+        twine check dist/*
+        python3 -m pip install "$( find dist -name '*.tar.gz' | head -1 )"
+
+    - name: Unit Tests
+      run: |
+        pytest PySquashfsImage/tests/test_*.py
diff --git a/PySquashfsImage/compressor.py b/PySquashfsImage/compressor.py
index 3e48766..de84296 100644
--- a/PySquashfsImage/compressor.py
+++ b/PySquashfsImage/compressor.py
@@ -30,6 +30,36 @@ def uncompress(self, src, size, outsize):
         return self._lib.decompress(src, False, outsize)
 
 
+class LZMACompressor(Compressor):
+    name = "lzma"
+
+    def __init__(self):
+        try:
+            import lzma
+        except ImportError:
+            from backports import lzma
+        self._lib = lzma
+
+    def uncompress(self, src, size, outsize):
+        # https://github.com/plougher/squashfs-tools/blob/a04910367d64a5220f623944e15be282647d77ba/squashfs-tools/
+        #   lzma_wrapper.c#L40
+        # res = LzmaCompress(dest + LZMA_HEADER_SIZE, &outlen, src, size, dest,
+        #                    &props_size, 5, block_size, 3, 0, 2, 32, 1);
+        # https://github.com/jljusten/LZMA-SDK/blob/781863cdf592da3e97420f50de5dac056ad352a5/C/LzmaLib.h#L96
+        # -> level=5, dictSize=block_size, lc=3, lp=0, pb=2, fb=32, numThreads=1
+        # https://github.com/plougher/squashfs-tools/blob/a04910367d64a5220f623944e15be282647d77ba/squashfs-tools/
+        #   lzma_wrapper.c#L30
+        # For some reason, squashfs does not store raw lzma but adds a custom header of 5 B and 8 B little-endian
+        # uncompressed size, which can be read with struct.unpack('<Q', src[5:5+8]))
+        LZMA_PROPS_SIZE = 5
+        LZMA_HEADER_SIZE = LZMA_PROPS_SIZE + 8
+        return self._lib.decompress(
+            src[LZMA_HEADER_SIZE:],
+            format=self._lib.FORMAT_RAW,
+            filters=[{"id": self._lib.FILTER_LZMA1, 'lc': 3, 'lp': 0, 'pb': 2}],
+        )
+
+
 class XZCompressor(Compressor):
     name = "xz"
 
@@ -48,11 +78,11 @@ class LZ4Compressor(Compressor):
     name = "lz4"
 
     def __init__(self):
-        import lz4.frame
-        self._lib = lz4.frame
+        import lz4.block
+        self._lib = lz4.block
 
     def uncompress(self, src, size, outsize):
-        return self._lib.decompress(src)
+        return self._lib.decompress(src, outsize)
 
 
 class ZSTDCompressor(Compressor):
@@ -69,6 +99,7 @@ def uncompress(self, src, size, outsize):
 compressors = {
     Compression.NO: Compressor,
     Compression.ZLIB: ZlibCompressor,
+    Compression.LZMA: LZMACompressor,
     Compression.LZO: LZOCompressor,
     Compression.XZ: XZCompressor,
     Compression.LZ4: LZ4Compressor,
diff --git a/PySquashfsImage/tests/test_library.py b/PySquashfsImage/tests/test_library.py
new file mode 100644
index 0000000..8506a67
--- /dev/null
+++ b/PySquashfsImage/tests/test_library.py
@@ -0,0 +1,38 @@
+import io
+import os
+import subprocess
+import tarfile
+import tempfile
+
+import pytest
+
+import PySquashfsImage
+
+
+def _createFile(tarArchive, name, contents):
+    tinfo = tarfile.TarInfo(name)
+    tinfo.size = len(contents)
+    tarArchive.addfile(tinfo, io.BytesIO(contents.encode()))
+
+
+@pytest.mark.parametrize("compression", ["", "gzip", "lz4", "lzma", "lzo", "xz", "zstd"])
+def test_compressions(compression):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tarPath = os.path.join(tmpdir, "foo.tar")
+        with tarfile.open(name=tarPath, mode='w:') as tarArchive:
+            _createFile(tarArchive, "foo", "bar")
+
+        squashfsPath = os.path.join(tmpdir, f"foo.{compression if compression else 'no-compression'}.squashfs")
+        compressionOptions = ["-comp", compression] if compression else ["-noI", "-noId", "-noD", "-noF", "-noX"]
+        process = subprocess.Popen(
+            ["sqfstar"] + compressionOptions + [squashfsPath], stdin=subprocess.PIPE, stdout=subprocess.PIPE
+        )
+        with open(tarPath, 'rb') as file:
+            process.communicate(file.read())
+
+        with open(squashfsPath, 'rb') as file, PySquashfsImage.SquashFsImage(file) as image:
+            entries = list(iter(image))
+            assert len(entries) == 2
+            assert entries[0].path == "/"
+            assert entries[1].path == "/foo"
+            assert image.read_file(entries[1].inode) == b"bar"