Use the SMT solver to convert symbolic to concrete value(s)

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - JoinBytes simplification rule
 - New simplification rule to help deal with abi.encodeWithSelector
 - More simplification rules for Props
+- Using the SMT solver to get a single concrete value for a symbolic expression
+  and continue running, whenever possible
 
 ## Fixed
 - We now try to simplify expressions fully before trying to cast them to a concrete value

src/EVM.hs

@@ -330,7 +330,6 @@ exec1 = do
     else do
       let ?op = getOpW8 vm.state
       let opName = getOpName vm.state
-
       case getOp (?op) of
 
         OpPush0 -> do
@@ -1380,10 +1379,10 @@ getCodeLocation :: VM t s -> CodeLocation
 getCodeLocation vm = (vm.state.contract, vm.state.pc)
 
 query :: Query t s -> EVM t s ()
-query = assign #result . Just . HandleEffect . Query
+query q = assign #result $ Just $ HandleEffect (Query q)
 
 choose :: Choose s -> EVM Symbolic s ()
-choose = assign #result . Just . HandleEffect . Choose
+choose c = assign #result $ Just $ HandleEffect (Choose c)
 
 -- | Construct RPC Query and halt execution until resolved
 fetchAccount :: VMOps t => Expr EAddr -> (Contract -> EVM t s ()) -> EVM t s ()
@@ -1403,11 +1402,11 @@ fetchAccount addr continue =
           Nothing -> do
             base <- use (#config % #baseState)
             assign (#result) . Just . HandleEffect . Query $
-              PleaseFetchContract a base
-                (\c -> do assign (#cache % #fetched % at a) (Just c)
-                          assign (#env % #contracts % at addr) (Just c)
-                          assign #result Nothing
-                          continue c)
+              PleaseFetchContract a base $ \c -> do
+                assign (#cache % #fetched % at a) (Just c)
+                assign (#env % #contracts % at addr) (Just c)
+                assign #result Nothing
+                continue c
       GVar _ -> internalError "Unexpected GVar"
 
 accessStorage
@@ -1446,13 +1445,11 @@ accessStorage addr slot continue = do
         else do
           modifying (#env % #contracts % ix addr % #storage) (writeStorage slot (Lit 0))
           continue $ Lit 0
-      mkQuery a s = query $
-        PleaseFetchSlot a s
-          (\x -> do
+      mkQuery a s = query $ PleaseFetchSlot a s $ \x -> do
               modifying (#cache % #fetched % ix a % #storage) (writeStorage (Lit s) (Lit x))
               modifying (#env % #contracts % ix (LitAddr a) % #storage) (writeStorage (Lit s) (Lit x))
               assign #result Nothing
-              continue (Lit x))
+              continue $ Lit x
 
 accessTStorage
   :: VMOps t => Expr EAddr
@@ -1585,24 +1582,29 @@ notStatic continue = do
 
 forceAddr :: VMOps t => Expr EWord -> String -> (Expr EAddr -> EVM t s ()) -> EVM t s ()
 forceAddr n msg continue = case wordToAddr n of
-  Nothing -> do
-    vm <- get
-    partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
+  Nothing -> oneSolution n $ \case
+    Just sol -> continue $ LitAddr (truncateToAddr sol)
+    Nothing -> fallback
   Just c -> continue c
+  where fallback = do
+          vm <- get
+          partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
 
 forceConcrete :: VMOps t => Expr EWord -> String -> (W256 -> EVM t s ()) -> EVM t s ()
 forceConcrete n msg continue = case maybeLitWordSimp n of
-  Nothing -> do
-    vm <- get
-    partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
+  Nothing -> oneSolution n $ maybe fallback continue
   Just c -> continue c
+  where fallback = do
+          vm <- get
+          partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
 
 forceConcreteAddr :: VMOps t => Expr EAddr -> String -> (Addr -> EVM t s ()) -> EVM t s ()
 forceConcreteAddr n msg continue = case maybeLitAddrSimp n of
-  Nothing -> do
-    vm <- get
-    partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
+  Nothing -> oneSolution (WAddr n) $ maybe fallback $ \c -> continue $ unsafeInto c
   Just c -> continue c
+  where fallback = do
+          vm <- get
+          partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) msg (wrap [n])
 
 forceConcreteAddr2 :: VMOps t => (Expr EAddr, Expr EAddr) -> String -> ((Addr, Addr) -> EVM t s ()) -> EVM t s ()
 forceConcreteAddr2 (n,m) msg continue = case (maybeLitAddrSimp n, maybeLitAddrSimp m) of
@@ -1699,8 +1701,13 @@ cheat gas (inOffset, inSize) (outOffset, outSize) xs = do
                     , context = newContext
                     }
   case maybeLitWordSimp abi of
-    Nothing -> partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) "symbolic cheatcode selector" (wrap [abi])
-    Just (unsafeInto -> abi') ->
+    Nothing -> oneSolution abi $ \case
+      Just concAbi -> runCheat concAbi input
+      Nothing -> partial $ UnexpectedSymbolicArg vm.state.pc (getOpName vm.state) "symbolic cheatcode selector" (wrap [abi])
+    Just concAbi -> runCheat concAbi input
+  where
+    runCheat abi input =  do
+      let abi' = unsafeInto abi
       case Map.lookup abi' cheatActions of
         Nothing -> vmError (BadCheatCode "Cannot understand cheatcode." abi')
         Just action -> action input
@@ -2933,7 +2940,7 @@ instance VMOps Symbolic where
   toGas _ = ()
   whenSymbolicElse a _ = a
 
-  partial e = assign #result (Just (Unfinished e))
+  partial e = assign #result $ Just (Unfinished e)
   branch cond continue = do
     loc <- codeloc
     pathconds <- use #constraints
@@ -2954,6 +2961,17 @@ instance VMOps Symbolic where
       choosePath loc Unknown =
         choose . PleaseChoosePath condSimp $ choosePath loc . Case
 
+  oneSolution ewordExpr continue = do
+    pathconds <- use #constraints
+    query $ PleaseGetSol ewordExpr pathconds $ \case
+      Just concVal -> do
+        assign #result Nothing
+        pushTo #constraints $ Expr.simplifyProp (ewordExpr .== (Lit concVal))
+        continue $ Just concVal
+      Nothing -> do
+        assign #result Nothing
+        continue Nothing
+
 instance VMOps Concrete where
   burn' n continue = do
     available <- use (#state % #gas)
@@ -3076,16 +3094,12 @@ instance VMOps Concrete where
     push (into vm.state.gas)
 
   enoughGas cost gasCap = cost <= gasCap
-
   subGas gasCap cost = gasCap - cost
-
   toGas = id
-
   whenSymbolicElse _ a = a
-
   partial _ = internalError "won't happen during concrete exec"
-
   branch (forceLit -> cond) continue = continue (cond > 0)
+  oneSolution _ _ = internalError "SMT solver should never be needed in concrete mode"
 
 -- Create symbolic VM from concrete VM
 symbolify :: VM Concrete s -> VM Symbolic s

src/EVM/Fetch.hs

@@ -27,6 +27,7 @@ import Numeric.Natural (Natural)
 import System.Environment (lookupEnv, getEnvironment)
 import System.Process
 import Control.Monad.IO.Class
+import Control.Monad (when)
 import EVM.Effects
 
 -- | Abstract representation of an RPC fetch request
@@ -212,6 +213,10 @@ oracle solvers info q = do
          -- Is is possible to satisfy the condition?
          continue <$> checkBranch solvers (branchcondition ./= (Lit 0)) pathconds
 
+    PleaseGetSol symAddr pathconditions continue -> do
+         let pathconds = foldl' PAnd (PBool True) pathconditions
+         continue <$> getSolution solvers symAddr pathconds
+
     PleaseFetchContract addr base continue -> do
       contract <- case info of
         Nothing -> let
@@ -239,6 +244,38 @@ oracle solvers info q = do
 
 type Fetcher t m s = App m => Query t s -> m (EVM t s ())
 
+getSolution :: forall m . App m => SolverGroup -> Expr EWord -> Prop -> m (Maybe W256)
+getSolution solvers symAddr pathconditions = do
+  conf <- readConfig
+  liftIO $ do
+    ret <- collectSolutions [] 1 pathconditions conf
+    case ret of
+      Nothing -> pure Nothing
+      Just r -> case length r of
+        0 -> pure Nothing
+        -- Temporary, a future improvement can deal with more than one solution
+        1 -> pure $ Just (r !! 0)
+        _ -> pure Nothing
+    where
+      collectSolutions :: [W256] -> Int -> Prop -> Config -> IO (Maybe [W256])
+      collectSolutions addrs maxSols conds conf = do
+        if (length addrs > maxSols) then pure Nothing
+        else
+          checkSat solvers (assertProps conf [(PEq (Var "addrQuery") symAddr) .&& conds]) >>= \case
+            Sat (SMTCex vars _ _ _ _ _)  -> case (Map.lookup (Var "addrQuery") vars) of
+              Just addr -> do
+                let newConds = PAnd conds (symAddr ./= (Lit addr))
+                when conf.debug $ putStrLn $ "Got one solution to symbolic query:" <> show addr <> " now have " <> show (length addrs + 1) <> " solution(s), max is: " <> show maxSols
+                collectSolutions (addr:addrs) maxSols newConds conf
+              _ -> internalError "No solution to symbolic query"
+            Unsat -> do
+              when conf.debug $ putStrLn "No more solution(s) to symbolic query."
+              pure $ Just addrs
+            -- Error or timeout, we need to be conservative
+            res -> do
+              when conf.debug $ putStrLn $ "Symbolic query result is neither SAT nor UNSAT:" <> show res
+              pure Nothing
+
 -- | Checks which branches are satisfiable, checking the pathconditions for consistency
 -- if the third argument is true.
 -- When in debug mode, we do not want to be able to navigate to dead paths,

src/EVM/Types.hs

@@ -594,6 +594,7 @@ data Query t s where
   PleaseFetchContract :: Addr -> BaseState -> (Contract -> EVM t s ()) -> Query t s
   PleaseFetchSlot     :: Addr -> W256 -> (W256 -> EVM t s ()) -> Query t s
   PleaseAskSMT        :: Expr EWord -> [Prop] -> (BranchCondition -> EVM Symbolic s ()) -> Query Symbolic s
+  PleaseGetSol        :: Expr EWord -> [Prop] -> (Maybe W256 -> EVM Symbolic s ()) -> Query Symbolic s
   PleaseDoFFI         :: [String] -> Map String String -> (ByteString -> EVM t s ()) -> Query t s
   PleaseReadEnv       :: String -> (String -> EVM t s ()) -> Query t s
 
@@ -617,6 +618,10 @@ instance Show (Query t s) where
       (("<EVM.Query: ask SMT about "
         ++ show condition ++ " in context "
         ++ show constraints ++ ">") ++)
+    PleaseGetSol expr constraints _ ->
+      (("<EVM.Query: ask SMT to get W256 for expression "
+        ++ show expr ++ " in context "
+        ++ show constraints ++ ">") ++)
     PleaseDoFFI cmd env _ ->
       (("<EVM.Query: do ffi: " ++ (show cmd) ++ " env: " ++ (show env)) ++)
     PleaseReadEnv variable _ ->
@@ -854,6 +859,7 @@ class VMOps (t :: VMType) where
 
   partial :: PartialExec -> EVM t s ()
   branch :: Expr EWord -> (Bool -> EVM t s ()) -> EVM t s ()
+  oneSolution :: Expr EWord -> (Maybe W256 -> EVM t s ()) -> EVM t s ()
 
 -- Bytecode Representations ------------------------------------------------------------------------