Simplify Exec.hs

2 months ago · 318a71ada1
parent 73819e31cd
commit 318a71ada1
1 changed files with 151 additions and 141 deletions
--- 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,17 +102,51 @@ execTxWith executeTx tx = do
        pure (vmResult, gasLeftBeforeTx - gasLeftAfterTx)
  where
  runFully = do
+    vmResult <- executeTx
+    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`.
+  handleErrorsAndConstruction vmResult vmBeforeTx = case (vmResult, tx.call) of
+    (Reversion, _) -> do
+      tracesBeforeVMReset <- gets (.traces)
+      codeContractBeforeVMReset <- gets (.state.codeContract)
+      calldataBeforeVMReset <- gets (.state.calldata)
+      callvalueBeforeVMReset <- gets (.state.callvalue)
+      -- If a transaction reverts reset VM to state before the transaction.
+      put vmBeforeTx
+      -- Undo reset of some of the VM state.
+      -- Otherwise we'd loose all information about the reverted transaction like
+      -- contract address, calldata, result and traces.
+      #result ?= vmResult
+      #state % #calldata .= calldataBeforeVMReset
+      #state % #callvalue .= callvalueBeforeVMReset
+      #traces .= tracesBeforeVMReset
+      #state % #codeContract .= codeContractBeforeVMReset
+    (VMFailure x, _) -> do
+      dapp <- asks (.dapp)
+      vm <- get
+      vmExcept (Just (dapp, vm)) x
+    (VMSuccess (ConcreteBuf bytecode'), SolCreate _) -> do
+      -- Handle contract creation.
+      #env % #contracts % at (LitAddr tx.dst) % _Just % #code .= InitCode mempty mempty
+      fromEVM $ replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bytecode'))
+      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)
+

-    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
+  -- information from the outside. We provide this information, and then
+  -- the execution is resumed.
+
  -- A previously unknown contract is required
-      Just q@(PleaseFetchContract addr _ continuation) -> do
+  handleQuery q@(PleaseFetchContract addr _ continuation) = do
    cacheRef <- asks (.fetchContractCache)
    cache <- liftIO $ readIORef cacheRef
    case Map.lookup addr cache of
@ -123,7 +157,8 @@ execTxWith executeTx tx = do
        put v'
      Nothing -> do
        logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
          Just rpcUrl -> do
            ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
            case ret of
@ -145,10 +180,9 @@ execTxWith executeTx tx = do
            -- 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
+  handleQuery q@(PleaseFetchSlot addr slot continuation) = do
    cacheRef <- asks (.fetchSlotCache)
    cache <- liftIO $ readIORef cacheRef
    case Map.lookup addr cache >>= Map.lookup slot of
@ -156,7 +190,8 @@ execTxWith executeTx tx = do
      Just Nothing -> fromEVM (continuation 0)
      Nothing -> do
        logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
          Just rpcUrl -> do
            ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
            case ret of
@ -175,47 +210,15 @@ execTxWith executeTx tx = 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
+  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)
-        runFully
-
-      -- No queries to answer, the tx is fully executed and the result is final
-      _ -> pure vmResult

-  -- | Handles reverts, failures and contract creations that might be the result
-  -- (`vmResult`) of executing transaction `tx`.
-  handleErrorsAndConstruction vmResult vmBeforeTx = case (vmResult, tx.call) of
-    (Reversion, _) -> do
-      tracesBeforeVMReset <- gets (.traces)
-      codeContractBeforeVMReset <- gets (.state.codeContract)
-      calldataBeforeVMReset <- gets (.state.calldata)
-      callvalueBeforeVMReset <- gets (.state.callvalue)
-      -- If a transaction reverts reset VM to state before the transaction.
-      put vmBeforeTx
-      -- Undo reset of some of the VM state.
-      -- Otherwise we'd loose all information about the reverted transaction like
-      -- contract address, calldata, result and traces.
-      #result ?= vmResult
-      #state % #calldata .= calldataBeforeVMReset
-      #state % #callvalue .= callvalueBeforeVMReset
-      #traces .= tracesBeforeVMReset
-      #state % #codeContract .= codeContractBeforeVMReset
-    (VMFailure x, _) -> do
-      dapp <- asks (.dapp)
-      vm <- get
-      vmExcept (Just (dapp, vm)) x
-    (VMSuccess (ConcreteBuf bytecode'), SolCreate _) -> do
-      -- Handle contract creation.
-      #env % #contracts % at (LitAddr tx.dst) % _Just % #code .= InitCode mempty mempty
-      fromEVM $ replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bytecode'))
-      modify' $ execState $ loadContract (LitAddr tx.dst)
-    _ -> pure ()
+  handleQuery (PleaseDoFFI [] _) = error "Malformed FFI call"

 logMsg :: (MonadIO m, MonadReader Env m) => String -> m ()
 logMsg msg = do
@ -262,9 +265,14 @@ execTxWithCov tx = do
    _ -> pure False

  pure (r, grew || grew')
-  where
-    -- the same as EVM.exec but collects coverage, will stop on a query
-    execCov env covContextRef = do
+
+-- | 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'
@ -288,18 +296,11 @@ execTxWithCov tx = 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
+    maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ createCoverageVec contract

    case maybeCovVec of
      Nothing -> pure ()
-          Just vec -> do
+      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.
@ -313,6 +314,15 @@ execTxWithCov tx = do
            _ ->
              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)