Cross-correlate events with the Focaccia tracing sequence

2 files changed, 54 insertions, 9 deletions

diff --git a/src/focaccia/deterministic.py b/src/focaccia/deterministic.py
index f920382..635a573 100644
--- a/src/focaccia/deterministic.py
+++ b/src/focaccia/deterministic.py
@@ -6,6 +6,7 @@ from typing import Union
 import brotli
 
 from .arch import Arch
+from .snapshot import ReadableProgramState
 
 try:
     import capnp
@@ -96,6 +97,12 @@ class Event:
         self.registers = registers
         self.mem_writes = memory_writes
 
+    def match(self, pc: int, target: ReadableProgramState) -> bool:
+        # TODO: match the rest of the state to be sure
+        if self.pc == pc:
+            return True
+        return False
+
     def __repr__(self) -> str:
         reg_repr = ''
         for reg, value in self.registers.items():
@@ -169,5 +176,17 @@ class DeterministicLog:
                           raw_event.event.which(),
                           registers, mem_writes)
             events.append(event)
-        return events
+
+        # deduplicate
+        deduped_events = []
+        for i in range(0, len(events), 2):
+            if events[i].event_type == 'syscall':
+                if events[i+1].pc == 0:
+                    deduped_events.append(events[i])
+                    break
+                if events[i+1].event_type != 'syscall':
+                    raise Exception(f'Event {events[i+1]} should follow {events[i]} but does not')
+                deduped_events.append(events[i+1])
+
+        return deduped_events
 
diff --git a/src/focaccia/symbolic.py b/src/focaccia/symbolic.py
index 81c78e6..fa20942 100644
--- a/src/focaccia/symbolic.py
+++ b/src/focaccia/symbolic.py
@@ -722,6 +722,9 @@ class SymbolicTracer:
         self.cross_validate = cross_validate
         self.target = SpeculativeTracer(self.create_debug_target())
 
+        self.nondet_events = self.env.detlog.events()
+        self.next_event: int | None = None
+
     def create_debug_target(self) -> LLDBConcreteTarget:
         binary = self.env.binary_name
         if self.remote is False:
@@ -779,9 +782,24 @@ class SymbolicTracer:
                                       f' mem[{hex(addr)}:{hex(addr+len(data))}] = {conc_data}.'
                                       f'\nFaulty transformation: {transform}')
 
-    def progress(self, new_pc, instruction: Instruction) -> int | None:
+    def is_stepping_instr(self, pc: int, instruction: Instruction):
+        if self.nondet_events:
+            if self.next_event and self.nondet_events[self.next_event].match(pc, self.target):
+                debug('Current instruction matches next event; stepping through it')
+                if self.next_event < len(self.nondet_events):
+                    self.next_event += 1
+                else:
+                    self.next_event = None
+
+                return True
+        else:
+            if self.target.arch.is_instr_syscall(str(instruction)):
+                return True
+        return False
+
+    def progress(self, new_pc, step: bool = False) -> int | None:
         self.target.speculate(new_pc)
-        if self.target.arch.is_instr_syscall(str(instruction)):
+        if step:
             self.target.progress_execution()
             if self.target.is_exited():
                 return None
@@ -801,13 +819,19 @@ class SymbolicTracer:
         if start_addr is not None:
             self.target.run_until(start_addr)
 
+        for i in range(len(self.nondet_events)):
+            if self.nondet_events[i].pc == self.target.read_pc():
+                self.next_event = i
+                debug(f'Starting from event {self.nondet_events[i]} onwards')
+                break
+
         ctx = DisassemblyContext(self.target)
         arch = ctx.arch
 
-        debug('Non-deterministic events handled:')
-        nondet_events = self.env.detlog.events()
-        for event in nondet_events:
-            debug(event)
+        if logger.isEnabledFor(logging.DEBUG):
+            debug('Tracing program with the following non-deterministic events')
+            for event in self.nondet_events:
+                debug(event)
 
         # Trace concolically
         strace: list[SymbolicTransform] = []
@@ -845,6 +869,8 @@ class SymbolicTracer:
                         continue
                     raise # forward exception
 
+            needs_step = self.is_stepping_instr(pc, instruction)
+
             # Run instruction
             conc_state = MiasmSymbolResolver(self.target, ctx.loc_db)
 
@@ -861,7 +887,7 @@ class SymbolicTracer:
                 new_pc = int(new_pc)
                 transform = SymbolicTransform(modified, [instruction], arch, pc, new_pc)
                 pred_regs, pred_mems = self.predict_next_state(instruction, transform)
-                self.progress(new_pc, instruction)
+                self.progress(new_pc, step=needs_step)
 
                 try:
                     self.validate(instruction, transform, pred_regs, pred_mems)
@@ -871,7 +897,7 @@ class SymbolicTracer:
                         continue
                     raise
             else:
-                new_pc = self.progress(new_pc, instruction)
+                new_pc = self.progress(new_pc, step=needs_step)
                 if new_pc is None:
                     continue # we're done
                 transform = SymbolicTransform(modified, [instruction], arch, pc, new_pc)