From faaf6b76567a23188eae331ff741ce57386a3693 Mon Sep 17 00:00:00 2001
From: Nicola Ruaro <ruaronicola@gmail.com>
Date: Wed, 15 Nov 2023 21:41:20 -0800
Subject: [PATCH] fix sha

---
 greed/TAC/special_ops.py  | 31 +++++++++++++------------------
 greed/solver/shortcuts.py | 21 +++++++++++++++++++--
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/greed/TAC/special_ops.py b/greed/TAC/special_ops.py
index b27ea49..f3cb16d 100644
--- a/greed/TAC/special_ops.py
+++ b/greed/TAC/special_ops.py
@@ -44,16 +44,16 @@ def handle(self, state: SymbolicEVMState):
         # calculate the possible solutions and apply the constraints.
         if options.GREEDY_SHA:
             log.debug(f"Using GREEDY_SHA strategy to try to resolve {new_sha.symbol.name}")
-            size_sol = state.solver.eval(self.size_val, raw=True)
+
+            size_sol = concretize(state, self.size_val)
+            offset_sol = concretize(state, self.offset_val)
             log.debug(f"    Size solution: {bv_unsigned_value(size_sol)}")
-            offset_sol = state.solver.eval(self.offset_val, raw=True)
             log.debug(f"    Offset solution: {bv_unsigned_value(offset_sol)}")
             
-            if not state.solver.is_formula_sat(NotEqual(self.size_val, size_sol)) and \
-               not state.solver.is_formula_sat(NotEqual(self.offset_val, offset_sol)):
-
-                # If size_sol is zero, return e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
-                if size_sol.value == 0:
+            if size_sol is None or offset_sol is None:
+                log.debug(f"    Cannot calculate concrete SHA3 for {new_sha.symbol.name} due to multiple size and offset solutions")
+            elif size_sol.value == 0:
+                    # If size_sol is zero, return e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
                     res = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
                     log.debug(f"    Calculated concrete SHA3 {res}")
 
@@ -66,14 +66,14 @@ def handle(self, state: SymbolicEVMState):
                     
                     # Just set the solution here
                     state.registers[self.res1_var] = BVV(int(res, 16), 256)
-                    return [state]
-
+            else:
                 # Else, get a solution for the input buffer
-                buffer_sol = state.solver.eval_memory_at(state.memory, offset_sol, 
-                                                                        size_sol, 
-                                                                        raw=True)
+                buffer = state.memory.readn(offset_sol, size_sol)
+                buffer_sol = concretize(state, buffer)
 
-                if not state.solver.is_formula_sat(NotEqual(state.memory.readn(offset_sol, size_sol), buffer_sol)):
+                if buffer_sol is None:
+                    log.debug(f"    Cannot calculate concrete SHA3 for {new_sha.symbol.name} due to multiple SHA solutions")
+                else:
                     # Everything has only one solution, we can calculate the SHA
                     keccak256 = sha3.keccak_256()
                     buffer_sol = bv_unsigned_value(buffer_sol).to_bytes(bv_unsigned_value(size_sol), 'big')
@@ -94,13 +94,8 @@ def handle(self, state: SymbolicEVMState):
                     
                     # Just set the solution here
                     state.registers[self.res1_var] = BVV(int(res,16),256)
-                else:
-                    log.debug(f"    Cannot calculate concrete SHA3 for {new_sha.symbol.name} due to multiple SHA solutions")
-            else:
-                log.debug(f"    Cannot calculate concrete SHA3 for {new_sha.symbol.name} due to multiple size and offset solutions")
 
         state.set_next_pc()
-
         return [state]
 
 
diff --git a/greed/solver/shortcuts.py b/greed/solver/shortcuts.py
index 32b7c8a..1173fa6 100644
--- a/greed/solver/shortcuts.py
+++ b/greed/solver/shortcuts.py
@@ -1,6 +1,6 @@
 __all__ = [
-    "ctx_or_symbolic", "BVSort", "BVV", "BVS", "Array", "If", "Equal", "NotEqual", "Or", "And", "Not",
-    "bv_unsigned_value", "get_bv_by_name", "is_concrete", "BV_Extract", "BV_Concat", "BV_Add", "BV_Sub",
+    "ctx_or_symbolic", "concretize", "BVSort", "BVV", "BVS", "Array", "If", "Equal", "NotEqual", "Or", "And",
+    "Not", "bv_unsigned_value", "get_bv_by_name", "is_concrete", "BV_Extract", "BV_Concat", "BV_Add", "BV_Sub",
     "BV_Mul", "BV_UDiv", "BV_SDiv", "BV_SMod", "BV_SRem", "BV_URem", "BV_Sign_Extend", "BV_Zero_Extend",
     "BV_UGE", "BV_ULE", "BV_UGT", "BV_ULT", "BV_SGE", "BV_SLE", "BV_SGT", "BV_SLT", "BV_And", "BV_Or",
     "BV_Xor", "BV_Not", "BV_Shl", "BV_Shr", "BV_Sar", "Array_Store", "Array_Select"
@@ -25,6 +25,23 @@ def ctx_or_symbolic(v, ctx, xid, nbits=256):
     return ctx[v]
 
 
+def concretize(state, val, force=False):
+    if is_concrete(val):
+        return val
+    else:
+        try:
+            val_sol = state.solver.eval(val, raw=True)
+            if state.solver.is_formula_true(Equal(val, val_sol)):
+                # one possible solution
+                return val_sol
+            elif force is True:
+                # multiple possible solutions
+                state.add_constraint(Equal(val, val_sol))
+                return val_sol
+        except:
+            return None
+
+
 # TYPES