diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs index 48f6e4ac6..24953a82f 100644 --- a/lib/Echidna/Transaction.hs +++ b/lib/Echidna/Transaction.hs @@ -3,12 +3,14 @@ module Echidna.Transaction where -import Control.Monad (join, when) +import Control.Monad (join, replicateM, when, zipWithM) import Control.Monad.IO.Class (MonadIO, liftIO) import Control.Monad.Random.Strict (MonadRandom, getRandomR, uniform) -import Control.Monad.Reader (MonadReader, ask) +import Control.Monad.Reader (MonadReader, ask, asks) import Control.Monad.State.Strict (MonadState, gets, modify', execState) import Data.ByteString qualified as BS +import Data.IORef (readIORef) +import Data.List.NonEmpty qualified as NE import Data.Map (Map, toList) import Data.Maybe (catMaybes) import Data.Set (Set) @@ -31,7 +33,7 @@ import Echidna.Types (fromEVM, Gas) import Echidna.Types.Config (Env(..), EConfig(..)) import Echidna.Types.Random import Echidna.Types.Signature - (SignatureMap, SolCall, ContractA) + (ContractA, SignatureMap, SolCall, SolCallPrototype) import Echidna.Types.Tx import Echidna.Types.World (World(..)) import Echidna.Types.Campaign @@ -62,14 +64,98 @@ genTx -> Map (Expr EAddr) Contract -> m Tx genTx world deployedContracts = do + contracts <- callableContracts world deployedContracts + (dstAddr, solCall) <- genRandomCall contracts + toTx world dstAddr solCall + +-- | Generate a 'Transaction' calling the function a prototype names, letting +-- the generator fill in whichever of its arguments the prototype left open. +-- +-- Falls back to a fully random transaction when nothing deployed exposes such +-- a function, so a prototype naming a function that isn't there costs +-- diversity rather than stalling the worker. +genTxFromPrototype + :: (MonadIO m, MonadRandom m, MonadState WorkerState m, MonadReader Env m) + => World + -> Map (Expr EAddr) Contract + -> SolCallPrototype + -> m Tx +genTxFromPrototype world deployedContracts prototype = do + contracts <- callableContracts world deployedContracts + (dstAddr, solCall) <- case matchingContracts prototype contracts of + [] -> genRandomCall contracts + candidates -> genPrototypeCall prototype candidates + toTx world dstAddr solCall + +-- | The deployed contracts Echidna knows how to build a call against, each +-- paired with the ABI resolved for it. +callableContracts + :: (MonadIO m, MonadRandom m, MonadReader Env m) + => World + -> Map (Expr EAddr) Contract + -> m [ContractA] +callableContracts world deployedContracts = do env <- ask - let txConf = env.cfg.txConf - genDict <- gets (.genDict) sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap + catMaybes <$> liftIO (mapM (toContractA env sigMap) (toList deployedContracts)) + where + toContractA :: Env -> SignatureMap -> (Expr EAddr, Contract) -> IO (Maybe ContractA) + toContractA env sigMap (addr, c) = + fmap (forceAddr addr,) . snd <$> lookupUsingCodehash env.codehashMap c env.dapp sigMap + +-- | The contracts exposing the function a prototype names, each cut down to +-- just the signatures matching it. +-- +-- Matching is on name and arity only: an argument the prototype leaves open +-- carries no type to compare against, so same-arity overloads all qualify and +-- one is picked at random. +matchingContracts :: SolCallPrototype -> [ContractA] -> [ContractA] +matchingContracts (name, args) contracts = + [ (addr, matchingSigs) + | (addr, sigs) <- contracts + , Just matchingSigs <- [NE.nonEmpty (NE.filter matches sigs)] + ] + where + matches (n, types) = n == name && length types == length args + +-- | Pick one of the given contracts and generate a call to a random function +-- of it. +genRandomCall + :: (MonadRandom m, MonadState WorkerState m) + => [ContractA] + -> m (Addr, SolCall) +genRandomCall contracts = do + genDict <- gets (.genDict) + (dstAddr, dstAbis) <- rElem' $ Set.fromList contracts + (dstAddr,) <$> genInteractionsM genDict dstAbis + +-- | Pick one of the contracts 'matchingContracts' selected and generate the +-- prototype's call against it, generating a value for every argument left open. +genPrototypeCall + :: (MonadRandom m, MonadState WorkerState m) + => SolCallPrototype + -> [ContractA] -- ^ From 'matchingContracts', so every signature matches + -> m (Addr, SolCall) +genPrototypeCall (name, args) candidates = do + genDict <- gets (.genDict) + (dstAddr, dstAbis) <- rElem' $ Set.fromList candidates + -- Only the argument types are taken from the signature; its name is the one + -- the prototype asked for. + (_, types) <- rElem dstAbis + vals <- zipWithM (\arg t -> maybe (genAbiValueM' genDict name 0 t) pure arg) args types + pure (dstAddr, (name, vals)) + +-- | Wrap a chosen call into a 'Tx', giving it a random sender, value and delay. +toTx + :: (MonadRandom m, MonadState WorkerState m, MonadReader Env m) + => World + -> Addr + -> SolCall + -> m Tx +toTx world dstAddr solCall = do + txConf <- asks (.cfg.txConf) + genDict <- gets (.genDict) sender <- rElem' world.senders - contractAList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts) - (dstAddr, dstAbis) <- rElem' $ Set.fromList $ catMaybes contractAList - solCall <- genInteractionsM genDict dstAbis value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues <*> genDelay txConf.maxBlockDelay genDict.dictValues @@ -81,10 +167,61 @@ genTx world deployedContracts = do , value = value , delay = level ts } + +-- | Maximum number of random transactions inserted between two consecutive +-- calls of a prioritized sequence. +maxInterleavedTxs :: Int +maxInterleavedTxs = 3 + +-- | Expand a prioritized sequence of prototypes into a transaction sequence of +-- the configured length. +-- +-- The prototype calls are generated in order, with up to 'maxInterleavedTxs' +-- random transactions between consecutive ones for diversity, and prefixed +-- with the start of a corpus entry so the sequence runs from a state the +-- campaign already reached. Worker 0 always takes the empty prefix, so the +-- prototype is exercised from the initial state too. The result is padded with +-- random transactions, or truncated, to respect @seqLen@. +genPrioritizedSeq + :: (MonadIO m, MonadRandom m, MonadState WorkerState m, MonadReader Env m) + => Map (Expr EAddr) Contract + -> [SolCallPrototype] + -> m [Tx] +genPrioritizedSeq deployedContracts prototypes = do + env <- ask + let world = env.world + seqLen = env.cfg.campaignConf.seqLen + + txs <- expand world prototypes + prefix <- corpusPrefix (seqLen - length txs) + + let combined = prefix ++ txs + padding = seqLen - length combined + if padding > 0 + then (combined ++) <$> replicateM padding (genTx world deployedContracts) + else pure $ take seqLen combined + where - toContractA :: Env -> SignatureMap -> (Expr EAddr, Contract) -> IO (Maybe ContractA) - toContractA env sigMap (addr, c) = - fmap (forceAddr addr,) . snd <$> lookupUsingCodehash env.codehashMap c env.dapp sigMap + expand _ [] = pure [] + expand world (p:ps) = do + tx <- genTxFromPrototype world deployedContracts p + case ps of + [] -> pure [tx] + _ -> do + n <- getRandomR (0, maxInterleavedTxs) + filler <- replicateM n (genTx world deployedContracts) + ((tx : filler) ++) <$> expand world ps + + -- The start of one corpus entry, at most @room@ transactions of it. + corpusPrefix room = do + workerId <- gets (.workerId) + corpus <- asks (.corpusRef) >>= liftIO . readIORef + if workerId == 0 || room <= 0 || Set.null corpus + then pure [] + else do + (_, corpusTxs) <- rElem' corpus + k <- getRandomR (0, min (length corpusTxs) room) + pure $ take k corpusTxs genDelay :: MonadRandom m => W256 -> Set W256 -> m W256 genDelay mv ds = diff --git a/lib/Echidna/Types/Campaign.hs b/lib/Echidna/Types/Campaign.hs index fef6d95d4..20e23fe39 100644 --- a/lib/Echidna/Types/Campaign.hs +++ b/lib/Echidna/Types/Campaign.hs @@ -14,6 +14,7 @@ import EVM.Solvers (Solver(..)) import Echidna.ABI (GenDict, emptyDict) import Echidna.Types import Echidna.Types.Coverage (CoverageFileType, CoverageMap) +import Echidna.Types.Signature (SolCallPrototype) import Echidna.Types.Tx (TxResult(..)) -- | Maximum number of functions a single worker samples at once. @@ -195,6 +196,10 @@ data WorkerState = WorkerState -- ^ Functions whose calls are sampled for return-value range and revert -- history, keyed by canonical signature (e.g. @"totalSupply()"@). Empty -- unless sampling was explicitly enabled for this worker. + , prioritizedSequences :: ![(Double, [SolCallPrototype])] + -- ^ Call sequences to bias generation towards, each with the probability + -- of being used in place of a corpus-mutated sequence. Empty unless + -- sequences were explicitly injected into this worker. } initialWorkerState :: WorkerState @@ -207,6 +212,7 @@ initialWorkerState = , totalGas = 0 , runningThreads = [] , sampledFunctions = Map.empty + , prioritizedSequences = [] } defaultTestLimit :: Int diff --git a/lib/Echidna/Types/InterWorker.hs b/lib/Echidna/Types/InterWorker.hs index e25ce6716..a5a164823 100644 --- a/lib/Echidna/Types/InterWorker.hs +++ b/lib/Echidna/Types/InterWorker.hs @@ -16,6 +16,7 @@ module Echidna.Types.InterWorker import Control.Concurrent.STM (TChan) import Data.Text (Text) +import Echidna.Types.Signature (SolCallPrototype) import Echidna.Types.Tx (Tx) import Echidna.Types.Worker (WorkerId) @@ -30,6 +31,11 @@ data FuzzerCmd -- Capped per worker by 'Echidna.Types.Campaign.maxSampledFunctions'. | ClearSampling -- ^ Forget every sampled function and its statistics. + | FuzzSequence [SolCallPrototype] Double + -- ^ Bias generation towards a sequence of calls, using it with the given + -- probability in place of a corpus-mutated one. + | ClearPrioritization + -- ^ Forget every prioritized sequence. deriving Show -- | A message every agent gets to see. diff --git a/lib/Echidna/Types/Signature.hs b/lib/Echidna/Types/Signature.hs index 5d7cd9cb0..fbc9de829 100644 --- a/lib/Echidna/Types/Signature.hs +++ b/lib/Echidna/Types/Signature.hs @@ -25,6 +25,10 @@ type SolSignature = (FunctionName, [AbiType]) -- A tuple for the name of the function and then any 'AbiValue' arguments passed (as a list). type SolCall = (FunctionName, [AbiValue]) +-- | A 'SolCall' with holes in it: the arguments given as 'Just' are fixed, and +-- the ones given as 'Nothing' are left for the fuzzer to fill in. +type SolCallPrototype = (FunctionName, [Maybe AbiValue]) + -- | A contract is just an address with an ABI (for our purposes). type ContractA = (Addr, NonEmpty SolSignature) diff --git a/lib/Echidna/Worker/Command.hs b/lib/Echidna/Worker/Command.hs index 8a94e72f2..f52be8d28 100644 --- a/lib/Echidna/Worker/Command.hs +++ b/lib/Echidna/Worker/Command.hs @@ -43,3 +43,11 @@ handleCmd (EnableSampling sig) = handleCmd ClearSampling = modify' $ \workerState -> workerState { sampledFunctions = Map.empty } + +handleCmd (FuzzSequence prototypes prob) = + modify' $ \workerState -> workerState + { prioritizedSequences = (prob, prototypes) : workerState.prioritizedSequences + } + +handleCmd ClearPrioritization = + modify' $ \workerState -> workerState { prioritizedSequences = [] } diff --git a/lib/Echidna/Worker/Fuzz.hs b/lib/Echidna/Worker/Fuzz.hs index 91e9ee606..636b3f697 100644 --- a/lib/Echidna/Worker/Fuzz.hs +++ b/lib/Echidna/Worker/Fuzz.hs @@ -6,11 +6,12 @@ module Echidna.Worker.Fuzz (runFuzzWorker) where import Control.Concurrent.STM (atomically, dupTChan) import Control.Monad (forM_, replicateM, void) import Control.Monad.Catch (MonadThrow) -import Control.Monad.Random.Strict (MonadRandom, evalRandT) +import Control.Monad.Random.Strict (MonadRandom, evalRandT, getRandom) import Control.Monad.Reader (MonadReader, ask, asks, liftIO) import Control.Monad.State.Strict (MonadIO, MonadState, StateT, gets, runStateT) import Control.Monad.Trans (lift) import Data.IORef (atomicModifyIORef', readIORef) +import Data.List.NonEmpty qualified as NE import Data.Map (Map) import System.Random (mkStdGen) @@ -23,6 +24,7 @@ import Echidna.Shrink (isShrinkable, shrinkWorkerTests) import Echidna.Transaction import Echidna.Types.Campaign import Echidna.Types.Config +import Echidna.Types.Random (rElem) import Echidna.Types.Test import Echidna.Types.Test qualified as Test import Echidna.Types.Tx (Tx) @@ -110,13 +112,32 @@ runFuzzWorker callback vm dict workerId initialCorpus testLimit = do -- workers will "drain" the work queue. shrink = shrinkWorkerTests workerId vm --- | Generate a new sequences of transactions, either using the corpus or with --- randomly created transactions +-- | Generate a new sequence of transactions: from one of the sequences +-- prioritized over the bus, or, failing that, the standard way. randseq :: (MonadRandom m, MonadReader Env m, MonadState WorkerState m, MonadIO m) => Map (Expr 'EAddr) Contract -> m [Tx] randseq deployedContracts = do + prioritized <- gets (.prioritizedSequences) + case NE.nonEmpty prioritized of + Nothing -> genStandardSeq deployedContracts + Just seqs -> do + -- Pick one of the prioritized sequences, then use it only with the + -- probability it was injected with. + (prob, prototypes) <- rElem seqs + roll <- getRandom + if roll <= prob + then genPrioritizedSeq deployedContracts prototypes + else genStandardSeq deployedContracts + +-- | Generate a new sequence of transactions, either using the corpus or with +-- randomly created transactions +genStandardSeq + :: (MonadRandom m, MonadReader Env m, MonadState WorkerState m, MonadIO m) + => Map (Expr 'EAddr) Contract + -> m [Tx] +genStandardSeq deployedContracts = do env <- ask let world = env.world