Prioritized and injected transaction sequences

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 =

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

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.

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)
 

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 = [] }

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