From 2a9439e72b75796f9eac53c648aa8e631fbe230f Mon Sep 17 00:00:00 2001 From: pyrco <105293448+pyrco@users.noreply.github.com> Date: Thu, 14 Nov 2024 16:35:56 +0100 Subject: [PATCH] Fix the mapping stream add() function (#59) The sorting should be only on the runlist index. Blanket sorting on all the tuple fields will fail as filehandles/BinaryIO is not sortable. --- dissect/util/stream.py | 4 +++- tests/test_stream.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/dissect/util/stream.py b/dissect/util/stream.py index bb6fb4e..3894875 100644 --- a/dissect/util/stream.py +++ b/dissect/util/stream.py @@ -281,9 +281,11 @@ def add(self, offset: int, size: int, fh: BinaryIO, file_offset: int = 0) -> Non size: The size that this mapped fh spans in the stream. fh: The file-like object to map. file_offset: The offset in the fh to start from. + + Note that there is no check on overlapping offsets and/or sizes. """ self._runs.append((offset, size, fh, file_offset)) - self._runs = sorted(self._runs) + self._runs = sorted(self._runs, key=lambda run: run[0]) self._buf = None self.size = self._runs[-1][0] + self._runs[-1][1] diff --git a/tests/test_stream.py b/tests/test_stream.py index 083c85f..bd18cbe 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -96,6 +96,20 @@ def test_mapping_stream() -> None: assert fh.read(1) == b"" +def test_mapping_stream_same_offset() -> None: + buffers = [ + io.BytesIO(b"\x01" * 512), + io.BytesIO(b"\x02" * 512), + ] + fh = stream.MappingStream() + # Add them in different orders to test if sorting works + fh.add(0, 1024, buffers[0]) + fh.add(0, 1024, buffers[1]) # This should not raise an exception in add() + + assert fh._runs[0][2] == buffers[0] + assert fh._runs[1][2] == buffers[1] + + def test_runlist_stream() -> None: buf = io.BytesIO(b"\x01" * 512 + b"\x02" * 512 + b"\x03" * 512) fh = stream.RunlistStream(buf, [(0, 32), (32, 16), (48, 48)], 1536, 16)