Reduce the number of calls to the SMT solver in EVM (#2411)

* extracted relevant code

* fix

* linted code

* fixed types

* fixes

* fix

* fixes

* fix

* regex fix

* fix code climate

* use solve_one_n_batched to reduce solver usage

* remove some useless calls to solve_one

* use solve_one_n_batched to show memory values

* used solve_one_n_batched in save_input_symbols

* clean-up

* clean-up

* Fix type hints

* Lint

Co-authored-by: Eric Hennenfent <eric.hennenfent@trailofbits.com>

manticore/core/smtlib/solver.py

@@ -22,7 +22,7 @@
 import shlex
 import time
 from functools import lru_cache
-from typing import Dict, Tuple, Sequence, Optional
+from typing import Any, Dict, Tuple, Sequence, Optional, List
 from subprocess import PIPE, Popen, check_output
 import re
 from . import operators as Operators
@@ -70,6 +70,9 @@ class SolverType(config.ConfigEnum):
 )
 
 # Regular expressions used by the solver
+RE_GET_EXPR_VALUE_ALL = re.compile(
+    "\(([a-zA-Z0-9_]*)[ \\n\\s]*(#b[0-1]*|#x[0-9a-fA-F]*|[\(]?_ bv[0-9]* [0-9]*|true|false)\\)"
+)
 RE_GET_EXPR_VALUE_FMT_BIN = re.compile(r"\(\((?P<expr>(.*))[ \n\s]*#b(?P<value>([0-1]*))\)\)")
 RE_GET_EXPR_VALUE_FMT_DEC = re.compile(r"\(\((?P<expr>(.*))\ \(_\ bv(?P<value>(\d*))\ \d*\)\)\)")
 RE_GET_EXPR_VALUE_FMT_HEX = re.compile(r"\(\((?P<expr>(.*))\ #x(?P<value>([0-9a-fA-F]*))\)\)")
@@ -81,6 +84,26 @@ class SolverType(config.ConfigEnum):
 SOLVER_STATS = {"unknown": 0, "timeout": 0}
 
 
+def _convert(v):
+    r = None
+    if v == "true":
+        r = True
+    elif v == "false":
+        r = False
+    elif v.startswith("#b"):
+        r = int(v[2:], 2)
+    elif v.startswith("#x"):
+        r = int(v[2:], 16)
+    elif v.startswith("_ bv"):
+        r = int(v[len("_ bv") : -len(" 256")], 10)
+    elif v.startswith("(_ bv"):
+        v = v[len("(_ bv") :]
+        r = int(v[: v.find(" ")], 10)
+
+    assert r is not None
+    return r
+
+
 class SingletonMixin(object):
     __singleton_instances: Dict[Tuple[int, int], "SingletonMixin"] = {}
 
@@ -376,6 +399,15 @@ def __getvalue_bool(self, expression_str):
         ret = self._smtlib.recv()
         return {"true": True, "false": False, "#b0": False, "#b1": True}[ret[2:-2].split(" ")[1]]
 
+    def __getvalue_all(
+        self, expressions_str: List[str], is_bv: List[bool]
+    ) -> Tuple[Dict[str, int], str]:
+        all_expressions_str = " ".join(expressions_str)
+        self._smtlib.send(f"(get-value ({all_expressions_str}))")
+        ret_solver = self._smtlib.recv()
+        return_values = re.findall(RE_GET_EXPR_VALUE_ALL, ret_solver)
+        return {value[0]: _convert(value[1]) for value in return_values}, ret_solver
+
     def _getvalue(self, expression) -> Union[int, bool, bytes]:
         """
         Ask the solver for one possible assignment for given expression using current set of constraints.
@@ -612,64 +644,95 @@ def _optimize_fancy(self, constraints: ConstraintSet, x: BitVec, goal: str, max_
             raise SolverError("Optimize failed")
 
     def get_value(self, constraints: ConstraintSet, *expressions):
+        values = self.get_value_in_batch(constraints, expressions)
+        if len(expressions) == 1:
+            return values[0]
+        else:
+            return values
+
+    def get_value_in_batch(self, constraints: ConstraintSet, expressions):
         """
         Ask the solver for one possible result of given expressions using
         given set of constraints.
         """
-        values = []
+        values: List[Any] = [None] * len(expressions)
         start = time.time()
         with constraints.related_to(*expressions) as temp_cs:
-            for expression in expressions:
+            vars: List[Any] = []
+            for idx, expression in enumerate(expressions):
                 if not issymbolic(expression):
-                    values.append(expression)
+                    values[idx] = expression
+                    vars.append(None)
                     continue
                 assert isinstance(expression, (Bool, BitVec, Array))
                 if isinstance(expression, Bool):
                     var = temp_cs.new_bool()
+                    vars.append(var)
+                    temp_cs.add(var == expression)
                 elif isinstance(expression, BitVec):
                     var = temp_cs.new_bitvec(expression.size)
+                    vars.append(var)
+                    temp_cs.add(var == expression)
                 elif isinstance(expression, Array):
                     var = []
-                    result = []
                     for i in range(expression.index_max):
                         subvar = temp_cs.new_bitvec(expression.value_bits)
                         var.append(subvar)
                         temp_cs.add(subvar == simplify(expression[i]))
-                    self._reset(temp_cs.to_string())
-                    if not self._is_sat():
-                        raise SolverError(
-                            "Solver could not find a value for expression under current constraint set"
-                        )
+                    vars.append(var)
 
-                    for i in range(expression.index_max):
-                        result.append(self.__getvalue_bv(var[i].name))
-                    values.append(bytes(result))
-                    if time.time() - start > consts.timeout:
-                        SOLVER_STATS["timeout"] += 1
-                        raise SolverError("Timeout")
-                    continue
-
-                temp_cs.add(var == expression)
+            self._reset(temp_cs.to_string())
+            if not self._is_sat():
+                raise SolverError(
+                    "Solver could not find a value for expression under current constraint set"
+                )
 
-                self._reset(temp_cs.to_string())
+            values_to_ask: List[str] = []
+            is_bv: List[bool] = []
+            for idx, expression in enumerate(expressions):
+                if not issymbolic(expression):
+                    continue
+                var = vars[idx]
+                if isinstance(expression, Bool):
+                    values_to_ask.append(var.name)
+                    is_bv.append(False)
+                if isinstance(expression, BitVec):
+                    values_to_ask.append(var.name)
+                    is_bv.append(True)
+                if isinstance(expression, Array):
+                    # result = []
+                    for i in range(expression.index_max):
+                        values_to_ask.append(var[i].name)
+                        is_bv.append(True)
 
-                if not self._is_sat():
-                    raise SolverError(
-                        "Solver could not find a value for expression under current constraint set"
-                    )
+            if values_to_ask == []:
+                return values
 
+            values_returned, sol = self.__getvalue_all(values_to_ask, is_bv)
+            for idx, expression in enumerate(expressions):
+                if not issymbolic(expression):
+                    continue
+                var = vars[idx]
                 if isinstance(expression, Bool):
-                    values.append(self.__getvalue_bool(var.name))
+                    values[idx] = values_returned[var.name]
                 if isinstance(expression, BitVec):
-                    values.append(self.__getvalue_bv(var.name))
+                    if var.name not in values_returned:
+                        logger.error(
+                            "var.name", var.name, "not in values_returned", values_returned
+                        )
+
+                    values[idx] = values_returned[var.name]
+                if isinstance(expression, Array):
+                    result = []
+                    for i in range(expression.index_max):
+                        result.append(values_returned[var[i].name])
+                    values[idx] = bytes(result)
+
             if time.time() - start > consts.timeout:
                 SOLVER_STATS["timeout"] += 1
                 raise SolverError("Timeout")
 
-        if len(expressions) == 1:
-            return values[0]
-        else:
-            return values
+        return values
 
 
 class Z3Solver(SMTLIBSolver):

manticore/core/state.py

@@ -1,7 +1,10 @@
 import copy
 import logging
 
+from typing import List, Tuple, Sequence
+
 from .smtlib import solver, Bool, issymbolic, BitVecConstant
+from .smtlib.expression import Expression
 from ..utils.event import Eventful
 from ..utils.helpers import PickleSerializer
 from ..utils import config
@@ -451,30 +454,41 @@ def solve_one(self, expr, constrain=False):
         """
         return self.solve_one_n(expr, constrain=constrain)[0]
 
-    def solve_one_n(self, *exprs, constrain=False):
+    def solve_one_n(self, *exprs: Expression, constrain: bool = False) -> List[int]:
         """
-        Concretize a symbolic :class:`~manticore.core.smtlib.expression.Expression` into
-        one solution.
+        Concretize a list of symbolic :class:`~manticore.core.smtlib.expression.Expression` into
+        a list of solutions.
 
         :param exprs: An iterable of manticore.core.smtlib.Expression
         :param bool constrain: If True, constrain expr to solved solution value
-        :return: Concrete value or a tuple of concrete values
-        :rtype: int
+        :return: List of concrete value or a tuple of concrete values
         """
-        values = []
-        for expr in exprs:
-            if not issymbolic(expr):
-                values.append(expr)
-            else:
-                expr = self.migrate_expression(expr)
-                value = self._solver.get_value(self._constraints, expr)
-                if constrain:
-                    self.constrain(expr == value)
-                # Include forgiveness here
-                if isinstance(value, bytearray):
-                    value = bytes(value)
-                values.append(value)
-        return values
+        return self.solve_one_n_batched(exprs, constrain)
+
+    def solve_one_n_batched(
+        self, exprs: Sequence[Expression], constrain: bool = False
+    ) -> List[int]:
+        """
+        Concretize a list of symbolic :class:`~manticore.core.smtlib.expression.Expression` into
+        a list of solutions.
+        :param exprs: An iterable of manticore.core.smtlib.Expression
+        :param bool constrain: If True, constrain expr to solved solution value
+        :return: List of concrete value or a tuple of concrete values
+        """
+        # Return ret instead of value, to allow the bytearray/bytes conversion
+        ret = []
+        exprs = [self.migrate_expression(x) for x in exprs]
+        values = self._solver.get_value_in_batch(self._constraints, exprs)
+        assert len(values) == len(exprs)
+        for idx, expr in enumerate(exprs):
+            value = values[idx]
+            if constrain:
+                self.constrain(expr == values[idx])
+            # Include forgiveness here
+            if isinstance(value, bytearray):
+                value = bytes(value)
+            ret.append(value)
+        return ret
 
     def solve_n(self, expr, nsolves):
         """

manticore/core/workspace.py

@@ -658,7 +658,7 @@ def save_constraints(testcase, state: StateBase):
     @staticmethod
     def save_input_symbols(testcase, state: StateBase):
         with testcase.open_stream("input") as f:
-            for symbol in state.input_symbols:
+            bufs = state.solve_one_n_batched(state.input_symbols)
+            for symbol, buf in zip(state.input_symbols, bufs):
                 # TODO - constrain=False here, so the extra migration shouldn't cause problems, right?
-                buf = state.solve_one(symbol)
                 f.write(f"{symbol.name}: {buf!r}\n")

manticore/platforms/evm.py

@@ -182,13 +182,24 @@ def concretize(self, state, constrain=False):
         :param state: a manticore state
         :param bool constrain: If True, constrain expr to concretized value
         """
-        conc_caller = state.solve_one(self.caller, constrain=constrain)
-        conc_address = state.solve_one(self.address, constrain=constrain)
-        conc_value = state.solve_one(self.value, constrain=constrain)
-        conc_gas = state.solve_one(self.gas, constrain=constrain)
-        conc_data = state.solve_one(self.data, constrain=constrain)
-        conc_return_data = state.solve_one(self.return_data, constrain=constrain)
-        conc_used_gas = state.solve_one(self.used_gas, constrain=constrain)
+        all_elems = [
+            self.caller,
+            self.address,
+            self.value,
+            self.gas,
+            self.data,
+            self._return_data,
+            self.used_gas,
+        ]
+        values = state.solve_one_n_batched(all_elems, constrain=constrain)
+        conc_caller = values[0]
+        conc_address = values[1]
+        conc_value = values[2]
+        conc_gas = values[3]
+        conc_data = values[4]
+        conc_return_data = values[5]
+        conc_used_gas = values[6]
+
         return Transaction(
             self.sort,
             conc_address,
@@ -291,9 +302,7 @@ def dump(self, stream, state, mevm, conc_tx=None):
                 )  # is this redundant since arguments are all concrete?
                 stream.write("Function call:\n")
                 stream.write("Constructor(")
-                stream.write(
-                    ",".join(map(repr, map(state.solve_one, arguments)))
-                )  # is this redundant since arguments are all concrete?
+                stream.write(",".join(map(repr, arguments)))
                 stream.write(") -> %s %s\n" % (self.result, flagged(is_argument_symbolic)))
 
         if self.sort == "CALL":
@@ -3435,14 +3444,18 @@ def dump(self, stream, state, mevm, message):
                 stream.write("Balance: %d %s\n" % (balance, flagged(is_balance_symbolic)))
 
             storage = blockchain.get_storage(account_address)
-            concrete_indexes = set()
-            for sindex in storage.written:
-                concrete_indexes.add(state.solve_one(sindex, constrain=True))
+            concrete_indexes = []
+            if len(storage.written) > 0:
+                concrete_indexes = state.solve_one_n_batched(storage.written, constrain=True)
+
+            concrete_values = []
+            if len(concrete_indexes) > 0:
+                concrete_values = state.solve_one_n_batched(concrete_indexes, constrain=True)
+
+            assert len(concrete_indexes) == len(concrete_values)
+            for index, value in zip(concrete_indexes, concrete_values):
+                stream.write(f"storage[{index:x}] = {value:x}\n")
 
-            for index in concrete_indexes:
-                stream.write(
-                    f"storage[{index:x}] = {state.solve_one(storage[index], constrain=True):x}\n"
-                )
             storage = blockchain.get_storage(account_address)
             stream.write("Storage: %s\n" % translate_to_smtlib(storage, use_bindings=False))