From 9e429da6f444286eff5f131406555d72477aa100 Mon Sep 17 00:00:00 2001
From: Kevin Schneider <kevin.schneider@tum.de>
Date: Tue, 25 Apr 2023 16:10:14 +0200
Subject: [PATCH] add unicorn dependency as submodule; remove modified tbs from
 cache during emulation

---
 .gitmodules                        | 3 +++
 emulation_worker/Cargo.lock        | 4 +---
 emulation_worker/Cargo.toml        | 4 +++-
 emulation_worker/src/hooks.rs      | 6 ++++++
 emulation_worker/src/hooks/util.rs | 8 ++++++++
 emulation_worker/unicorn           | 1 +
 6 files changed, 22 insertions(+), 4 deletions(-)
 create mode 160000 emulation_worker/unicorn

diff --git a/.gitmodules b/.gitmodules
index a1a8b17..9b861d2 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,3 +1,6 @@
 [submodule "qemu"]
 	path = qemu
 	url = https://github.com/Fraunhofer-AISEC/archie-qemu.git
+[submodule "emulation_worker/unicorn"]
+	path = emulation_worker/unicorn
+	url = https://github.com/unicorn-engine/unicorn.git
diff --git a/emulation_worker/Cargo.lock b/emulation_worker/Cargo.lock
index 8eaec9d..f52daec 100644
--- a/emulation_worker/Cargo.lock
+++ b/emulation_worker/Cargo.lock
@@ -440,9 +440,7 @@ checksum = "84a22b9f218b40614adcb3f4ff08b703773ad44fa9423e4e0d346d5db86e4ebc"
 
 [[package]]
 name = "unicorn-engine"
-version = "2.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f3b881bfd9837ff4f62e81a1e64b40a584604375ae0a73d0d5f09b7a72350b96"
+version = "2.0.2"
 dependencies = [
  "bitflags",
  "cc",
diff --git a/emulation_worker/Cargo.toml b/emulation_worker/Cargo.toml
index 8b89aaf..96d2a00 100644
--- a/emulation_worker/Cargo.toml
+++ b/emulation_worker/Cargo.toml
@@ -2,6 +2,8 @@
 name = "emulation_worker"
 version = "0.1.0"
 edition = "2021"
+authors = ["Kevin Schneider"]
+license = "Apache-2.0"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [lib]
@@ -10,7 +12,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = { version = "0.17.3", features = ["extension-module"] }
-unicorn-engine = "2.0.0"
+unicorn-engine = { path = "./unicorn" }
 num = "0.4.0"
 priority-queue = "1.3.1"
 capstone = "0.11.0"
diff --git a/emulation_worker/src/hooks.rs b/emulation_worker/src/hooks.rs
index a469900..866367a 100644
--- a/emulation_worker/src/hooks.rs
+++ b/emulation_worker/src/hooks.rs
@@ -212,6 +212,12 @@ fn fault_hook_cb(uc: &mut Unicorn<'_, ()>, address: u64, _size: u32, state: &Arc
                 apply_model(&data, fault).to_bytes_le().as_slice(),
             )
             .expect("failed writing fault data to memory");
+            if matches!(fault.kind, FaultType::Instruction) {
+                // We need to remove the tb containing the modified instructions from the cache
+                // since they might not have any effect otherwise
+                uc.ctl_remove_cache(fault.address, fault.address + fault_size as u64)
+                    .unwrap();
+            }
             dump_memory(
                 uc,
                 fault.address,
diff --git a/emulation_worker/src/hooks/util.rs b/emulation_worker/src/hooks/util.rs
index d45c253..b5b5234 100644
--- a/emulation_worker/src/hooks/util.rs
+++ b/emulation_worker/src/hooks/util.rs
@@ -95,6 +95,14 @@ pub fn undo_faults(
         FaultType::Data | FaultType::Instruction => {
             uc.mem_write(fault.address, prefault_data.to_bytes_le().as_slice())
                 .expect("failed restoring memory value");
+            if matches!(fault.kind, FaultType::Instruction) {
+                //uc.ctl_arg_2(UC_CTL_TB_REMOVE_CACHE | UC_CTL_IO_WRITE, unsafe { std::mem::transmute::<u64, *mut c_void>(address) }, unsafe { std::mem::transmute::<u64, *mut c_void>(fault.address + prefault_data.to_bytes_le().len() as u64) }).unwrap();
+                uc.ctl_remove_cache(
+                    fault.address,
+                    fault.address + prefault_data.to_bytes_le().len() as u64,
+                )
+                .unwrap();
+            }
         }
     }
 
diff --git a/emulation_worker/unicorn b/emulation_worker/unicorn
new file mode 160000
index 0000000..bde3cd7
--- /dev/null
+++ b/emulation_worker/unicorn
@@ -0,0 +1 @@
+Subproject commit bde3cd7dae991caadf38243771f70ab2c6e1a2b2