PE: Read the ARM64 and ARMNT exception directory

cle/backends/pe/pe.py

@@ -470,7 +470,13 @@ def _handle_exports(self, coff_symbol_types: dict[int, set[SymbolType]]):
                         self.deps.append(forwardlib)
 
     def _handle_seh(self):
-        if hasattr(self._pe, "DIRECTORY_ENTRY_EXCEPTION"):
+        assert self._pe.FILE_HEADER is not None
+        machine = self._pe.FILE_HEADER.Machine
+        if machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARM64"]:
+            self._handle_seh_arm(arm64=True)
+        elif machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARMNT"]:
+            self._handle_seh_arm(arm64=False)
+        elif hasattr(self._pe, "DIRECTORY_ENTRY_EXCEPTION"):
             for entry in self._pe.DIRECTORY_ENTRY_EXCEPTION:
                 self.function_hints.append(
                     FunctionHint(
@@ -480,6 +486,54 @@ def _handle_seh(self):
                     )
                 )
 
+    def _handle_seh_arm(self, arm64: bool):
+        """
+        Read the ARM64 and ARMNT exception directory, which pefile parses for x86-64 and Itanium
+        only. Each entry is two words: the function's start RVA, then either the RVA of an .xdata
+        record or unwind data packed into the word itself, distinguished by its low two bits.
+        Described in Microsoft's ARM and ARM64 exception handling references.
+        """
+        exc_dd = self._meta_dd("IMAGE_DIRECTORY_ENTRY_EXCEPTION")
+        if exc_dd is None:
+            return
+        try:
+            table = self._pe.get_data(exc_dd.VirtualAddress, exc_dd.Size)
+        except pefile.PEFormatError:
+            log.warning("PE exception directory lies outside the image")
+            return
+
+        instruction_unit = 4 if arm64 else 2
+        for offset in range(0, len(table) - 7, 8):
+            begin, unwind_data = struct.unpack_from("<II", table, offset)
+            if begin == 0 and unwind_data == 0:
+                continue
+            flag = unwind_data & 3
+            if flag == 0:
+                header = self._pe.get_dword_at_rva(unwind_data)
+                if header is None:
+                    continue
+                # ARMNT marks a function fragment with bit 22 of the .xdata header. ARM64 has no
+                # such bit; there a fragment only ever appears in the packed form below.
+                if not arm64 and header & (1 << 22):
+                    continue
+                length = (header & 0x3FFFF) * instruction_unit
+            elif flag == 1:
+                length = ((unwind_data >> 2) & 0x7FF) * instruction_unit
+            else:
+                # Flag 2 describes a piece of a function that begins elsewhere, so its start
+                # address is not a function entry. Flag 3 is reserved.
+                continue
+            # Bit 0 of an ARMNT start address marks Thumb code. CLE names an ARM function by the
+            # address with that bit set, as an ELF symbol table does, because it is what selects
+            # the decoder; the function's first instruction is at the address without it.
+            self.function_hints.append(
+                FunctionHint(
+                    begin + self.linked_base,
+                    length,
+                    FunctionHintSource.EH_FRAME,
+                )
+            )
+
     def _parse_meta_regions(self):
         """
         Walk pefile's parsed data directories and build meta_regions describing the locations and layouts of PE
@@ -507,8 +561,14 @@ def _meta_pe_context(self) -> tuple[pefile.PE, int, bool, int]:
         ptr_size = 8 if is_64 else 4
         return pe, base, is_64, ptr_size
 
-    def _meta_dd(self, name: str) -> pefile.Structure | None:
-        """Return a data directory entry if it has a nonzero VirtualAddress and Size, else None."""
+    def _meta_dd(self, name: str):
+        """
+        Return a data directory entry if it has a nonzero VirtualAddress and Size, else None.
+
+        The return type is inferred rather than declared: pefile fills a data directory entry in from the format
+        string it parsed, so the pefile.Structure this used to promise declares neither VirtualAddress nor Size and
+        hides both from every caller.
+        """
         idx = pefile.DIRECTORY_ENTRY[name]
         dd = self._pe.OPTIONAL_HEADER.DATA_DIRECTORY[idx]
         if dd.VirtualAddress and dd.Size:

tests/test_pe_function_hints.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import os
+import unittest
+
+import pefile
+
+import cle
+from cle.backends.backend import FunctionHintSource
+
+TEST_BASE = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.path.join("..", "..", "binaries"))
+
+ARM64_PE = os.path.join(TEST_BASE, "tests", "aarch64", "windows", "pe_reloc_arm64.exe")
+ARMNT_PE = os.path.join(TEST_BASE, "tests", "armel", "windows", "pe_reloc_armnt.exe")
+AMD64_PE = os.path.join(TEST_BASE, "tests", "x86_64", "windows", "sioctl.sys")
+
+
+def load(path):
+    return cle.Loader(path, auto_load_libs=False, main_opts={"backend": "pe"}).main_object
+
+
+def exception_directory_hints(obj):
+    return sorted((h.addr, h.size) for h in obj.function_hints if h.source == FunctionHintSource.EH_FRAME)
+
+
+# pylint: disable=no-self-use
+class TestPEFunctionHints(unittest.TestCase):
+    """
+    Function hints taken from the exception directory, whose entries have a different shape on
+    each architecture.
+    """
+
+    def test_aarch64(self):
+        obj = load(ARM64_PE)
+        base = obj.linked_base
+        # the last entry uses the packed form, the first two point at an .xdata record
+        assert exception_directory_hints(obj) == [
+            (base + 0x1000, 0x138),
+            (base + 0x1138, 0x64),
+            (base + 0x119C, 0x1C),
+        ]
+
+    def test_armnt(self):
+        obj = load(ARMNT_PE)
+        base = obj.linked_base
+        # every function in the image is Thumb, so every start address carries the Thumb bit
+        assert exception_directory_hints(obj) == [
+            (base + 0x1001, 0xE8),
+            (base + 0x10E9, 0x3C),
+            (base + 0x1125, 0x1A),
+            (base + 0x113F, 0xC),
+            (base + 0x114B, 0x10),
+            (base + 0x115B, 0xC),
+        ]
+
+    def test_armnt_matches_the_function_pointer_table(self):
+        obj = load(ARMNT_PE)
+        base = obj.linked_base
+        # the three entries of the function pointer table name three of the same functions the
+        # exception directory does, Thumb bit and all
+        table = [obj.memory.unpack_word(0x2000 + i * 4, size=4) for i in range(3)]
+        assert sorted(table) == [base + 0x113F, base + 0x114B, base + 0x115B]
+        hint_addrs = {addr for addr, _ in exception_directory_hints(obj)}
+        assert hint_addrs.issuperset(table)
+
+    def test_x86_64_matches_pefile(self):
+        obj = load(AMD64_PE)
+        pe = pefile.PE(AMD64_PE, fast_load=True)
+        pe.parse_data_directories(directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXCEPTION"]])
+        entries = getattr(pe, "DIRECTORY_ENTRY_EXCEPTION")  # pefile only sets this for x86-64 and Itanium
+        expected = sorted(
+            (obj.linked_base + e.struct.BeginAddress, e.struct.EndAddress - e.struct.BeginAddress) for e in entries
+        )
+        assert exception_directory_hints(obj) == expected
+
+
+if __name__ == "__main__":
+    unittest.main()