Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Code cleanup #906

Merged
merged 8 commits into from
Jan 24, 2023
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 34 additions & 20 deletions lib/Echidna.hs
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
module Echidna where

import Control.Monad.Catch (MonadThrow(..))
import Data.HashMap.Strict (toList)
import Data.Map.Strict (keys)
import Data.List (nub, find)
import Data.HashMap.Strict qualified as HM
import Data.List (find)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import System.FilePath ((</>))

import EVM
import EVM.ABI (AbiValue(AbiAddress))
Expand Down Expand Up @@ -34,14 +36,17 @@ import Echidna.RPC (loadEtheno, extractFromEtheno)
-- * A VM with the contract deployed and ready for testing
-- * A World with all the required data for generating random transctions
-- * A list of Echidna tests to check
-- * A prepopulated dictionary (if any)
-- * A prepopulated dictionary
-- * A list of transaction sequences to initialize the corpus
prepareContract :: EConfig -> NE.NonEmpty FilePath -> Maybe ContractName -> Seed
-> IO (VM, SourceCache, [SolcContract], World, [EchidnaTest], Maybe GenDict, [[Tx]])
-> IO (VM, SourceCache, [SolcContract], World, [EchidnaTest], GenDict, [[Tx]])
prepareContract cfg fs c g = do
ctxs1 <- loadTxs (fmap (++ "/reproducers/") cd)
ctxs2 <- loadTxs (fmap (++ "/coverage/") cd)
let ctxs = ctxs1 ++ ctxs2
ctxs <- case cfg._cConf._corpusDir of
Nothing -> pure []
Just dir -> do
ctxs1 <- loadTxs (dir </> "reproducers")
ctxs2 <- loadTxs (dir </> "coverage")
pure (ctxs1 ++ ctxs2)

let solConf = cfg._sConf

Expand All @@ -51,27 +56,36 @@ prepareContract cfg fs c g = do

-- run processors
si <- runSlither (NE.head fs) solConf._cryticArgs
case find (< minSupportedSolcVersion) $ solcVersions si of
case find (< minSupportedSolcVersion) si.solcVersions of
Just outdatedVersion -> throwM $ OutdatedSolcVersion outdatedVersion
Nothing -> return ()

-- load tests
let (v, w, ts) = prepareForTest solConf p c si
let (vm, world, ts) = prepareForTest solConf p c si

-- get signatures
let sigs = nub $ concatMap (NE.toList . snd) (toList $ w._highSignatureMap)
let sigs = Set.fromList $ concatMap NE.toList (HM.elems world.highSignatureMap)

let ads = addresses solConf
let ads' = AbiAddress <$> keys v._env._contracts
let constants' = enhanceConstants si ++ timeConstants ++ extremeConstants ++ NE.toList ads ++ ads'
let ads' = AbiAddress <$> Map.keys vm._env._contracts
let constants' = Set.fromList $ enhanceConstants si ++
timeConstants ++
extremeConstants ++
Set.toList ads ++
ads'

-- load transactions from init sequence (if any)
es' <- maybe (return []) loadEtheno it
let txs = ctxs ++ maybe [] (const [extractFromEtheno es' sigs]) it
ethenoCorpus <-
case cfg._sConf._initialize of
Nothing -> pure []
Just fp -> do
es' <- loadEtheno fp
pure [extractFromEtheno es' sigs]

let corp = ctxs ++ ethenoCorpus

-- start ui and run tests
let sc = selectSourceCache c scs
pure (v, sc, cs, w, ts, Just $ mkGenDict df constants' [] g (returnTypes cs), txs)
where cd = cfg._cConf._corpusDir
df = cfg._cConf._dictFreq
it = cfg._sConf._initialize

let dict = mkGenDict cfg._cConf._dictFreq constants' Set.empty g (returnTypes cs)

pure (vm, sc, cs, world, ts, dict, corp)
43 changes: 20 additions & 23 deletions lib/Echidna/ABI.hs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ module Echidna.ABI where

import Control.Lens
import Control.Monad (join, liftM2, liftM3, foldM, replicateM)
import Control.Monad.Random.Strict (MonadRandom, getRandom, getRandoms, getRandomR, uniformMay)
import Control.Monad.Random.Strict (MonadRandom, getRandom, getRandoms, getRandomR)
import Control.Monad.Random.Strict qualified as R
import Data.Binary.Put (runPut, putWord32be)
import Data.BinaryWord (unsignedWord)
Expand All @@ -20,12 +20,11 @@ import Data.Foldable (toList)
import Data.Hashable (Hashable(..))
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as M
import Data.HashSet (HashSet, fromList, union)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.List (intercalate)
import Data.List.NonEmpty qualified as NE
import Data.Maybe (fromMaybe, catMaybes, mapMaybe)
import Data.Maybe (fromMaybe, catMaybes)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding (encodeUtf8)
Expand Down Expand Up @@ -104,9 +103,9 @@ hashSig = abiKeccak . TE.encodeUtf8
-- | Configuration necessary for generating new 'SolCalls'. Don't construct this by hand! Use 'mkConf'.
data GenDict = GenDict { _pSynthA :: Float
-- ^ Fraction of time to use dictionary vs. synthesize
, _constants :: HashMap AbiType (HashSet AbiValue)
, _constants :: HashMap AbiType (Set AbiValue)
-- ^ Constants to use, sorted by type
, _wholeCalls :: HashMap SolSignature (HashSet SolCall)
, _wholeCalls :: HashMap SolSignature (Set SolCall)
-- ^ Whole calls to use, sorted by type
, _defSeed :: Int
-- ^ Default seed to use if one is not provided in EConfig
Expand All @@ -118,14 +117,14 @@ data GenDict = GenDict { _pSynthA :: Float

makeLenses 'GenDict

hashMapBy :: (Hashable k, Hashable a, Eq k, Ord a) => (a -> k) -> [a] -> HashMap k (HashSet a)
hashMapBy f = M.fromListWith union . fmap (\v -> (f v, fromList [v]))
hashMapBy :: (Hashable k, Hashable a, Eq k, Ord a) => (a -> k) -> Set a -> HashMap k (Set a)
hashMapBy f = M.fromListWith Set.union . fmap (\v -> (f v, Set.singleton v)) . Set.toList

gaddCalls :: [SolCall] -> GenDict -> GenDict
gaddCalls :: Set SolCall -> GenDict -> GenDict
gaddCalls c = wholeCalls <>~ hashMapBy (fmap $ fmap abiValueType) c

defaultDict :: GenDict
defaultDict = mkGenDict 0 [] [] 0 (const Nothing)
defaultDict = mkGenDict 0 Set.empty Set.empty 0 (const Nothing)

deriving anyclass instance Hashable AbiType
deriving anyclass instance Hashable AbiValue
Expand All @@ -134,17 +133,17 @@ deriving anyclass instance Hashable Addr
-- | Construct a 'GenDict' from some dictionaries, a 'Float', a default seed, and a typing rule for
-- return values
mkGenDict :: Float -- ^ Percentage of time to mutate instead of synthesize. Should be in [0,1]
-> [AbiValue] -- ^ A list of 'AbiValue' constants to use during dictionary-based generation
-> [SolCall] -- ^ A list of complete 'SolCall's to mutate
-> Set AbiValue -- ^ A list of 'AbiValue' constants to use during dictionary-based generation
-> Set SolCall -- ^ A list of complete 'SolCall's to mutate
-> Int -- ^ A default seed
-> (Text -> Maybe AbiType)
-- ^ A return value typing rule
-> GenDict
mkGenDict p vs cs s tr =
GenDict p (hashMapBy abiValueType vs) (hashMapBy (fmap $ fmap abiValueType) cs) s tr (mkDictValues vs)

mkDictValues :: [AbiValue] -> Set W256
mkDictValues vs = Set.fromList $ mapMaybe fromValue vs
mkDictValues :: Set AbiValue -> Set W256
mkDictValues = Set.foldl' (\acc e -> maybe acc (`Set.insert` acc) (fromValue e)) Set.empty
where fromValue (AbiUInt _ n) = Just (fromIntegral n)
fromValue (AbiInt _ n) = Just (fromIntegral n)
fromValue _ = Nothing
Expand All @@ -167,10 +166,6 @@ genAbiValue = genAbiValueM defaultDict
genAbiCall :: MonadRandom m => SolSignature -> m SolCall
genAbiCall = traverse $ traverse genAbiValue

-- | Synthesize a random 'SolCall' given a list of 'SolSignature's (effectively, an ABI). Doesn't use a dictionary.
genInteractions :: MonadRandom m => NE.NonEmpty SolSignature -> m SolCall
genInteractions l = genAbiCall =<< rElem l

-- Mutation helper functions

-- | Given an 'Integral' number n, get a random number in [0,2n].
Expand Down Expand Up @@ -252,7 +247,7 @@ shrinkAbiValue :: MonadRandom m => AbiValue -> m AbiValue
shrinkAbiValue (AbiUInt n m) = AbiUInt n <$> shrinkInt m
shrinkAbiValue (AbiInt n m) = AbiInt n <$> shrinkInt m
shrinkAbiValue (AbiAddress 0) = pure $ AbiAddress 0
shrinkAbiValue (AbiAddress _) = rElem $ NE.fromList [AbiAddress 0, AbiAddress 0xdeadbeef]
shrinkAbiValue (AbiAddress _) = rElem' $ Set.fromList [AbiAddress 0, AbiAddress 0xdeadbeef]
shrinkAbiValue (AbiBool _) = pure $ AbiBool False
shrinkAbiValue (AbiBytes n b) = AbiBytes n <$> addNulls b
shrinkAbiValue (AbiBytesDynamic b) = fmap AbiBytesDynamic $ addNulls =<< shrinkBS b
Expand Down Expand Up @@ -309,16 +304,18 @@ mutateAbiCall = traverse f
-- @a@ from a GenDict, return a generator that takes an @a@ and either synthesizes new @b@s with the
-- provided generator or uses the 'GenDict' dictionary (when available).
genWithDict :: (Eq a, Hashable a, MonadRandom m)
=> GenDict -> HashMap a [b] -> (a -> m b) -> a -> m b
=> GenDict -> HashMap a (Set b) -> (a -> m b) -> a -> m b
genWithDict genDict m g t = do
r <- getRandom
let maybeValM = if genDict ^. pSynthA >= r then fromDict else pure Nothing
fromDict = uniformMay (M.lookupDefault [] t m)
let maybeValM = if genDict._pSynthA >= r then fromDict else pure Nothing
fromDict = case M.lookup t m of
Nothing -> pure Nothing
Just cs -> Just <$> rElem' cs
fromMaybe <$> g t <*> maybeValM

-- | Synthesize a random 'AbiValue' given its 'AbiType'. Requires a dictionary.
genAbiValueM :: MonadRandom m => GenDict -> AbiType -> m AbiValue
genAbiValueM genDict = genWithDict genDict (toList <$> genDict ^. constants) $ \case
genAbiValueM genDict = genWithDict genDict genDict._constants $ \case
ggrieco-tob marked this conversation as resolved.
Show resolved Hide resolved
(AbiUIntType n) -> fixAbiUInt n . fromInteger <$> getRandomUint n
(AbiIntType n) -> fixAbiInt n . fromInteger <$> getRandomInt n
AbiAddressType -> AbiAddress . fromInteger <$> getRandomR (0, 2 ^ (160 :: Integer) - 1)
Expand All @@ -337,7 +334,7 @@ genAbiValueM genDict = genWithDict genDict (toList <$> genDict ^. constants) $ \
genAbiCallM :: MonadRandom m => GenDict -> SolSignature -> m SolCall
genAbiCallM genDict abi = do
solCall <- genWithDict genDict
(toList <$> genDict ^. wholeCalls)
genDict._wholeCalls
(traverse $ traverse (genAbiValueM genDict))
abi
mutateAbiCall solCall
Expand Down
Loading