Merge pull request #635 from onekey-sec/634-fix-zstd-infinite-loop

fix(handlers): fix infinite loop in zstd handler.
onekey-sec · Jul 31, 2023 · 1facc7d · 1facc7d
2 parents 3733fb3 + 87ccbea
commit 1facc7d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/tests/integration/compression/zstd/__input__/truncated.dos.zstd b/tests/integration/compression/zstd/__input__/truncated.dos.zstd
diff --git a/tests/integration/compression/zstd/__output__/truncated.dos.zstd_extract/.gitkeep b/tests/integration/compression/zstd/__output__/truncated.dos.zstd_extract/.gitkeep
diff --git a/unblob/handlers/compression/zstd.py b/unblob/handlers/compression/zstd.py
@@ -58,9 +58,11 @@ def calculate_chunk(self, file: File, start_offset: int) -> Optional[ValidChunk]
 
         last_block = False
         while not last_block:
-            block_header = int.from_bytes(
-                file.read(BLOCK_HEADER_LEN), byteorder="little"
-            )
+            block_header_val = file.read(BLOCK_HEADER_LEN)
+            # EOF
+            if not block_header_val:
+                raise InvalidInputFormat("Premature end of ZSTD stream.")
+            block_header = int.from_bytes(block_header_val, byteorder="little")
             last_block = block_header >> 0 & 0b1
             block_type = block_header >> 1 & 0b11