Merge pull request #65 from mrexodia/exception-hooks

Implement exception hooks

src/dumpulator/__init__.py

@@ -1,2 +1,2 @@
-from .dumpulator import Dumpulator
+from .dumpulator import Dumpulator, ExceptionType, MemoryViolation, ExceptionInfo
 from .ntsyscalls import syscall

src/dumpulator/dumpulator.py

@@ -3,7 +3,7 @@
 import sys
 import traceback
 from enum import Enum
-from typing import List, Union, NamedTuple
+from typing import List, Union, NamedTuple, Callable
 import inspect
 from collections import OrderedDict
 from dataclasses import dataclass, field
@@ -39,22 +39,41 @@ class ExceptionType(Enum):
  ContextSwitch = 3
  Terminate = 4
 
+class MemoryViolation(Enum):
+ Unknown = 0
+ ReadUnmapped = 1
+ WriteUnmapped = 2
+ ExecuteUnmapped = 3
+ ReadProtect = 4
+ WriteProtect = 5
+ ExecuteProtect = 6
+ ReadUnaligned = 7
+ WriteUnaligned = 8
+ ExecuteUnaligned = 9
+
 @dataclass
 class ExceptionInfo:
  type: ExceptionType = ExceptionType.NoException
- memory_access: int = 0 # refers to UC_MEM_* values
+ # type == ExceptionType.Memory
+ memory_violation: MemoryViolation = MemoryViolation.Unknown
  memory_address: int = 0
  memory_size: int = 0
  memory_value: int = 0
+ # type == ExceptionType.Interrupt
  interrupt_number: int = 0
+
+ # Internal state
+ _handling: bool = False
+
+@dataclass
+class UnicornExceptionInfo(ExceptionInfo):
+ final: bool = False
  code_hook_h: Optional[int] = None # represents a `unicorn.uc_hook_h` value (from uc.hook_add)
  context: Optional[unicorn.UcContext] = None
  tb_start: int = 0
  tb_size: int = 0
  tb_icount: int = 0
  step_count: int = 0
- final: bool = False
- handling: bool = False
 
  def __str__(self):
  return f"{self.type}, ({hex(self.tb_start)}, {hex(self.tb_size)}, {self.tb_icount})"
@@ -305,8 +324,9 @@ def __init__(self, minidump_file, *, trace=False, quiet=False, thread_id=None, d
  self.kill_me = None
  self.exit_code = None
  self.exports = self._all_exports()
- self.exception = ExceptionInfo()
- self.last_exception: Optional[ExceptionInfo] = None
+ self._exception = UnicornExceptionInfo()
+ self._last_exception: Optional[UnicornExceptionInfo] = None
+ self._exception_hook: Optional[Callable[[ExceptionInfo], Optional[int]]] = None
  if not self._quiet:
  print("Memory map:")
  self.print_memory()
@@ -896,15 +916,28 @@ def allocate(self, size, page_align=False):
  self.memory.commit(self.memory.align_page(ptr), self.memory.align_page(size))
  return ptr
 
- def handle_exception(self):
- assert not self.exception.handling
- self.exception.handling = True
+ def set_exception_hook(self, exception_hook: Optional[Callable[[ExceptionInfo], Optional[int]]]):
+ previous_hook = self._exception_hook
+ self._exception_hook = exception_hook
+ return previous_hook
 
- if self.exception.type == ExceptionType.ContextSwitch:
+ def handle_exception(self):
+ assert not self._exception._handling
+ self._exception._handling = True
+
+ if self._exception_hook is not None:
+ hook_result = self._exception_hook(self._exception)
+ if hook_result is not None:
+ # Clear the pending exception
+ self._last_exception = self._exception
+ self._exception = UnicornExceptionInfo()
+ return hook_result
+
+ if self._exception.type == ExceptionType.ContextSwitch:
  self.info(f"context switch, cip: {hex(self.regs.cip)}")
  # Clear the pending exception
- self.last_exception = self.exception
- self.exception = ExceptionInfo()
+ self._last_exception = self._exception
+ self._exception = UnicornExceptionInfo()
  # NOTE: the context has already been restored using context_restore in the caller
  return self.regs.cip
 
@@ -961,22 +994,23 @@ def handle_exception(self):
  context_ex.XState.Offset = 0xF0 if self._x64 else 0x20
  context_ex.XState.Length = 0x160 if self._x64 else 0x140
  record = record_type()
- if self.exception.type == ExceptionType.Memory:
- record.ExceptionCode = 0xC0000005
+ alignment_violations = [MemoryViolation.ReadUnaligned, MemoryViolation.WriteUnaligned, MemoryViolation.ExecuteUnaligned]
+ if self._exception.type == ExceptionType.Memory and self._exception.memory_violation not in alignment_violations:
+ record.ExceptionCode = STATUS_ACCESS_VIOLATION
  record.ExceptionFlags = 0
  record.ExceptionAddress = self.regs.cip
  record.NumberParameters = 2
  types = {
- UC_MEM_READ_UNMAPPED: EXCEPTION_READ_FAULT,
- UC_MEM_WRITE_UNMAPPED: EXCEPTION_WRITE_FAULT,
- UC_MEM_FETCH_UNMAPPED: EXCEPTION_READ_FAULT,
- UC_MEM_READ_PROT: EXCEPTION_READ_FAULT,
- UC_MEM_WRITE_PROT: EXCEPTION_WRITE_FAULT,
- UC_MEM_FETCH_PROT: EXCEPTION_EXECUTE_FAULT,
+ MemoryViolation.ReadUnmapped: EXCEPTION_READ_FAULT,
+ MemoryViolation.WriteUnmapped: EXCEPTION_WRITE_FAULT,
+ MemoryViolation.ExecuteUnmapped: EXCEPTION_READ_FAULT,
+ MemoryViolation.ReadProtect: EXCEPTION_READ_FAULT,
+ MemoryViolation.WriteProtect: EXCEPTION_WRITE_FAULT,
+ MemoryViolation.ExecuteProtect: EXCEPTION_EXECUTE_FAULT,
  }
- record.ExceptionInformation[0] = types[self.exception.memory_access]
- record.ExceptionInformation[1] = self.exception.memory_address
- elif self.exception.type == ExceptionType.Interrupt and self.exception.interrupt_number == 3:
+ record.ExceptionInformation[0] = types[self._exception.memory_violation]
+ record.ExceptionInformation[1] = self._exception.memory_address
+ elif self._exception.type == ExceptionType.Interrupt and self._exception.interrupt_number == 3:
  if self._x64:
  context.Rip -= 1 # TODO: long int3 and prefixes
  record.ExceptionCode = 0x80000003
@@ -990,11 +1024,11 @@ def handle_exception(self):
  record.ExceptionAddress = context.Eip
  record.NumberParameters = 1
  else:
- raise NotImplementedError(f"{self.exception}") # TODO: implement
+ raise NotImplementedError(f"{self._exception}") # TODO: implement
 
  # Clear the pending exception
- self.last_exception = self.exception
- self.exception = ExceptionInfo()
+ self._last_exception = self._exception
+ self._exception = UnicornExceptionInfo()
 
  def write_stack(cur_ptr: int, data: bytes):
  self.write(cur_ptr, data)
@@ -1024,19 +1058,19 @@ def write_stack(cur_ptr: int, data: bytes):
 
  def start(self, begin, end=0xffffffffffffffff, count=0) -> None:
  # Clear exceptions before starting
- self.exception = ExceptionInfo()
+ self._exception = UnicornExceptionInfo()
  emu_begin = begin
  emu_until = end
  emu_count = count
  while True:
  try:
- if self.exception.type != ExceptionType.NoException:
- if self.exception.final:
+ if self._exception.type != ExceptionType.NoException:
+ if self._exception.final:
  # Restore the context (unicorn might mess with it before stopping)
- if self.exception.context is not None:
- self._uc.context_restore(self.exception.context)
+ if self._exception.context is not None:
+ self._uc.context_restore(self._exception.context)
 
- if self.exception.type == ExceptionType.Terminate:
+ if self._exception.type == ExceptionType.Terminate:
  if self.exit_code is not None:
  self.info(f"exit code: {hex(self.exit_code)}")
  break
@@ -1051,20 +1085,20 @@ def start(self, begin, end=0xffffffffffffffff, count=0) -> None:
  emu_count = 0
  else:
  # If this happens there was an error restarting simulation
- assert self.exception.step_count == 0
+ assert self._exception.step_count == 0
 
  # Hook should be installed at this point
- assert self.exception.code_hook_h is not None
+ assert self._exception.code_hook_h is not None
 
  # Restore the context (unicorn might mess with it before stopping)
- assert self.exception.context is not None
- self._uc.context_restore(self.exception.context)
+ assert self._exception.context is not None
+ self._uc.context_restore(self._exception.context)
 
  # Restart emulation
  self.info(f"restarting emulation to handle exception...")
  emu_begin = self.regs.cip
  emu_until = 0xffffffffffffffff
- emu_count = self.exception.tb_icount + 1
+ emu_count = self._exception.tb_icount + 1
 
  self.info(f"emu_start({hex(emu_begin)}, {hex(emu_until)}, {emu_count})")
  self.kill_me = None
@@ -1076,7 +1110,7 @@ def start(self, begin, end=0xffffffffffffffff, count=0) -> None:
  except UcError as err:
  if self.kill_me is not None and type(self.kill_me) is not UcError:
  raise self.kill_me
- if self.exception.type != ExceptionType.NoException:
+ if self._exception.type != ExceptionType.NoException:
  # Handle the exception outside of the except handler
  continue
  else:
@@ -1232,7 +1266,7 @@ def load_dll(self, file_name: str, file_data: bytes):
 def _hook_code_exception(uc: Uc, address, size, dp: Dumpulator):
  try:
  dp.info(f"exception step: {hex(address)}[{size}]")
- ex = dp.exception
+ ex = dp._exception
  ex.step_count += 1
  if ex.step_count >= ex.tb_icount:
  raise Exception("Stepped past the basic block without reaching exception")
@@ -1246,18 +1280,27 @@ def _hook_mem(uc: Uc, access, address, size, value, dp: Dumpulator):
  return True
 
  fetch_accesses = [UC_MEM_FETCH, UC_MEM_FETCH_PROT, UC_MEM_FETCH_UNMAPPED]
- if access == UC_MEM_FETCH_UNMAPPED and address >= FORCE_KILL_ADDR - 0x10 and address <= FORCE_KILL_ADDR + 0x10 and dp.kill_me is not None:
+ if access == UC_MEM_FETCH_UNMAPPED and FORCE_KILL_ADDR - 0x10 <= address <= FORCE_KILL_ADDR + 0x10 and dp.kill_me is not None:
  dp.error(f"forced exit memory operation {access} of {hex(address)}[{hex(size)}] = {hex(value)}")
  return False
- if dp.exception.final and access in fetch_accesses:
+ if dp._exception.final and access in fetch_accesses:
  dp.info(f"fetch from {hex(address)}[{size}] already reported")
  return False
  # TODO: figure out why when you start executing at 0 this callback is triggered more than once
  try:
+ violation = {
+ UC_MEM_READ_UNMAPPED: MemoryViolation.ReadUnmapped,
+ UC_MEM_WRITE_UNMAPPED: MemoryViolation.WriteUnmapped,
+ UC_MEM_FETCH_UNMAPPED: MemoryViolation.ExecuteUnmapped,
+ UC_MEM_READ_PROT: MemoryViolation.ReadProtect,
+ UC_MEM_WRITE_PROT: MemoryViolation.WriteProtect,
+ UC_MEM_FETCH_PROT: MemoryViolation.ExecuteProtect,
+ }.get(access, MemoryViolation.Unknown)
+ assert violation != MemoryViolation.Unknown, f"Unexpected memory access {access}"
  # Extract exception information
- exception = ExceptionInfo()
+ exception = UnicornExceptionInfo()
  exception.type = ExceptionType.Memory
- exception.memory_access = access
+ exception.memory_violation = violation
  exception.memory_address = address
  exception.memory_size = size
  exception.memory_value = value
@@ -1269,7 +1312,7 @@ def _hook_mem(uc: Uc, access, address, size, value, dp: Dumpulator):
  exception.tb_icount = tb.icount
 
  # Print exception info
- final = dp.trace or dp.exception.code_hook_h is not None
+ final = dp.trace or dp._exception.code_hook_h is not None
  info = "final" if final else "initial"
  if access == UC_MEM_READ_UNMAPPED:
  dp.error(f"{info} unmapped read from {hex(address)}[{hex(size)}], cip = {hex(dp.regs.cip)}, exception: {exception}")
@@ -1295,25 +1338,25 @@ def _hook_mem(uc: Uc, access, address, size, value, dp: Dumpulator):
  if final:
  # Make sure this is the same exception we expect
  if not dp.trace:
- assert access == dp.exception.memory_access
- assert address == dp.exception.memory_address
- assert size == dp.exception.memory_size
- assert value == dp.exception.memory_value
+ assert violation == dp._exception.memory_violation
+ assert address == dp._exception.memory_address
+ assert size == dp._exception.memory_size
+ assert value == dp._exception.memory_value
 
  # Delete the code hook
- uc.hook_del(dp.exception.code_hook_h)
- dp.exception.code_hook_h = None
+ uc.hook_del(dp._exception.code_hook_h)
+ dp._exception.code_hook_h = None
 
  # At this point we know for sure the context is correct so we can report the exception
- dp.exception = exception
- dp.exception.final = True
+ dp._exception = exception
+ dp._exception.final = True
 
  # Stop emulation (we resume it on KiUserExceptionDispatcher later)
  dp.stop()
  return False
 
  # There should not be an exception active
- assert dp.exception.type == ExceptionType.NoException
+ assert dp._exception.type == ExceptionType.NoException
 
  # Remove the translation block cache for this block
  # Without doing this single stepping the block won't work
@@ -1325,7 +1368,7 @@ def _hook_mem(uc: Uc, access, address, size, value, dp: Dumpulator):
  exception.code_hook_h = uc.hook_add(UC_HOOK_CODE, _hook_code_exception, user_data=dp)
 
  # Store the exception info
- dp.exception = exception
+ dp._exception = exception
 
  # Stop emulation (we resume execution later)
  dp.stop()
@@ -1452,7 +1495,7 @@ def _arg_type_string(arg):
 def _hook_interrupt(uc: Uc, number, dp: Dumpulator):
  try:
  # Extract exception information
- exception = ExceptionInfo()
+ exception = UnicornExceptionInfo()
  exception.type = ExceptionType.Interrupt
  exception.interrupt_number = number
  exception.context = uc.context_save()
@@ -1470,11 +1513,11 @@ def _hook_interrupt(uc: Uc, number, dp: Dumpulator):
  dp.error(f"interrupt {number} ({description}), cip = {hex(dp.regs.cip)}, cs = {hex(dp.regs.cs)}")
 
  # There should not be an exception active
- assert dp.exception.type == ExceptionType.NoException
+ assert dp._exception.type == ExceptionType.NoException
 
  # At this point we know for sure the context is correct so we can report the exception
- dp.exception = exception
- dp.exception.final = True
+ dp._exception = exception
+ dp._exception.final = True
  except AssertionError as err:
  traceback.print_exc()
  raise err
@@ -1560,7 +1603,7 @@ def syscall_arg(index):
  status = syscall_impl(dp, *args)
  if isinstance(status, ExceptionInfo):
  print("context switch, stopping emulation")
- dp.exception = status
+ dp._exception = status
  raise dp.raise_kill(UcError(UC_ERR_EXCEPTION)) from None
  else:
  dp.info(f"status = {hex(status)}")
@@ -1610,11 +1653,11 @@ def _hook_invalid(uc: Uc, dp: Dumpulator):
  instr = next(dp.cs.disasm(code, address, 1))
  if _emulate_unsupported_instruction(dp, instr):
  # Resume execution with a context switch
- assert dp.exception.type == ExceptionType.NoException
- exception = ExceptionInfo()
+ assert dp._exception.type == ExceptionType.NoException
+ exception = UnicornExceptionInfo()
  exception.type = ExceptionType.ContextSwitch
  exception.final = True
- dp.exception = exception
+ dp._exception = exception
  return False # NOTE: returning True would stop emulation
  except StopIteration:
  pass # Unsupported instruction

src/dumpulator/native.py

@@ -8,6 +8,7 @@
 # NTSTATUS
 STATUS_SUCCESS = 0
 STATUS_NOT_IMPLEMENTED = 0xC0000002
+STATUS_ACCESS_VIOLATION = 0xC0000005
 STATUS_INVALID_HANDLE = 0xC0000008
 STATUS_NO_SUCH_FILE = 0xC000000F
 STATUS_ACCESS_DENIED = 0xC0000022