Treat only last occurrence of trace terminating block as a stop point (…

…angr#2992)

in SimEngineUnicorn

angr/engines/unicorn.py

@@ -169,6 +169,7 @@ def process_successors(self, successors, **kwargs):
             return super().process_successors(successors, **kwargs)
 
         extra_stop_points = kwargs.get('extra_stop_points', None)
+        last_block_details = kwargs.get('last_block_details', None)
         step = kwargs.get('step', None)
         if extra_stop_points is None:
             extra_stop_points = set(self.project._sim_procedures)
@@ -210,6 +211,8 @@ def process_successors(self, successors, **kwargs):
 
         try:
             state.unicorn.set_stops(extra_stop_points)
+            if last_block_details is not None:
+                state.unicorn.set_last_block_details(last_block_details)
             state.unicorn.set_tracking(track_bbls=o.UNICORN_TRACK_BBL_ADDRS in state.options,
                                        track_stack=o.UNICORN_TRACK_STACK_POINTERS in state.options)
             state.unicorn.hook()

angr/exploration_techniques/tracer.py

@@ -173,6 +173,10 @@ def __init__(self,
         # whether we should follow the trace
         self._no_follow = self._trace is None
 
+        # Keep track of count of termination point
+        self._last_block_total_count = self._trace.count(self._trace[-1])
+        self._last_block_seen_count = 0
+
         # sanity check: copy_states must be enabled in Permissive mode since we may need to backtrack from a previous
         # state.
         if self._mode == TracingMode.Permissive and not self._copy_states:
@@ -343,8 +347,17 @@ def step_state(self, simgr, state, **kwargs):
                 self.project.hook(state.addr, RepHook(insn.mnemonic.split(" ")[1]).run, length=insn.size)
 
         # perform the step. ask qemu to stop at the termination point.
-        stops = set(kwargs.pop('extra_stop_points', ())) | {self._trace[-1]}
-        succs_dict = simgr.step_state(state, extra_stop_points=stops, **kwargs)
+        # if termination point occurs multiple times in trace, pass details to SimEngineUnicorn's native interface so
+        # that it can stop at last block
+        if self._last_block_total_count > 1:
+            stops = set(kwargs.pop('extra_stop_points', ()))
+            last_block_details = {"addr": self._trace[-1], "tot_count": self._last_block_total_count,
+                                  "curr_count": self._last_block_seen_count}
+        else:
+            stops = set(kwargs.pop('extra_stop_points', ())) | {self._trace[-1]}
+            last_block_details = None
+
+        succs_dict = simgr.step_state(state, extra_stop_points=stops, last_block_details=last_block_details, **kwargs)
         if None not in succs_dict and simgr.errored:
             raise simgr.errored[-1].error
         sat_succs = succs_dict[None]  # satisfiable states
@@ -473,6 +486,8 @@ def _update_state_tracking(self, state: 'angr.SimState'):
         sync = state.globals['sync_idx']
         timer = state.globals['sync_timer']
 
+        self._last_block_seen_count += state.history.recent_bbl_addrs.count(self._trace[-1])
+
         if state.history.recent_block_count > 1:
             # multiple blocks were executed this step. they should follow the trace *perfectly*
             # or else something is up

angr/state_plugins/unicorn_engine.py

@@ -364,6 +364,7 @@ def _setup_prototype_explicit(handle, func, restype, *argtypes):
         _setup_prototype(h, 'syscall_count', ctypes.c_uint64, state_t)
         _setup_prototype(h, 'step', ctypes.c_uint64, state_t)
         _setup_prototype(h, 'activate_page', None, state_t, ctypes.c_uint64, ctypes.c_void_p, ctypes.c_void_p)
+        _setup_prototype(h, 'set_last_block_details', None, state_t, ctypes.c_uint64, ctypes.c_int64, ctypes.c_int64)
         _setup_prototype(h, 'set_stops', None, state_t, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64))
         _setup_prototype(h, 'cache_page', ctypes.c_bool, state_t, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_char_p, ctypes.c_uint64)
         _setup_prototype(h, 'uncache_pages_touching_region', None, state_t, ctypes.c_uint64, ctypes.c_uint64)
@@ -675,6 +676,9 @@ def _setup_unicorn(self):
         if self.state.arch.uc_mode is None:
             raise SimUnicornUnsupport("unsupported architecture %r" % self.state.arch)
 
+    def set_last_block_details(self, details):
+        _UC_NATIVE.set_last_block_details(self._uc_state, details["addr"], details["curr_count"], details["tot_count"])
+
     def set_stops(self, stop_points):
         _UC_NATIVE.set_stops(self._uc_state,
             ctypes.c_uint64(len(stop_points)),

native/angr_native.def

@@ -14,6 +14,7 @@ EXPORTS
   simunicorn_syscall_count
   simunicorn_step
   simunicorn_activate_page
+  simunicorn_set_last_block_details
   simunicorn_set_stops
   simunicorn_set_map_callback
   simunicorn_cache_page

native/sim_unicorn.cpp

@@ -52,6 +52,9 @@ State::State(uc_engine *_uc, uint64_t cache_key):uc(_uc) {
 	mode = *((uc_mode*)((uc_arch*)uc + 1));
 	curr_block_details.reset();
 	symbolic_read_in_progress = false;
+	trace_last_block_addr = 0;
+	trace_last_block_tot_count = -1;
+	trace_last_block_curr_count = -1;
 }
 
 /*
@@ -204,6 +207,11 @@ void State::step(address_t current_address, int32_t size, bool check_stop_points
 		if (stop_point != stop_points.end() && *stop_point < current_address + real_size) {
 			stop(STOP_STOPPOINT);
 		}
+		else if ((trace_last_block_tot_count > 0) && (trace_last_block_addr >= current_address) &&
+		  (trace_last_block_addr < current_address + real_size) && (trace_last_block_curr_count == trace_last_block_tot_count - 1)) {
+			// Executing last block in trace. Stop.
+			stop(STOP_STOPPOINT);
+		}
 	}
 }
 
@@ -250,6 +258,9 @@ void State::commit() {
 		}
 		blocks_with_symbolic_instrs.emplace_back(curr_block_details);
 	}
+	if (curr_block_details.block_addr == trace_last_block_addr) {
+		trace_last_block_curr_count += 1;
+	}
 	// Clear all block level taint status trackers and symbolic instruction list
 	block_symbolic_registers.clear();
 	block_concrete_registers.clear();
@@ -388,6 +399,13 @@ mem_update_t *State::sync() {
 	return mem_updates_head;
 }
 
+void State::set_last_block_details(address_t block_addr, int64_t curr_count, int64_t tot_count) {
+	trace_last_block_addr = block_addr;
+	trace_last_block_curr_count = curr_count;
+	trace_last_block_tot_count = tot_count;
+	return;
+}
+
 void State::set_stops(uint64_t count, address_t *stops) {
 	stop_points.clear();
 	for (auto i = 0; i < count; i++) {
@@ -2384,6 +2402,11 @@ uint64_t simunicorn_step(State *state) {
 	return state->cur_steps;
 }
 
+extern "C"
+void simunicorn_set_last_block_details(State *state, address_t block_addr, uint64_t curr_count, uint64_t total_count) {
+	state->set_last_block_details(block_addr, curr_count, total_count);
+}
+
 extern "C"
 void simunicorn_set_stops(State *state, uint64_t count, uint64_t *stops)
 {

native/sim_unicorn.hpp

@@ -518,6 +518,10 @@ class State {
 	//std::map<uint64_t, taint_t *> active_pages;
 	std::set<uint64_t> stop_points;
 
+	address_t trace_last_block_addr;
+	int64_t trace_last_block_curr_count;
+	int64_t trace_last_block_tot_count;
+
 	address_t taint_engine_next_instr_address, taint_engine_stop_mem_read_instruction;
 	uint32_t taint_engine_stop_mem_read_size;
 	bool symbolic_read_in_progress;
@@ -734,6 +738,11 @@ class State {
 		 */
 		mem_update_t *sync();
 
+		/*
+		 * set details of the last block of trace
+		 */
+		void set_last_block_details(address_t block_addr, int64_t curr_count, int64_t tot_count);
+
 		/*
 		 * set a list of stops to stop execution at
 		 */