Prune unnecessary transaction sequences from corpus (#625)

* prune sequences WIP

* keep track of total pruning

* do it only every 3 minutes

* resolve TODOs, change logging format, add config option, move state into struct

* comments and docs

* Prune in parallel

* add comments

* Update fuzzing/fuzzer.go

Co-authored-by: anishnaik <[email protected]>

* suggestions from @anishnaik

* go fmt

* move corpus_pruner and remove dependency on fuzzer

* 'values' -> 'corpus items'

---------

Co-authored-by: anishnaik <[email protected]>

Author: samalws-tob

docs/src/project_configuration/fuzzing_config.md

@@ -53,6 +53,14 @@ The fuzzing configuration defines the parameters for the fuzzing campaign.
   Enabling coverage allows for improved code exploration.
 - **Default**: `true`
 
+### `pruneFrequency`
+
+- **Type**: Integer
+- **Description**: Determines how often, in minutes, the corpus should be pruned to remove unnecessary members.
+  Setting `pruneFrequency` to 0 disables pruning.
+  `pruneFrequency` only matters if `coverageEnabled` is set to true; otherwise, no pruning will occur.
+- **Default**: `5`
+
 ### `corpusDirectory`
 
 - **Type**: String

docs/src/static/function_level_testing_medusa.json

@@ -5,6 +5,7 @@
     "timeout": 0,
     "testLimit": 1000,
     "callSequenceLength": 1,
+    "pruneFrequency": 5,
     "corpusDirectory": "",
     "coverageEnabled": true,
     "targetContracts": ["TestDepositContract"],

docs/src/static/medusa.json

@@ -6,6 +6,7 @@
     "testLimit": 0,
     "shrinkLimit": 5000,
     "callSequenceLength": 100,
+    "pruneFrequency": 5,
     "corpusDirectory": "",
     "coverageEnabled": true,
     "coverageFormats": ["html", "lcov"],

fuzzing/config/config.go

@@ -55,6 +55,11 @@ type FuzzingConfig struct {
 	// CallSequenceLength describes the maximum length a transaction sequence can be generated as.
 	CallSequenceLength int `json:"callSequenceLength"`
 
+	// PruneFrequncy determines how often, in minutes, the corpus should be pruned to remove unnecessary members.
+	// Setting PruneFrequency to 0 disables pruning.
+	// PruneFrequency only matters if CoverageEnabled is set to true; otherwise, no pruning will occur.
+	PruneFrequency uint64 `json:"pruneFrequency"`
+
 	// CorpusDirectory describes the name for the folder that will hold the corpus and the coverage files. If empty,
 	// the in-memory corpus will be used, but not flush to disk.
 	CorpusDirectory string `json:"corpusDirectory"`

fuzzing/config/config_defaults.go

@@ -45,6 +45,7 @@ func GetDefaultProjectConfig(platform string) (*ProjectConfig, error) {
 			TestLimit:               0,
 			ShrinkLimit:             5_000,
 			CallSequenceLength:      100,
+			PruneFrequency:          5,
 			TargetContracts:         []string{},
 			TargetContractsBalances: []*ContractBalance{},
 			PredeployedContracts:    map[string]string{},

fuzzing/corpus/corpus.go

@@ -1,7 +1,10 @@
 package corpus
 
 import (
+	"math/rand"
+
 	"bytes"
+	"context"
 	"fmt"
 	"math/big"
 	"os"
@@ -455,14 +458,13 @@ func (c *Corpus) AddTestResultCallSequence(callSequence calls.CallSequence, muta
 	return c.addCallSequence(c.testResultSequenceFiles, callSequence, false, mutationChooserWeight, flushImmediately)
 }
 
-// CheckSequenceCoverageAndUpdate checks if the most recent call executed in the provided call sequence achieved
-// coverage the Corpus did not with any of its call sequences. If it did, the call sequence is added to the corpus
-// and the Corpus coverage maps are updated accordingly.
-// Returns an error if one occurs.
-func (c *Corpus) CheckSequenceCoverageAndUpdate(callSequence calls.CallSequence, mutationChooserWeight *big.Int, flushImmediately bool) error {
+// checkSequenceCoverageAndUpdate checks if the most recent call executed in the provided call sequence achieved
+// coverage the not already included in coverageMaps. If it did, coverageMaps is updated accordingly.
+// Returns a boolean indicating whether any change happened, and an error if one occurs.
+func checkSequenceCoverageAndUpdate(callSequence calls.CallSequence, coverageMaps *coverage.CoverageMaps) (bool, error) {
 	// If we have coverage-guided fuzzing disabled or no calls in our sequence, there is nothing to do.
 	if len(callSequence) == 0 {
-		return nil
+		return false, nil
 	}
 
 	// Obtain our coverage maps for our last call.
@@ -473,14 +475,22 @@ func (c *Corpus) CheckSequenceCoverageAndUpdate(callSequence calls.CallSequence,
 
 	// If we have none, because a coverage tracer wasn't attached when processing this call, we can stop.
 	if lastMessageCoverageMaps == nil {
-		return nil
+		return false, nil
 	}
 
 	// Memory optimization: Remove them from the results now that we obtained them, to free memory later.
 	coverage.RemoveCoverageTracerResults(lastMessageResult)
 
 	// Merge the coverage maps into our total coverage maps and check if we had an update.
-	coverageUpdated, err := c.coverageMaps.Update(lastMessageCoverageMaps)
+	return coverageMaps.Update(lastMessageCoverageMaps)
+}
+
+// CheckSequenceCoverageAndUpdate checks if the most recent call executed in the provided call sequence achieved
+// coverage the Corpus did not with any of its call sequences. If it did, the call sequence is added to the corpus
+// and the Corpus coverage maps are updated accordingly.
+// Returns an error if one occurs.
+func (c *Corpus) CheckSequenceCoverageAndUpdate(callSequence calls.CallSequence, mutationChooserWeight *big.Int, flushImmediately bool) error {
+	coverageUpdated, err := checkSequenceCoverageAndUpdate(callSequence, c.coverageMaps)
 	if err != nil {
 		return err
 	}
@@ -551,3 +561,77 @@ func (c *Corpus) Flush() error {
 
 	return nil
 }
+
+// PruneSequences removes unnecessary entries from the corpus. It does this by:
+//   - Initialize a blank coverage map tmpMap
+//   - Grab all sequences in the corpus
+//   - Randomize the order
+//   - For each transaction, see whether it adds anything new to tmpMap.
+//     If it does, add the new coverage and continue.
+//     If it doesn't, remove it from the corpus.
+//
+// By doing this, we hope to find a smaller set of txn sequences that still preserves our current coverage.
+// PruneSequences takes a chain.TestChain parameter used to run transactions.
+// It returns an int indicating the number of sequences removed from the corpus, and an error if any occurred.
+func (c *Corpus) PruneSequences(ctx context.Context, chain *chain.TestChain) (int, error) {
+	chainOriginalIndex := uint64(len(chain.CommittedBlocks()))
+	tmpMap := coverage.NewCoverageMaps()
+
+	c.callSequencesLock.Lock()
+	seqs := make([]calls.CallSequence, len(c.mutationTargetSequenceChooser.Choices))
+	for i, seq := range c.mutationTargetSequenceChooser.Choices {
+		seqCloned, err := seq.Data.Clone()
+		if err != nil {
+			c.callSequencesLock.Unlock()
+			return 0, err
+		}
+		seqs[i] = seqCloned
+	}
+	c.callSequencesLock.Unlock()
+	// We don't need to lock during the next part as long as the ordering of Choices doesn't change.
+	// New items could get added in the meantime, but older items won't be touched.
+
+	toRemove := map[int]bool{}
+
+	// Iterate seqs in a random order
+	for _, i := range rand.Perm(len(seqs)) {
+		if utils.CheckContextDone(ctx) {
+			return 0, nil
+		}
+
+		seq := seqs[i]
+
+		fetchElementFunc := func(currentIndex int) (*calls.CallSequenceElement, error) {
+			if currentIndex >= len(seq) {
+				return nil, nil
+			}
+			return seq[currentIndex], nil
+		}
+
+		// Never quit early
+		executionCheckFunc := func(currentlyExecutedSequence calls.CallSequence) (bool, error) { return false, nil }
+
+		seq, err := calls.ExecuteCallSequenceIteratively(chain, fetchElementFunc, executionCheckFunc)
+		if err != nil {
+			return 0, err
+		}
+
+		coverageUpdated, err := checkSequenceCoverageAndUpdate(seq, tmpMap)
+		if err != nil {
+			return 0, err
+		}
+
+		if !coverageUpdated {
+			// No new coverage was added. We can remove this from the corpus.
+			toRemove[i] = true
+		}
+
+		err = chain.RevertToBlockIndex(chainOriginalIndex)
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	c.mutationTargetSequenceChooser.RemoveChoices(toRemove)
+	return len(toRemove), nil
+}

fuzzing/corpus/corpus_pruner.go

@@ -0,0 +1,113 @@
+package corpus
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/crytic/medusa/chain"
+	"github.com/crytic/medusa/fuzzing/coverage"
+	"github.com/crytic/medusa/logging"
+	"github.com/crytic/medusa/utils"
+)
+
+// CorpusPruner is a job that runs every `PruneFrequency` minutes.
+// It removes unnecessary items from the corpus by calling `Corpus.PruneSequences`.
+type CorpusPruner struct {
+	// enabled determines if the pruner is enabled
+	enabled bool
+
+	// corpus is the corpus to be pruned
+	corpus *Corpus
+
+	// logger is used to log when pruning and on error
+	logger *logging.Logger
+
+	// ctx is the CorpusPruner's context which can be used to cancel the pruner
+	ctx context.Context
+
+	// pruneFrequency determines how often, in minutes, the pruning should occur
+	pruneFrequency uint64
+
+	// totalCorpusPruned counts the total number of sequences pruned so far
+	totalCorpusPruned int
+
+	// chain is the test chain used during pruning
+	chain *chain.TestChain
+}
+
+// NewCorpusPruner creates a new CorpusPruner.
+func NewCorpusPruner(enabled bool, pruneFrequency uint64, logger *logging.Logger) *CorpusPruner {
+	if !enabled {
+		return &CorpusPruner{}
+	}
+	return &CorpusPruner{
+		enabled:        enabled,
+		pruneFrequency: pruneFrequency,
+		logger:         logger,
+	}
+}
+
+// pruneCorpus is a wrapper around Corpus.PruneSequences that adds timing, logging, and updating totalCorpusPruned.
+// It is used by mainLoop.
+func (cp *CorpusPruner) pruneCorpus() error {
+	start := time.Now() // We'll track how long pruning takes
+	n, err := cp.corpus.PruneSequences(cp.ctx, cp.chain)
+	// PruneSequences takes a while, so ctx could've finished in the meantime.
+	// If it did, we skip the log message.
+	if err != nil || utils.CheckContextDone(cp.ctx) {
+		return err
+	}
+	cp.totalCorpusPruned += n
+	cp.logger.Info(fmt.Sprintf("Pruned %d corpus items in %v. Total pruned this run: %d", n, time.Since(start), cp.totalCorpusPruned))
+	return nil
+}
+
+// mainLoop calls pruneCorpus every `pruneFrequency` minutes.
+// It runs infinitely until ctx.Done is triggered.
+func (cp *CorpusPruner) mainLoop() {
+	defer cp.chain.Close()
+	ticker := time.NewTicker(time.Duration(cp.pruneFrequency) * time.Minute)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-cp.ctx.Done():
+			return
+		case <-ticker.C:
+			err := cp.pruneCorpus()
+			if err != nil {
+				cp.logger.Error("Corpus pruner encountered an error", err)
+				return
+			}
+		}
+	}
+}
+
+// Start takes a context, a corpus to prune, and a base chain in a setup state ready for testing.
+// It clones the base chain, then prunes the corpus every `PruneFrequency` minutes.
+// This runs until ctx cancels the operation.
+// Returns an error if one occurred.
+func (cp *CorpusPruner) Start(ctx context.Context, corpus *Corpus, baseTestChain *chain.TestChain) error {
+	if !cp.enabled {
+		return nil
+	}
+
+	// Clone our chain, attaching a tracer.
+	clonedChain, err := baseTestChain.Clone(func(initializedChain *chain.TestChain) error {
+		initializedChain.AddTracer(coverage.NewCoverageTracer().NativeTracer(), true, false)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	cp.chain = clonedChain
+
+	// Write our params to the struct so we don't have to pass them all over the place as function args.
+	cp.ctx = ctx
+	cp.corpus = corpus
+
+	// Start up the main loop in a goroutine.
+	go cp.mainLoop()
+
+	return nil
+}

fuzzing/fuzzer.go

@@ -85,6 +85,9 @@ type Fuzzer struct {
 	// revertReporter tracks per-function reversion metrics, if enabled
 	revertReporter *reverts.RevertReporter
 
+	// corpusPruner is a service that will prune the corpus at a given frequency to reduce corpus size and memory overhead.
+	corpusPruner *corpus.CorpusPruner
+
 	// randomProvider describes the provider used to generate random values in the Fuzzer. All other random providers
 	// used by the Fuzzer's subcomponents are derived from this one.
 	randomProvider *rand.Rand
@@ -174,6 +177,10 @@ func NewFuzzer(config config.ProjectConfig) (*Fuzzer, error) {
 		return nil, err
 	}
 
+	// Create the corpus pruner.
+	pruneEnabled := config.Fuzzing.CoverageEnabled && config.Fuzzing.PruneFrequency > 0
+	corpusPruner := corpus.NewCorpusPruner(pruneEnabled, config.Fuzzing.PruneFrequency, logger)
+
 	// Create and return our fuzzing instance.
 	fuzzer := &Fuzzer{
 		config:              config,
@@ -184,6 +191,7 @@ func NewFuzzer(config config.ProjectConfig) (*Fuzzer, error) {
 		testCases:           make([]TestCase, 0),
 		testCasesFinished:   make(map[string]TestCase),
 		revertReporter:      revertReporter,
+		corpusPruner:        corpusPruner,
 		Hooks: FuzzerHooks{
 			NewCallSequenceGeneratorConfigFunc: defaultCallSequenceGeneratorConfigFunc,
 			NewShrinkingValueMutatorFunc:       defaultShrinkingValueMutatorFunc,
@@ -853,6 +861,13 @@ func (f *Fuzzer) Start() error {
 		)
 	}
 
+	// Start the corpus pruner.
+	err = f.corpusPruner.Start(f.ctx, f.corpus, baseTestChain)
+	if err != nil {
+		f.logger.Error("Error starting corpus pruner", err)
+		return err
+	}
+
 	// Log the start of our fuzzing campaign.
 	f.logger.Info("Fuzzing with ", colors.Bold, f.config.Fuzzing.Workers, colors.Reset, " workers")