Save corpus while fuzzing; add TestSimplified event; change event system

lib/Echidna/Async.hs

@@ -0,0 +1,39 @@
+module Echidna.Async where
+
+import Control.Concurrent.Thread.Group (forkIO)
+import Control.Monad (void)
+import Control.Monad.Reader (MonadReader, asks, ask)
+import Control.Monad.State.Strict (MonadState, gets)
+import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.IORef (atomicModifyIORef', readIORef)
+import Data.Time (LocalTime)
+
+import Echidna.Types.Campaign (CampaignEvent, WorkerState(..))
+import Echidna.Types.Config (Env(..))
+import Echidna.Utility (getTimestamp)
+
+spawnThread :: Env -> IO a -> IO ()
+spawnThread env io = void $ forkIO env.threadGroup io
+
+addEventHandler
+  :: (MonadReader Env m, MonadIO m)
+  => ((Int, LocalTime, CampaignEvent) -> IO ())
+  -> m ()
+addEventHandler f = do
+  handlersRef <- asks (.eventHandlers)
+  liftIO $ atomicModifyIORef' handlersRef (\l -> (f:l, ()))
+
+pushEvent
+  :: (MonadReader Env m, MonadState WorkerState m, MonadIO m)
+  => CampaignEvent
+  -> m ()
+pushEvent event = do
+  workerId <- gets (.workerId)
+  env <- ask
+  liftIO $ pushEventIO env workerId event
+
+pushEventIO :: Env -> Int -> CampaignEvent -> IO ()
+pushEventIO env workerId event = do
+  time <- liftIO getTimestamp
+  handlers <- readIORef env.eventHandlers
+  mapM_ (\f -> spawnThread env $ f (workerId, time, event)) handlers

lib/Echidna/Campaign.hs

@@ -4,7 +4,6 @@ module Echidna.Campaign where
 
 import Optics.Core hiding ((|>))
 
-import Control.Concurrent (writeChan)
 import Control.DeepSeq (force)
 import Control.Monad (replicateM, when, void, forM_)
 import Control.Monad.Catch (MonadThrow(..))
@@ -29,6 +28,7 @@ import EVM.ABI (getAbi, AbiType(AbiAddressType), AbiValue(AbiAddress))
 import EVM.Types hiding (Env, Frame(state))
 
 import Echidna.ABI
+import Echidna.Async (pushEvent)
 import Echidna.Exec
 import Echidna.Events (extractEvents)
 import Echidna.Mutator.Corpus
@@ -46,7 +46,6 @@ import Echidna.Types.Test
 import Echidna.Types.Test qualified as Test
 import Echidna.Types.Tx (TxCall(..), Tx(..), call)
 import Echidna.Types.World (World)
-import Echidna.Utility (getTimestamp)
 
 instance MonadThrow m => MonadThrow (RandT g m) where
   throwM = lift . throwM
@@ -216,7 +215,7 @@ callseq vm txSeq = do
 
     cov <- liftIO . readIORef =<< asks (.coverageRef)
     points <- liftIO $ scoveragePoints cov
-    pushEvent (NewCoverage points (length cov) newSize)
+    pushEvent (NewCoverage points (length cov) newSize (fst <$> results))
 
   modify' $ \workerState ->
 
@@ -392,13 +391,3 @@ updateTest vmForShrink (vm, xs) test = do
       -- but requires passing `vmForShrink` and feels a bit wrong.
       shrinkTest vmForShrink test
     _ -> pure Nothing
-
-pushEvent
-  :: (MonadReader Env m, MonadState WorkerState m, MonadIO m)
-  => CampaignEvent
-  -> m ()
-pushEvent event = do
-  workerId <- gets (.workerId)
-  time <- liftIO getTimestamp
-  chan <- asks (.eventQueue)
-  liftIO $ writeChan chan (workerId, time, event)

lib/Echidna/Output/Corpus.hs

@@ -1,5 +1,9 @@
 module Echidna.Output.Corpus where
 
+import Control.Exception (handle, IOException)
+import Control.Monad (unless)
+import Control.Monad.Reader (MonadReader, ask)
+import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Extra (unlessM)
 import Data.Aeson (ToJSON(..), decodeStrict, encodeFile)
 import Data.ByteString qualified as BS
@@ -8,13 +12,17 @@ import Data.Maybe (catMaybes)
 import System.Directory (createDirectoryIfMissing, makeRelativeToCurrentDirectory, doesFileExist)
 import System.FilePath ((</>), (<.>))
 
+import Echidna.Async (addEventHandler, pushEventIO)
+import Echidna.Types.Campaign (CampaignEvent(..), CampaignConf(..))
+import Echidna.Types.Config (Env(..), EConfig(..))
+import Echidna.Types.Test (EchidnaTest(..))
 import Echidna.Types.Tx (Tx)
 import Echidna.Utility (listDirectory, withCurrentDirectory)
 
-saveTxs :: FilePath -> [[Tx]] -> IO ()
-saveTxs dir = mapM_ saveTxSeq where
+saveTxs :: FilePath -> String -> [[Tx]] -> IO ()
+saveTxs dir prefix = mapM_ saveTxSeq where
   saveTxSeq txSeq = do
-    let file = dir </> (show . hash . show) txSeq <.> "txt"
+    let file = dir </> prefix ++ (show . abs . hash . show) txSeq <.> "txt"
     unlessM (doesFileExist file) $ encodeFile file (toJSON txSeq)
 
 loadTxs :: FilePath -> IO [[Tx]]
@@ -26,3 +34,21 @@ loadTxs dir = do
   putStrLn ("Loaded " ++ show (length txSeqs) ++ " transaction sequences from " ++ dir)
   pure txSeqs
   where readCall f = decodeStrict <$> BS.readFile f
+
+-- setup a handler to save to corpus in the background while tests are running
+setupCorpusSaver :: (MonadReader Env m, MonadIO m) => m ()
+setupCorpusSaver = do
+  env <- ask
+  maybe (pure ()) (addEventHandler . saveEvent env) env.cfg.campaignConf.corpusDir
+  where
+    saveEvent env dir (workerId, _, event) = maybe (pure ()) (saveFile workerId env dir) $ getEventInfo event
+
+    getEventInfo (TestFalsified test) = Just ("reproducers", "unshrunk-", test.reproducer)
+    getEventInfo (TestOptimized test) = Just ("reproducers", "", test.reproducer)
+    getEventInfo (TestShrunk    test) = Just ("reproducers", "", test.reproducer)
+    getEventInfo (NewCoverage _ _ _ txs) = Just ("coverage", "", txs)
+    getEventInfo _ = Nothing
+
+    saveFile workerId env dir (subdir, prefix, txs) = unless (null txs) $ handle (exceptionHandler workerId env) $ saveTxs (dir </> subdir) prefix [txs]
+
+    exceptionHandler workerId env (e :: IOException) = pushEventIO env workerId . HandlerFailed $ "Problem while writing to file: " ++ show e

lib/Echidna/Shrink.hs

@@ -2,34 +2,36 @@ module Echidna.Shrink (shrinkTest) where
 
 import Control.Monad ((<=<))
 import Control.Monad.Catch (MonadThrow)
+import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Random.Strict (MonadRandom, getRandomR, uniform)
 import Control.Monad.Reader.Class (MonadReader (ask), asks)
-import Control.Monad.State.Strict (MonadIO)
+import Control.Monad.State.Strict (MonadState)
 import Data.Set qualified as Set
 import Data.List qualified as List
 
 import EVM.Types (VM)
 
+import Echidna.Async (pushEvent)
 import Echidna.Events (extractEvents)
 import Echidna.Exec
 import Echidna.Transaction
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Test (TestValue(..), EchidnaTest(..), TestState(..), isOptimizationTest)
 import Echidna.Types.Tx (Tx(..))
 import Echidna.Types.Config
-import Echidna.Types.Campaign (CampaignConf(..))
+import Echidna.Types.Campaign (CampaignConf(..), CampaignEvent(..), WorkerState(..))
 import Echidna.Test (getResultFromVM, checkETest)
 
 shrinkTest
-  :: (MonadIO m, MonadThrow m, MonadRandom m, MonadReader Env m)
+  :: (MonadIO m, MonadThrow m, MonadRandom m, MonadReader Env m, MonadState WorkerState m)
   => VM
   -> EchidnaTest
   -> m (Maybe EchidnaTest)
 shrinkTest vm test = do
   env <- ask
   case test.state of
     Large i | i >= env.cfg.campaignConf.shrinkLimit && not (isOptimizationTest test) ->
-      pure $ Just test { state = Solved }
+      solvedEvent $ test { state = Solved }
     Large i ->
       if length test.reproducer > 1 || any canShrinkTx test.reproducer then do
         maybeShrunk <- shrinkSeq vm (checkETest test) test.value test.reproducer
@@ -43,11 +45,13 @@ shrinkTest vm test = do
           Nothing ->
             -- No success with shrinking this time, just bump trials
             Just test { state = Large (i + 1) }
+      else if isOptimizationTest test then
+        pure $ Just test { state = Large (i + 1) }
       else
-        pure $ Just test { state = if isOptimizationTest test
-                                 then Large (i + 1)
-                                 else Solved }
+        solvedEvent $ test { state = Solved }
     _ -> pure Nothing
+  where
+    solvedEvent test' = pushEvent (TestShrunk test') >> pure (Just test')
 
 -- | Given a call sequence that solves some Echidna test, try to randomly
 -- generate a smaller one that still solves that test.

lib/Echidna/Types/Campaign.hs

@@ -44,8 +44,12 @@ data CampaignConf = CampaignConf
 data CampaignEvent
   = TestFalsified !EchidnaTest
   | TestOptimized !EchidnaTest
-  | NewCoverage !Int !Int !Int
+  | TestShrunk !EchidnaTest
+  | NewCoverage !Int !Int !Int [Tx]
   | TxSequenceReplayed !Int !Int
+  | HandlerFailed !String
+  -- ^ Error occurred while handling another event
+  -- (e.g. failed to write coverage to a file)
   | WorkerStopped WorkerStopReason
   -- ^ This is a terminal event. Worker exits and won't push any events after
   -- this one
@@ -62,21 +66,19 @@ data WorkerStopReason
 ppCampaignEvent :: CampaignEvent -> String
 ppCampaignEvent = \case
   TestFalsified test ->
-    let name = case test.testType of
-                 PropertyTest n _ -> n
-                 AssertionTest _ n _ -> encodeSig n
-                 CallTest n _ -> n
-                 _ -> error "impossible"
-    in "Test " <> T.unpack name <> " falsified!"
+    "Test " <> T.unpack (showTest test) <> " falsified!"
   TestOptimized test ->
     let name = case test.testType of OptimizationTest n _ -> n; _ -> error "fixme"
     in "New maximum value of " <> T.unpack name <> ": " <> show test.value
-  NewCoverage points codehashes corpus ->
+  TestShrunk test ->
+    "Test " <> T.unpack (showTest test) <> " shrunk."
+  NewCoverage points codehashes corpus _ ->
     "New coverage: " <> show points <> " instr, "
       <> show codehashes <> " contracts, "
       <> show corpus <> " seqs in corpus"
   TxSequenceReplayed current total ->
     "Sequence replayed from corpus (" <> show current <> "/" <> show total <> ")"
+  HandlerFailed s -> "Error while handling event: " ++ s
   WorkerStopped TestLimitReached ->
     "Test limit reached. Stopping."
   WorkerStopped TimeLimitReached ->
@@ -89,6 +91,12 @@ ppCampaignEvent = \case
     "Crashed:\n\n" <>
     e <>
     "\n\nPlease report it to https://github.com/crytic/echidna/issues"
+  where
+    showTest test = case test.testType of
+      PropertyTest n _ -> n
+      AssertionTest _ n _ -> encodeSig n
+      CallTest n _ -> n
+      _ -> error "impossible"
 
 -- | The state of a fuzzing campaign.
 data WorkerState = WorkerState

lib/Echidna/Types/Config.hs

@@ -1,6 +1,6 @@
 module Echidna.Types.Config where
 
-import Control.Concurrent (Chan)
+import Control.Concurrent.Thread.Group (ThreadGroup)
 import Data.Aeson.Key (Key)
 import Data.IORef (IORef)
 import Data.Map (Map)
@@ -63,9 +63,11 @@ data Env = Env
   { cfg :: EConfig
   , dapp :: DappInfo
 
-  -- | Shared between all workers. Events are fairly rare so contention is
-  -- minimal.
-  , eventQueue :: Chan (Int, LocalTime, CampaignEvent)
+  , eventHandlers :: IORef [(Int, LocalTime, CampaignEvent) -> IO ()]
+
+  -- mainly for handling events, but can be used for any purpose
+  -- `wait` is called on this group before echidna closes
+  , threadGroup :: ThreadGroup
 
   , testsRef :: IORef [EchidnaTest]
   , coverageRef :: IORef CoverageMap

lib/Echidna/UI.hs

@@ -35,6 +35,7 @@ import UnliftIO.Concurrent hiding (killThread, threadDelay)
 import EVM.Types (Addr, Contract, VM, W256)
 
 import Echidna.ABI
+import Echidna.Async (addEventHandler, pushEventIO)
 import Echidna.Campaign (runWorker)
 import Echidna.Output.JSON qualified
 import Echidna.Types.Campaign
@@ -87,16 +88,13 @@ ui vm world dict initialCorpus = do
   workers <- forM (zip corpusChunks [0..(nworkers-1)]) $
     uncurry (spawnWorker env perWorkerTestLimit)
 
-  -- A var used to block and wait for listener to finish
-  listenerStopVar <- newEmptyMVar
-
   case effectiveMode of
 #ifdef INTERACTIVE_UI
     Interactive -> do
       -- Channel to push events to update UI
       uiChannel <- liftIO $ newBChan 1000
       let forwardEvent = writeBChan uiChannel . WorkerEvent
-      liftIO $ spawnListener env forwardEvent nworkers listenerStopVar
+      addEventHandler forwardEvent
 
       ticker <- liftIO . forkIO . forever $ do
         threadDelay 200_000 -- 200 ms
@@ -144,9 +142,6 @@ ui vm world dict initialCorpus = do
       -- Exited from the UI, stop the workers, not needed anymore
       stopWorkers workers
 
-      -- wait for all events to be processed
-      takeMVar listenerStopVar
-
       liftIO $ killThread ticker
 
       states <- workerStates workers
@@ -164,7 +159,7 @@ ui vm world dict initialCorpus = do
         installHandler sig (Catch $ stopWorkers workers) Nothing
 #endif
       let forwardEvent = putStrLn . ppLogLine
-      liftIO $ spawnListener env forwardEvent nworkers listenerStopVar
+      addEventHandler forwardEvent
 
       let printStatus = do
             states <- liftIO $ workerStates workers
@@ -177,9 +172,6 @@ ui vm world dict initialCorpus = do
         threadDelay 3_000_000 -- 3 seconds
         printStatus
 
-      -- wait for all events to be processed
-      takeMVar listenerStopVar
-
       liftIO $ killThread ticker
 
       -- print final status regardless the last scheduled update
@@ -216,37 +208,14 @@ ui vm world dict initialCorpus = do
         , Handler $ \(e :: SomeException)  -> pure $ Crashed (show e)
         ]
 
-      time <- liftIO getTimestamp
-      writeChan env.eventQueue (workerId, time, WorkerStopped stopReason)
+      liftIO $ pushEventIO env workerId (WorkerStopped stopReason)
 
     pure (threadId, stateRef)
 
   -- | Get a snapshot of all worker states
   workerStates workers =
     forM workers $ \(_, stateRef) -> readIORef stateRef
 
--- | Listener reads events and forwards all of them to the UI using the
--- 'forwardEvent' function. It exits after receiving all 'WorkerStopped'
--- events and sets the passed 'MVar' so the parent thread can block on listener
--- until all workers are done.
-spawnListener
-  :: Env
-  -> ((Int, LocalTime, CampaignEvent) -> IO ())
-  -- ^ a function that forwards event to the UI
-  -> Int     -- ^ number of workers
-  -> MVar () -- ^ use to join this thread
-  -> IO ()
-spawnListener env forwardEvent nworkers stopVar =
-  void $ forkFinally (loop nworkers) (const $ putMVar stopVar ())
-  where
-  loop !workersAlive =
-    when (workersAlive > 0) $ do
-      event <- readChan env.eventQueue
-      forwardEvent event
-      case event of
-        (_, _, WorkerStopped _) -> loop (workersAlive - 1)
-        _                       -> loop workersAlive
-
 #ifdef INTERACTIVE_UI
  -- | Order the workers to stop immediately
 stopWorkers :: MonadIO m => [(ThreadId, a)] -> m ()
@@ -282,7 +251,7 @@ monitor = do
         modify' $ \state -> state { workerEvents = state.workerEvents |> event }
 
         case campaignEvent of
-          NewCoverage coverage numCodehashes size ->
+          NewCoverage coverage numCodehashes size _ ->
             modify' $ \state ->
               state { coverage = max state.coverage coverage -- max not really needed
                     , corpusSize = size

package.yaml

@@ -38,6 +38,7 @@ dependencies:
   - semver
   - split
   - text
+  - threads
   - transformers
   - time
   - unliftio