Fix unlimited memory usage by bounding dynamic dict growth [Vibecoded]

lib/Echidna.hs

@@ -100,6 +100,9 @@ prepareContract cfg solFiles buildOutput selectedContract seed = do
                      Set.empty
                      seed
                      (returnTypes contracts)
+                     env.cfg.campaignConf.dictDynamicConstantsLimit
+                     env.cfg.campaignConf.dictDynamicValuesLimit
+                     env.cfg.campaignConf.dictDynamicCallsLimit
                      nonViewPureSigs
   pure (vm, env, dict)
 

lib/Echidna/ABI.hs

@@ -17,7 +17,9 @@ import Data.List.NonEmpty (NonEmpty)
 import Data.List.NonEmpty qualified as NE
 import Data.Map (Map)
 import Data.Map qualified as Map
-import Data.Maybe (fromMaybe, catMaybes)
+import Data.Maybe (catMaybes)
+import Data.Sequence (Seq, (|>))
+import Data.Sequence qualified as Seq
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Text (Text)
@@ -111,20 +113,74 @@ hashSig = abiKeccak . TE.encodeUtf8
 data GenDict = GenDict
   { pSynthA    :: Float
     -- ^ Fraction of time to use dictionary vs. synthesize
-  , constants  :: !(Map AbiType (Set AbiValue))
-    -- ^ Constants to use, sorted by type
-  , wholeCalls :: !(Map SolSignature (Set SolCall))
-    -- ^ Whole calls to use, sorted by type
+  , staticConstants  :: !(Map AbiType (Set AbiValue))
+    -- ^ Static constants to use, sorted by type
+  , dynamicConstants :: !(Map AbiType (BoundedSet AbiValue))
+    -- ^ Runtime-mined constants to use, sorted by type
+  , staticWholeCalls :: !(Map SolSignature (Set SolCall))
+    -- ^ Static whole calls to use, sorted by type
+  , dynamicWholeCalls :: !(Map SolSignature (BoundedSet SolCall))
+    -- ^ Runtime-mined whole calls to use, sorted by type
   , defSeed    :: Int
     -- ^ Default seed to use if one is not provided in EConfig
   , rTypes     :: Text -> Maybe AbiType
     -- ^ Return types of any methods we scrape return values from
-  , dictValues :: !(Set W256)
-    -- ^ A set of int/uint constants for better performance
+  , staticDictValues :: !(Set W256)
+    -- ^ Static int/uint constants for better performance
+  , dynamicDictValues :: !(BoundedSet W256)
+    -- ^ Runtime-mined int/uint constants for better performance
+  , dynamicConstantsLimit :: !Int
+    -- ^ Maximum runtime-mined constants to keep per ABI type
+  , dynamicValuesLimit :: !Int
+    -- ^ Maximum runtime-mined int/uint constants to keep
+  , dynamicCallsLimit :: !Int
+    -- ^ Maximum runtime-mined whole calls to keep per signature
   , callbackSigs    :: ![SolSignature]
     -- ^ A list of callback signatures (for generating random callbacks)
   }
 
+data BoundedSet a = BoundedSet
+  { boundedOrder :: !(Seq a)
+  , boundedMembers :: !(Set a)
+  }
+
+emptyBoundedSet :: BoundedSet a
+emptyBoundedSet = BoundedSet mempty Set.empty
+
+boundedSize :: BoundedSet a -> Int
+boundedSize = Set.size . (.boundedMembers)
+
+boundedInsert :: Ord a => Int -> Set a -> a -> BoundedSet a -> BoundedSet a
+boundedInsert limit staticMembers value bounded
+  | limit <= 0 = emptyBoundedSet
+  | value `Set.member` staticMembers = bounded
+  | value `Set.member` bounded.boundedMembers = bounded
+  | otherwise = trimBounded limit $
+      BoundedSet (bounded.boundedOrder |> value) (Set.insert value bounded.boundedMembers)
+
+boundedInsertSet :: Ord a => Int -> Set a -> Set a -> BoundedSet a -> BoundedSet a
+boundedInsertSet limit staticMembers values bounded =
+  Set.foldl' (\acc value -> boundedInsert limit staticMembers value acc) bounded values
+
+trimBounded :: Ord a => Int -> BoundedSet a -> BoundedSet a
+trimBounded limit bounded
+  | boundedSize bounded <= limit = bounded
+  | otherwise =
+      case Seq.viewl bounded.boundedOrder of
+        Seq.EmptyL -> emptyBoundedSet
+        value Seq.:< rest -> trimBounded limit $ BoundedSet rest (Set.delete value bounded.boundedMembers)
+
+rElemStaticDynamic :: (MonadRandom m, Ord a) => Set a -> BoundedSet a -> m (Maybe a)
+rElemStaticDynamic staticMembers dynamicMembers =
+  case Set.size staticMembers + boundedSize dynamicMembers of
+    0 -> pure Nothing
+    total -> do
+      idx <- getRandomR (0, total - 1)
+      pure . Just $
+        if idx < Set.size staticMembers
+          then Set.elemAt idx staticMembers
+          else Seq.index dynamicMembers.boundedOrder (idx - Set.size staticMembers)
+
 hashMapBy
   :: (Ord k, Eq k, Ord a)
   => (a -> k)
@@ -134,7 +190,32 @@ hashMapBy f = Map.fromListWith Set.union . fmap (\v -> (f v, Set.singleton v)) .
 
 gaddCalls :: Set SolCall -> GenDict -> GenDict
 gaddCalls calls dict =
-  dict { wholeCalls = dict.wholeCalls <> hashMapBy (fmap $ fmap abiValueType) calls }
+  dict { dynamicWholeCalls = Map.foldlWithKey' insertCalls dict.dynamicWholeCalls callsBySig }
+  where
+    callsBySig = hashMapBy (fmap $ fmap abiValueType) calls
+    insertCalls acc sig callsForSig =
+      let staticCalls = Map.findWithDefault Set.empty sig dict.staticWholeCalls
+          existingCalls = Map.findWithDefault emptyBoundedSet sig acc
+          updatedCalls = boundedInsertSet dict.dynamicCallsLimit staticCalls callsForSig existingCalls
+      in if boundedSize updatedCalls == 0
+         then Map.delete sig acc
+         else Map.insert sig updatedCalls acc
+
+addDynamicConstants :: Map AbiType (Set AbiValue) -> GenDict -> GenDict
+addDynamicConstants additions dict =
+  dict
+    { dynamicConstants = Map.foldlWithKey' insertConstants dict.dynamicConstants additions
+    , dynamicDictValues = boundedInsertSet dict.dynamicValuesLimit dict.staticDictValues dynamicValues dict.dynamicDictValues
+    }
+  where
+    insertConstants acc abiType values =
+      let staticValues = Map.findWithDefault Set.empty abiType dict.staticConstants
+          existingValues = Map.findWithDefault emptyBoundedSet abiType acc
+          updatedValues = boundedInsertSet dict.dynamicConstantsLimit staticValues values existingValues
+      in if boundedSize updatedValues == 0
+         then Map.delete abiType acc
+         else Map.insert abiType updatedValues acc
+    dynamicValues = mkDictValues $ Set.unions $ Map.elems additions
 
 -- | Construct a 'GenDict' from some dictionaries, a 'Float', a default seed,
 -- and a typing rule for return values
@@ -144,18 +225,27 @@ mkGenDict
   -> Set SolCall  -- ^ A list of complete 'SolCall's to mutate
   -> Int          -- ^ A default seed
   -> (Text -> Maybe AbiType) -- ^ A return value typing rule
+  -> Int
+  -> Int
+  -> Int
   -> [SolSignature]
   -> GenDict
-mkGenDict mutationChance abiValues solCalls seed typingRule =
+mkGenDict mutationChance abiValues solCalls seed typingRule constantsLimit valuesLimit callsLimit =
   GenDict mutationChance
           (hashMapBy abiValueType abiValues)
+          Map.empty
           (hashMapBy (fmap $ fmap abiValueType) solCalls)
+          Map.empty
           seed
           typingRule
           (mkDictValues abiValues)
+          emptyBoundedSet
+          constantsLimit
+          valuesLimit
+          callsLimit
 
 emptyDict :: GenDict
-emptyDict = mkGenDict 0 Set.empty Set.empty 0 (const Nothing) []
+emptyDict = mkGenDict 0 Set.empty Set.empty 0 (const Nothing) 0 0 0 []
 
 mkDictValues :: Set AbiValue -> Set W256
 mkDictValues =
@@ -361,19 +451,32 @@ 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, Ord a, MonadRandom m)
+  :: (MonadRandom m)
   => GenDict
-  -> Map a (Set b)
+  -> (a -> m (Maybe b))
   -> (a -> m b)
   -> a
   -> m b
-genWithDict genDict m g t = do
+genWithDict genDict fromDict g t = do
   r <- getRandom
-  let maybeValM = if genDict.pSynthA >= r then fromDict else pure Nothing
-      fromDict = case Map.lookup t m of
-                   Nothing -> pure Nothing
-                   Just cs -> Just <$> rElem' cs
-  fromMaybe <$> g t <*> maybeValM
+  maybeVal <- if genDict.pSynthA >= r then fromDict t else pure Nothing
+  maybe (g t) pure maybeVal
+
+constantFromDict :: MonadRandom m => GenDict -> AbiType -> m (Maybe AbiValue)
+constantFromDict genDict abiType =
+  rElemStaticDynamic
+    (Map.findWithDefault Set.empty abiType genDict.staticConstants)
+    (Map.findWithDefault emptyBoundedSet abiType genDict.dynamicConstants)
+
+wholeCallFromDict :: MonadRandom m => GenDict -> SolSignature -> m (Maybe SolCall)
+wholeCallFromDict genDict sig =
+  rElemStaticDynamic
+    (Map.findWithDefault Set.empty sig genDict.staticWholeCalls)
+    (Map.findWithDefault emptyBoundedSet sig genDict.dynamicWholeCalls)
+
+dictValueFromDict :: MonadRandom m => GenDict -> m (Maybe W256)
+dictValueFromDict genDict =
+  rElemStaticDynamic genDict.staticDictValues genDict.dynamicDictValues
 
 -- | A small number of dummy addresses
 pregenAdds :: [Addr]
@@ -418,15 +521,15 @@ genAbiValueM' genDict funcName depth t =
         AbiTupleType v        -> AbiTuple <$> traverse (genAbiValueM' genDict funcName (depth + 1)) v
         AbiFunctionType       -> AbiFunction <$> rElem (NE.fromList pregenAdds)
                                              <*> (FunctionSelector <$> getRandom)
-  in genWithDict genDict genDict.constants go t
+  in genWithDict genDict (constantFromDict genDict) go t
 
 -- | Given a 'SolSignature', generate a random 'SolCall' with that signature,
 -- possibly with a dictionary.
 genAbiCallM :: MonadRandom m => GenDict -> SolSignature -> m SolCall
 genAbiCallM genDict (name, types) = do
   let genVals = zipWithM (flip (genAbiValueM' genDict name)) types (repeat 0)
   solCall <- genWithDict genDict
-                         genDict.wholeCalls
+                         (wholeCallFromDict genDict)
                          (const ((name,) <$> genVals))
                          (name, types)
   mutateAbiCall solCall

lib/Echidna/Campaign.hs

@@ -500,12 +500,8 @@ callseq vm txSeq = do
       eventDiffs = extractEventValues env.dapp vm vm'
       -- union the return results with the new addresses
       additions = Map.unionsWith Set.union [resultMap, eventDiffs, diffs]
-      -- append to the constants dictionary
-      updatedDict = workerState.genDict
-        { constants = Map.unionWith Set.union workerState.genDict.constants additions
-        , dictValues = Set.union (mkDictValues $ Set.unions $ Map.elems additions)
-                                 workerState.genDict.dictValues
-        }
+      -- append to the runtime-mined constants dictionary
+      updatedDict = addDynamicConstants additions workerState.genDict
 
     -- Update the worker state
     in workerState

lib/Echidna/Config.hs

@@ -72,6 +72,12 @@ instance FromJSON EConfigWithUsage where
           fail $ show k <> ": value does not fit in 256 bits"
         else
           pure $ fromIntegral value
+      nonNegative k def = do
+        value <- v ..:? k ..!= def
+        if value < 0 then
+          lift $ fail $ show k <> ": value must be non-negative"
+        else
+          pure value
 
       txConfParser = TxConf
         <$> v ..:? "propMaxGas" ..!= maxGasPerBlock
@@ -96,6 +102,9 @@ instance FromJSON EConfigWithUsage where
         <*> (v ..:? "coverage" <&> \case Just False -> Nothing;  _ -> Just mempty)
         <*> v ..:? "seed"
         <*> v ..:? "dictFreq" ..!= 0.40
+        <*> nonNegative "dictDynamicConstantsLimit" defaultDictDynamicConstantsLimit
+        <*> nonNegative "dictDynamicValuesLimit" defaultDictDynamicValuesLimit
+        <*> nonNegative "dictDynamicCallsLimit" defaultDictDynamicCallsLimit
         <*> v ..:? "corpusDir" ..!= Nothing
         <*> v ..:? "coverageDir" ..!= Nothing
         <*> v ..:? "mutConsts" ..!= defaultMutationConsts

lib/Echidna/Transaction.hs

@@ -11,7 +11,6 @@ import Control.Monad.State.Strict (MonadState, gets, modify', execState)
 import Data.ByteString qualified as BS
 import Data.Map (Map, toList)
 import Data.Maybe (catMaybes)
-import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Vector qualified as V
 import Optics.Core
@@ -70,9 +69,9 @@ genTx world deployedContracts = do
   contractAList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
   (dstAddr, dstAbis) <- rElem' $ Set.fromList $ catMaybes contractAList
   solCall <- genInteractionsM genDict dstAbis
-  value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall
-  ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues
-            <*> genDelay txConf.maxBlockDelay genDict.dictValues
+  value <- genValue txConf.maxValue genDict world.payableSigs solCall
+  ts <- (,) <$> genDelay txConf.maxTimeDelay genDict
+            <*> genDelay txConf.maxBlockDelay genDict
   pure $ Tx { call = SolCall solCall
             , src = sender
             , dst = dstAddr
@@ -86,20 +85,20 @@ genTx world deployedContracts = do
     toContractA env sigMap (addr, c) =
       fmap (forceAddr addr,) . snd <$> lookupUsingCodehash env.codehashMap c env.dapp sigMap
 
-genDelay :: MonadRandom m => W256 -> Set W256 -> m W256
-genDelay mv ds =
+genDelay :: MonadRandom m => W256 -> GenDict -> m W256
+genDelay mv genDict =
   join $ oftenUsually fromDict randValue
   where randValue = fromIntegral <$> getRandomR (0 :: Integer, fromIntegral mv)
-        fromDict = (`mod` (mv + 1)) <$> rElem' ds
+        fromDict = maybe randValue (pure . (`mod` (mv + 1))) =<< dictValueFromDict genDict
 
 genValue
   :: MonadRandom m
   => W256
-  -> Set W256
+  -> GenDict
   -> [FunctionSelector]
   -> SolCall
   -> m W256
-genValue mv ds ps sc =
+genValue mv genDict ps sc =
   if sig `elem` ps then
     join $ oftenUsually fromDict randValue
   else
@@ -108,7 +107,7 @@ genValue mv ds ps sc =
   where
     randValue = fromIntegral <$> getRandomR (0 :: Integer, fromIntegral mv)
     sig = (hashSig . encodeSig . signatureCall) sc
-    fromDict = (`mod` (mv + 1)) <$> rElem' ds
+    fromDict = maybe randValue (pure . (`mod` (mv + 1))) =<< dictValueFromDict genDict
 
 -- | Check if a 'Transaction' is as \"small\" (simple) as possible (using ad-hoc heuristics).
 canShrinkTx :: Tx -> Bool

lib/Echidna/Types/Campaign.hs

@@ -29,6 +29,12 @@ data CampaignConf = CampaignConf
     -- ^ Seed used for the generation of random transactions
   , dictFreq           :: Float
     -- ^ Frequency for the use of dictionary values in the random transactions
+  , dictDynamicConstantsLimit :: Int
+    -- ^ Maximum runtime-mined constants to keep per ABI type
+  , dictDynamicValuesLimit :: Int
+    -- ^ Maximum runtime-mined int/uint constants to keep
+  , dictDynamicCallsLimit :: Int
+    -- ^ Maximum runtime-mined whole calls to keep per signature
   , corpusDir          :: Maybe FilePath
     -- ^ Directory to load and save lists of transactions
   , coverageDir        :: Maybe FilePath
@@ -108,6 +114,15 @@ defaultSequenceLength = 100
 defaultShrinkLimit :: Int
 defaultShrinkLimit = 5000
 
+defaultDictDynamicConstantsLimit :: Int
+defaultDictDynamicConstantsLimit = 4096
+
+defaultDictDynamicValuesLimit :: Int
+defaultDictDynamicValuesLimit = 8192
+
+defaultDictDynamicCallsLimit :: Int
+defaultDictDynamicCallsLimit = 1024
+
 defaultSymExecTimeout :: Int
 defaultSymExecTimeout = 30
 

src/test/Spec.hs

@@ -6,6 +6,7 @@ import Tests.Cheat (cheatTests)
 import Tests.Compile (compilationTests)
 import Tests.Config (configTests)
 import Tests.Coverage (coverageTests)
+import Tests.Dict (dictTests)
 import Tests.Encoding (encodingJSONTests)
 import Tests.Foundry (foundryTests)
 import Tests.FoundryTestGen (foundryTestGenTests)
@@ -21,6 +22,7 @@ main :: IO ()
 main = withCurrentDirectory "./tests/solidity" . defaultMain $
          testGroup "Echidna"
            [ configTests
+           , dictTests
            , compilationTests
            , seedTests
            , integrationTests