evm: aggressively check & migrate expressions into current ConstraintSet in case they are global/external

manticore/core/smtlib/constraints.py

@@ -1,4 +1,4 @@
-
+import itertools
 from .expression import BitVecVariable, BoolVariable, ArrayVariable, Array, Bool, BitVec, BoolConstant, ArrayProxy, BoolEq, Variable, Constant
 from .visitors import GetDeclarations, TranslatorSmtlib, get_variables, simplify, replace, translate_to_smtlib
 import logging
@@ -17,16 +17,18 @@ def __init__(self):
         self._constraints = list()
         self._parent = None
         self._sid = 0
+        self._declarations = {}
         self._child = None
 
     def __reduce__(self):
-        return (self.__class__, (), {'_parent': self._parent, '_constraints': self._constraints, '_sid': self._sid})
+        return (self.__class__, (), {'_parent': self._parent, '_constraints': self._constraints, '_sid': self._sid, '_declarations': self._declarations})
 
     def __enter__(self):
         assert self._child is None
         self._child = self.__class__()
         self._child._parent = self
         self._child._sid = self._sid
+        self._child._declarations = dict(self._declarations)
         return self._child
 
     def __exit__(self, ty, value, traceback):
@@ -159,8 +161,26 @@ def to_string(self, related_to=None, replace_constants=True):
             constraint_str = translator.pop()
         return result
 
+    def _declare(self, var):
+        ''' Declare the variable `var` '''
+        if var.name in self._declarations:
+            raise ValueError('Variable already declared')
+        self._declarations[var.name] = var
+        return var
+
+    def get_variable(self, name):
+        ''' Returns the variable declared under name or None if it does not exists '''
+        if name not in self._declarations:
+            return None
+        return self._declarations[name]
+
+    def get_declared_variables(self):
+        ''' Returns the variable expressions of this constraint set '''
+        return self._declarations.values()
+
     @property
     def declarations(self):
+        ''' Returns the variable expressions of this constraint set '''
         declarations = GetDeclarations()
         for a in self.constraints:
             try:
@@ -191,68 +211,134 @@ def __str__(self):
         ''' Returns a smtlib representation of the current state '''
         return self.to_string()
 
-    def _get_new_name(self, name='VAR'):
+    def _make_unique_name(self, name='VAR'):
         ''' Makes an uniq variable name'''
-        return '%s_%d' % (name, self._get_sid())
-
-    def migrate(self, expression, name=None, bindings=None):
-        ''' Migrate an expression created for a different constraint set
+        # the while loop is necessary because appending the result of _get_sid()
+        # is not guaranteed to make a unique name on the first try; a colliding
+        # name could have been added previously
+        while name in self._declarations:
+            name = '%s_%d' % (name, self._get_sid())
+        return name
+
+    def is_declared(self, expression_var):
+        ''' True if expression_var is declared in this constraint set '''
+        if not isinstance(expression_var, Variable):
+            raise ValueError("Expression must be a Variable")
+        return any(expression_var is x for x in self.get_declared_variables())
+
+    def migrate(self, expression, name_migration_map=None):
+        ''' Migrate an expression created for a different constraint set to self.
             Returns an expression that can be used with this constraintSet
-        '''
-        # Simply check there are no name overlappings
-        if bindings is None:
-            bindings = {}
-        variables = get_variables(expression)
-        for var in variables:
-            if name is None:
-                name = self._get_new_name(var.name + '_migrated')
-
-            if var in bindings:
-                continue
 
-            if isinstance(var, Bool):
-                new_var = self.new_bool(name=name)
-            elif isinstance(var, BitVec):
-                new_var = self.new_bitvec(var.size, name=name)
-            elif isinstance(var, Array):
-                new_var = self.new_array(index_max=var.index_max, index_bits=var.index_bits, value_bits=var.value_bits, name=name)
-            else:
-                raise NotImplemented("Unknown type {}".format(type(var)))
+            All the foreign variables used in the expression are replaced by
+            variables of this constraint set. If the variable was replaced before
+            the replacement is taken from the provided migration map.
+
+            The migration mapping is updated with new replacements.
 
-            bindings[var] = new_var
+            :param expression: the potentially foreign expression
+            :param name_migration_map: mapping of already migrated variables. maps from string name of foreign variable to its currently existing migrated string name. this is updated during this migration.
+            :return: a migrated expresion where all the variables are local. name_migration_map is updated
 
-        migrated_expression = replace(expression, bindings)
+        '''
+        if name_migration_map is None:
+            name_migration_map = {}
+
+        #  name_migration_map -> object_migration_map
+        #  Based on the name mapping in name_migration_map build an object to
+        #  object mapping to be used in the replacing of variables
+        #  inv: object_migration_map's keys should ALWAYS be external/foreign
+        #  expressions, and its values should ALWAYS be internal/local expressions
+        object_migration_map = {}
+
+        #List of foreign vars used in expression
+        foreign_vars = itertools.filterfalse(self.is_declared, get_variables(expression))
+        for foreign_var in foreign_vars:
+            # If a variable with the same name was previously migrated
+            if foreign_var.name in name_migration_map:
+                migrated_name = name_migration_map[foreign_var.name]
+                native_var = self.get_variable(migrated_name)
+                assert native_var is not None, "name_migration_map contains a variable that does not exist in this ConstraintSet"
+                object_migration_map[foreign_var] = native_var
+            else:
+                # foreign_var was not found in the local declared variables nor
+                # any variable with the same name was previously migrated
+                # lets make a new uniq internal name for it
+                migrated_name = foreign_var.name
+                if migrated_name in self._declarations:
+                    migrated_name = self._make_unique_name(foreign_var.name + '_migrated')
+                # Create and declare a new variable of given type
+                if isinstance(foreign_var, Bool):
+                    new_var = self.new_bool(name=migrated_name)
+                elif isinstance(foreign_var, BitVec):
+                    new_var = self.new_bitvec(expression_var.size, name=migrated_name)
+                elif isinstance(foreign_var, Array):
+                    # Note that we are discarding the ArrayProxy encapsulation
+                    new_var = self.new_array(index_max=foreign_var.index_max, index_bits=foreign_var.index_bits, value_bits=foreign_var.value_bits, name=migrated_name).array
+                else:
+                    raise NotImplemented("Unknown expression type {} encountered during expression migration".format(type(var)))
+                # Update the var to var mapping
+                object_migration_map[foreign_var] = new_var
+                # Update the name to name mapping
+                name_migration_map[foreign_var.name] = new_var.name
+
+        #  Actually replace each appearence of migrated variables by the new ones
+        migrated_expression = replace(expression, object_migration_map)
         return migrated_expression
 
-    def new_bool(self, name='B', taint=frozenset()):
+    def new_bool(self, name=None, taint=frozenset(), avoid_collisions=False):
         ''' Declares a free symbolic boolean in the constraint store
             :param name: try to assign name to internal variable representation,
                          if not uniq a numeric nonce will be appended
+            :param avoid_collisions: potentially avoid_collisions the variable to avoid name colisions if True
             :return: a fresh BoolVariable
         '''
-        name = self._get_new_name(name)
-        return BoolVariable(name, taint=taint)
-
-    def new_bitvec(self, size, name='V', taint=frozenset()):
+        if name is None:
+            name = 'B'
+            avoid_collisions = True
+        if avoid_collisions:
+            name = self._make_unique_name(name)
+        if not avoid_collisions and name in self._declarations:
+            raise ValueError("Name already used")
+        var = BoolVariable(name, taint=taint)
+        return self._declare(var)
+
+    def new_bitvec(self, size, name=None, taint=frozenset(), avoid_collisions=False):
         ''' Declares a free symbolic bitvector in the constraint store
             :param size: size in bits for the bitvector
             :param name: try to assign name to internal variable representation,
                          if not uniq a numeric nonce will be appended
+            :param avoid_collisions: potentially avoid_collisions the variable to avoid name colisions if True
             :return: a fresh BitVecVariable
         '''
         if not (size == 1 or size % 8 == 0):
             raise Exception('Invalid bitvec size %s' % size)
-        name = self._get_new_name(name)
-        return BitVecVariable(size, name, taint=taint)
-
-    def new_array(self, index_bits=32, name='A', index_max=None, value_bits=8, taint=frozenset()):
+        if name is None:
+            name = 'BV'
+            avoid_collisions = True
+        if avoid_collisions:
+            name = self._make_unique_name(name)
+        if not avoid_collisions and name in self._declarations:
+            raise ValueError("Name already used")
+        var = BitVecVariable(size, name, taint=taint)
+        return self._declare(var)
+
+    def new_array(self, index_bits=32, name=None, index_max=None, value_bits=8, taint=frozenset(), avoid_collisions=False):
         ''' Declares a free symbolic array of value_bits long bitvectors in the constraint store.
             :param index_bits: size in bits for the array indexes one of [32, 64]
             :param value_bits: size in bits for the array values
             :param name: try to assign name to internal variable representation,
                          if not uniq a numeric nonce will be appended
             :param index_max: upper limit for indexes on ths array (#FIXME)
+            :param avoid_collisions: potentially avoid_collisions the variable to avoid name colisions if True
             :return: a fresh ArrayProxy
         '''
-        name = self._get_new_name(name)
-        return ArrayProxy(ArrayVariable(index_bits, index_max, value_bits, name, taint=taint))
+        if name is None:
+            name = 'A'
+            avoid_collisions = True
+        if avoid_collisions:
+            name = self._make_unique_name(name)
+        if not avoid_collisions and name in self._declarations:
+            raise ValueError("Name already used")
+        var = self._declare(ArrayVariable(index_bits, index_max, value_bits, name, taint=taint))
+        return ArrayProxy(var)

manticore/core/smtlib/visitors.py

@@ -704,6 +704,8 @@ def visit_Variable(self, expression):
 
 
 def replace(expression, bindings):
+    if not bindings:
+        return expression
     visitor = Replace(bindings)
     visitor.visit(expression, use_fixed_point=True)
     result_expression = visitor.result