Skip to content

Commit

Permalink
Save corpus and reproducers continuously
Browse files Browse the repository at this point in the history
  • Loading branch information
arcz committed Jan 18, 2024
1 parent 988bda7 commit 6a438b9
Show file tree
Hide file tree
Showing 8 changed files with 165 additions and 86 deletions.
59 changes: 49 additions & 10 deletions lib/Echidna/Campaign.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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(..))
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 ->

Expand Down Expand Up @@ -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 _ ->
Expand All @@ -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
38 changes: 37 additions & 1 deletion lib/Echidna/Output/Corpus.hs
Original file line number Diff line number Diff line change
@@ -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]]
Expand All @@ -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)
27 changes: 17 additions & 10 deletions lib/Echidna/Server.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 ()
Expand Down
38 changes: 25 additions & 13 deletions lib/Echidna/Types/Campaign.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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 ]

Expand All @@ -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 ->
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion lib/Echidna/Types/Config.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 6a438b9

Please sign in to comment.