From 5368716e424c2e950b6e0fae3a6eaf1c5d574587 Mon Sep 17 00:00:00 2001
From: Mark Mossberg <mark.mossberg@gmail.com>
Date: Mon, 20 Mar 2017 17:56:25 -0400
Subject: [PATCH] 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 |  9 ++++-----
 manticore/core/state.py    | 11 ++++++++++-
 tests/test_state.py        | 16 ++++++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/manticore/core/executor.py b/manticore/core/executor.py
index 546c280..869afb7 100644
--- a/manticore/core/executor.py
+++ b/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'):
diff --git a/manticore/core/state.py b/manticore/core/state.py
index 8223f88..5e82dab 100644
--- a/manticore/core/state.py
+++ b/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)
\ No newline at end of file
+        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
diff --git a/tests/test_state.py b/tests/test_state.py
index 62d3277..87eff6b 100644
--- a/tests/test_state.py
+++ b/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)
+