From 9ea6fadb829e9741cca0b6e0c5d9ee78577d2fa4 Mon Sep 17 00:00:00 2001 From: Marek Mahut Date: Wed, 5 Aug 2026 12:39:30 +0200 Subject: [PATCH] fix: auto-detect image compression in StorageMux flashing --- .../jumpstarter_driver_opendal/client.py | 6 +- .../jumpstarter_driver_opendal/driver.py | 5 +- .../jumpstarter_driver_opendal/driver_test.py | 61 +++++++++++++++++++ .../jumpstarter/jumpstarter/common/storage.py | 6 +- .../jumpstarter/common/storage_test.py | 42 +++++++++++++ 5 files changed, 117 insertions(+), 3 deletions(-) create mode 100644 python/packages/jumpstarter/jumpstarter/common/storage_test.py diff --git a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py index 6aa654c2e..b57e1eb5a 100644 --- a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py +++ b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py @@ -786,7 +786,11 @@ def flash( operator: Operator | None = None, compression: Compression | None = None, ): - """Flash image to DUT""" + """Flash image to DUT + + gzip, xz, bz2 and zstd compressed images are detected from their + file signature and decompressed transparently on the exporter. + """ if target is not None: raise ArgumentError(f"target is not supported for StorageMuxFlasherClient, {target} provided") diff --git a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py index e488cb3cc..0d833aa55 100644 --- a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py +++ b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver.py @@ -17,6 +17,7 @@ from .adapter import AsyncFileStream from .common import Capability, HashAlgo, Metadata, Mode, PresignedRequest from jumpstarter.driver import Driver, export +from jumpstarter.streams.encoding import AutoDecompressIterator @dataclass(kw_only=True) @@ -390,7 +391,9 @@ async def off(self): async def write(self, src: str): async with await FileWriteStream.from_path(self.file.name) as stream: async with self.resource(src) as res: - async for chunk in res: + # match write_to_storage_device: compressed images are + # detected by file signature and decompressed transparently + async for chunk in AutoDecompressIterator(source=res): await stream.send(chunk) @export diff --git a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py index 3f730ff0d..06137c360 100644 --- a/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py +++ b/python/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/driver_test.py @@ -1,5 +1,9 @@ +import bz2 +import gzip import hashlib +import lzma import os +import sys from contextlib import contextmanager from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path @@ -11,6 +15,11 @@ import pytest from opendal import Operator +if sys.version_info >= (3, 14): + from compression import zstd +else: + from backports import zstd + from .common import PresignedRequest from .driver import MockFlasher, MockStorageMux, MockStorageMuxFlasher, Opendal from jumpstarter.client.core import DriverError @@ -177,6 +186,58 @@ def test_driver_mock_storage_mux_flasher(tmp_path): assert (tmp_path / "dump.img").read_bytes() == b"hello" +@pytest.mark.parametrize( + "compress", + [gzip.compress, lambda data: lzma.compress(data, format=lzma.FORMAT_XZ), bz2.compress, zstd.compress], + ids=["gzip", "xz", "bz2", "zstd"], +) +def test_driver_mock_storage_mux_flasher_auto_decompress(tmp_path, compress): + original = b"hello compressed world" * 1024 + with serve(MockStorageMuxFlasher()) as flasher: + (tmp_path / "disk.img").write_bytes(compress(original)) + + flasher.flash(tmp_path / "disk.img") + flasher.dump(tmp_path / "dump.img") + + assert (tmp_path / "dump.img").read_bytes() == original + + +def test_driver_mock_storage_mux_flasher_http_auto_decompress(tmp_path): + """Flashing a compressed image from a direct HTTP URL must auto-decompress (issue #54).""" + original = b"hello compressed world" * 1024 + compressed = lzma.compress(original, format=lzma.FORMAT_XZ) + + class CompressedHandler(BaseHTTPRequestHandler): + def do_HEAD(self): + self.send_response(200) + self.send_header("content-length", str(len(compressed))) + self.end_headers() + + def do_GET(self): + self.send_response(200) + self.send_header("content-length", str(len(compressed))) + self.end_headers() + self.wfile.write(compressed) + + def log_message(self, format, *args): + pass + + with serve(MockStorageMuxFlasher()) as flasher: + server = HTTPServer(("127.0.0.1", 0), CompressedHandler) + port = server.server_address[1] + server_thread = Thread(target=server.serve_forever) + server_thread.daemon = True + server_thread.start() + try: + flasher.flash(f"http://127.0.0.1:{port}/image.raw.xz") + flasher.dump(tmp_path / "dump.img") + + assert (tmp_path / "dump.img").read_bytes() == original + finally: + server.shutdown() + server.server_close() + + def test_drivers_mock_storage_mux_fs(monkeypatch: pytest.MonkeyPatch): with serve(MockStorageMux()) as client: with TemporaryDirectory() as tempdir: diff --git a/python/packages/jumpstarter/jumpstarter/common/storage.py b/python/packages/jumpstarter/jumpstarter/common/storage.py index 9c8679612..8f3082aac 100644 --- a/python/packages/jumpstarter/jumpstarter/common/storage.py +++ b/python/packages/jumpstarter/jumpstarter/common/storage.py @@ -7,6 +7,8 @@ from anyio.abc import AnyByteStream from anyio.streams.file import FileReadStream, FileWriteStream +from jumpstarter.streams.encoding import AutoDecompressIterator + async def wait_for_storage_device( # noqa: C901 storage_device: str | os.PathLike, @@ -68,7 +70,9 @@ async def write_to_storage_device( async with FileWriteStream(file) as stream: total_bytes = 0 next_print = 0 - async for chunk in resource: + # gzip/xz/bz2/zstd images are detected by file signature and + # decompressed transparently; uncompressed data passes through + async for chunk in AutoDecompressIterator(source=resource): await stream.send(chunk) if logger: total_bytes += len(chunk) diff --git a/python/packages/jumpstarter/jumpstarter/common/storage_test.py b/python/packages/jumpstarter/jumpstarter/common/storage_test.py new file mode 100644 index 000000000..ee8d620b0 --- /dev/null +++ b/python/packages/jumpstarter/jumpstarter/common/storage_test.py @@ -0,0 +1,42 @@ +import bz2 +import gzip +import lzma +import sys + +import pytest + +if sys.version_info >= (3, 14): + from compression import zstd +else: + from backports import zstd + +from .storage import write_to_storage_device + +pytestmark = pytest.mark.anyio + + +async def _chunks(data: bytes, chunk_size: int = 16): + for i in range(0, len(data), chunk_size): + yield data[i : i + chunk_size] + + +@pytest.mark.parametrize( + "compress", + [ + lambda data: data, + gzip.compress, + lambda data: lzma.compress(data, format=lzma.FORMAT_XZ), + bz2.compress, + zstd.compress, + ], + ids=["raw", "gzip", "xz", "bz2", "zstd"], +) +async def test_write_to_storage_device_auto_decompress(tmp_path, compress): + original = b"jumpstarter" * 1024 + device = tmp_path / "device" + # simulate a present storage device: wait_for_storage_device requires a nonzero size + device.write_bytes(b"\x00" * len(original)) + + await write_to_storage_device(device, _chunks(compress(original)), leeway=0) + + assert device.read_bytes() == original