From 87ccbeaa46833638a26d6108b2b9af54ea66c5c8 Mon Sep 17 00:00:00 2001 From: Quentin Kaiser Date: Mon, 31 Jul 2023 09:34:23 +0200 Subject: [PATCH] fix(handlers): fix infinite loop in zstd handler. Given an invalid or truncated zstd file, unblob/handlers/compression/zstd.py, in calculate_chunk(), did not check for EOF in the while loop. After reaching EOF, it always "reads" a block of type 0 and of size 0. This is a valid block type (i.e. no raised exception) but not an "end block" type, which is the only other condition to exit the loop. Therefore, it kept looping actively and never returned. Since we're reading from an mmap'ed file, reading from the file after reaching EOF does not trigger an exception. Fixed by checking if the read value is not empty and yielding an InvalidInputFormat exception otherwise. An integration test case reproducing the infinite loop has been added to zstd test cases (truncated.dos.zstd). --- .../compression/zstd/__input__/truncated.dos.zstd | 3 +++ .../zstd/__output__/truncated.dos.zstd_extract/.gitkeep | 0 unblob/handlers/compression/zstd.py | 8 +++++--- 3 files changed, 8 insertions(+), 3 deletions(-) create mode 100644 tests/integration/compression/zstd/__input__/truncated.dos.zstd create mode 100644 tests/integration/compression/zstd/__output__/truncated.dos.zstd_extract/.gitkeep diff --git a/tests/integration/compression/zstd/__input__/truncated.dos.zstd b/tests/integration/compression/zstd/__input__/truncated.dos.zstd new file mode 100644 index 0000000000..1719097383 --- /dev/null +++ b/tests/integration/compression/zstd/__input__/truncated.dos.zstd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:052ba0688dc57289ad3bb5863ac952b4ff8d91156e3a3a274e60ef2ad1746015 +size 9 diff --git a/tests/integration/compression/zstd/__output__/truncated.dos.zstd_extract/.gitkeep b/tests/integration/compression/zstd/__output__/truncated.dos.zstd_extract/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/unblob/handlers/compression/zstd.py b/unblob/handlers/compression/zstd.py index 3afe3281bd..1bd59d3d62 100644 --- 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