fix #38: better thumb mode handling

Author: DennyDai

src/patcherex2/components/binary_analyzers/angr.py

@@ -2,11 +2,12 @@
 
 import logging
 import traceback
+from bisect import bisect_right
 
 import angr
 from archinfo import ArchARM
 
-from .binary_analyzer import BinaryAnalyzer
+from .binary_analyzer import BinaryAnalyzer, UnknownInstructionModeError
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +21,7 @@ def __init__(self, binary_path: str, **kwargs) -> None:
         self._p = None
         self._cfg = None
         self._load_base = None
+        self._mapping_symbols = None
 
     @property
     def load_base(self) -> int:
@@ -124,6 +126,8 @@ def _flatten(n):
                             for instr_addr in instr_addrs
                         ],
                     }
+        except UnknownInstructionModeError:
+            raise
         except Exception:  # noqa: BLE001
             logger.error(
                 f"angr RegionIdentifier failed for function containing {hex(addr)}, falling back to use cfg nodes\n{traceback.format_exc()}"
@@ -208,16 +212,60 @@ def get_function(self, name_or_addr: int | str) -> dict[str, int] | None:
         else:
             raise TypeError(f"Invalid type for name_or_addr: {type(name_or_addr)}")
 
-    def is_thumb(self, addr: int) -> bool:
+    @property
+    def mapping_symbols(self) -> list[tuple[int, str]]:
+        if self._mapping_symbols is None:
+            self._mapping_symbols = sorted(
+                (symbol.rebased_addr, symbol.name[:2])
+                for symbol in self.p.loader.main_object.symbols
+                if symbol.name and symbol.name[:2] in {"$a", "$d", "$t"}
+            )
+        return self._mapping_symbols
+
+    def _mapping_thumb_mode(self, addr: int) -> tuple[bool, bool | None]:
+        index = bisect_right(self.mapping_symbols, (addr, "\uffff")) - 1
+        if index < 0:
+            return False, None
+        _, kind = self.mapping_symbols[index]
+        if kind == "$t":
+            return True, True
+        if kind == "$a":
+            return True, False
+        return True, None
+
+    def _cfg_thumb_mode(self, addr: int) -> bool | None:
+        candidates = (addr, addr + 1) if addr % 2 == 0 else (addr,)
+        modes = {
+            node.thumb
+            for node in self.cfg.model.nodes()
+            if any(candidate in node.instruction_addrs for candidate in candidates)
+        }
+        if len(modes) != 1:
+            return None
+        return modes.pop()
+
+    def thumb_mode(self, addr: int) -> bool | None:
+        """Return the ARM instruction mode, or None when it is unknown."""
         if not isinstance(self.p.arch, ArchARM):
             return False
+
         addr = self.denormalize_addr(addr)
+        mapping_found, mapping_mode = self._mapping_thumb_mode(addr)
+        cfg_mode = self._cfg_thumb_mode(addr)
 
-        for node in self.cfg.model.nodes():
-            if addr in node.instruction_addrs:
-                return node.thumb
-        if addr % 2 == 0:
-            return self.is_thumb(self.normalize_addr(addr + 1))
-        else:
-            logger.error(f"Cannot find a block containing address {hex(addr)}")
-            return False
+        if not mapping_found:
+            return cfg_mode
+        if cfg_mode is None:
+            return mapping_mode
+        if mapping_mode is not None and mapping_mode == cfg_mode:
+            return mapping_mode
+        return None
+
+    def is_thumb(self, addr: int) -> bool:
+        mode = self.thumb_mode(addr)
+        if mode is None:
+            addr = self.denormalize_addr(addr)
+            raise UnknownInstructionModeError(
+                f"Cannot determine ARM instruction mode at {hex(addr)}"
+            )
+        return mode

src/patcherex2/components/binary_analyzers/binary_analyzer.py

@@ -1,2 +1,6 @@
+class UnknownInstructionModeError(RuntimeError):
+    """Raised when an instruction's architecture mode cannot be determined."""
+
+
 class BinaryAnalyzer:
     pass

src/patcherex2/components/binary_analyzers/ghidra.py

@@ -3,7 +3,7 @@
 import logging
 import tempfile
 
-from .binary_analyzer import BinaryAnalyzer
+from .binary_analyzer import BinaryAnalyzer, UnknownInstructionModeError
 
 logger = logging.getLogger(__name__)
 
@@ -41,11 +41,10 @@ def __init__(self, binary_path: str, language: str | None = None, **kwargs):
             self.flatapi = FlatProgramAPI(self.currentProgram)
 
             import ghidra
+            import ghidra.program.model.block
 
             self.ghidra = ghidra
-            self.bbm = self.ghidra.program.model.block.BasicBlockModel(
-                self.currentProgram
-            )
+            self.bbm = ghidra.program.model.block.BasicBlockModel(self.currentProgram)
         except BaseException:
             self.shutdown()
             raise
@@ -78,59 +77,70 @@ def shutdown(self):
                     if temp_proj_dir_ctx is not None:
                         temp_proj_dir_ctx.cleanup()
 
-    def normalize_addr(self, addr):
-        addr = addr.getOffset()
+    @property
+    def load_base(self) -> int:
+        return self.currentProgram.getImageBase().getOffset()
+
+    def normalize_addr(self, addr: int) -> int:
         if self.currentProgram.getRelocationTable().isRelocatable():
-            addr -= self.currentProgram.getImageBase().getOffset()
+            addr -= self.load_base
         return addr
 
-    def denormalize_addr(self, addr):
+    def denormalize_addr(self, addr: int) -> int:
         if self.currentProgram.getRelocationTable().isRelocatable():
-            addr += self.currentProgram.getImageBase().getOffset()
-        return self.flatapi.toAddr(hex(addr))
+            addr += self.load_base
+        return addr
+
+    def _normalize_ghidra_addr(self, addr) -> int:
+        return self.normalize_addr(addr.getOffset())
+
+    def _to_ghidra_addr(self, addr: int):
+        return self.flatapi.toAddr(hex(self.denormalize_addr(addr)))
 
     def mem_addr_to_file_offset(self, addr: int) -> int:
-        addr = self.denormalize_addr(addr)
+        ghidra_addr = self._to_ghidra_addr(addr)
         try:
             return (
                 self.currentProgram.getMemory()
-                .getAddressSourceInfo(addr)
+                .getAddressSourceInfo(ghidra_addr)
                 .getFileOffset()
             )
         except Exception:  # noqa: BLE001
             raise ValueError("Can't get file offset for addr") from None
 
     def get_basic_block(self, addr: int) -> dict[str, int | list[int]]:
-        logger.info(f"getting basic block at 0x{addr} with ghidra")
-        addr = self.denormalize_addr(addr)
+        logger.info(f"getting basic block at {hex(addr)} with ghidra")
+        ghidra_addr = self._to_ghidra_addr(addr)
 
         block = self.bbm.getFirstCodeBlockContaining(
-            addr, self.ghidra.util.task.TaskMonitor.DUMMY
+            ghidra_addr, self.ghidra.util.task.TaskMonitor.DUMMY
         )
         if block is None:
-            raise ValueError(f"Cannot find block containing address 0x{addr}")
+            raise ValueError(f"Cannot find block containing address {hex(addr)}")
         instrs = []
         ii = self.currentProgram.getListing().getInstructions(block, True)
         for i in ii:
-            instrs.append(self.normalize_addr(i.getAddress()))
+            instrs.append(self._normalize_ghidra_addr(i.getAddress()))
         return {
-            "start": self.normalize_addr(block.getMinAddress()),
-            "end": self.normalize_addr(block.getMinAddress()) + block.getNumAddresses(),
+            "start": self._normalize_ghidra_addr(block.getMinAddress()),
+            "end": self._normalize_ghidra_addr(block.getMinAddress())
+            + block.getNumAddresses(),
             "size": block.getNumAddresses(),
             "instruction_addrs": instrs,
         }
 
     def get_instr_bytes_at(self, addr: int, num_instr=1):
-        addr = self.denormalize_addr(addr)
-        instr = self.currentProgram.getListing().getInstructionContaining(addr)
+        ghidra_addr = self._to_ghidra_addr(addr)
+        instr = self.currentProgram.getListing().getInstructionContaining(ghidra_addr)
         if instr is None:
             return None
-        b = instr.getBytes()
+        b = bytes(instr.getBytes())
         for _i in range(1, num_instr):
             instr = instr.getNext()
-            b = b"".join([b, instr.getBytes()])
+            b = b"".join([b, bytes(instr.getBytes())])
         logger.info(
-            f"got instr bytes of length {len(b)} for {num_instr} instrs at 0x{addr} with ghidra"
+            f"got instr bytes of length {len(b)} for {num_instr} instrs at "
+            f"{hex(addr)} with ghidra"
         )
         return b
 
@@ -143,7 +153,7 @@ def get_unused_funcs(self) -> list[dict[str, int]]:
                 b = f.getBody()
                 unused_funcs.append(
                     {
-                        "addr": self.normalize_addr(b.getMinAddress()),
+                        "addr": self._normalize_ghidra_addr(b.getMinAddress()),
                         "size": b.getNumAddresses(),
                     }
                 )
@@ -162,15 +172,16 @@ def get_all_symbols(self) -> dict[str, int]:
         for f in fi:
             if f.getName() in symbols:
                 continue
-            symbols[f.getName()] = self.normalize_addr(f.getEntryPoint())
+            symbols[f.getName()] = self._normalize_ghidra_addr(f.getEntryPoint())
             if self.is_thumb(symbols[f.getName()]):
                 symbols[f.getName()] += 1
         return symbols
 
     def get_function(self, name_or_addr: int | str) -> dict[str, int] | None:
         if isinstance(name_or_addr, int):
-            name_or_addr = self.denormalize_addr(name_or_addr)
-            func = self.currentProgram.getListing().getFunctionContaining(name_or_addr)
+            func = self.currentProgram.getListing().getFunctionContaining(
+                self._to_ghidra_addr(name_or_addr)
+            )
             if func is None:
                 return None
         elif isinstance(name_or_addr, str):
@@ -183,16 +194,33 @@ def get_function(self, name_or_addr: int | str) -> dict[str, int] | None:
 
         b = func.getBody()
         return {
-            "addr": self.normalize_addr(b.getMinAddress()),
+            "addr": self._normalize_ghidra_addr(b.getMinAddress()),
             "size": b.getNumAddresses(),
         }
 
-    def is_thumb(self, addr: int) -> bool:
-        addr = self.denormalize_addr(addr)
-        r = self.currentProgram.getRegister("TMode")
-        if r is None:
+    def thumb_mode(self, addr: int) -> bool | None:
+        """Return the Ghidra ARM mode, or None when it is unknown."""
+        register = self.currentProgram.getRegister("TMode")
+        if register is None:
             return False
-        v = self.currentProgram.getProgramContext().getRegisterValue(r, addr)
-        t = v.unsignedValueIgnoreMask.intValue() == 1
-        logger.info(f"address 0x{addr} {'is' if t else 'is not'} thumb from ghidra")
-        return t
+
+        value = self.currentProgram.getProgramContext().getRegisterValue(
+            register, self._to_ghidra_addr(addr)
+        )
+        if value is None or not value.hasValue():
+            logger.info(f"address {hex(addr)} has no TMode value in ghidra")
+            return None
+
+        is_thumb = value.unsignedValueIgnoreMask.intValue() == 1
+        logger.info(
+            f"address {hex(addr)} {'is' if is_thumb else 'is not'} thumb from ghidra"
+        )
+        return is_thumb
+
+    def is_thumb(self, addr: int) -> bool:
+        mode = self.thumb_mode(addr)
+        if mode is None:
+            raise UnknownInstructionModeError(
+                f"Cannot determine ARM instruction mode at {hex(addr)}"
+            )
+        return mode

tests/test_angr_thumb.py

@@ -0,0 +1,89 @@
+from types import SimpleNamespace
+
+import pytest
+from archinfo import ArchARM, Endness
+
+from patcherex2.components.binary_analyzers.angr import Angr
+from patcherex2.components.binary_analyzers.binary_analyzer import (
+    UnknownInstructionModeError,
+)
+
+
+class FakeSymbol:
+    def __init__(self, name: str, address: int) -> None:
+        self.name = name
+        self.rebased_addr = address
+
+
+class FakeNode:
+    def __init__(self, address: int, thumb: bool) -> None:
+        self.instruction_addrs = {address}
+        self.thumb = thumb
+
+
+def make_analyzer(
+    mapping_symbols: list[tuple[str, int]], cfg_nodes: list[tuple[int, bool]]
+) -> Angr:
+    main_object = SimpleNamespace(
+        mapped_base=0,
+        pic=False,
+        symbols=[FakeSymbol(name, address) for name, address in mapping_symbols],
+    )
+    analyzer = Angr.__new__(Angr)
+    analyzer._p = SimpleNamespace(
+        arch=ArchARM(Endness.LE),
+        loader=SimpleNamespace(main_object=main_object),
+    )
+    analyzer._cfg = SimpleNamespace(
+        model=SimpleNamespace(
+            nodes=lambda: [FakeNode(address, thumb) for address, thumb in cfg_nodes]
+        )
+    )
+    analyzer._load_base = 0
+    analyzer._mapping_symbols = None
+    return analyzer
+
+
+@pytest.mark.parametrize(
+    ("mapping_symbols", "cfg_nodes", "expected"),
+    [
+        (["$t"], [(0x101, True)], True),
+        (["$a"], [(0x101, False)], False),
+        (["$t"], [], True),
+        (["$a"], [], False),
+        ([], [(0x101, True)], True),
+        ([], [(0x101, False)], False),
+        (["$t"], [(0x101, False)], None),
+        (["$a"], [(0x101, True)], None),
+        (["$d"], [(0x101, True)], None),
+        (["$d"], [], None),
+        ([], [], None),
+    ],
+)
+def test_thumb_mode_combines_mapping_symbols_and_cfg(
+    mapping_symbols, cfg_nodes, expected
+):
+    analyzer = make_analyzer([(name, 0x100) for name in mapping_symbols], cfg_nodes)
+
+    assert analyzer.thumb_mode(0x100) is expected
+
+
+def test_thumb_mode_uses_mapping_symbol_ranges():
+    analyzer = make_analyzer([("$t", 0x100), ("$a", 0x200)], [])
+
+    assert analyzer.thumb_mode(0x1FF) is True
+    assert analyzer.thumb_mode(0x200) is False
+
+
+def test_is_thumb_rejects_unknown_mode():
+    analyzer = make_analyzer([("$t", 0x100)], [(0x101, False)])
+
+    with pytest.raises(UnknownInstructionModeError, match="0x100"):
+        analyzer.is_thumb(0x100)
+
+
+def test_get_instr_bytes_at_rejects_unknown_mode():
+    analyzer = make_analyzer([], [])
+
+    with pytest.raises(UnknownInstructionModeError, match="0x100"):
+        analyzer.get_instr_bytes_at(0x100)