Skip to content

Commit 318a71a

Browse files
samalws-tobsamalws-tob
committed
Simplify Exec.hs
1 parent 73819e3 commit 318a71a

File tree

1 file changed

+151
-141
lines changed

1 file changed

+151
-141
lines changed

Diff for: ‎lib/Echidna/Exec.hs

+151-141
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import Control.Monad.Reader (MonadReader, ask, asks)
1414
import Control.Monad.ST (ST, stToIO, RealWorld)
1515
import Data.Bits
1616
import Data.ByteString qualified as BS
17-
import Data.IORef (readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
17+
import Data.IORef (IORef, readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
1818
import Data.Map qualified as Map
1919
import Data.Maybe (fromMaybe, fromJust)
2020
import Data.Text qualified as T
@@ -102,91 +102,8 @@ execTxWith executeTx tx = do
102102
pure (vmResult, gasLeftBeforeTx - gasLeftAfterTx)
103103
where
104104
runFully = do
105-
config <- asks (.cfg)
106-
-- TODO: Is the latest block a good default? It makes fuzzing hard to reproduce. Rethink this.
107-
let rpcBlock = maybe EVM.Fetch.Latest (EVM.Fetch.BlockNumber . fromIntegral) config.rpcBlock
108-
109105
vmResult <- executeTx
110-
-- For queries, we halt execution because the VM needs some additional
111-
-- information from the outside. We provide this information and resume
112-
-- the execution by recursively calling `runFully`.
113-
case getQuery vmResult of
114-
-- A previously unknown contract is required
115-
Just q@(PleaseFetchContract addr _ continuation) -> do
116-
cacheRef <- asks (.fetchContractCache)
117-
cache <- liftIO $ readIORef cacheRef
118-
case Map.lookup addr cache of
119-
Just (Just contract) -> fromEVM (continuation contract)
120-
Just Nothing -> do
121-
v <- get
122-
v' <- liftIO $ stToIO $ execStateT (continuation emptyAccount) v
123-
put v'
124-
Nothing -> do
125-
logMsg $ "INFO: Performing RPC: " <> show q
126-
case config.rpcUrl of
127-
Just rpcUrl -> do
128-
ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
129-
case ret of
130-
-- TODO: fix hevm to not return an empty contract in case of an error
131-
Just contract | contract.code /= RuntimeCode (ConcreteRuntimeCode "") -> do
132-
fromEVM (continuation contract)
133-
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr (Just contract) cache
134-
_ -> do
135-
-- TODO: better error reporting in HEVM, when intermittent
136-
-- network error then retry
137-
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
138-
logMsg $ "ERROR: Failed to fetch contract: " <> show q
139-
-- TODO: How should we fail here? It could be a network error,
140-
-- RPC server returning junk etc.
141-
fromEVM (continuation emptyAccount)
142-
Nothing -> do
143-
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
144-
logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
145-
-- TODO: How should we fail here? RPC is not configured but VM
146-
-- wants to fetch
147-
fromEVM (continuation emptyAccount)
148-
runFully -- resume execution
149-
150-
-- A previously unknown slot is required
151-
Just q@(PleaseFetchSlot addr slot continuation) -> do
152-
cacheRef <- asks (.fetchSlotCache)
153-
cache <- liftIO $ readIORef cacheRef
154-
case Map.lookup addr cache >>= Map.lookup slot of
155-
Just (Just value) -> fromEVM (continuation value)
156-
Just Nothing -> fromEVM (continuation 0)
157-
Nothing -> do
158-
logMsg $ "INFO: Performing RPC: " <> show q
159-
case config.rpcUrl of
160-
Just rpcUrl -> do
161-
ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
162-
case ret of
163-
Just value -> do
164-
fromEVM (continuation value)
165-
liftIO $ atomicWriteIORef cacheRef $
166-
Map.insertWith Map.union addr (Map.singleton slot (Just value)) cache
167-
Nothing -> do
168-
-- TODO: How should we fail here? It could be a network error,
169-
-- RPC server returning junk etc.
170-
logMsg $ "ERROR: Failed to fetch slot: " <> show q
171-
liftIO $ atomicWriteIORef cacheRef $
172-
Map.insertWith Map.union addr (Map.singleton slot Nothing) cache
173-
fromEVM (continuation 0)
174-
Nothing -> do
175-
logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
176-
-- Use the zero slot
177-
fromEVM (continuation 0)
178-
runFully -- resume execution
179-
180-
-- Execute a FFI call
181-
Just (PleaseDoFFI (cmd : args) continuation) -> do
182-
(_, stdout, _) <- liftIO $ readProcessWithExitCode cmd args ""
183-
let encodedResponse = encodeAbiValue $
184-
AbiTuple (V.fromList [AbiBytesDynamic . hexText . T.pack $ stdout])
185-
fromEVM (continuation encodedResponse)
186-
runFully
187-
188-
-- No queries to answer, the tx is fully executed and the result is final
189-
_ -> pure vmResult
106+
maybe (pure vmResult) (\q -> handleQuery q >> runFully) (getQuery vmResult)
190107

191108
-- | Handles reverts, failures and contract creations that might be the result
192109
-- (`vmResult`) of executing transaction `tx`.
@@ -217,6 +134,92 @@ execTxWith executeTx tx = do
217134
modify' $ execState $ loadContract (LitAddr tx.dst)
218135
_ -> pure ()
219136

137+
getRpcInfo = do
138+
config <- asks (.cfg)
139+
-- TODO: Is the latest block a good default? It makes fuzzing hard to reproduce. Rethink this.
140+
let rpcBlock = maybe EVM.Fetch.Latest (EVM.Fetch.BlockNumber . fromIntegral) config.rpcBlock
141+
return (config.rpcUrl, rpcBlock)
142+
143+
144+
-- For queries, we halt execution because the VM needs some additional
145+
-- information from the outside. We provide this information, and then
146+
-- the execution is resumed.
147+
148+
-- A previously unknown contract is required
149+
handleQuery q@(PleaseFetchContract addr _ continuation) = do
150+
cacheRef <- asks (.fetchContractCache)
151+
cache <- liftIO $ readIORef cacheRef
152+
case Map.lookup addr cache of
153+
Just (Just contract) -> fromEVM (continuation contract)
154+
Just Nothing -> do
155+
v <- get
156+
v' <- liftIO $ stToIO $ execStateT (continuation emptyAccount) v
157+
put v'
158+
Nothing -> do
159+
logMsg $ "INFO: Performing RPC: " <> show q
160+
(maybeRpcUrl, rpcBlock) <- getRpcInfo
161+
case maybeRpcUrl of
162+
Just rpcUrl -> do
163+
ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
164+
case ret of
165+
-- TODO: fix hevm to not return an empty contract in case of an error
166+
Just contract | contract.code /= RuntimeCode (ConcreteRuntimeCode "") -> do
167+
fromEVM (continuation contract)
168+
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr (Just contract) cache
169+
_ -> do
170+
-- TODO: better error reporting in HEVM, when intermittent
171+
-- network error then retry
172+
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
173+
logMsg $ "ERROR: Failed to fetch contract: " <> show q
174+
-- TODO: How should we fail here? It could be a network error,
175+
-- RPC server returning junk etc.
176+
fromEVM (continuation emptyAccount)
177+
Nothing -> do
178+
liftIO $ atomicWriteIORef cacheRef $ Map.insert addr Nothing cache
179+
logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
180+
-- TODO: How should we fail here? RPC is not configured but VM
181+
-- wants to fetch
182+
fromEVM (continuation emptyAccount)
183+
184+
-- A previously unknown slot is required
185+
handleQuery q@(PleaseFetchSlot addr slot continuation) = do
186+
cacheRef <- asks (.fetchSlotCache)
187+
cache <- liftIO $ readIORef cacheRef
188+
case Map.lookup addr cache >>= Map.lookup slot of
189+
Just (Just value) -> fromEVM (continuation value)
190+
Just Nothing -> fromEVM (continuation 0)
191+
Nothing -> do
192+
logMsg $ "INFO: Performing RPC: " <> show q
193+
(maybeRpcUrl, rpcBlock) <- getRpcInfo
194+
case maybeRpcUrl of
195+
Just rpcUrl -> do
196+
ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
197+
case ret of
198+
Just value -> do
199+
fromEVM (continuation value)
200+
liftIO $ atomicWriteIORef cacheRef $
201+
Map.insertWith Map.union addr (Map.singleton slot (Just value)) cache
202+
Nothing -> do
203+
-- TODO: How should we fail here? It could be a network error,
204+
-- RPC server returning junk etc.
205+
logMsg $ "ERROR: Failed to fetch slot: " <> show q
206+
liftIO $ atomicWriteIORef cacheRef $
207+
Map.insertWith Map.union addr (Map.singleton slot Nothing) cache
208+
fromEVM (continuation 0)
209+
Nothing -> do
210+
logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
211+
-- Use the zero slot
212+
fromEVM (continuation 0)
213+
214+
-- Execute a FFI call
215+
handleQuery (PleaseDoFFI (cmd : args) continuation) = do
216+
(_, stdout, _) <- liftIO $ readProcessWithExitCode cmd args ""
217+
let encodedResponse = encodeAbiValue $
218+
AbiTuple (V.fromList [AbiBytesDynamic . hexText . T.pack $ stdout])
219+
fromEVM (continuation encodedResponse)
220+
221+
handleQuery (PleaseDoFFI [] _) = error "Malformed FFI call"
222+
220223
logMsg :: (MonadIO m, MonadReader Env m) => String -> m ()
221224
logMsg msg = do
222225
cfg <- asks (.cfg)
@@ -262,63 +265,70 @@ execTxWithCov tx = do
262265
_ -> pure False
263266

264267
pure (r, grew || grew')
268+
269+
-- | The same as EVM.exec but collects coverage, will stop on a query
270+
execCov
271+
:: (MonadIO m, MonadState (VM Concrete RealWorld) m, MonadThrow m)
272+
=> Env
273+
-> IORef (Bool, Maybe (VMut.IOVector CoverageInfo, Int))
274+
-> m (VMResult Concrete RealWorld)
275+
execCov env covContextRef = do
276+
vm <- get
277+
(r, vm') <- liftIO $ loop vm
278+
put vm'
279+
pure r
265280
where
266-
-- the same as EVM.exec but collects coverage, will stop on a query
267-
execCov env covContextRef = do
268-
vm <- get
269-
(r, vm') <- liftIO $ loop vm
270-
put vm'
271-
pure r
272-
where
273-
-- | Repeatedly exec a step and add coverage until we have an end result
274-
loop :: VM Concrete RealWorld -> IO (VMResult Concrete RealWorld, VM Concrete RealWorld)
275-
loop !vm = case vm.result of
276-
Nothing -> do
277-
addCoverage vm
278-
stepVM vm >>= loop
279-
Just r -> pure (r, vm)
280-
281-
-- | Execute one instruction on the EVM
282-
stepVM :: VM Concrete RealWorld -> IO (VM Concrete RealWorld)
283-
stepVM = stToIO . execStateT exec1
284-
285-
-- | Add current location to the CoverageMap
286-
addCoverage :: VM Concrete RealWorld -> IO ()
287-
addCoverage !vm = do
288-
let (pc, opIx, depth) = currentCovLoc vm
289-
contract = currentContract vm
290-
291-
maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
292-
let size = BS.length . forceBuf . fromJust . view bytecode $ contract
293-
if size == 0 then pure Nothing else do
294-
-- IO for making a new vec
295-
vec <- VMut.new size
296-
-- We use -1 for opIx to indicate that the location was not covered
297-
forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
298-
pure $ Just vec
299-
300-
case maybeCovVec of
301-
Nothing -> pure ()
302-
Just vec -> do
303-
-- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
304-
-- we observed this in some real-world scenarios. This is likely a
305-
-- bug in another place, investigate.
306-
-- ... this should be fixed now, since we use `codeContract` instead
307-
-- of `contract` for everything; it may be safe to remove this check.
308-
when (pc < VMut.length vec) $
309-
VMut.read vec pc >>= \case
310-
(_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
311-
VMut.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
312-
writeIORef covContextRef (True, Just (vec, pc))
313-
_ ->
314-
modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
315-
316-
-- | Get the VM's current execution location
317-
currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
318-
319-
-- | Get the current contract being executed
320-
currentContract vm = fromMaybe (error "no contract information on coverage") $
321-
vm ^? #env % #contracts % at vm.state.codeContract % _Just
281+
-- | Repeatedly exec a step and add coverage until we have an end result
282+
loop :: VM Concrete RealWorld -> IO (VMResult Concrete RealWorld, VM Concrete RealWorld)
283+
loop !vm = case vm.result of
284+
Nothing -> do
285+
addCoverage vm
286+
stepVM vm >>= loop
287+
Just r -> pure (r, vm)
288+
289+
-- | Execute one instruction on the EVM
290+
stepVM :: VM Concrete RealWorld -> IO (VM Concrete RealWorld)
291+
stepVM = stToIO . execStateT exec1
292+
293+
-- | Add current location to the CoverageMap
294+
addCoverage :: VM Concrete RealWorld -> IO ()
295+
addCoverage !vm = do
296+
let (pc, opIx, depth) = currentCovLoc vm
297+
contract = currentContract vm
298+
299+
maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ createCoverageVec contract
300+
301+
case maybeCovVec of
302+
Nothing -> pure ()
303+
Just vec ->
304+
-- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
305+
-- we observed this in some real-world scenarios. This is likely a
306+
-- bug in another place, investigate.
307+
-- ... this should be fixed now, since we use `codeContract` instead
308+
-- of `contract` for everything; it may be safe to remove this check.
309+
when (pc < VMut.length vec) $
310+
VMut.read vec pc >>= \case
311+
(_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
312+
VMut.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
313+
writeIORef covContextRef (True, Just (vec, pc))
314+
_ ->
315+
modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
316+
317+
createCoverageVec contract = do
318+
let size = BS.length . forceBuf . fromJust . view bytecode $ contract
319+
if size == 0 then pure Nothing else do
320+
-- IO for making a new vec
321+
vec <- VMut.new size
322+
-- We use -1 for opIx to indicate that the location was not covered
323+
forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
324+
pure $ Just vec
325+
326+
-- | Get the VM's current execution location
327+
currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
328+
329+
-- | Get the current contract being executed
330+
currentContract vm = fromMaybe (error "no contract information on coverage") $
331+
vm ^? #env % #contracts % at vm.state.codeContract % _Just
322332

323333
initialVM :: Bool -> ST s (VM Concrete s)
324334
initialVM ffi = do

0 commit comments

Comments
 (0)