Save corpus and reproducers continuously (#1167)

Author: arcz

lib/Echidna/Campaign.hs

@@ -3,7 +3,7 @@
 
 module Echidna.Campaign where
 
-import Control.Concurrent (writeChan)
+import Control.Concurrent
 import Control.DeepSeq (force)
 import Control.Monad (replicateM, when, void, forM_)
 import Control.Monad.Catch (MonadThrow(..))
@@ -22,6 +22,7 @@ import Data.Maybe (isJust, mapMaybe, fromMaybe)
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Text (Text)
+import Data.Time (LocalTime)
 import System.Random (mkStdGen)
 
 import EVM (cheatCode)
@@ -67,7 +68,7 @@ replayCorpus
 replayCorpus vm txSeqs =
   forM_ (zip [1..] txSeqs) $ \(i, txSeq) -> do
     _ <- callseq vm txSeq
-    pushEvent (TxSequenceReplayed i (length txSeqs))
+    pushWorkerEvent (TxSequenceReplayed i (length txSeqs))
 
 -- | Run a fuzzing campaign given an initial universe state, some tests, and an
 -- optional dictionary to generate calls with. Return the 'Campaign' state once
@@ -206,7 +207,11 @@ callseq vm txSeq = do
 
     cov <- liftIO . readIORef =<< asks (.coverageRef)
     points <- liftIO $ scoveragePoints cov
-    pushEvent (NewCoverage points (length cov) newSize)
+    pushWorkerEvent NewCoverage { points
+                                , numCodehashes = length cov
+                                , corpusSize = newSize
+                                , transactions = fst <$> results
+                                }
 
   modify' $ \workerState ->
 
@@ -368,10 +373,10 @@ updateTest vmForShrink (vm, xs) test = do
         test' = updateOpenTest test xs (testValue, vm', results)
       case test'.state of
         Large _ -> do
-          pushEvent (TestFalsified test')
+          pushWorkerEvent (TestFalsified test')
           pure (Just test')
         _ | test'.value > test.value -> do
-          pushEvent (TestOptimized test')
+          pushWorkerEvent (TestOptimized test')
           pure (Just test')
         _ -> pure Nothing
     Large _ ->
@@ -381,12 +386,46 @@ updateTest vmForShrink (vm, xs) test = do
       shrinkTest vmForShrink test
     _ -> pure Nothing
 
-pushEvent
+pushWorkerEvent
   :: (MonadReader Env m, MonadState WorkerState m, MonadIO m)
-  => CampaignEvent
+  => WorkerEvent
   -> m ()
-pushEvent event = do
+pushWorkerEvent event = do
   workerId <- gets (.workerId)
+  env <- ask
+  liftIO $ pushCampaignEvent env (WorkerEvent workerId event)
+
+pushCampaignEvent :: Env -> CampaignEvent -> IO ()
+pushCampaignEvent env event = do
   time <- liftIO getTimestamp
-  chan <- asks (.eventQueue)
-  liftIO $ writeChan chan (workerId, time, event)
+  writeChan env.eventQueue (time, event)
+
+-- | Listener reads events and runs the given 'handler' function. It exits after
+-- receiving all 'WorkerStopped' events and sets the returned 'MVar' so the
+-- parent thread can safely block on listener until all events are processed.
+--
+-- NOTE: because the 'Failure' event does not come from a specific fuzzing worker
+-- it is possible that a listener won't process it if emitted after all workers
+-- are stopped. This is quite unlikely and non-critical but should be addressed
+-- in the long term.
+spawnListener
+  :: (MonadReader Env m, MonadIO m)
+  => ((LocalTime, CampaignEvent) -> IO ())
+  -- ^ a function that handles the events
+  -> m (MVar ())
+spawnListener handler = do
+  cfg <- asks (.cfg)
+  let nworkers = fromMaybe 1 cfg.campaignConf.workers
+  eventQueue <- asks (.eventQueue)
+  chan <- liftIO $ dupChan eventQueue
+  stopVar <- liftIO newEmptyMVar
+  liftIO $ void $ forkFinally (loop chan nworkers) (const $ putMVar stopVar ())
+  pure stopVar
+  where
+  loop chan !workersAlive =
+    when (workersAlive > 0) $ do
+      event <- readChan chan
+      handler event
+      case event of
+        (_, WorkerEvent _ (WorkerStopped _)) -> loop chan (workersAlive - 1)
+        _                                    -> loop chan workersAlive

lib/Echidna/Output/Corpus.hs

@@ -1,20 +1,28 @@
 module Echidna.Output.Corpus where
 
+import Control.Exception (IOException, handle)
+import Control.Monad (unless)
 import Control.Monad.Extra (unlessM)
 import Data.Aeson (ToJSON(..), decodeStrict, encodeFile)
 import Data.ByteString qualified as BS
 import Data.Hashable (hash)
 import Data.Maybe (catMaybes)
+import Data.Time (LocalTime)
 import System.Directory (createDirectoryIfMissing, makeRelativeToCurrentDirectory, doesFileExist)
 import System.FilePath ((</>), (<.>))
 
+import Echidna.Campaign (pushCampaignEvent)
+import Echidna.Types.Config
+import Echidna.Types.Campaign
+import Echidna.Types.Test (EchidnaTest(..))
 import Echidna.Types.Tx (Tx)
 import Echidna.Utility (listDirectory, withCurrentDirectory)
 
 saveTxs :: FilePath -> [[Tx]] -> IO ()
 saveTxs dir = mapM_ saveTxSeq where
   saveTxSeq txSeq = do
-    let file = dir </> (show . hash . show) txSeq <.> "txt"
+    createDirectoryIfMissing True dir
+    let file = dir </> (show . abs . hash . show) txSeq <.> "txt"
     unlessM (doesFileExist file) $ encodeFile file (toJSON txSeq)
 
 loadTxs :: FilePath -> IO [[Tx]]
@@ -26,3 +34,31 @@ loadTxs dir = do
   putStrLn ("Loaded " ++ show (length txSeqs) ++ " transaction sequences from " ++ dir)
   pure txSeqs
   where readCall f = decodeStrict <$> BS.readFile f
+
+-- Save corpus/reproducers transactions based on an event
+saveCorpusEvent :: Env -> (LocalTime, CampaignEvent) -> IO ()
+saveCorpusEvent env (_time, campaignEvent) = do
+  case env.cfg.campaignConf.corpusDir of
+    Just corpusDir -> saveEvent corpusDir campaignEvent
+    Nothing -> pure ()
+  where
+    saveEvent dir (WorkerEvent _workerId event) =
+      maybe (pure ()) (saveFile dir) $ getEventInfo event
+    saveEvent _ _ = pure ()
+
+    getEventInfo = \case
+      -- TODO: We save intermediate reproducers in separate directories.
+      -- This is to because there can be a lot of them and we want to skip
+      -- loading those on startup. Ideally, we should override the same file
+      -- with a better version of a reproducer, this is smaller or more optimized.
+      TestFalsified test -> Just ("reproducers-unshrunk", test.reproducer)
+      TestOptimized test -> Just ("reproducers-optimizations", test.reproducer)
+      NewCoverage { transactions } -> Just ("coverage", transactions)
+      _ -> Nothing
+
+    saveFile dir (subdir, txs) =
+      unless (null txs) $
+        handle exceptionHandler $ saveTxs (dir </> subdir) [txs]
+
+    exceptionHandler (e :: IOException) =
+      pushCampaignEvent env (Failure $ "Problem while writing to file: " ++ show e)

lib/Echidna/Server.hs

@@ -10,17 +10,21 @@ import Data.Word (Word16)
 import Network.Wai.EventSource (ServerEvent(..), eventSourceAppIO)
 import Network.Wai.Handler.Warp (run)
 
-import Echidna.Types.Campaign (CampaignEvent (..))
+import Echidna.Types.Campaign
 import Echidna.Types.Config (Env(..))
 
-newtype SSE = SSE (Int, LocalTime, CampaignEvent)
+newtype SSE = SSE (LocalTime, CampaignEvent)
 
 instance ToJSON SSE where
-  toJSON (SSE (workerId, time, event)) =
+  toJSON (SSE (time, WorkerEvent workerId event)) =
     object [ "worker" .= workerId
            , "timestamp" .= time
            , "data" .= event
            ]
+  toJSON (SSE (time, Failure reason)) =
+    object [ "timestamp" .= time
+           , "data" .= reason
+           ]
 
 runSSEServer :: MVar () -> Env -> Word16 -> Int -> IO ()
 runSSEServer serverStopVar env port nworkers = do
@@ -32,15 +36,18 @@ runSSEServer serverStopVar env port nworkers = do
         if aliveNow == 0 then
           pure CloseEvent
         else do
-          event@(_, _, campaignEvent) <- readChan sseChan
+          event@(_, campaignEvent) <- readChan sseChan
           let eventName = \case
-                TestFalsified _ -> "test_falsified"
-                TestOptimized _ -> "test_optimized"
-                NewCoverage {} -> "new_coverage"
-                TxSequenceReplayed _ _ -> "tx_sequence_replayed"
-                WorkerStopped _ -> "worker_stopped"
+                WorkerEvent _ workerEvent ->
+                  case workerEvent of
+                    TestFalsified _ -> "test_falsified"
+                    TestOptimized _ -> "test_optimized"
+                    NewCoverage {} -> "new_coverage"
+                    TxSequenceReplayed _ _ -> "tx_sequence_replayed"
+                    WorkerStopped _ -> "worker_stopped"
+                Failure _err -> "failure"
           case campaignEvent of
-            WorkerStopped _ -> do
+            WorkerEvent _ (WorkerStopped _) -> do
               aliveAfter <- atomicModifyIORef' aliveRef (\n -> (n-1, n-1))
               when (aliveAfter == 0) $ putMVar serverStopVar ()
             _ -> pure ()

lib/Echidna/Types/Campaign.hs

@@ -45,22 +45,28 @@ data CampaignConf = CampaignConf
     -- ^ Server-Sent Events HTTP port number, if missing server is not ran
   }
 
+type WorkerId = Int
+
 data CampaignEvent
+  = WorkerEvent WorkerId WorkerEvent
+  | Failure String
+
+data WorkerEvent
   = TestFalsified !EchidnaTest
   | TestOptimized !EchidnaTest
-  | NewCoverage !Int !Int !Int
+  | NewCoverage { points :: !Int, numCodehashes :: !Int, corpusSize :: !Int, transactions :: [Tx] }
   | TxSequenceReplayed !Int !Int
   | WorkerStopped WorkerStopReason
   -- ^ This is a terminal event. Worker exits and won't push any events after
   -- this one
   deriving Show
 
-instance ToJSON CampaignEvent where
+instance ToJSON WorkerEvent where
   toJSON = \case
     TestFalsified test -> toJSON test
     TestOptimized test -> toJSON test
-    NewCoverage coverage numContracts corpusSize ->
-      object [ "coverage" .= coverage, "contracts" .= numContracts, "corpus_size" .= corpusSize]
+    NewCoverage { points, numCodehashes, corpusSize } ->
+      object [ "coverage" .= points, "contracts" .= numCodehashes, "corpus_size" .= corpusSize]
     TxSequenceReplayed current total -> object [ "current" .= current, "total" .= total ]
     WorkerStopped reason -> object [ "reason" .= show reason ]
 
@@ -74,20 +80,20 @@ data WorkerStopReason
 
 ppCampaignEvent :: CampaignEvent -> String
 ppCampaignEvent = \case
+  WorkerEvent _ e -> ppWorkerEvent e
+  Failure err -> err
+
+ppWorkerEvent :: WorkerEvent -> String
+ppWorkerEvent = \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 ->
+  NewCoverage { points, numCodehashes, corpusSize } ->
     "New coverage: " <> show points <> " instr, "
-      <> show codehashes <> " contracts, "
-      <> show corpus <> " seqs in corpus"
+      <> show numCodehashes <> " contracts, "
+      <> show corpusSize <> " seqs in corpus"
   TxSequenceReplayed current total ->
     "Sequence replayed from corpus (" <> show current <> "/" <> show total <> ")"
   WorkerStopped TestLimitReached ->
@@ -102,6 +108,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

@@ -65,7 +65,7 @@ data Env = Env
 
   -- | Shared between all workers. Events are fairly rare so contention is
   -- minimal.
-  , eventQueue :: Chan (Int, LocalTime, CampaignEvent)
+  , eventQueue :: Chan (LocalTime, CampaignEvent)
 
   , testsRef :: IORef [EchidnaTest]
   , coverageRef :: IORef CoverageMap