diff --git a/ext4/volume.py b/ext4/volume.py index 8b76610..6785b14 100644 --- a/ext4/volume.py +++ b/ext4/volume.py @@ -211,13 +211,36 @@ def seek(self, offset: int, mode: int = io.SEEK_SET) -> int: def read(self, size: int) -> bytes: _ = self.stream.seek(self.offset + self.cursor) - data = self.stream.read(size) + remaining = size + data = bytearray() + while remaining > 0: + chunk = self.stream.read(remaining) + if not chunk: + break + + data.extend(chunk) + remaining -= len(chunk) + + data = data[:size] self.cursor += len(data) - return data + return bytes(data) def peek(self, size: int) -> bytes: _ = self.stream.seek(self.offset + self.cursor) - return self.stream.peek(size) + remaining = size + data = bytearray() + while remaining > 0: + chunk = self.stream.peek(remaining) + if not chunk: + break + + data.extend(chunk) + remaining -= len(chunk) + _ = self.stream.seek(len(chunk), io.SEEK_CUR) + + data = data[:size] + _ = self.stream.seek(self.offset + self.cursor) + return bytes(data) def tell(self) -> int: return self.cursor diff --git a/pyproject.toml b/pyproject.toml index 2c2daf8..1adf11e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ext4" -version = "1.4" +version = "1.4.1" authors = [ { name="Eeems", email="eeems@eeems.email" }, ] diff --git a/test.py b/test.py index 43c360c..8964222 100644 --- a/test.py +++ b/test.py @@ -11,6 +11,48 @@ ) import ext4 +from ext4._compat import PeekableStream + + +class _ShortReadStream: + """Returns all but the last byte to test Volume.read() loops on short reads.""" + + def __init__(self, stream: PeekableStream) -> None: + self._stream: PeekableStream = stream + + def read(self, size: int | None = -1, /) -> bytes: + if size is None or size < 0: + data = self._stream.read(size) + if len(data) <= 1: + return data + + _ = self._stream.seek(-1, os.SEEK_CUR) + return data[:-1] + + if size <= 1: + return self._stream.read(size) + + return self._stream.read(size - 1) + + def peek(self, size: int = 0, /) -> bytes: + if size < 0: + data = self._stream.peek(size) + if len(data) <= 1: + return data + + return data[:-1] + + if size <= 1: + return self._stream.peek(size) + + return self._stream.peek(size - 1) + + def seek(self, offset: int, whence: int = os.SEEK_SET, /) -> int: + return self._stream.seek(offset, whence) + + def tell(self) -> int: + return self._stream.tell() + FAILED = False @@ -180,6 +222,22 @@ def test_root_inode(volume: ext4.Volume) -> None: _assert("isinstance(inode, ext4.SymbolicLink)") _assert('inode.readlink() == b"test.txt"', inode.readlink) + try: + volume = ext4.Volume(_ShortReadStream(f), offset=offset) + + except Exception: + FAILED = True # pyright: ignore[reportConstantRedefinition] + traceback.print_exc() + continue + + _assert("volume.superblock is not None") + test_root_inode(volume) + inode = cast(ext4.File, volume.inode_at("/test.txt")) + _assert("isinstance(inode, ext4.File)") + data = inode.open().read() + _assert("len(data) > 0") + _assert("data.startswith(b'hello world')") + img_file = "test_htree.ext4" print(f"Testing image: {img_file}") with open(img_file, "rb") as f: @@ -346,5 +404,6 @@ def test_root_inode(volume: ext4.Volume) -> None: if inode_no is not None: inode = volume.inodes[inode_no] _assert("isinstance(inode, ext4.File)", lambda: inode) + if FAILED: sys.exit(1)