Fix state.branches (#74)

* init

* rm old code

pretty sure it's wrong

* Clean

* rm

* Rename record_fork, move to ConcretizeRegister exception handler

It doesn't need to be in fork() because it's only relevant for when
we fork due to PC.

* Add test for record_branches

* Add back record_branches which got removed in rebase

Rebase went oddly because State got moved to a new file

manticore/core/executor.py

@@ -43,7 +43,7 @@ from math import ceil, log
 
 from ..utils.nointerrupt import DelayedKeyboardInterrupt
 from .cpu.abstractcpu import ConcretizeRegister, ConcretizeMemory, \
-        InvalidPCException, IgnoreAPI
+        InvalidPCException, IgnoreAPI, SymbolicPCException
 from .memory import MemoryException, SymbolicMemoryException
 from .smtlib import solver, Expression, Operators, SolverException, Array, BitVec, Bool, ConstraintSet
 from ..utils.event import Signal
@@ -293,10 +293,6 @@ class Executor(object):
         # get last executed instruction
         last_cpu, last_pc = state.last_pc
 
-        try:
-            state.branches[(last_pc, state.cpu.PC)] += 1
-        except KeyError:
-            state.branches[(last_pc, state.cpu.PC)] = 1
         item = (last_pc, state.cpu.PC)
         assert not issymbolic(last_pc)
         assert not issymbolic(state.cpu.PC)
@@ -734,6 +730,9 @@ class Executor(object):
                     assert len(vals) > 0, "It should be at least one solution"
                     logger.debug("%s. Possible values are: %s", e.message, ["0x%016x"%x for x in vals])
 
+                    if isinstance(e, SymbolicPCException):
+                        current_state.record_branches(vals)
+
                     def setstate(state, val):
                         # XXX This only applies to ARM since x86 has no conditional instructions
                         if hasattr(state.cpu, 'forceNextInstruction'):

manticore/core/state.py

@@ -201,4 +201,13 @@ class State(object):
 
     def solve_n(self, expr, nsolves=1, policy='minmax'):
         # type: (Expression, int) -> list
-        return self._solver.get_all_values(self.constraints, expr, nsolves, silent=True)
+        return self._solver.get_all_values(self.constraints, expr, nsolves, silent=True)
+
+    def record_branches(self, targets):
+        _, branch = self.last_pc
+        for target in targets:
+            key = (branch, target)
+            try:
+                self.branches[key] += 1
+            except KeyError:
+                self.branches[key] = 1

tests/test_state.py

@@ -100,3 +100,19 @@ class StateTest(unittest.TestCase):
         with self.assertRaises(Exception):
             expr = self.state.new_symbolic_value(length)
 
+    def test_record_branches(self):
+        branch = 0x80488bb
+        target = 0x8048997
+        fallthrough = 0x80488c1
+        self.state.last_pc = (0, branch)
+
+        self.state.record_branches([target, fallthrough])
+
+        self.assertEqual(self.state.branches[(branch, target)], 1)
+        self.assertEqual(self.state.branches[(branch, fallthrough)], 1)
+
+        self.state.record_branches([target, fallthrough])
+
+        self.assertEqual(self.state.branches[(branch, target)], 2)
+        self.assertEqual(self.state.branches[(branch, fallthrough)], 2)
+