Collect coverage during init (#1304)

Co-authored-by: Emilio López <emilio.lopez@trailofbits.com>

lib/Echidna.hs

@@ -118,7 +118,8 @@ mkEnv cfg buildOutput tests world slitherInfo = do
   codehashMap <- newIORef mempty
   chainId <- maybe (pure Nothing) EVM.Fetch.fetchChainIdFrom cfg.rpcUrl
   eventQueue <- newChan
-  coverageRef <- newIORef mempty
+  coverageRefInit <- newIORef mempty
+  coverageRefRuntime <- newIORef mempty
   corpusRef <- newIORef mempty
   testRefs <- traverse newIORef tests
   (contractCache, slotCache) <- Onchain.loadRpcCache cfg
@@ -127,6 +128,6 @@ mkEnv cfg buildOutput tests world slitherInfo = do
   -- TODO put in real path
   let dapp = dappInfo "/" buildOutput
   pure $ Env { cfg, dapp, codehashMap, fetchContractCache, fetchSlotCache
-             , chainId, eventQueue, coverageRef, corpusRef, testRefs, world
+             , chainId, eventQueue, coverageRefInit, coverageRefRuntime, corpusRef, testRefs, world
              , slitherInfo
              }

lib/Echidna/Campaign.hs

@@ -43,7 +43,7 @@ import Echidna.Transaction
 import Echidna.Types (Gas)
 import Echidna.Types.Campaign
 import Echidna.Types.Corpus (Corpus, corpusSize)
-import Echidna.Types.Coverage (scoveragePoints)
+import Echidna.Types.Coverage (coverageStats)
 import Echidna.Types.Config
 import Echidna.Types.Signature (FunctionName)
 import Echidna.Types.Test
@@ -353,10 +353,9 @@ callseq vm txSeq = do
       let !corp' = force $ addToCorpus (ncallseqs + 1) results corp
       in (corp', corpusSize corp')
 
-    cov <- liftIO . readIORef =<< asks (.coverageRef)
-    points <- liftIO $ scoveragePoints cov
+    (points, numCodehashes) <- liftIO $ coverageStats env.coverageRefInit env.coverageRefRuntime
     pushWorkerEvent NewCoverage { points
-                                , numCodehashes = length cov
+                                , numCodehashes
                                 , corpusSize = newSize
                                 , transactions = fst <$> results
                                 }

lib/Echidna/Deploy.hs

@@ -3,20 +3,22 @@ module Echidna.Deploy where
 import Control.Monad (foldM)
 import Control.Monad.Catch (MonadThrow(..), throwM)
 import Control.Monad.Reader (MonadReader, asks)
-import Control.Monad.State.Strict (MonadIO)
+import Control.Monad.State.Strict (MonadIO, runStateT)
 import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
 import Data.ByteString.Base16 qualified as BS16 (decode)
 import Data.Either (fromRight)
+import Data.Maybe (isJust)
 import Data.Text (Text, unlines)
 import Data.Text.Encoding (encodeUtf8)
 
 import EVM.Solidity
 import EVM.Types hiding (Env)
 
-import Echidna.Exec (execTx)
+import Echidna.Exec (execTx, execTxWithCov)
 import Echidna.Events (extractEvents)
-import Echidna.Types.Config (Env(..))
+import Echidna.Types.Campaign (CampaignConf(..))
+import Echidna.Types.Config (Env(..), EConfig(..))
 import Echidna.Types.Solidity (SolException(..))
 import Echidna.Types.Tx (createTx, unlimitedGasPerBlock)
 import Control.Monad.ST (RealWorld)
@@ -50,8 +52,9 @@ deployBytecodes'
 deployBytecodes' cs src initialVM = foldM deployOne initialVM cs
   where
   deployOne vm (dst, bytecode) = do
-    (_, vm') <-
-      execTx vm $ createTx (bytecode <> zeros) src dst unlimitedGasPerBlock (0, 0)
+    coverageEnabled <- asks (isJust . (.cfg.campaignConf.knownCoverage))
+    let deployTx = createTx (bytecode <> zeros) src dst unlimitedGasPerBlock (0, 0)
+    vm' <- if coverageEnabled then snd <$> runStateT (execTxWithCov deployTx) vm else snd <$> execTx vm deployTx
     case vm'.result of
       Just (VMSuccess _) -> pure vm'
       _ -> do

lib/Echidna/Exec.hs

@@ -287,9 +287,15 @@ execTxWithCov tx = do
       addCoverage !vm = do
         let (pc, opIx, depth) = currentCovLoc vm
             contract = currentContract vm
-
-        maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
+            covRef = case contract.code of
+              InitCode _ _ -> env.coverageRefInit
+              _ -> env.coverageRefRuntime
+
+        maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp covRef $ do
+          let
+            size = case contract.code of
+              InitCode b _ -> BS.length b
+              _ -> BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
             vec <- VMut.new size

lib/Echidna/Output/JSON.hs

@@ -20,7 +20,7 @@ import Echidna.Events (Events, extractEvents)
 import Echidna.Types (Gas)
 import Echidna.Types.Campaign (WorkerState(..))
 import Echidna.Types.Config (Env(..))
-import Echidna.Types.Coverage (CoverageInfo)
+import Echidna.Types.Coverage (CoverageInfo, mergeCoverageMaps)
 import Echidna.Types.Test qualified as T
 import Echidna.Types.Test (EchidnaTest(..))
 import Echidna.Types.Tx (Tx(..), TxCall(..))
@@ -101,7 +101,7 @@ instance ToJSON Transaction where
 encodeCampaign :: Env -> [WorkerState] -> IO L.ByteString
 encodeCampaign env workerStates = do
   tests <- traverse readIORef env.testRefs
-  frozenCov <- mapM VU.freeze =<< readIORef env.coverageRef
+  frozenCov <- mergeCoverageMaps env.dapp env.coverageRefInit env.coverageRefRuntime
   -- TODO: this is ugly, refactor seed to live in Env
   let worker0 = Prelude.head workerStates
   pure $ encode Campaign

lib/Echidna/Output/Source.hs

@@ -7,7 +7,6 @@ import Prelude hiding (writeFile)
 import Control.Monad (unless)
 import Data.ByteString qualified as BS
 import Data.Foldable
-import Data.IORef (readIORef)
 import Data.List (nub, sort)
 import Data.Maybe (fromMaybe, mapMaybe)
 import Data.Map (Map)
@@ -19,7 +18,7 @@ import Data.Text qualified as T
 import Data.Text.Encoding (decodeUtf8)
 import Data.Text.IO (writeFile)
 import Data.Vector qualified as V
-import Data.Vector.Unboxed.Mutable qualified as VU
+import Data.Vector.Unboxed qualified as VU
 import HTMLEntities.Text qualified as HTML
 import System.Directory (createDirectoryIfMissing)
 import System.FilePath ((</>))
@@ -30,7 +29,7 @@ import EVM.Solidity (SourceCache(..), SrcMap, SolcContract(..))
 
 import Echidna.Types.Campaign (CampaignConf(..))
 import Echidna.Types.Config (Env(..), EConfig(..))
-import Echidna.Types.Coverage (OpIx, unpackTxResults, CoverageMap, CoverageFileType (..))
+import Echidna.Types.Coverage (OpIx, unpackTxResults, FrozenCoverageMap, CoverageFileType (..), mergeCoverageMaps)
 import Echidna.Types.Tx (TxResult(..))
 import Echidna.SourceAnalysis.Slither (AssertLocation(..), assertLocationList, SlitherInfo(..))
 
@@ -43,7 +42,7 @@ saveCoverages
   -> IO ()
 saveCoverages env seed d sc cs = do
   let fileTypes = env.cfg.campaignConf.coverageFormats
-  coverage <- readIORef env.coverageRef
+  coverage <- mergeCoverageMaps env.dapp env.coverageRefInit env.coverageRefRuntime
   mapM_ (\ty -> saveCoverage ty seed d sc cs coverage) fileTypes
 
 saveCoverage
@@ -52,12 +51,12 @@ saveCoverage
   -> FilePath
   -> SourceCache
   -> [SolcContract]
-  -> CoverageMap
+  -> FrozenCoverageMap
   -> IO ()
 saveCoverage fileType seed d sc cs covMap = do
   let extension = coverageFileExtension fileType
       fn = d </> "covered." <> show seed <> extension
-  cc <- ppCoveredCode fileType sc cs covMap
+      cc = ppCoveredCode fileType sc cs covMap
   createDirectoryIfMissing True d
   writeFile fn cc
 
@@ -67,12 +66,12 @@ coverageFileExtension Html = ".html"
 coverageFileExtension Txt = ".txt"
 
 -- | Pretty-print the covered code
-ppCoveredCode :: CoverageFileType -> SourceCache -> [SolcContract] -> CoverageMap -> IO Text
-ppCoveredCode fileType sc cs s | null s = pure "Coverage map is empty"
-  | otherwise = do
-  -- List of covered lines during the fuzzing campaign
-  covLines <- srcMapCov sc s cs
+ppCoveredCode :: CoverageFileType -> SourceCache -> [SolcContract] -> FrozenCoverageMap -> Text
+ppCoveredCode fileType sc cs s | null s = "Coverage map is empty"
+  | otherwise =
   let
+    -- List of covered lines during the fuzzing campaign
+    covLines = srcMapCov sc s cs
     -- Collect all the possible lines from all the files
     allFiles = (\(path, src) -> (path, V.fromList (decodeUtf8 <$> BS.split 0xa src))) <$> Map.elems sc.files
     -- Excludes lines such as comments or blanks
@@ -102,7 +101,7 @@ ppCoveredCode fileType sc cs s | null s = pure "Coverage map is empty"
       Html -> "<code>" : ls ++ ["", "</code>","<br />"]
       Txt  -> ls
     -- ^ Alter file contents, in the case of html encasing it in <code> and adding a line break
-  pure $ topHeader <> T.unlines (map ppFile allFiles)
+  in topHeader <> T.unlines (map ppFile allFiles)
 
 -- | Mark one particular line, from a list of lines, keeping the order of them
 markLines :: CoverageFileType -> V.Vector Text -> S.Set Int -> Map Int [TxResult] -> V.Vector Text
@@ -148,11 +147,11 @@ getMarker ErrorOutOfGas = 'o'
 getMarker _             = 'e'
 
 -- | Given a source cache, a coverage map, a contract returns a list of covered lines
-srcMapCov :: SourceCache -> CoverageMap -> [SolcContract] -> IO (Map FilePath (Map Int [TxResult]))
-srcMapCov sc covMap contracts = do
-  Map.unionsWith Map.union <$> mapM linesCovered contracts
+srcMapCov :: SourceCache -> FrozenCoverageMap -> [SolcContract] -> Map FilePath (Map Int [TxResult])
+srcMapCov sc covMap contracts =
+  Map.unionsWith Map.union $ linesCovered <$> contracts
   where
-  linesCovered :: SolcContract -> IO (Map FilePath (Map Int [TxResult]))
+  linesCovered :: SolcContract -> Map FilePath (Map Int [TxResult])
   linesCovered c =
     case Map.lookup c.runtimeCodehash covMap of
       Just vec -> VU.foldl' (\acc covInfo -> case covInfo of
@@ -197,11 +196,11 @@ checkAssertionsCoverage
   -> Env
   -> IO ()
 checkAssertionsCoverage sc env = do
+  covMap <- mergeCoverageMaps env.dapp env.coverageRefInit env.coverageRefRuntime
   let
     cs = Map.elems env.dapp.solcByName
     asserts = maybe [] (concatMap assertLocationList . Map.elems . (.asserts)) env.slitherInfo
-  covMap <- readIORef env.coverageRef
-  covLines <- srcMapCov sc covMap cs
+    covLines = srcMapCov sc covMap cs
   mapM_ (checkAssertionReached covLines) asserts
 
 -- | Helper function for `checkAssertionsCoverage` which checks a single assertion

lib/Echidna/Solidity.hs

@@ -7,6 +7,7 @@ import Control.Monad.Catch (MonadThrow(..))
 import Control.Monad.Extra (whenM)
 import Control.Monad.Reader (ReaderT(runReaderT))
 import Control.Monad.ST (stToIO, RealWorld)
+import Control.Monad.State (runStateT)
 import Data.Foldable (toList)
 import Data.List (find, partition, isSuffixOf, (\\))
 import Data.List.NonEmpty (NonEmpty((:|)))
@@ -40,9 +41,10 @@ import Echidna.ABI
 import Echidna.Deploy (deployContracts, deployBytecodes)
 import Echidna.Etheno (loadEthenoBatch)
 import Echidna.Events (extractEvents)
-import Echidna.Exec (execTx, initialVM)
+import Echidna.Exec (execTx, execTxWithCov, initialVM)
 import Echidna.SourceAnalysis.Slither
 import Echidna.Test (createTests, isAssertionMode, isPropertyMode, isDapptestMode)
+import Echidna.Types.Campaign (CampaignConf(..))
 import Echidna.Types.Config (EConfig(..), Env(..))
 import Echidna.Types.Signature
   (ContractName, SolSignature, SignatureMap, FunctionName)
@@ -200,14 +202,18 @@ loadSpecified env mainContract cs = do
     vm2 <- deployBytecodes solConf.deployBytecodes solConf.deployer vm1
 
     -- main contract deployment
-    let deployment = execTx vm2 $ createTxWithValue
-                                    mainContract.creationCode
-                                    solConf.deployer
-                                    solConf.contractAddr
-                                    unlimitedGasPerBlock
-                                    (fromIntegral solConf.balanceContract)
-                                    (0, 0)
-    (_, vm3) <- deployment
+    let
+      coverageEnabled = isJust env.cfg.campaignConf.knownCoverage
+      deployTx = createTxWithValue
+                    mainContract.creationCode
+                    solConf.deployer
+                    solConf.contractAddr
+                    unlimitedGasPerBlock
+                    (fromIntegral solConf.balanceContract)
+                    (0, 0)
+      deployment = if coverageEnabled then snd <$> runStateT (execTxWithCov deployTx) vm2 else snd <$> execTx vm2 deployTx
+
+    vm3 <- deployment
     when (isNothing $ currentContract vm3) $
       throwM $ DeploymentFailed solConf.contractAddr $ T.unlines $ extractEvents True env.dapp vm3
 

lib/Echidna/Types/Config.hs

@@ -71,7 +71,8 @@ data Env = Env
   , eventQueue :: Chan (LocalTime, CampaignEvent)
 
   , testRefs :: [IORef EchidnaTest]
-  , coverageRef :: IORef CoverageMap
+  , coverageRefInit :: IORef CoverageMap
+  , coverageRefRuntime :: IORef CoverageMap
   , corpusRef :: IORef Corpus
 
   , slitherInfo :: Maybe SlitherInfo

lib/Echidna/Types/Coverage.hs

@@ -1,14 +1,20 @@
 module Echidna.Types.Coverage where
 
+import Control.Monad ((>=>))
 import Data.Aeson (ToJSON(toJSON), FromJSON(parseJSON), withText)
 import Data.Bits (testBit)
+import Data.IORef (IORef, readIORef)
 import Data.List (foldl')
 import Data.Map qualified as Map
 import Data.Map.Strict (Map)
+import Data.Set qualified as Set
 import Data.Text (toLower)
 import Data.Vector.Unboxed.Mutable (IOVector)
-import Data.Vector.Unboxed.Mutable qualified as V
+import Data.Vector.Unboxed.Mutable qualified as VM
+import Data.Vector.Unboxed qualified as V
 import Data.Word (Word64)
+import EVM.Dapp (DappInfo(..))
+import EVM.Solidity (SolcContract(..))
 import EVM.Types (W256)
 
 import Echidna.Types.Tx (TxResult)
@@ -17,6 +23,10 @@ import Echidna.Types.Tx (TxResult)
 -- Indexed by contracts' compile-time codehash; see `CodehashMap`.
 type CoverageMap = Map W256 (IOVector CoverageInfo)
 
+-- | CoverageMap, but using Vectors instead of IOVectors.
+-- IO is not required to access this map's members.
+type FrozenCoverageMap = Map W256 (V.Vector CoverageInfo)
+
 -- | Basic coverage information
 type CoverageInfo = (OpIx, StackDepths, TxResults)
 
@@ -29,12 +39,42 @@ type StackDepths = Word64
 -- | Packed TxResults used for coverage, corresponding bits are set
 type TxResults = Word64
 
+-- | Given the CoverageMaps used for contract init and runtime, produce a single combined coverage map
+-- with op indices from init correctly shifted over (see srcMapForOpLocation in Echidna.Output.Source).
+-- Takes IORef CoverageMap because this is how they are stored in the Env.
+mergeCoverageMaps :: DappInfo -> IORef CoverageMap -> IORef CoverageMap -> IO FrozenCoverageMap
+mergeCoverageMaps dapp initMap runtimeMap = mergeFrozenCoverageMaps dapp <$> freeze initMap <*> freeze runtimeMap
+  where freeze = readIORef >=> mapM V.freeze
+
+-- | Given the FrozenCoverageMaps used for contract init and runtime, produce a single combined coverage map
+-- with op indices from init correctly shifted over (see srcMapForOpLocation in Echidna.Output.Source).
+-- Helper function for mergeCoverageMaps.
+mergeFrozenCoverageMaps :: DappInfo -> FrozenCoverageMap -> FrozenCoverageMap -> FrozenCoverageMap
+mergeFrozenCoverageMaps dapp initMap runtimeMap = Map.unionWith (<>) runtimeMap initMap'
+  where
+    initMap' = Map.mapWithKey modifyInitMapEntry initMap
+    -- eta reduced, second argument is a vec
+    modifyInitMapEntry hash = V.map $ modifyCoverageInfo $ getOpOffset hash
+    modifyCoverageInfo toAdd (op, x, y) = (op + toAdd, x, y)
+    getOpOffset hash = maybe 0 (length . (.runtimeSrcmap) . snd) $ Map.lookup hash dapp.solcByHash
+
+-- | Given the CoverageMaps used for contract init and runtime,
+-- return the point coverage and the number of unique contracts hit.
+-- Takes IORef CoverageMap because this is how they are stored in the Env.
+coverageStats :: IORef CoverageMap -> IORef CoverageMap -> IO (Int, Int)
+coverageStats initRef runtimeRef = do
+  initMap <- readIORef initRef
+  runtimeMap <- readIORef runtimeRef
+  pointsInit <- scoveragePoints initMap
+  pointsRuntime <- scoveragePoints runtimeMap
+  pure (pointsInit + pointsRuntime, length $ Set.fromList $ Map.keys initMap ++ Map.keys runtimeMap)
+
 -- | Given good point coverage, count the number of unique points but
 -- only considering the different instruction PCs (discarding the TxResult).
 -- This is useful for reporting a coverage measure to the user
 scoveragePoints :: CoverageMap -> IO Int
 scoveragePoints cm = do
-  sum <$> mapM (V.foldl' countCovered 0) (Map.elems cm)
+  sum <$> mapM (VM.foldl' countCovered 0) (Map.elems cm)
 
 countCovered :: Int -> CoverageInfo -> Int
 countCovered acc (opIx,_,_) = if opIx == -1 then acc else acc + 1

lib/Echidna/UI.hs

@@ -39,7 +39,7 @@ import Echidna.Server (runSSEServer)
 import Echidna.Types.Campaign
 import Echidna.Types.Config
 import Echidna.Types.Corpus qualified as Corpus
-import Echidna.Types.Coverage (scoveragePoints)
+import Echidna.Types.Coverage (coverageStats)
 import Echidna.Types.Test (EchidnaTest(..), didFail, isOptimizationTest)
 import Echidna.Types.Tx (Tx)
 import Echidna.UI.Report
@@ -339,7 +339,7 @@ statusLine
   -> IO String
 statusLine env states = do
   tests <- traverse readIORef env.testRefs
-  points <- scoveragePoints =<< readIORef env.coverageRef
+  (points, _) <- coverageStats env.coverageRefInit env.coverageRefRuntime
   corpus <- readIORef env.corpusRef
   let totalCalls = sum ((.ncalls) <$> states)
   pure $ "tests: " <> show (length $ filter didFail tests) <> "/" <> show (length tests)