error handling

Author: evakhoni

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/driver.py

@@ -51,7 +51,7 @@ async def flash(self, source, partition: str | None = None):
         async with await FileWriteStream.from_path(self.parent.validate_partition(partition)) as stream:
             async with self.resource(source) as res:
                 # Wrap with auto-decompression to handle .gz, .xz, .bz2, .zstd files
-                async for chunk in AutoDecompressIterator(source=res.__aiter__()):
+                async for chunk in AutoDecompressIterator(source=res):
                     await stream.send(chunk)
 
     @export

packages/jumpstarter/jumpstarter/streams/encoding.py

@@ -162,10 +162,26 @@ class AutoDecompressIterator(AsyncIterator[bytes]):
 
     source: AsyncIterator[bytes]
     _decompressor: Any = field(init=False, default=None)
+    _compression: Compression | None = field(init=False, default=None)
     _detected: bool = field(init=False, default=False)
     _buffer: bytes = field(init=False, default=b"")
     _exhausted: bool = field(init=False, default=False)
 
+    def _call_decompressor(self, method_name: str, *args) -> bytes:
+        """Call decompressor method with error handling.
+
+        Args:
+            method_name: decompressor method to call
+            *args: Arguments to the method
+        """
+        try:
+            method = getattr(self._decompressor, method_name)
+            return method(*args)
+        except (zlib.error, lzma.LZMAError, OSError, zstd.ZstdError) as e:
+            raise RuntimeError(
+                f"Failed to decompress {self._compression}: {e}"
+            ) from e
+
     async def _detect_compression(self) -> None:
         """Read enough bytes to detect compression format."""
         # Buffer data until we have enough for detection
@@ -180,6 +196,7 @@ async def _detect_compression(self) -> None:
         # Detect compression from buffered data
         compression = detect_compression_from_signature(self._buffer)
         if compression is not None:
+            self._compression = compression
             self._decompressor = create_decompressor(compression)
 
         self._detected = True
@@ -194,7 +211,7 @@ async def __anext__(self) -> bytes:
             data = self._buffer
             self._buffer = b""
             if self._decompressor is not None:
-                return self._decompressor.decompress(data)
+                return self._call_decompressor("decompress", data)
             return data
 
         # Stream exhausted
@@ -208,14 +225,14 @@ async def __anext__(self) -> bytes:
             self._exhausted = True
             # Flush any remaining data from decompressor (gzip needs this)
             if self._decompressor is not None and hasattr(self._decompressor, "flush"):
-                remaining = self._decompressor.flush()
+                remaining = self._call_decompressor("flush")
                 self._decompressor = None
                 if remaining:
                     return remaining
             raise
 
         if self._decompressor is not None:
-            return self._decompressor.decompress(chunk)
+            return self._call_decompressor("decompress", chunk)
         return chunk
 
     def __aiter__(self) -> AsyncIterator[bytes]:

packages/jumpstarter/jumpstarter/streams/encoding_test.py

@@ -164,3 +164,13 @@ async def test_large_data(self):
         original = b"x" * 1024 * 1024  # 1MB of data
         compressed = gzip.compress(original)
         await self._decompress_and_check(compressed, original, chunk_size=65536)
+
+    async def test_corrupted_gzip(self):
+        """Corrupted gzip data should raise RuntimeError with clear message."""
+        # Create fake gzip data: valid signature but corrupted payload
+        corrupted = b"\x1f\x8b\x08" + b"corrupted data here"
+
+        with pytest.raises(RuntimeError, match=r"Failed to decompress gzip:.*"):
+            chunks = []
+            async for chunk in AutoDecompressIterator(source=self._async_iter_from_bytes(corrupted, 16)):
+                chunks.append(chunk)