Pack and unpack words struct cannot describe

cle/memory.py

@@ -5,12 +5,43 @@
 import struct
 from collections.abc import Iterator
 from mmap import mmap
-from typing import Any, cast
+from typing import Any, Literal, cast
 
 import archinfo
 
 __all__ = ("ClemoryBase", "Clemory", "ClemoryView", "ClemoryTranslator", "UninitializedClemory")
 
+# The word sizes struct has an integer format character for. Every other width - the 3-byte word of
+# a 24-bit architecture, or anything wider than 8 bytes - has to be composed from its bytes.
+_STRUCT_WORD_SIZES = frozenset((1, 2, 4, 8))
+
+
+def _byteorder(endness: archinfo.Endness) -> Literal["little", "big"]:
+    """
+    Translate an :class:`archinfo.Endness` into a byte order accepted by ``int.from_bytes``.
+    """
+    if endness == archinfo.Endness.BE:
+        return "big"
+    if endness == archinfo.Endness.LE:
+        return "little"
+    raise ValueError(f"Unsupported endness value {endness}.")
+
+
+def _classify_struct_error(error: struct.error, fmt: str, available: int, addr: int) -> Exception:
+    """
+    Decide what an access that raised `error` should report: a ``KeyError`` if it ran off the end of
+    its backer, and the error itself otherwise.
+
+    Return the exception rather than raising it so the caller can raise it outside its own handler,
+    where a short access does not carry the struct error as its context.
+    """
+    try:
+        fmt_size = struct.calcsize(fmt)
+    except struct.error:
+        # `fmt` itself is malformed, so this is not a question of how much room is left
+        return error
+    return KeyError(addr) if available < fmt_size else error
+
 
 class ClemoryBase:
     """
@@ -60,15 +91,17 @@ def unpack(self, addr: int, fmt: str) -> tuple[Any, ...]:
         try:
             return struct.unpack_from(fmt, backer, addr - start)
         except struct.error as e:
-            if len(backer) - (addr - start) >= struct.calcsize(fmt):
-                raise e
-            raise KeyError(addr)  # pylint: disable=raise-missing-from
+            error = _classify_struct_error(e, fmt, len(backer) - (addr - start), addr)
+        raise error
 
     def unpack_word(
         self, addr: int, size: int | None = None, signed: bool = False, endness: archinfo.Endness | None = None
     ) -> int:
         """
-        Use the ``struct`` module to unpack a single integer from the address `addr`.
+        Unpack a single integer from the address `addr`.
+
+        Any positive size works. Widths ``struct`` has no format character for are read as bytes and
+        recombined.
 
         You may override any of the attributes of the word being extracted:
 
@@ -77,26 +110,19 @@ def unpack_word(
         :param bool signed: Whether the data should be extracted signed/unsigned. Default unsigned
         :param archinfo.Endness endness: The endian to use in packing/unpacking. Defaults to memory endness
         """
-        if size is not None and size > 8:
-            # support larger wordsizes via recursive algorithm
-            subsize = size >> 1
-            if size != subsize << 1:
-                raise ValueError("Cannot unpack non-power-of-two sizes")
-
-            if endness is None:
-                endness = self._arch.memory_endness
-            if endness == archinfo.Endness.BE:
-                lo_off, hi_off = subsize, 0
-            elif endness == archinfo.Endness.LE:
-                lo_off, hi_off = 0, subsize
-            else:
-                raise ValueError(f"Unsupported endness value {endness}.")
+        word_size: int = self._arch.bytes if size is None else size
+        word_endness: archinfo.Endness = self._arch.memory_endness if endness is None else endness
+        if word_size <= 0:
+            raise ValueError(f"Invalid size: {word_size}")
 
-            lo = self.unpack_word(addr + lo_off, size=subsize, signed=False, endness=endness)
-            hi = self.unpack_word(addr + hi_off, size=subsize, signed=signed, endness=endness)
-            return (hi << (subsize << 3)) | lo
+        if word_size not in _STRUCT_WORD_SIZES:
+            data = self.load(addr, word_size)
+            if len(data) != word_size:
+                # `load` stops at the first unmapped byte instead of raising
+                raise KeyError(addr)
+            return int.from_bytes(data, _byteorder(word_endness), signed=signed)
 
-        return self.unpack(addr, self._arch.struct_fmt(size=size, signed=signed, endness=endness))[0]
+        return self.unpack(addr, self._arch.struct_fmt(size=word_size, signed=signed, endness=word_endness))[0]
 
     def load_null_terminated_bytes(self, addr: int, max_size: int = 4096) -> bytes:
         """
@@ -130,9 +156,8 @@ def pack(self, addr: int, fmt: str, *data):
         try:
             return struct.pack_into(fmt, backer, addr - start, *data)
         except struct.error as e:
-            if len(backer) - (addr - start) >= struct.calcsize(fmt):
-                raise e
-            raise KeyError(addr)  # pylint: disable=raise-missing-from
+            error = _classify_struct_error(e, fmt, len(backer) - (addr - start), addr)
+        raise error
 
     def pack_word(
         self,
@@ -143,7 +168,11 @@ def pack_word(
         endness: archinfo.Endness | None = None,
     ):
         """
-        Use the ``struct`` module to pack a single integer `data` into memory at the address `addr`.
+        Pack a single integer `data` into memory at the address `addr`.
+
+        Any positive size works. Widths ``struct`` has no format character for are written as bytes.
+        The whole word has to be backed; a write that would run off the end raises ``KeyError`` and
+        leaves memory alone.
 
         You may override any of the attributes of the word being packed:
 
@@ -152,9 +181,22 @@ def pack_word(
         :param bool signed: Whether the data should be extracted signed/unsigned. Default unsigned
         :param archinfo.Endness endness: The endian to use in packing/unpacking. Defaults to memory endness
         """
+        word_size: int = self._arch.bytes if size is None else size
+        word_endness: archinfo.Endness = self._arch.memory_endness if endness is None else endness
+        if word_size <= 0:
+            raise ValueError(f"Invalid size: {word_size}")
+
         if not signed:
-            data &= (1 << (size * 8 if size is not None else self._arch.bits)) - 1
-        return self.pack(addr, self._arch.struct_fmt(size=size, signed=signed, endness=endness), data)
+            data &= (1 << (word_size * 8)) - 1
+
+        if word_size not in _STRUCT_WORD_SIZES:
+            if len(self.load(addr, word_size)) != word_size:
+                # `store` writes the bytes that fit before it reports the overrun, so check the whole
+                # word is backed before writing any of it
+                raise KeyError(addr)
+            return self.store(addr, data.to_bytes(word_size, _byteorder(word_endness), signed=signed))
+
+        return self.pack(addr, self._arch.struct_fmt(size=word_size, signed=signed, endness=word_endness), data)
 
     def read(self, nbytes: int):
         """
@@ -786,9 +828,8 @@ def unpack(self, addr, fmt):
                 try:
                     return struct.unpack_from(fmt, data, addr - start)
                 except struct.error as ex:
-                    if len(data) - (addr - start) >= struct.calcsize(fmt):
-                        raise ex
-                    raise KeyError(addr) from ex
+                    error = _classify_struct_error(ex, fmt, len(data) - (addr - start), addr)
+                raise error
 
         idx = bisect.bisect_right(self._flattened_backers, addr, key=lambda x: x[0])
         if idx > 0:
@@ -803,9 +844,8 @@ def unpack(self, addr, fmt):
             self._last_backer_pos = idx
             return v
         except struct.error as ex:
-            if len(data) - (addr - start) >= struct.calcsize(fmt):
-                raise ex
-            raise KeyError(addr) from ex
+            error = _classify_struct_error(ex, fmt, len(data) - (addr - start), addr)
+        raise error
 
     def _flatten_backers(self):
         for start, backer in self._clemory.backers():

tests/test_clemory.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import struct
 import timeit
 
 import cffi
+import pytest
 
 import cle
 
@@ -115,6 +117,47 @@ def test_clemory_contains():
     assert clemory.consecutive is True
 
 
+def test_clemory_malformed_format():
+    # a format string struct cannot parse is a format error, not an out-of-bounds access, and
+    # classifying it must not raise struct.error out of the handler doing the classification
+    clemory = cle.Clemory(None, root=True)
+    clemory.add_backer(0, b"A" * 4)
+
+    with pytest.raises(struct.error) as excinfo:
+        clemory.unpack(0, "<Z")
+    assert excinfo.value.__context__ is None
+
+    with pytest.raises(struct.error) as excinfo:
+        clemory.pack(0, "<Z", 1)
+    assert excinfo.value.__context__ is None
+
+    # a format struct can parse that runs off the end of the backer is still a KeyError
+    with pytest.raises(KeyError):
+        clemory.unpack(2, "<I")
+    with pytest.raises(KeyError):
+        clemory.pack(2, "<I", 1)
+
+
+def test_clemory_read_only_view_malformed_format():
+    clemory = cle.Clemory(None, root=True)
+    clemory.add_backer(0, b"A" * 4)
+    view = cle.ClemoryReadOnlyView(None, clemory)
+
+    with pytest.raises(struct.error) as excinfo:
+        view.unpack(0, "<Z")
+    assert excinfo.value.__context__ is None
+    with pytest.raises(KeyError):
+        view.unpack(2, "<I")
+
+    # and again now that the backer the reads above found is cached, which is a separate handler
+    assert view.unpack(0, "<H") == (0x4141,)
+    with pytest.raises(struct.error) as excinfo:
+        view.unpack(0, "<Z")
+    assert excinfo.value.__context__ is None
+    with pytest.raises(KeyError):
+        view.unpack(2, "<I")
+
+
 def main():
     g = globals()
     for func_name, func in g.items():

tests/test_tls_resiliency.py

@@ -4,8 +4,15 @@
 import unittest
 from unittest import TestCase
 
+import archinfo
+
 import cle
 
+try:
+    import pypcode
+except ImportError:
+    pypcode = None
+
 test_location = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.path.join("..", "..", "binaries", "tests"))
 
 
@@ -18,6 +25,33 @@ def test_tls_pe_incorrect_tls_data_start():
         th = path_ld.tls.new_thread()
         assert th is not None
 
+    @staticmethod
+    @unittest.skipIf(pypcode is None, "pypcode not installed")
+    def test_tls_24bit_arch():
+        # isqrt_atmega128.o is an ATmega128 object whose e_flags name the extended-address AVR
+        # variant, which Ghidra's AVR8 opinion file maps to avr8:LE:16:extended -- a language whose
+        # word is three bytes wide despite the 16 in its name. cle's own opinion matching compares
+        # the opinion's secondary constraint against e_type instead of e_flags, so it picks
+        # avr8:LE:16:default for every EM_AVR ELF; name the language here rather than wait for that
+        # to be fixed. Setting up the ELF TLS header writes the DTV pointer one word at a time.
+        p = os.path.join(test_location, "avr", "isqrt_atmega128.o")
+        ld = cle.Loader(
+            p,
+            main_opts={"arch": archinfo.ArchPcode("avr8:LE:16:extended")},
+            auto_load_libs=False,
+        )
+        arch = ld.main_object.arch
+        assert arch.bits == 24
+
+        thread = ld.tls.new_thread()
+        assert thread is not None
+
+        # the DTV pointer is the write that fails, so read it back rather than trusting the load
+        elf_tls = arch.elf_tls
+        assert elf_tls is not None and elf_tls.dtv_offsets
+        for offset in elf_tls.dtv_offsets:
+            assert thread.memory.unpack_word(offset + thread.tcb_offset) == thread.dtv_offset
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_unpackword.py

@@ -3,6 +3,7 @@
 from io import BytesIO
 
 import archinfo
+import pytest
 
 import cle
 
@@ -70,5 +71,46 @@ def test_unpackword():
     assert ymmword == 0xFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFD0F0E0D0C0B0A09080706050403020137 - 2**256
 
 
+def test_word_sizes_struct_cannot_express():
+    # struct has integer format characters for 1, 2, 4 and 8 bytes only. Every other width - the
+    # 3-byte word of a 24-bit architecture among them - is composed from its bytes.
+    clemory = cle.Clemory(archinfo.ArchX86(), root=True)
+    clemory.add_backer(0, bytes(32))
+
+    clemory.pack_word(0, 0x123456, size=3)
+    assert clemory.load(0, 4) == b"\x56\x34\x12\x00"
+    assert clemory.unpack_word(0, 3) == 0x123456
+
+    clemory.pack_word(4, 0x123456, size=3, endness=archinfo.Endness.BE)
+    assert clemory.load(4, 3) == b"\x12\x34\x56"
+    assert clemory.unpack_word(4, 3, endness=archinfo.Endness.BE) == 0x123456
+
+    clemory.pack_word(8, -2, size=3, signed=True)
+    assert clemory.load(8, 3) == b"\xfe\xff\xff"
+    assert clemory.unpack_word(8, 3, signed=True) == -2
+    assert clemory.unpack_word(8, 3) == 0xFFFFFE
+
+    # the same goes for anything wider than 8 bytes, power of two or not
+    clemory.pack_word(12, 0x0102030405060708090A, size=10)
+    assert clemory.unpack_word(12, 10) == 0x0102030405060708090A
+    clemory.pack_word(12, 0x0102030405060708090A0B0C0D0E0F10, size=16)
+    assert clemory.unpack_word(12, 16) == 0x0102030405060708090A0B0C0D0E0F10
+
+
+def test_word_off_the_end_of_a_backer():
+    clemory = cle.Clemory(archinfo.ArchX86(), root=True)
+    clemory.add_backer(0, bytes(2))
+
+    with pytest.raises(KeyError):
+        clemory.unpack_word(0, 3)
+
+    # a write that does not fit leaves memory alone rather than storing the bytes that do fit
+    with pytest.raises(KeyError):
+        clemory.pack_word(0, 0x123456, size=3)
+    assert clemory.load(0, 2) == b"\x00\x00"
+
+
 if __name__ == "__main__":
     test_unpackword()
+    test_word_sizes_struct_cannot_express()
+    test_word_off_the_end_of_a_backer()