From 318a71ada1dd34402c1fa452f631bdf4644a8d18 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Wed, 11 Sep 2024 14:10:24 -0400
Subject: [PATCH] Simplify Exec.hs

---
 lib/Echidna/Exec.hs | 292 +++++++++++++++++++++++---------------------
 1 file changed, 151 insertions(+), 141 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index f76691897..0ef9f5cf2 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -14,7 +14,7 @@ import Control.Monad.Reader (MonadReader, ask, asks)
 import Control.Monad.ST (ST, stToIO, RealWorld)
 import Data.Bits
 import Data.ByteString qualified as BS
-import Data.IORef (readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
+import Data.IORef (IORef, readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe, fromJust)
 import Data.Text qualified as T
@@ -102,91 +102,8 @@ execTxWith executeTx tx = do
         pure (vmResult, gasLeftBeforeTx - gasLeftAfterTx)
   where
   runFully = do
-    config <- asks (.cfg)
-    -- TODO: Is the latest block a good default? It makes fuzzing hard to reproduce. Rethink this.
-    let rpcBlock = maybe EVM.Fetch.Latest (EVM.Fetch.BlockNumber . fromIntegral) config.rpcBlock
-
     vmResult <- executeTx
-    -- For queries, we halt execution because the VM needs some additional
-    -- information from the outside. We provide this information and resume
-    -- the execution by recursively calling `runFully`.
-    case getQuery vmResult of
-      -- A previously unknown contract is required
-      Just q@(PleaseFetchContract addr _ continuation) -> do
-        cacheRef <- asks (.fetchContractCache)
-        cache <- liftIO $ readIORef cacheRef
-        case Map.lookup addr cache of
-          Just (Just contract) -> fromEVM (continuation contract)
-          Just Nothing -> do
-            v <- get
-            v' <- liftIO $ stToIO $ execStateT (continuation emptyAccount) v
-            put v'
-          Nothing -> do
-            logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
-              Just rpcUrl -> do
-                ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
-                case ret of
-                  -- TODO: fix hevm to not return an empty contract in case of an error
-                  Just contract | contract.code /= RuntimeCode (ConcreteRuntimeCode "") -> do
-                    fromEVM (continuation contract)
-                    liftIO $ atomicWriteIORef cacheRef $ Map.insert addr (Just contract) cache
-                  _ -> do
-                    -- TODO: better error reporting in HEVM, when intermittent
-                    -- network error then retry
-                    liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
-                    logMsg $ "ERROR: Failed to fetch contract: " <> show q
-                    -- TODO: How should we fail here? It could be a network error,
-                    -- RPC server returning junk etc.
-                    fromEVM (continuation emptyAccount)
-              Nothing -> do
-                liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
-                logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
-                -- TODO: How should we fail here? RPC is not configured but VM
-                -- wants to fetch
-                fromEVM (continuation emptyAccount)
-        runFully -- resume execution
-
-      -- A previously unknown slot is required
-      Just q@(PleaseFetchSlot addr slot continuation) -> do
-        cacheRef <- asks (.fetchSlotCache)
-        cache <- liftIO $ readIORef cacheRef
-        case Map.lookup addr cache >>= Map.lookup slot of
-          Just (Just value) -> fromEVM (continuation value)
-          Just Nothing -> fromEVM (continuation 0)
-          Nothing -> do
-            logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
-              Just rpcUrl -> do
-                ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
-                case ret of
-                  Just value -> do
-                    fromEVM (continuation value)
-                    liftIO $ atomicWriteIORef cacheRef $
-                      Map.insertWith Map.union addr (Map.singleton slot (Just value)) cache
-                  Nothing -> do
-                    -- TODO: How should we fail here? It could be a network error,
-                    -- RPC server returning junk etc.
-                    logMsg $ "ERROR: Failed to fetch slot: " <> show q
-                    liftIO $ atomicWriteIORef cacheRef $
-                      Map.insertWith Map.union addr (Map.singleton slot Nothing) cache
-                    fromEVM (continuation 0)
-              Nothing -> do
-                logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
-                -- Use the zero slot
-                fromEVM (continuation 0)
-        runFully -- resume execution
-
-      -- Execute a FFI call
-      Just (PleaseDoFFI (cmd : args) continuation) -> do
-        (_, stdout, _) <- liftIO $ readProcessWithExitCode cmd args ""
-        let encodedResponse = encodeAbiValue $
-              AbiTuple (V.fromList [AbiBytesDynamic . hexText . T.pack $ stdout])
-        fromEVM (continuation encodedResponse)
-        runFully
-
-      -- No queries to answer, the tx is fully executed and the result is final
-      _ -> pure vmResult
+    maybe (pure vmResult) (\q -> handleQuery q >> runFully) (getQuery vmResult)
 
   -- | Handles reverts, failures and contract creations that might be the result
   -- (`vmResult`) of executing transaction `tx`.
@@ -217,6 +134,92 @@ execTxWith executeTx tx = do
       modify' $ execState $ loadContract (LitAddr tx.dst)
     _ -> pure ()
 
+  getRpcInfo = do
+    config <- asks (.cfg)
+    -- TODO: Is the latest block a good default? It makes fuzzing hard to reproduce. Rethink this.
+    let rpcBlock = maybe EVM.Fetch.Latest (EVM.Fetch.BlockNumber . fromIntegral) config.rpcBlock
+    return (config.rpcUrl, rpcBlock)
+
+
+  -- For queries, we halt execution because the VM needs some additional
+  -- information from the outside. We provide this information, and then
+  -- the execution is resumed.
+
+  -- A previously unknown contract is required
+  handleQuery q@(PleaseFetchContract addr _ continuation) = do
+    cacheRef <- asks (.fetchContractCache)
+    cache <- liftIO $ readIORef cacheRef
+    case Map.lookup addr cache of
+      Just (Just contract) -> fromEVM (continuation contract)
+      Just Nothing -> do
+        v <- get
+        v' <- liftIO $ stToIO $ execStateT (continuation emptyAccount) v
+        put v'
+      Nothing -> do
+        logMsg $ "INFO: Performing RPC: " <> show q
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
+          Just rpcUrl -> do
+            ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
+            case ret of
+              -- TODO: fix hevm to not return an empty contract in case of an error
+              Just contract | contract.code /= RuntimeCode (ConcreteRuntimeCode "") -> do
+                fromEVM (continuation contract)
+                liftIO $ atomicWriteIORef cacheRef $ Map.insert addr (Just contract) cache
+              _ -> do
+                -- TODO: better error reporting in HEVM, when intermittent
+                -- network error then retry
+                liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
+                logMsg $ "ERROR: Failed to fetch contract: " <> show q
+                -- TODO: How should we fail here? It could be a network error,
+                -- RPC server returning junk etc.
+                fromEVM (continuation emptyAccount)
+          Nothing -> do
+            liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
+            logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
+            -- TODO: How should we fail here? RPC is not configured but VM
+            -- wants to fetch
+            fromEVM (continuation emptyAccount)
+
+  -- A previously unknown slot is required
+  handleQuery q@(PleaseFetchSlot addr slot continuation) = do
+    cacheRef <- asks (.fetchSlotCache)
+    cache <- liftIO $ readIORef cacheRef
+    case Map.lookup addr cache >>= Map.lookup slot of
+      Just (Just value) -> fromEVM (continuation value)
+      Just Nothing -> fromEVM (continuation 0)
+      Nothing -> do
+        logMsg $ "INFO: Performing RPC: " <> show q
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
+          Just rpcUrl -> do
+            ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
+            case ret of
+              Just value -> do
+                fromEVM (continuation value)
+                liftIO $ atomicWriteIORef cacheRef $
+                  Map.insertWith Map.union addr (Map.singleton slot (Just value)) cache
+              Nothing -> do
+                -- TODO: How should we fail here? It could be a network error,
+                -- RPC server returning junk etc.
+                logMsg $ "ERROR: Failed to fetch slot: " <> show q
+                liftIO $ atomicWriteIORef cacheRef $
+                  Map.insertWith Map.union addr (Map.singleton slot Nothing) cache
+                fromEVM (continuation 0)
+          Nothing -> do
+            logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
+            -- Use the zero slot
+            fromEVM (continuation 0)
+
+  -- Execute a FFI call
+  handleQuery (PleaseDoFFI (cmd : args) continuation) = do
+    (_, stdout, _) <- liftIO $ readProcessWithExitCode cmd args ""
+    let encodedResponse = encodeAbiValue $
+          AbiTuple (V.fromList [AbiBytesDynamic . hexText . T.pack $ stdout])
+    fromEVM (continuation encodedResponse)
+
+  handleQuery (PleaseDoFFI [] _) = error "Malformed FFI call"
+
 logMsg :: (MonadIO m, MonadReader Env m) => String -> m ()
 logMsg msg = do
   cfg <- asks (.cfg)
@@ -262,63 +265,70 @@ execTxWithCov tx = do
     _ -> pure False
 
   pure (r, grew || grew')
+
+-- | The same as EVM.exec but collects coverage, will stop on a query
+execCov
+  :: (MonadIO m, MonadState (VM Concrete RealWorld) m, MonadThrow m)
+  => Env
+  -> IORef (Bool, Maybe (VMut.IOVector CoverageInfo, Int))
+  -> m (VMResult Concrete RealWorld)
+execCov env covContextRef = do
+  vm <- get
+  (r, vm') <- liftIO $ loop vm
+  put vm'
+  pure r
   where
-    -- the same as EVM.exec but collects coverage, will stop on a query
-    execCov env covContextRef = do
-      vm <- get
-      (r, vm') <- liftIO $ loop vm
-      put vm'
-      pure r
-      where
-      -- | Repeatedly exec a step and add coverage until we have an end result
-      loop :: VM Concrete RealWorld -> IO (VMResult Concrete RealWorld, VM Concrete RealWorld)
-      loop !vm = case vm.result of
-        Nothing -> do
-          addCoverage vm
-          stepVM vm >>= loop
-        Just r -> pure (r, vm)
-
-      -- | Execute one instruction on the EVM
-      stepVM :: VM Concrete RealWorld -> IO (VM Concrete RealWorld)
-      stepVM = stToIO . execStateT exec1
-
-      -- | Add current location to the CoverageMap
-      addCoverage :: VM Concrete RealWorld -> IO ()
-      addCoverage !vm = do
-        let (pc, opIx, depth) = currentCovLoc vm
-            contract = currentContract vm
-
-        maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
-          if size == 0 then pure Nothing else do
-            -- IO for making a new vec
-            vec <- VMut.new size
-            -- We use -1 for opIx to indicate that the location was not covered
-            forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
-            pure $ Just vec
-
-        case maybeCovVec of
-          Nothing -> pure ()
-          Just vec -> do
-            -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
-            -- we observed this in some real-world scenarios. This is likely a
-            -- bug in another place, investigate.
-            -- ... this should be fixed now, since we use `codeContract` instead
-            -- of `contract` for everything; it may be safe to remove this check.
-            when (pc < VMut.length vec) $
-              VMut.read vec pc >>= \case
-                (_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
-                  VMut.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
-                  writeIORef covContextRef (True, Just (vec, pc))
-                _ ->
-                  modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
-
-      -- | Get the VM's current execution location
-      currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
-
-      -- | Get the current contract being executed
-      currentContract vm = fromMaybe (error "no contract information on coverage") $
-        vm ^? #env % #contracts % at vm.state.codeContract % _Just
+  -- | Repeatedly exec a step and add coverage until we have an end result
+  loop :: VM Concrete RealWorld -> IO (VMResult Concrete RealWorld, VM Concrete RealWorld)
+  loop !vm = case vm.result of
+    Nothing -> do
+      addCoverage vm
+      stepVM vm >>= loop
+    Just r -> pure (r, vm)
+
+  -- | Execute one instruction on the EVM
+  stepVM :: VM Concrete RealWorld -> IO (VM Concrete RealWorld)
+  stepVM = stToIO . execStateT exec1
+
+  -- | Add current location to the CoverageMap
+  addCoverage :: VM Concrete RealWorld -> IO ()
+  addCoverage !vm = do
+    let (pc, opIx, depth) = currentCovLoc vm
+        contract = currentContract vm
+
+    maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ createCoverageVec contract
+
+    case maybeCovVec of
+      Nothing -> pure ()
+      Just vec ->
+        -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
+        -- we observed this in some real-world scenarios. This is likely a
+        -- bug in another place, investigate.
+        -- ... this should be fixed now, since we use `codeContract` instead
+        -- of `contract` for everything; it may be safe to remove this check.
+        when (pc < VMut.length vec) $
+          VMut.read vec pc >>= \case
+            (_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
+              VMut.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
+              writeIORef covContextRef (True, Just (vec, pc))
+            _ ->
+              modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
+
+  createCoverageVec contract = do
+    let size = BS.length . forceBuf . fromJust . view bytecode $ contract
+    if size == 0 then pure Nothing else do
+      -- IO for making a new vec
+      vec <- VMut.new size
+      -- We use -1 for opIx to indicate that the location was not covered
+      forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
+      pure $ Just vec
+
+  -- | Get the VM's current execution location
+  currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
+
+  -- | Get the current contract being executed
+  currentContract vm = fromMaybe (error "no contract information on coverage") $
+    vm ^? #env % #contracts % at vm.state.codeContract % _Just
 
 initialVM :: Bool -> ST s (VM Concrete s)
 initialVM ffi = do