Merge pull request #25 from icicle-emu/memory-capacity

Add support for dynamically updating the memory capacity

.cargo/config.toml

@@ -0,0 +1,2 @@
+[env]
+GHIDRA_SRC = { value = "python/icicle", relative = true }

.github/workflows/CI.yml

@@ -12,14 +12,14 @@ jobs:
       PYTHONUNBUFFERED: '1'
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           submodules: 'true'
       
       - name: Python environment
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
 
       - name: Install Rust toolchain
         uses: dtolnay/rust-toolchain@stable
@@ -42,9 +42,9 @@ jobs:
           python tests/invalid.py
     
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-windows
           path: dist
 
   macos:
@@ -56,14 +56,14 @@ jobs:
       PYTHONUNBUFFERED: '1'
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           submodules: 'true'
       
       - name: Python environment
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
 
       - name: Install Rust toolchain
         uses: dtolnay/rust-toolchain@stable
@@ -76,7 +76,7 @@ jobs:
           DEVELOPER_DIR: /Applications/Xcode.app/Contents/Developer
           MACOSX_DEPLOYMENT_TARGET: '10.9'
           ARCHFLAGS: -arch x86_64 -arch arm64
-          PYO3_CROSS_LIB_DIR: /Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.8/lib
+          PYO3_CROSS_LIB_DIR: /Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.9/lib
         run: |
           pip install -r requirements.txt
           python setup.py bdist_wheel --py-limited-api=cp37
@@ -90,22 +90,22 @@ jobs:
           python tests/invalid.py
     
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-macos
           path: dist
 
   linux:
     # Skip building pull requests from the same repository
     if: ${{ github.event_name == 'push' || (github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository) }}
-    runs-on: ubuntu-latest
-    container: quay.io/pypa/manylinux2014_x86_64
+    runs-on: ubuntu-24.04
+    container: quay.io/pypa/manylinux_2_28_x86_64
     env:
       # Disable output buffering in an attempt to get readable errors
       PYTHONUNBUFFERED: '1'
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           submodules: 'true'
       
@@ -116,7 +116,7 @@ jobs:
           export PATH="$PATH:$HOME/.cargo/bin"
           export PATH="/opt/python/cp38-cp38/bin:$PATH"
           pip install -r requirements.txt
-          python setup.py bdist_wheel --py-limited-api=cp37 --plat-name manylinux2014_x86_64
+          python setup.py bdist_wheel --py-limited-api=cp37 --plat-name manylinux_2_28_x86_64
           auditwheel show dist/*.whl
           pip install --force-reinstall dist/*.whl
           python -c "import icicle"
@@ -129,33 +129,32 @@ jobs:
           python tests/invalid.py
     
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-linux
           path: dist
 
   release:
     if: ${{ startsWith(github.ref, 'refs/tags/') }}
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: [windows, macos, linux]
     permissions:
       contents: write
       discussions: write
+      id-token: write
     steps:
       - name: Download wheels
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
         with:
-          name: wheels
+          pattern: wheels-*
+          merge-multiple: true
           path: dist
 
       - name: Publish to PyPI
-        uses: pypa/gh-action-pypi-publish@b7f401de30cb6434a1e19f805ff006643653240e # v1.8.10
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
-          packages-dir: dist/
+        uses: pypa/gh-action-pypi-publish@67339c736fd9354cd4f8cb0b744f2b82a74b5c70 # v1.12.13
 
       - name: Release
-        uses: softprops/action-gh-release@c9b46fe7aad9f02afd89b12450b780f52dacfb2d # master 2023-03-26
+        uses: softprops/action-gh-release@e7a8f85e1c67a31e6ed99a94b41bd0b71bbee6b8 # v2.0.9
         with:
           generate_release_notes: true
           files: dist/*

Cargo.toml

@@ -1,24 +1,34 @@
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[workspace]
+members = ["tests"]
+exclude = ["icicle-emu"]
 
 [package]
 name = "icicle-python"
-version = "0.0.2"
+version = "0.0.10"
 edition = "2021"
 
 [lib]
 name = "icicle"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "lib"]
+
+[workspace.dependencies]
+pyo3 = { version = "0.23.3", features = ["indexmap", "abi3-py37"] }
 
 [dependencies]
+pyo3 = { workspace = true }
 icicle-cpu = { path = "icicle-emu/icicle-cpu" }
 icicle-vm = { path = "icicle-emu/icicle-vm" }
 pcode = { path = "icicle-emu/sleigh/pcode" }
 sleigh-runtime = { path = "icicle-emu/sleigh/sleigh-runtime" }
-indexmap = "1.9.3"
-pyo3 = { version = "0.19.2", features = ["extension-module", "indexmap", "abi3-py37"] }
+indexmap = "2.2.6"
 target-lexicon = "0.12.7"
 tracing = "*"
 tracing-subscriber = "0.3.17"
 
 [build-dependencies]
-pyo3-build-config = "0.19.2"
+pyo3-build-config = "0.23.3"
+
+# Reference: https://github.com/PyO3/pyo3/issues/340
+[features]
+extension-module = ["pyo3/extension-module"]
+default = ["extension-module"]

python/icicle/Ghidra/Processors/x86/data/languages/ia.sinc

@@ -1145,78 +1145,78 @@ macro ptr8(r,x) {
 
 macro push22(x) {
   mysave:2 = x;
-  SP = SP -2;
-  tmp:$(SIZE) = segment(SS,SP);
+  tmp:$(SIZE) = segment(SS,SP-2);
   *:2 tmp = mysave;
+  SP = SP-2;
 }
 
 macro push24(x) {
   mysave:4 = x;
-  SP = SP-4;
-  tmp:$(SIZE) = segment(SS,SP);
+  tmp:$(SIZE) = segment(SS,SP-4);
   *:4 tmp = mysave;
+  SP = SP-4;
 }
 
 macro push28(x) {
   mysave:8 = x;
-  SP = SP-8;
-  tmp:$(SIZE) = segment(SS,SP);
+  tmp:$(SIZE) = segment(SS,SP-8);
   *:8 tmp = mysave;
+  SP = SP-8;
 }
 
 macro push42(x) {
   mysave:2 = x;
-  $(STACKPTR) = $(STACKPTR) - 2;
-  *:2 $(STACKPTR) = mysave;
+  *:2 ($(STACKPTR)-2) = mysave;
+  $(STACKPTR) = $(STACKPTR)-2;
 }
 
 macro push44(x) {
   mysave:4 = x;
-  $(STACKPTR) = $(STACKPTR) - 4;
-  *:4 $(STACKPTR) = mysave;
+  *:4 ($(STACKPTR)-4) = mysave;
+  $(STACKPTR) = $(STACKPTR)-4;
 }
 
 macro pushseg44(x) {
   mysave:2 = x;
-  $(STACKPTR) = $(STACKPTR) - 4;
-  *:2 $(STACKPTR) = mysave;
+  *:2 ($(STACKPTR)-4) = mysave;
+  $(STACKPTR) = $(STACKPTR)-4;
 }
 
 macro push48(x) {
   mysave:8 = x;
-  $(STACKPTR) = $(STACKPTR) - 8;
-  *:8 $(STACKPTR) = mysave;
+  *:8 ($(STACKPTR)-8) = mysave;
+  $(STACKPTR) = $(STACKPTR)-8;
 }
 
 @ifdef IA64
 macro push82(x) {
   mysave:2 = x;
-  $(STACKPTR) = $(STACKPTR) - 2;
-  *:2 $(STACKPTR) = mysave;
+  *:2 ($(STACKPTR)-2) = mysave;
+  $(STACKPTR) = $(STACKPTR)-2;
 }
 
 macro push84(x) {
   mysave:4 = x;
-  $(STACKPTR) = $(STACKPTR) - 4;
-  *:4 $(STACKPTR) = mysave;
+  *:4 ($(STACKPTR)-4) = mysave;
+  $(STACKPTR) = $(STACKPTR)-4;
 }
 
 macro push88(x) {
   mysave:8 = x;
-  $(STACKPTR) = $(STACKPTR) - 8;
-  *:8 $(STACKPTR) = mysave;
+  *:8 ($(STACKPTR)-8) = mysave;
+  $(STACKPTR) = $(STACKPTR)-8;
 }
 
 macro pushseg82(x) {
   mysave:2 = x;
-  $(STACKPTR) = $(STACKPTR) - 2;
-  *:2 $(STACKPTR) = mysave;
+  *:2 ($(STACKPTR)-2) = mysave;
+  $(STACKPTR) = $(STACKPTR)-2;
 }
 
 macro pushseg88(x) {
   mysave:2 = x;
-  $(STACKPTR) = $(STACKPTR) - 8;
-  *:2 $(STACKPTR) = mysave;
+  *:2 ($(STACKPTR)-8) = mysave;
+  $(STACKPTR) = $(STACKPTR)-8;
 }
 @endif
 

python/icicle/__init__.py

@@ -1,16 +1,163 @@
-from .icicle import *
+from typing import List, Dict, Tuple
+from enum import Enum
 
-class MemoryError(Exception):
-    def __init__(self, message, code):
+class MemoryProtection(Enum):
+    NoAccess = ...
+    ReadOnly = ...
+    ReadWrite = ...
+    ExecuteOnly = ...
+    ExecuteRead = ...
+    ExecuteReadWrite = ...
+
+class MemoryExceptionCode(Enum):
+    Unallocated = ...
+    Unmapped = ...
+    UnmappedRegisters = ...
+    Uninitialized = ...
+    ReadViolation = ...
+    WriteViolation = ...
+    ExecViolation = ...
+    ReadWatch = ...
+    WriteWatch = ...
+    Unaligned = ...
+    OutOfMemory = ...
+    SelfModifyingCode = ...
+    AddressOverflow = ...
+    Unknown = ...
+
+class RunStatus(Enum):
+    Running = ...
+    InstructionLimit = ...
+    Breakpoint = ...
+    Interrupted = ...
+    Halt = ...
+    Killed = ...
+    Deadlock = ...
+    OutOfMemory = ...
+    Unimplemented = ...
+    UnhandledException = ...
+
+class ExceptionCode(Enum):
+    NoException = ...
+    InstructionLimit = ...
+    Halt = ...
+    Sleep = ...
+    Syscall = ...
+    CpuStateChanged = ...
+    DivisionException = ...
+    ReadUnmapped = ...
+    ReadPerm = ...
+    ReadUnaligned = ...
+    ReadWatch = ...
+    ReadUninitialized = ...
+    WriteUnmapped = ...
+    WritePerm = ...
+    WriteWatch = ...
+    WriteUnaligned = ...
+    ExecViolation = ...
+    SelfModifyingCode = ...
+    OutOfMemory = ...
+    AddressOverflow = ...
+    InvalidInstruction = ...
+    UnknownInterrupt = ...
+    UnknownCpuID = ...
+    InvalidOpSize = ...
+    InvalidFloatSize = ...
+    CodeNotTranslated = ...
+    ShadowStackOverflow = ...
+    ShadowStackInvalid = ...
+    InvalidTarget = ...
+    UnimplementedOp = ...
+    ExternalAddr = ...
+    Environment = ...
+    JitError = ...
+    InternalError = ...
+    UnmappedRegister = ...
+    UnknownError = ...
+
+class Icicle:
+    def __init__(self, architecture: str, *,
+                 jit = True,
+                 jit_mem = True,
+                 shadow_stack = False,
+                 recompilation = True,
+                 track_uninitialized = False,
+                 optimize_instructions = True,
+                 optimize_block = False,
+                 tracing = False,
+                 ) -> None: ...
+
+    @property
+    def architecture(self) -> str: ...
+
+    @property
+    def exception_code(self) -> ExceptionCode: ...
+
+    @property
+    def exception_value(self) -> int: ...
+
+    icount: int
+
+    icount_limit: int
+
+    pc: int
+
+    sp: int
+
+    """
+    Physical memory capacity in pages (adjust when seeing OutOfMemory)
+    The default limit is set so that the maximum corresponds to ~400 MB of host memory.
+    """
+    mem_capacity: int
+
+    # TODO: API to get memory information?
+
+    def mem_map(self, address: int, size: int, protection: MemoryProtection): ...
+
+    def mem_unmap(self, address: int, size: int): ...
+
+    def mem_protect(self, address: int, size: int, protection: MemoryProtection): ...
+
+    def mem_read(self, address: int, size: int) -> bytes: ...
+
+    def mem_write(self, address: int, data: bytes) -> None: ...
+
+    def reg_list(self) -> Dict[str, Tuple[int, int]]: ...
+
+    def reg_offset(self, name: str) -> int: ...
+
+    def reg_size(self, name: str) -> int: ...
+
+    def reg_read(self, name: str) -> int: ...
+
+    def reg_write(self, name: str, value: int) -> None: ...
+
+    def reset(self): ...
+
+    def run(self) -> RunStatus: ...
+
+    def run_until(self, address: int) -> RunStatus: ...
+
+    def step(self, count: int) -> RunStatus: ...
+
+    def add_breakpoint(self, address: int) -> bool: ...
+
+    def remove_breakpoint(self, address: int) -> bool: ...
+
+def architectures() -> List[str]: ...
+
+class MemoryException(Exception):
+    def __init__(self, message: str, code: MemoryExceptionCode):
         super().__init__(message)
         self.code = code
+
     def __str__(self):
-        return f"{super().__str__()}: {self.code}"
+        return f"{super().__str__()} ({self.code})"
 
 def __ghidra_init():
     import os
     for path in __path__ + [os.getenv("GHIDRA_SRC", ".")]:
-        processors_dir = os.path.join(path, "Ghidra/Processors")
+        processors_dir = os.path.join(path, "Ghidra", "Processors")
         if os.path.isdir(processors_dir):
             os.putenv("GHIDRA_SRC", path)
             os.environ["GHIDRA_SRC"] = path
@@ -18,3 +165,6 @@ def __ghidra_init():
     raise FileNotFoundError("Ghidra processor definitions not found")
 
 __ghidra_init()
+
+# NOTE: This overrides the stubs at runtime with the actual implementation
+from .icicle import *

python/icicle/icicle.pyi

@@ -1,135 +0,0 @@
-from typing import List, Dict, Tuple
-from enum import Enum
-
-class MemoryProtection(Enum):
-    NoAccess = ...
-    ReadOnly = ...
-    ReadWrite = ...
-    ExecuteOnly = ...
-    ExecuteRead = ...
-    ExecuteReadWrite = ...
-
-class MemoryErrorCode(Enum):
-    Unallocated = ...
-    Unmapped = ...
-    Uninitialized = ...
-    ReadViolation = ...
-    WriteViolation = ...
-    ExecViolation = ...
-    ReadWatch = ...
-    WriteWatch = ...
-    Unaligned = ...
-    OutOfMemory = ...
-    SelfModifyingCode = ...
-    AddressOverflow = ...
-    Unknown = ...
-
-class RunStatus(Enum):
-    Running = ...
-    InstructionLimit = ...
-    Breakpoint = ...
-    Interrupted = ...
-    Halt = ...
-    Killed = ...
-    Deadlock = ...
-    OutOfMemory = ...
-    Unimplemented = ...
-    UnhandledException = ...
-
-class ExceptionCode(Enum):
-    NoException = ...
-    InstructionLimit = ...
-    Halt = ...
-    Sleep = ...
-    Syscall = ...
-    CpuStateChanged = ...
-    DivideByZero = ...
-    ReadUnmapped = ...
-    ReadPerm = ...
-    ReadUnaligned = ...
-    ReadWatch = ...
-    ReadUninitialized = ...
-    WriteUnmapped = ...
-    WritePerm = ...
-    WriteWatch = ...
-    WriteUnaligned = ...
-    ExecViolation = ...
-    SelfModifyingCode = ...
-    OutOfMemory = ...
-    AddressOverflow = ...
-    InvalidInstruction = ...
-    UnknownInterrupt = ...
-    UnknownCpuID = ...
-    InvalidOpSize = ...
-    InvalidFloatSize = ...
-    CodeNotTranslated = ...
-    ShadowStackOverflow = ...
-    ShadowStackInvalid = ...
-    InvalidTarget = ...
-    UnimplementedOp = ...
-    ExternalAddr = ...
-    Environment = ...
-    JitError = ...
-    InternalError = ...
-    UnknownError = ...
-
-class MemoryError(Exception):
-    code: MemoryErrorCode
-
-class Icicle:
-    def __init__(self, architecture: str, *,
-                 jit = True,
-                 jit_mem = True,
-                 shadow_stack = True,
-                 recompilation = True,
-                 track_uninitialized = False,
-                 optimize_instructions = True,
-                 optimize_block = True,
-                 tracing = False,
-                 ) -> None: ...
-
-    @property
-    def exception_code(self) -> ExceptionCode: ...
-
-    @property
-    def exception_value(self) -> int: ...
-
-    icount: int
-
-    icount_limit: int
-
-    # TODO: API to get memory information?
-
-    def mem_map(self, address: int, size: int, protection: MemoryProtection): ...
-
-    def mem_unmap(self, address: int, size: int): ...
-
-    def mem_protect(self, address: int, size: int, protection: MemoryProtection): ...
-
-    def mem_read(self, address: int, size: int) -> bytes: ...
-
-    def mem_write(self, address: int, data: bytes) -> None: ...
-
-    def reg_list(self) -> Dict[str, Tuple[int, int]]: ...
-
-    def reg_offset(self, name: str) -> int: ...
-
-    def reg_size(self, name: str) -> int: ...
-
-    def reg_read(self, name: str) -> int: ...
-
-    def reg_write(self, name: str, value: int) -> None: ...
-
-    def reset(self): ...
-
-    def run(self) -> RunStatus: ...
-
-    def run_until(self, address: int) -> RunStatus: ...
-
-    def step(self, count: int) -> RunStatus: ...
-
-    def add_breakpoint(self, address: int) -> bool: ...
-
-    def remove_breakpoint(self, address: int) -> bool: ...
-
-def architectures() -> List[str]: ...

src/lib.rs

@@ -1,22 +1,38 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
 use icicle_cpu::mem::{Mapping, MemError, perm};
-use icicle_cpu::{ExceptionCode, VmExit};
+use icicle_cpu::{Cpu, ExceptionCode, ValueSource, VmExit};
 use pyo3::prelude::*;
 use icicle_vm;
-use icicle_vm::linux::LinuxCpu;
 use pyo3::exceptions::*;
 use target_lexicon;
 use indexmap::IndexMap;
+use target_lexicon::Architecture;
 use sleigh_runtime::NamedRegister;
 
 // References:
 // - https://pyo3.rs/main/conversions/tables
 // - https://pyo3.rs/main/class
 
-#[pyclass(module = "icicle")]
-#[derive(Clone)]
-enum MemoryProtection {
+struct X86FlagsRegHandler {
+    pub eflags: pcode::VarNode,
+}
+
+impl icicle_cpu::RegHandler for X86FlagsRegHandler {
+    fn read(&mut self, cpu: &mut Cpu) {
+        let eflags = icicle_vm::x86::eflags(cpu);
+        cpu.write_var::<u32>(self.eflags, eflags);
+    }
+
+    fn write(&mut self, cpu: &mut Cpu) {
+        let eflags = cpu.read_var::<u32>(self.eflags);
+        icicle_vm::x86::set_eflags(cpu, eflags);
+    }
+}
+
+#[pyclass(eq, eq_int, module = "icicle")]
+#[derive(Clone, Debug, PartialEq)]
+pub enum MemoryProtection {
     NoAccess,
     ReadOnly,
     ReadWrite,
@@ -25,9 +41,9 @@ enum MemoryProtection {
     ExecuteReadWrite,
 }
 
-#[pyclass(module = "icicle")]
-#[derive(Clone)]
-enum RunStatus {
+#[pyclass(eq, eq_int, module = "icicle")]
+#[derive(Clone, Debug, PartialEq)]
+pub enum RunStatus {
     /// The VM is still running.
     Running,
 
@@ -59,10 +75,12 @@ enum RunStatus {
     UnhandledException,
 }
 
-#[pyclass(module = "icicle")]
-enum MemoryErrorCode {
+#[pyclass(eq, eq_int, module = "icicle")]
+#[derive(Clone, Debug, PartialEq)]
+pub enum MemoryExceptionCode {
     Unallocated,
     Unmapped,
+    UnmappedRegister,
     Uninitialized,
     ReadViolation,
     WriteViolation,
@@ -76,29 +94,30 @@ enum MemoryErrorCode {
     Unknown,
 }
 
-impl From<MemError> for MemoryErrorCode {
+impl From<MemError> for MemoryExceptionCode {
     fn from(value: MemError) -> Self {
         match value {
-            MemError::Unallocated => MemoryErrorCode::Unallocated,
-            MemError::Unmapped => MemoryErrorCode::Unmapped,
-            MemError::Uninitalized => MemoryErrorCode::Uninitialized,
-            MemError::ReadViolation => MemoryErrorCode::ReadViolation,
-            MemError::WriteViolation => MemoryErrorCode::WriteViolation,
-            MemError::ExecViolation => MemoryErrorCode::ExecViolation,
-            MemError::ReadWatch => MemoryErrorCode::ReadWatch,
-            MemError::WriteWatch => MemoryErrorCode::WriteWatch,
-            MemError::Unaligned => MemoryErrorCode::Unaligned,
-            MemError::OutOfMemory => MemoryErrorCode::OutOfMemory,
-            MemError::SelfModifyingCode => MemoryErrorCode::SelfModifyingCode,
-            MemError::AddressOverflow => MemoryErrorCode::AddressOverflow,
-            MemError::Unknown => MemoryErrorCode::Unknown,
+            MemError::Unallocated => MemoryExceptionCode::Unallocated,
+            MemError::Unmapped => MemoryExceptionCode::Unmapped,
+            MemError::UnmappedRegister => MemoryExceptionCode::UnmappedRegister,
+            MemError::Uninitalized => MemoryExceptionCode::Uninitialized,
+            MemError::ReadViolation => MemoryExceptionCode::ReadViolation,
+            MemError::WriteViolation => MemoryExceptionCode::WriteViolation,
+            MemError::ExecViolation => MemoryExceptionCode::ExecViolation,
+            MemError::ReadWatch => MemoryExceptionCode::ReadWatch,
+            MemError::WriteWatch => MemoryExceptionCode::WriteWatch,
+            MemError::Unaligned => MemoryExceptionCode::Unaligned,
+            MemError::OutOfMemory => MemoryExceptionCode::OutOfMemory,
+            MemError::SelfModifyingCode => MemoryExceptionCode::SelfModifyingCode,
+            MemError::AddressOverflow => MemoryExceptionCode::AddressOverflow,
+            MemError::Unknown => MemoryExceptionCode::Unknown,
         }
     }
 }
 
-#[pyclass(module = "icicle", name = "ExceptionCode")]
-#[derive(Clone)]
-enum ExceptionCodePy {
+#[pyclass(eq, eq_int, module = "icicle", name = "ExceptionCode")]
+#[derive(Clone, Debug, PartialEq)]
+pub enum ExceptionCodePy {
     NoException = 0x0000,
 
     InstructionLimit = 0x0001,
@@ -108,7 +127,7 @@ enum ExceptionCodePy {
 
     Syscall = 0x0101,
     CpuStateChanged = 0x0102,
-    DivideByZero = 0x0103,
+    DivisionException = 0x0103,
 
     ReadUnmapped = 0x0201,
     ReadPerm = 0x0202,
@@ -144,6 +163,7 @@ enum ExceptionCodePy {
 
     JitError = 0x3001,
     InternalError = 0x3002,
+    UnmappedRegister = 0x3003,
 
     UnknownError,
 }
@@ -158,7 +178,7 @@ impl From<ExceptionCode> for ExceptionCodePy {
             ExceptionCode::SoftwareBreakpoint => ExceptionCodePy::SoftwareBreakpoint,
             ExceptionCode::Syscall => ExceptionCodePy::Syscall,
             ExceptionCode::CpuStateChanged => ExceptionCodePy::CpuStateChanged,
-            ExceptionCode::DivideByZero => ExceptionCodePy::DivideByZero,
+            ExceptionCode::DivisionException => ExceptionCodePy::DivisionException,
             ExceptionCode::ReadUnmapped => ExceptionCodePy::ReadUnmapped,
             ExceptionCode::ReadPerm => ExceptionCodePy::ReadPerm,
             ExceptionCode::ReadUnaligned => ExceptionCodePy::ReadUnaligned,
@@ -187,17 +207,18 @@ impl From<ExceptionCode> for ExceptionCodePy {
             ExceptionCode::Environment => ExceptionCodePy::Environment,
             ExceptionCode::JitError => ExceptionCodePy::JitError,
             ExceptionCode::InternalError => ExceptionCodePy::InternalError,
+            ExceptionCode::UnmappedRegister => ExceptionCodePy::UnmappedRegister,
             ExceptionCode::UnknownError => ExceptionCodePy::UnknownError,
         }
     }
 }
 
 #[allow(non_snake_case)]
-fn raise_MemoryError(message: String, e: MemError) -> PyErr {
+fn raise_MemoryException(message: String, e: MemError) -> PyErr {
     Python::with_gil(|py| {
         let icicle = py.import("icicle").unwrap();
-        let exception = icicle.getattr("MemoryError").unwrap();
-        let args = (message, MemoryErrorCode::from(e));
+        let exception = icicle.getattr("MemoryException").unwrap();
+        let args = (message, MemoryExceptionCode::from(e));
         let inst = exception.call1(args).unwrap();
         PyErr::from_value(inst)
     })
@@ -215,13 +236,14 @@ fn convert_protection(protection: MemoryProtection) -> u8 {
 }
 
 #[pyclass(unsendable, module = "icicle")]
-struct Icicle {
+pub struct Icicle {
+    architecture: String,
     vm: icicle_vm::Vm,
     regs: HashMap<String, NamedRegister>,
 }
 
 fn reg_find<'a>(i: &'a Icicle, name: &str) -> PyResult<&'a NamedRegister> {
-    let sleigh = i.vm.cpu.sleigh();
+    let sleigh = &i.vm.cpu.arch.sleigh;
     match sleigh.get_reg(name) {
         None => {
             i.regs.get(name.to_lowercase().as_str())
@@ -236,48 +258,91 @@ fn reg_find<'a>(i: &'a Icicle, name: &str) -> PyResult<&'a NamedRegister> {
 #[pymethods]
 impl Icicle {
     #[getter]
-    fn get_icount_limit(&mut self) -> u64 {
+    pub fn get_icount_limit(&mut self) -> u64 {
         self.vm.icount_limit
     }
 
     #[setter]
-    fn set_icount_limit(&mut self, value: u64) {
+    pub fn set_icount_limit(&mut self, value: u64) {
         self.vm.icount_limit = value;
     }
 
     #[getter]
-    fn get_icount(&mut self) -> u64 {
-        return self.vm.cpu.icount;
+    pub fn get_icount(&mut self) -> u64 {
+        self.vm.cpu.icount
     }
 
     #[setter]
-    fn set_icount(&mut self, value: u64) {
+    pub fn set_icount(&mut self, value: u64) {
         self.vm.cpu.icount = value;
     }
 
     #[getter]
-    fn get_exception_code(&self) -> ExceptionCodePy {
+    pub fn get_exception_code(&self) -> ExceptionCodePy {
         ExceptionCode::from_u32(self.vm.cpu.exception.code).into()
     }
 
     #[getter]
-    fn get_exception_value(&self) -> u64 {
+    pub fn get_exception_value(&self) -> u64 {
         self.vm.cpu.exception.value
     }
 
+    #[getter]
+    pub fn get_architecture(&self) -> String {
+        self.architecture.to_string()
+    }
+
+    #[getter]
+    pub fn get_pc(&self) -> u64 {
+        self.vm.cpu.read_pc()
+    }
+
+    #[setter]
+    pub fn set_pc(&mut self, address: u64) {
+        self.vm.cpu.write_pc(address)
+    }
+
+    #[getter]
+    pub fn get_sp(&mut self) -> u64 {
+        self.vm.cpu.read_reg(self.vm.cpu.arch.reg_sp)
+    }
+
+    #[setter]
+    pub fn set_sp(&mut self, address: u64) {
+        self.vm.cpu.write_reg(self.vm.cpu.arch.reg_sp, address)
+    }
+
+    #[getter]
+    pub fn get_mem_capacity(&self) -> usize {
+        self.vm.cpu.mem.capacity()
+    }
+
+    #[setter]
+    pub fn set_mem_capacity(&mut self, capacity: usize) -> PyResult<()> {
+        if self.vm.cpu.mem.set_capacity(capacity) {
+            return Ok(());
+        }
+        Err(
+            raise_MemoryException(
+                format!("Reducing memory capacity is not supported"),
+                MemError::Unknown,
+            )
+        )
+    }
+
     #[new]
     #[pyo3(signature = (
         architecture,
         jit = true,
         jit_mem = true,
-        shadow_stack = true,
+        shadow_stack = false,
         recompilation = true,
         track_uninitialized = false,
         optimize_instructions = true,
-        optimize_block = true,
+        optimize_block = false,
         tracing = false,
     ))]
-    fn new(
+    pub fn new(
         architecture: String,
         jit: bool,
         jit_mem: bool,
@@ -304,7 +369,7 @@ impl Icicle {
             }
         }
 
-        // Setup the CPU state for the target triple
+        // Set up the CPU state for the target triple
         let mut config = icicle_vm::cpu::Config::from_target_triple(
             format!("{architecture}-none").as_str()
         );
@@ -323,45 +388,57 @@ impl Icicle {
         config.optimize_instructions = optimize_instructions;
         config.optimize_block = optimize_block;
 
-        let vm = icicle_vm::build(&config)
+        let mut vm = icicle_vm::build(&config)
             .map_err(|e| {
                 PyException::new_err(format!("VM build error: {e}"))
             })?;
 
         // Populate the lowercase register map
         let mut regs = HashMap::new();
-        let sleigh = vm.cpu.sleigh();
+        let sleigh = &vm.cpu.arch.sleigh;
         for reg in &sleigh.named_registers {
             let name = sleigh.get_str(reg.name);
             regs.insert(name.to_lowercase(), reg.clone());
         }
 
+        // Special handling for x86 flags
+        match config.triple.architecture {
+            Architecture::X86_32(_) | Architecture::X86_64 | Architecture::X86_64h => {
+                let eflags = sleigh.get_reg("eflags").unwrap().var;
+                let reg_handler = X86FlagsRegHandler { eflags };
+                vm.cpu.add_reg_handler(eflags.id, Box::new(reg_handler));
+            }
+            _ => {}
+        }
+
         Ok(Icicle {
+            architecture,
             vm,
             regs,
         })
     }
 
-    fn __str__(&mut self) -> String {
+    pub fn __str__(&mut self) -> String {
         let arch = &self.vm.cpu.arch;
         let endianness = if arch.sleigh.big_endian {
             "big endian"
         } else {
             "little endian"
         };
-        format!("Icicle VM for {0:?} ({endianness})", arch.triple.architecture)
+        format!("Icicle VM for {0:?} ({endianness})", self.architecture)
     }
 
-    fn mem_map(&mut self, address: u64, size: u64, protection: MemoryProtection) -> PyResult<()> {
+    pub fn mem_map(&mut self, address: u64, size: u64, protection: MemoryProtection) -> PyResult<()> {
+        let init_perm = if self.vm.cpu.mem.track_uninitialized { perm::NONE } else { perm::INIT };
         let mapping = Mapping {
-            perm: convert_protection(protection),
+            perm: convert_protection(protection) | init_perm,
             value: 0,
         };
         if self.vm.cpu.mem.map_memory_len(address, size, mapping) {
             Ok(())
         } else {
             Err(
-                raise_MemoryError(
+                raise_MemoryException(
                     format!("Failed to map memory {address:X}[{size:X}]"),
                     MemError::Unknown,
                 )
@@ -369,12 +446,12 @@ impl Icicle {
         }
     }
 
-    fn mem_unmap(&mut self, address: u64, size: u64) -> PyResult<()> {
+    pub fn mem_unmap(&mut self, address: u64, size: u64) -> PyResult<()> {
         if self.vm.cpu.mem.unmap_memory_len(address, size) {
             Ok(())
         } else {
             Err(
-                raise_MemoryError(
+                raise_MemoryException(
                     format!("Failed to unmap memory {address:X}[{size:X}]"),
                     MemError::Unknown,
                 )
@@ -382,10 +459,10 @@ impl Icicle {
         }
     }
 
-    fn mem_protect(&mut self, address: u64, size: usize, protection: MemoryProtection) -> PyResult<()> {
+    pub fn mem_protect(&mut self, address: u64, size: usize, protection: MemoryProtection) -> PyResult<()> {
         self.vm.cpu.mem.update_perm(address, size as u64, convert_protection(protection))
             .map_err(|e| {
-                raise_MemoryError(
+                raise_MemoryException(
                     format!("Failed to protect memory {address:X}[{size:X}]"),
                     e,
                 )
@@ -393,7 +470,7 @@ impl Icicle {
         Ok(())
     }
 
-    fn mem_read(&mut self, address: u64, size: usize) -> PyResult<Cow<[u8]>> {
+    pub fn mem_read(&mut self, address: u64, size: usize) -> PyResult<Cow<[u8]>> {
         // Allocate a buffer
         let mut buffer = Vec::with_capacity(size);
         buffer.resize(size, 0);
@@ -401,56 +478,62 @@ impl Icicle {
         // Read the memory
         self.vm.cpu.mem.read_bytes(address, &mut buffer[..], perm::NONE)
             .map_err(|e| {
-                raise_MemoryError(
+                raise_MemoryException(
                     format!("Failed to read memory {address:X}[{size:X}]"),
                     e,
                 )
             })?;
-        return Ok(Cow::Owned(buffer));
+        Ok(Cow::Owned(buffer))
     }
 
-    fn mem_write(&mut self, address: u64, data: Vec<u8>) -> PyResult<()> {
+    pub fn mem_write(&mut self, address: u64, data: Vec<u8>) -> PyResult<()> {
         let size = data.len();
         self.vm.cpu.mem.write_bytes(address, &data[..], perm::NONE)
             .map_err(|e| {
-                raise_MemoryError(
+                raise_MemoryException(
                     format!("Failed to write memory {address:X}[{size:X}]"),
                     e,
                 )
             })
     }
 
-    fn reg_list(&self) -> PyResult<IndexMap<String, (u32, u8)>> {
+    pub fn reg_list(&self) -> PyResult<IndexMap<String, (u32, u8)>> {
         let mut result = IndexMap::new();
-        let sleigh = self.vm.cpu.sleigh();
+        let sleigh = &self.vm.cpu.arch.sleigh;
         for reg in &sleigh.named_registers {
             let name = sleigh.get_str(reg.name);
             result.insert(name.to_string(), (reg.offset, reg.var.size));
         }
-        return Ok(result);
+        Ok(result)
     }
 
-    fn reg_offset(&self, name: &str) -> PyResult<u32> {
+    pub fn reg_offset(&self, name: &str) -> PyResult<u32> {
         Ok(reg_find(self, name)?.offset)
     }
 
-    fn reg_size(&self, name: &str) -> PyResult<u8> {
+    pub fn reg_size(&self, name: &str) -> PyResult<u8> {
         Ok(reg_find(self, name)?.var.size)
     }
 
-    fn reg_read(&mut self, name: &str) -> PyResult<u64> {
+    pub fn reg_read(&mut self, name: &str) -> PyResult<u64> {
         Ok(self.vm.cpu.read_reg(reg_find(self, name)?.var))
     }
 
-    fn reg_write(&mut self, name: &str, value: u64) -> PyResult<()> {
-        Ok(self.vm.cpu.write_reg(reg_find(self, name)?.var, value))
+    pub fn reg_write(&mut self, name: &str, value: u64) -> PyResult<()> {
+        let var = reg_find(self, name)?.var;
+        if var == self.vm.cpu.arch.reg_pc {
+            self.vm.cpu.write_pc(value);
+        } else {
+            self.vm.cpu.write_reg(var, value);
+        }
+        Ok(())
     }
 
-    fn reset(&mut self) {
+    pub fn reset(&mut self) {
         self.vm.reset();
     }
 
-    fn run(&mut self) -> RunStatus {
+    pub fn run(&mut self) -> RunStatus {
         match self.vm.run() {
             VmExit::Running => RunStatus::Running,
             VmExit::InstructionLimit => RunStatus::InstructionLimit,
@@ -465,7 +548,7 @@ impl Icicle {
         }
     }
 
-    fn run_until(&mut self, address: u64) -> RunStatus {
+    pub fn run_until(&mut self, address: u64) -> RunStatus {
         let breakpoint_added = self.vm.add_breakpoint(address);
         let status = self.run();
         if breakpoint_added {
@@ -474,7 +557,7 @@ impl Icicle {
         status
     }
 
-    fn step(&mut self, count: u64) -> RunStatus {
+    pub fn step(&mut self, count: u64) -> RunStatus {
         let old_limit = self.vm.icount_limit;
         self.vm.icount_limit = self.vm.cpu.icount.saturating_add(count);
         let status = self.run();
@@ -482,17 +565,17 @@ impl Icicle {
         status
     }
 
-    fn add_breakpoint(&mut self, address: u64) -> bool {
+    pub fn add_breakpoint(&mut self, address: u64) -> bool {
         self.vm.add_breakpoint(address)
     }
 
-    fn remove_breakpoint(&mut self, address: u64) -> bool {
+    pub fn remove_breakpoint(&mut self, address: u64) -> bool {
         self.vm.remove_breakpoint(address)
     }
 }
 
 #[pyfunction]
-fn architectures() -> PyResult<Vec<&'static str>> {
+pub fn architectures() -> PyResult<Vec<&'static str>> {
     Ok(vec![
         "i686",
         "x86_64",
@@ -504,11 +587,11 @@ fn architectures() -> PyResult<Vec<&'static str>> {
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
 #[pymodule]
-fn icicle(_: Python<'_>, m: &PyModule) -> PyResult<()> {
+fn icicle(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(architectures, m)?)?;
     m.add_class::<Icicle>()?;
     m.add_class::<MemoryProtection>()?;
-    m.add_class::<MemoryErrorCode>()?;
+    m.add_class::<MemoryExceptionCode>()?;
     m.add_class::<RunStatus>()?;
     m.add_class::<ExceptionCodePy>()?;
     Ok(())

src/main.rs

@@ -1,379 +0,0 @@
-use std::borrow::Cow;
-use std::collections::HashMap;
-use icicle_cpu::mem::{Mapping, perm};
-use icicle_cpu::{VmExit};
-use icicle_vm;
-use icicle_vm::linux::LinuxCpu;
-use target_lexicon;
-use indexmap::IndexMap;
-use sleigh_runtime::NamedRegister;
-
-#[allow(dead_code)]
-#[derive(Clone)]
-enum MemoryProtection {
-    NoAccess,
-    ReadOnly,
-    ReadWrite,
-    ExecuteOnly,
-    ExecuteRead,
-    ExecuteReadWrite,
-}
-
-struct Icicle {
-    vm: icicle_vm::Vm,
-    regs: HashMap<String, NamedRegister>,
-}
-
-fn convert_protection(protection: MemoryProtection) -> u8 {
-    match protection {
-        MemoryProtection::NoAccess => perm::NONE,
-        MemoryProtection::ReadOnly => perm::READ,
-        MemoryProtection::ReadWrite => perm::READ | perm::WRITE,
-        MemoryProtection::ExecuteOnly => perm::EXEC,
-        MemoryProtection::ExecuteRead => perm::EXEC | perm::READ,
-        MemoryProtection::ExecuteReadWrite => perm::EXEC | perm::READ | perm::WRITE,
-    }
-}
-
-fn reg_find<'a>(i: &'a Icicle, name: &str) -> Result<&'a NamedRegister, String> {
-    let sleigh = i.vm.cpu.sleigh();
-    match sleigh.get_reg(name) {
-        None => {
-            i.regs.get(name.to_lowercase().as_str())
-                .ok_or(
-                    format!("Register not found: {name}")
-                )
-        }
-        Some(r) => Ok(r),
-    }
-}
-
-impl Icicle {
-    #[allow(dead_code)]
-    fn get_icount_limit(&mut self) -> u64 {
-        self.vm.icount_limit
-    }
-
-    #[allow(dead_code)]
-    fn set_icount_limit(&mut self, value: u64) {
-        self.vm.icount_limit = value;
-    }
-
-    #[allow(dead_code)]
-    fn get_icount(&mut self) -> u64 {
-        return self.vm.cpu.icount;
-    }
-
-    #[allow(dead_code)]
-    fn set_icount(&mut self, value: u64) {
-        self.vm.cpu.icount = value;
-    }
-
-    fn new(
-        architecture: String,
-        jit: bool,
-        jit_mem: bool,
-        shadow_stack: bool,
-        recompilation: bool,
-        track_uninitialized: bool,
-        optimize_instructions: bool,
-        optimize_block: bool,
-        tracing: bool,
-    ) -> Result<Self, String> {
-        // Prevent mixing '_' and '-'
-        if architecture.split("-").count() != 1 {
-            return Err(
-                format!("Bad architecture format: {architecture}")
-            );
-        }
-
-        // TODO: support instantiating this multiple times
-        if tracing {
-            tracing_subscriber::fmt()
-                .with_max_level(tracing::Level::DEBUG)
-                .with_target(false)
-                .init();
-        }
-
-        // Setup the CPU state for the target triple
-        let mut config = icicle_vm::cpu::Config::from_target_triple(
-            format!("{architecture}-none").as_str()
-        );
-        if config.triple.architecture == target_lexicon::Architecture::Unknown {
-            return Err(
-                format!("Unknown architecture: {architecture}")
-            );
-        }
-
-        // Configuration
-        config.enable_jit = jit;
-        config.enable_jit_mem = jit_mem;
-        config.enable_shadow_stack = shadow_stack;
-        config.enable_recompilation = recompilation;
-        config.track_uninitialized = track_uninitialized;
-        config.optimize_instructions = optimize_instructions;
-        config.optimize_block = optimize_block;
-
-        let vm = icicle_vm::build(&config)
-            .map_err(|e| {
-                format!("VM build error: {e}")
-            })?;
-
-        // Populate the lowercase register map
-        let mut regs = HashMap::new();
-        let sleigh = vm.cpu.sleigh();
-        for reg in &sleigh.named_registers {
-            let name = sleigh.get_str(reg.name);
-            regs.insert(name.to_lowercase(), reg.clone());
-        }
-
-        Ok(Icicle {
-            vm,
-            regs,
-        })
-    }
-
-    fn __str__(&mut self) -> String {
-        let arch = &self.vm.cpu.arch;
-        let endianness = if arch.sleigh.big_endian {
-            "big endian"
-        } else {
-            "little endian"
-        };
-        format!("Icicle VM for {0:?} ({endianness})", arch.triple.architecture)
-    }
-
-    fn mem_map(&mut self, address: u64, size: u64, protection: MemoryProtection) -> Result<(), String> {
-        let mapping = Mapping {
-            perm: convert_protection(protection),
-            value: 0,
-        };
-        if self.vm.cpu.mem.map_memory_len(address, size, mapping) {
-            Ok(())
-        } else {
-            Err(
-                format!("Failed to map memory {address:X}[{size:X}]")
-            )
-        }
-    }
-
-    #[allow(dead_code)]
-    fn mem_unmap(&mut self, address: u64, size: u64) -> Result<(), String> {
-        if self.vm.cpu.mem.unmap_memory_len(address, size) {
-            Ok(())
-        } else {
-            Err(
-                format!("Failed to unmap memory {address:X}[{size:X}]")
-            )
-        }
-    }
-
-    fn mem_protect(&mut self, address: u64, size: usize, protection: MemoryProtection) -> Result<(), String> {
-        self.vm.cpu.mem.update_perm(address, size as u64, convert_protection(protection))
-            .map_err(|_| {
-                format!("Failed to protect memory {address:X}[{size:X}]")
-            })?;
-        Ok(())
-    }
-
-    #[allow(dead_code)]
-    fn mem_read(&mut self, address: u64, size: usize) -> Result<Cow<[u8]>, String> {
-        // Allocate a buffer
-        let mut buffer = Vec::with_capacity(size);
-        buffer.resize(size, 0);
-
-        // Read the memory
-        match self.vm.cpu.mem.read_bytes(address, &mut buffer[..], perm::NONE) {
-            Ok(_) => Ok(Cow::Owned(buffer)),
-            Err(_) => Err(format!("Failed to read memory {address:X}[{size:X}]"))
-        }
-    }
-
-    fn mem_write(&mut self, address: u64, data: Vec<u8>) -> Result<(), String> {
-        let size = data.len();
-        match self.vm.cpu.mem.write_bytes(address, &data[..], perm::NONE) {
-            Ok(_) => Ok(()),
-            Err(_) => Err(format!("Failed to write memory {address:X}[{size:X}]"))
-        }
-    }
-
-    #[allow(dead_code)]
-    fn reg_list(&self) -> IndexMap<String, (u32, u8)> {
-        let mut result = IndexMap::new();
-        let sleigh = self.vm.cpu.sleigh();
-        for reg in &sleigh.named_registers {
-            let name = sleigh.get_str(reg.name);
-            result.insert(name.to_string(), (reg.offset, reg.var.size));
-        }
-        return result;
-    }
-
-    #[allow(dead_code)]
-    fn reg_offset(&self, name: &str) -> Result<u32, String> {
-        Ok(reg_find(self, name)?.offset)
-    }
-
-    #[allow(dead_code)]
-    fn reg_size(&self, name: &str) -> Result<u8, String> {
-        Ok(reg_find(self, name)?.var.size)
-    }
-
-    fn reg_read(&mut self, name: &str) -> Result<u64, String> {
-        Ok(self.vm.cpu.read_reg(reg_find(self, name)?.var))
-    }
-
-    fn reg_write(&mut self, name: &str, value: u64) -> Result<(), String> {
-        Ok(self.vm.cpu.write_reg(reg_find(self, name)?.var, value))
-    }
-
-    #[allow(dead_code)]
-    fn reset(&mut self) {
-        self.vm.reset();
-    }
-
-    fn run(&mut self) -> VmExit {
-        self.vm.run()
-    }
-
-    #[allow(dead_code)]
-    fn run_until(&mut self, address: u64) -> VmExit {
-        let breakpoint_added = self.vm.add_breakpoint(address);
-        let status = self.run();
-        if breakpoint_added {
-            self.vm.remove_breakpoint(address);
-        }
-        status
-    }
-
-    #[allow(dead_code)]
-    fn step(&mut self, count: u64) -> VmExit {
-        let old_limit = self.vm.icount_limit;
-        self.vm.icount_limit = self.vm.cpu.icount.saturating_add(count);
-        let status = self.run();
-        self.vm.icount_limit = old_limit;
-        status
-    }
-
-    #[allow(dead_code)]
-    fn add_breakpoint(&mut self, address: u64) -> bool {
-        self.vm.add_breakpoint(address)
-    }
-
-    #[allow(dead_code)]
-    fn remove_breakpoint(&mut self, address: u64) -> bool {
-        self.vm.remove_breakpoint(address)
-    }
-}
-
-#[allow(dead_code)]
-fn architectures() -> Result<Vec<&'static str>, String> {
-    Ok(vec![
-        "i686",
-        "x86_64",
-        "aarch64",
-    ])
-}
-
-fn nx_start() {
-    let mut vm = Icicle::new(
-        "x86_64".to_string(),
-        false,
-        true,
-        true,
-        true,
-        false,
-        true,
-        true,
-        false,
-    ).unwrap();
-    let page = 0x10000;
-    vm.mem_map(page, 0x1000, MemoryProtection::ReadOnly).unwrap();
-    // <non-executable memory> inc eax; ret
-    vm.mem_write(page, b"\xFF\xC0\xC3".to_vec()).unwrap();
-    vm.reg_write("rip", page).unwrap();
-    let status = vm.run();
-    println!("status: {:?}", status);
-    println!("rip: {:#x}", vm.reg_read("rip").unwrap());
-}
-
-fn nx_middle() {
-    let mut vm = Icicle::new(
-        "x86_64".to_string(),
-        false,
-        true,
-        true,
-        true,
-        false,
-        true,
-        true,
-        false,
-    ).unwrap();
-    let page = 0x10000;
-    vm.mem_map(page, 0x2000, MemoryProtection::ExecuteRead).unwrap();
-    vm.mem_protect(page + 0x1000, 0x1000, MemoryProtection::ReadOnly).unwrap();
-    // inc eax; inc eax; <transition to non-executable region>; ret
-    let rip = page + 0x1000 - 2;
-    vm.mem_write(rip, b"\xFF\xC0\xFF\xC0\xC3".to_vec()).unwrap();
-    vm.reg_write("rip", rip).unwrap();
-    let status = vm.run();
-    println!("status: {:?}", status);
-    println!("rip: {:#x}", vm.reg_read("rip").unwrap());
-    println!("rax: {:#x}", vm.reg_read("rax").unwrap());
-}
-
-fn inv_start() {
-    let mut vm = Icicle::new(
-        "x86_64".to_string(),
-        false,
-        true,
-        true,
-        true,
-        false,
-        true,
-        true,
-        false,
-    ).unwrap();
-    let page = 0x10000;
-    vm.mem_map(page, 0x1000, MemoryProtection::ExecuteRead).unwrap();
-    // <invalid>; ret
-    vm.mem_write(page, b"\xFF\xFF\xC3".to_vec()).unwrap();
-    vm.reg_write("rip", page).unwrap();
-    let status = vm.run();
-    println!("status: {:?}", status);
-    println!("rip: {:#x}", vm.reg_read("rip").unwrap());
-    println!("rax: {:#x}", vm.reg_read("rax").unwrap());
-}
-
-fn inv_middle() {
-    let mut vm = Icicle::new(
-        "x86_64".to_string(),
-        false,
-        true,
-        true,
-        true,
-        false,
-        true,
-        true,
-        false,
-    ).unwrap();
-    let page = 0x10000;
-    vm.mem_map(page, 0x1000, MemoryProtection::ExecuteRead).unwrap();
-    // inc eax; <invalid>; ret
-    vm.mem_write(page, b"\xFF\xC0\xFF\xFF\xC3".to_vec()).unwrap();
-    vm.reg_write("rip", page).unwrap();
-    let status = vm.run();
-    println!("status: {:?}", status);
-    println!("rip: {:#x}", vm.reg_read("rip").unwrap());
-    println!("rax: {:#x}", vm.reg_read("rax").unwrap());
-}
-
-fn main() {
-    println!("=== NX (block start) ===");
-    nx_start();
-    println!("=== NX (block middle) ===");
-    nx_middle();
-    println!("=== Invalid instruction (block start) ===");
-    inv_start();
-    println!("=== Invalid instruction (block middle) ===");
-    inv_middle();
-}

tests/.cargo/config.toml

@@ -0,0 +1,2 @@
+[env]
+GHIDRA_SRC = { value = "../python/icicle", relative = true }

tests/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "tests"
+edition = "2021"
+
+[[bin]]
+name = "tests"
+path = "tests.rs"
+
+[dependencies]
+icicle-python = { path = "..", default-features = false }
+pyo3 = { workspace = true, features = ["auto-initialize"] }

tests/blockopt.py

@@ -0,0 +1,36 @@
+from icicle import *
+
+def main():
+    emu = Icicle("x86_64", jit=True, optimize_block=True, tracing=True)
+
+    instructions = bytes.fromhex("41 C1 EA 07 41 83 E2 1F 74 08 44 89 D0 48 89 54 C6 08 49 83 C1 04 4C 89 0E 4C 89 C9 44 8B 11 44 89 D0 F7 D0 49 89 C9 A8 03 0F 84 88 F6 FF FF EB 4C".replace(" ", ""))
+
+    addr = 0x140001A73
+    heap = 0x71000
+    emu.mem_map(heap, 0x1000, MemoryProtection.ReadWrite)
+    emu.mem_write(heap + 4, 0x1337.to_bytes(4, "little"))
+    emu.mem_map(addr & ~0xFFF, 0x1000, MemoryProtection.ExecuteRead)
+    emu.mem_write(addr, instructions)
+    emu.reg_write("r9", heap)
+    emu.reg_write("r10", 0x13)
+    emu.reg_write("rip", addr)
+    emu.reg_write("rsi", heap + 0x100)
+
+    for i in range(11):
+        rip = emu.reg_read("rip")
+        rcx = emu.reg_read("rcx")
+        r9 = emu.reg_read("r9")
+        print(f"[{i}] RIP: {hex(rip)}, RCX: {hex(rcx)}, R9: {hex(r9)}")
+
+        if rip == 0x140001A8F:
+            emu.reg_write("r9", 0x13370900)
+
+        emu.step(1)
+
+        if rip == 0x140001A9A:
+            assert rcx == r9, f"expected rcx({hex(rcx)}) == r9({hex(r9)})"
+
+    print("Everything works!")
+
+if __name__ == "__main__":
+    main()

tests/example.py

@@ -49,12 +49,12 @@ def old_test():
     vm.mem_map(addr, 0x1000, icicle.MemoryProtection.NoAccess)
     try:
         vm.mem_protect(addr, 0x2000, icicle.MemoryProtection.ExecuteRead)
-    except icicle.MemoryError as x:
-        print("MemoryError working!")
+    except icicle.MemoryException as x:
+        print("MemoryException working!")
         message = x.args[0]
         print(message, x.code)
     except Exception as x:
-        print(x, type(x), isinstance(x, icicle.MemoryError))
+        print(x, type(x), isinstance(x, icicle.MemoryException))
     vm.mem_protect(addr, 0x1000, icicle.MemoryProtection.ExecuteRead)
     vm.mem_write(addr, b"\x90\x90\x90")
     data = vm.mem_read(addr, 4)

tests/tests.rs

@@ -0,0 +1,397 @@
+#![allow(unused)]
+
+use icicle::*;
+use pyo3::PyResult;
+use std::process::exit;
+
+// NOTE: https://github.com/rust-lang/rust-analyzer/issues/18752
+#[test]
+fn example() -> PyResult<()> {
+    Err(pyo3::exceptions::PyException::new_err("test"))
+}
+
+fn new_vm(jit: bool) -> PyResult<Icicle> {
+    Icicle::new(
+        "x86_64".to_string(),
+        jit,
+        true,
+        false,
+        true,
+        false,
+        true,
+        false,
+        false,
+    )
+}
+
+fn new_trace_vm(jit: bool) -> PyResult<Icicle> {
+    Icicle::new(
+        "x86_64".to_string(),
+        jit,
+        true,
+        false,
+        true,
+        false,
+        true,
+        false,
+        true,
+    )
+}
+
+fn nx_start() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    let page = 0x10000;
+    vm.mem_map(page, 0x1000, MemoryProtection::ReadOnly)?;
+    // <non-executable memory> inc eax; ret
+    vm.mem_write(page, b"\xFF\xC0\xC3".to_vec())?;
+    vm.reg_write("rip", page)?;
+    let status = vm.run();
+    println!("status: {:?}", status);
+    println!("rip: {:#x}", vm.reg_read("rip")?);
+
+    Ok(())
+}
+
+fn nx_middle() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    let page = 0x10000;
+    vm.mem_map(page, 0x2000, MemoryProtection::ExecuteRead)?;
+    vm.mem_protect(page + 0x1000, 0x1000, MemoryProtection::ReadOnly)?;
+    // inc eax; inc eax; <transition to non-executable region>; ret
+    let rip = page + 0x1000 - 2;
+    vm.mem_write(rip, b"\xFF\xC0\xFF\xC0\xC3".to_vec())?;
+    vm.reg_write("rip", rip)?;
+    let status = vm.run();
+    println!("status: {:?}", status);
+    println!("rip: {:#x}", vm.reg_read("rip")?);
+    println!("rax: {:#x}", vm.reg_read("rax")?);
+
+    Ok(())
+}
+
+fn inv_start() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    let page = 0x10000;
+    vm.mem_map(page, 0x1000, MemoryProtection::ExecuteRead)?;
+    // <invalid>; ret
+    vm.mem_write(page, b"\xFF\xFF\xC3".to_vec())?;
+    vm.reg_write("rip", page)?;
+    let status = vm.run();
+    println!("status: {:?}", status);
+    println!("rip: {:#x}", vm.reg_read("rip")?);
+    println!("rax: {:#x}", vm.reg_read("rax")?);
+
+    Ok(())
+}
+
+fn inv_middle() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    let page = 0x10000;
+    vm.mem_map(page, 0x1000, MemoryProtection::ExecuteRead)?;
+    // inc eax; <invalid>; ret
+    vm.mem_write(page, b"\xFF\xC0\xFF\xFF\xC3".to_vec())?;
+    vm.reg_write("rip", page)?;
+    let status = vm.run();
+    println!("status: {:?}", status);
+    println!("rip: {:#x}", vm.reg_read("rip")?);
+    println!("rax: {:#x}", vm.reg_read("rax")?);
+
+    Ok(())
+}
+
+fn block_optimization() -> PyResult<()> {
+    let mut vm: Icicle = Icicle::new(
+        "x86_64".to_string(),
+        true,
+        true,
+        true,
+        true,
+        false,
+        true,
+        true,
+        false,
+    )?;
+
+    // Memory setup
+    let addr: u64 = 0x140001A73;
+    let heap: u64 = 0x71000;
+    vm.mem_map(heap, 0x1000, MemoryProtection::ReadWrite)?;
+    vm.mem_write(heap + 4, b"\x37\x13\x00\x00".to_vec())?;
+    vm.mem_map(addr & !0xFFF, 0x1000, MemoryProtection::ExecuteRead)?;
+    vm.mem_write(addr, b"\x41\xc1\xea\x07\x41\x83\xe2\x1f\x74\x08\x44\x89\xd0\x48\x89\x54\xc6\x08\x49\x83\xc1\x04\x4c\x89\x0e\x4c\x89\xc9\x44\x8b\x11\x44\x89\xd0\xf7\xd0\x49\x89\xc9\xa8\x03\x0f\x84\x88\xf6\xff\xff\xeb\x4c\xcc".to_vec())?;
+
+    // Register setup
+    vm.reg_write("r9", heap)?;
+    vm.reg_write("r10", 0x13)?;
+    vm.reg_write("rip", addr)?;
+    vm.reg_write("rsi", heap + 0x100)?;
+
+    // Step through instructions
+    for i in 0..11 {
+        let rip = vm.reg_read("rip")?;
+        let rcx = vm.reg_read("rcx")?;
+        let r9 = vm.reg_read("r9")?;
+
+        println!("[{}] RIP: {:#x}, RCX: {:#x}, R9: {:#x}", i, rip, rcx, r9);
+
+        if rip == 0x140001A8F {
+            vm.reg_write("r9", 0x13370900)?;
+        }
+
+        vm.step(1);
+
+        if rip == 0x140001A9A {
+            if rcx != r9 {
+                println!("[BUG] expected rcx({:#x}) == r9({:#x})", rcx, r9);
+            } else {
+                println!("Everything works!");
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn rewind() -> PyResult<()> {
+    let mut vm = new_vm(true)?;
+
+    vm.mem_map(0x100, 0x20, MemoryProtection::ExecuteRead)?;
+    vm.mem_map(0x200, 0x20, MemoryProtection::ReadOnly)?;
+
+    vm.mem_write(0x100, b"\x55\xCC".to_vec())?; // push rbp
+    vm.reg_write("rbp", 0xF00)?;
+    vm.reg_write("rsp", 0x210)?;
+    vm.reg_write("rip", 0x100)?;
+    let status = vm.step(1);
+    println!("run status      : {:?}", status);
+    println!("exception code  : {:?}", vm.get_exception_code());
+    println!("exception value : {:#x}", vm.get_exception_value());
+    println!("stack pointer   : {:#x}", vm.reg_read("rsp")?);
+
+    Ok(())
+}
+
+fn execute_uninitialized() -> PyResult<()> {
+    let mut vm = Icicle::new(
+        "x86_64".to_string(),
+        true,
+        true,
+        false,
+        true,
+        true, // NOTE: setting this to true is not properly supported
+        true,
+        false,
+        false,
+    )?;
+
+    // \x48\x8d\x05\x01\x00\x00\x00\x90\x8a\x18\x90
+
+    vm.mem_map(0x100, 0x20, MemoryProtection::ExecuteOnly)?;
+    vm.mem_write(0x100, b"\x90\xFF\xC0".to_vec())?; // inc eax
+    vm.reg_write("rip", 0x100)?;
+    {
+        println!("[pre1] icount: {}", vm.get_icount());
+        let status = vm.step(2);
+        // NOTE: the real reason is that INIT is not set
+        println!("run status      : {:?}", status);
+        println!("exception code  : {:?}", vm.get_exception_code());
+        println!("exception value : {:#x}", vm.get_exception_value());
+        println!("rax             : {:#x}", vm.reg_read("rax")?);
+    }
+
+    {
+        println!("[pre2] icount: {}", vm.get_icount());
+        let status = vm.step(1);
+        // NOTE: the real reason is that INIT is not set
+        println!("run status      : {:?}", status);
+        println!("exception code  : {:?}", vm.get_exception_code());
+        println!("exception value : {:#x}", vm.get_exception_value());
+        println!("rax             : {:#x}", vm.reg_read("rax")?);
+        println!("[post] icount: {}", vm.get_icount());
+    }
+
+    // TODO: status is now UnhandledException, should be InstructionLimit
+    // on the next stpe it should be UnhandledException -> ExecViolation
+
+    Ok(())
+}
+
+fn execute_only() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+
+    vm.mem_map(0x100, 0x20, MemoryProtection::ExecuteOnly)?;
+    /*
+    0x100: lea rax, [rip]
+    0x107: nop
+    0x108: mov bl, byte ptr [rax]
+    0x10A: int3
+    */
+    vm.mem_write(
+        0x100,
+        b"\x48\x8d\x05\x00\x00\x00\x00\x90\x8a\x18\xCC".to_vec(),
+    )?; // nop
+    vm.reg_write("rip", 0x100)?;
+    vm.step(2);
+    let status = vm.step(1);
+    // NOTE: the real reason is that INIT is not set
+    println!("run status      : {:?}", status);
+    println!("exception code  : {:?}", vm.get_exception_code());
+    println!("exception value : {:#x}", vm.get_exception_value());
+    println!("bl: {:#x}", vm.reg_read("bl")?);
+    println!("rip: {:#x}", vm.reg_read("rip")?);
+
+    Ok(())
+}
+
+fn self_modifying() -> PyResult<()> {
+    // TODO: add a self-modifying code check (where the previously-executed code is written to)
+
+    Ok(())
+}
+
+fn step_modify_rip() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    vm.mem_map(0x100, 0x20, MemoryProtection::ExecuteRead)?;
+
+    // 0x100:  48 01 d8                add    rax,rbx
+    // 0x103:  48 83 e9 05             sub    rcx,0x5
+    // 0x107:  48 89 d9                mov    rcx,rbx
+    // 0x10a:  90                      nop
+    // 0x10b:  90                      nop
+    vm.mem_write(
+        0x100,
+        b"\x48\x01\xD8\x48\x83\xE9\x05\x48\x89\xD9\x90\x90\xCC".to_vec(),
+    )?;
+
+    vm.reg_write("rax", 0xF00)?;
+    vm.reg_write("rbx", 0x210)?;
+    vm.reg_write("rip", 0x100)?;
+
+    println!("starting run at {:#x}", vm.reg_read("rip")?);
+    let mut status = vm.step(1);
+
+    println!(
+        "ending run at {:#x} (status: {:?})",
+        vm.reg_read("rip")?,
+        status
+    );
+    vm.reg_write("rip", 0x100)?;
+    //vm.write_pc(0x100);
+    //println!("pc: {:#x}", vm.read_pc());
+    println!("rip rewritten {:#x}", vm.reg_read("rip")?);
+    status = vm.step(1);
+    println!(
+        "ending run at {:#x} (status: {:?})",
+        vm.reg_read("rip")?,
+        status
+    );
+
+    Ok(())
+}
+
+fn eflags_reconstruction() -> PyResult<()> {
+    let mut vm = new_vm(false)?;
+    vm.mem_map(0x100, 0x20, MemoryProtection::ExecuteRead)?;
+
+    vm.mem_write(0x100, b"\x48\x01\xD8\xCC".to_vec())?;
+    vm.reg_write("rax", 0x7FFFFFFFFFFFFFFF)?;
+    vm.reg_write("rbx", 0x1)?;
+
+    let of_mask = (1 << 11) as u64;
+
+    {
+        let eflags = vm.reg_read("eflags")?;
+        let of = vm.reg_read("OF")?;
+        let of_set = (eflags & of_mask) == of_mask;
+        println!("[pre] eflags: {:#x}, OF: {:#x} == {}", eflags, of, of_set);
+    }
+
+    vm.set_pc(0x100);
+    let status = vm.step(1);
+    println!("run status: {:?}", status);
+
+    {
+        let eflags = vm.reg_read("eflags")?;
+        let rflags = vm.reg_read("rflags")?;
+        let of = vm.reg_read("OF")?;
+        let of_set = (eflags & of_mask) == of_mask;
+        println!(
+            "[post] eflags: {:#x} == {:#x}, OF: {:#x} == {}",
+            eflags, rflags, of, of_set
+        );
+    }
+
+    {
+        vm.reg_write("OF", 0)?;
+        let eflags = vm.reg_read("eflags")?;
+        let of = vm.reg_read("OF")?;
+        let of_set = (eflags >> 11) & 1;
+        println!("[OF=0] eflags: {:#x}, OF: {:#x} == {}", eflags, of, of_set);
+    }
+
+    {
+        let mut eflags = vm.reg_read("eflags")?;
+        eflags |= of_mask;
+        vm.reg_write("rflags", eflags)?;
+        let of = vm.reg_read("OF")?;
+        let of_set = (eflags >> 11) & 1;
+        println!(
+            "[rflags|={:#x}] eflags: {:#x}, OF: {:#x} == {}",
+            of_mask, eflags, of, of_set
+        );
+    }
+
+    Ok(())
+}
+
+fn main() {
+    // Make sure the GHIDRA_SRC environment variable is valid
+    match std::env::var("GHIDRA_SRC") {
+        Ok(ghidra_src) => {
+            // Make sure the directory $GHIDRA_SRC/Ghidra/Processors exists
+            if !std::path::Path::new(&ghidra_src)
+                .join("Ghidra")
+                .join("Processors")
+                .exists()
+            {
+                println!("GHIDRA_SRC environment variable invalid!");
+                exit(1);
+            }
+            println!("GHIDRA_SRC: {}", ghidra_src);
+        }
+        Err(_) => {
+            println!("GHIDRA_SRC environment variable not set!");
+            exit(1);
+        }
+    }
+
+    let tests: Vec<(&str, fn() -> PyResult<()>)> = vec![
+        //("NX (block start)", nx_start),
+        //("NX (block middle)", nx_middle),
+        //("Invalid instruction (block start)", inv_start),
+        //("Invalid instruction (block middle)", inv_middle),
+        //("Block optimization bug", block_optimization),
+        //("Rewind", rewind),
+        //("Execute only", execute_only),
+        ("Execute uninitialized", execute_uninitialized),
+        //("Step modify rip", step_modify_rip),
+        //("EFlags reconstruction", eflags_reconstruction),
+    ];
+
+    let mut success = 0;
+    for (name, f) in tests.iter() {
+        println!("=== {} ===", name);
+        match f() {
+            Ok(_) => {
+                success += 1;
+                println!("[OK]");
+            }
+            Err(e) => {
+                println!("[ERROR] {}", e);
+            }
+        }
+    }
+
+    println!("{}/{} tests passed", success, tests.len());
+    exit(if success == tests.len() { 0 } else { 1 });
+}