Capture return values during call sequence execution (#533)

* add primitive return types to value set during call sequence execution

* improve value generation of fixed byte arrays

* remove unit test because it is hard to test this properly

* fix lint and lower chance of generating random address, bytes, or string

* fix bug with byte array input generation

fuzzing/calls/call_sequence.go

@@ -219,6 +219,29 @@ func (cse *CallSequenceElement) Method() (*abi.Method, error) {
 	return method, err
 }
 
+// DecodedReturnValues returns the Go-equivalent decoded return values for the CallSequenceElement's return data
+func (cse *CallSequenceElement) DecodedReturnValues() ([]any, error) {
+	// First, retrieve the method that was called by the call sequence element
+	method, err := cse.Method()
+	if err != nil {
+		return nil, err
+	}
+
+	// Retrieve the ABI-encoded return data
+	encodedReturnData := cse.ChainReference.Block.MessageResults[cse.ChainReference.TransactionIndex].ExecutionResult.ReturnData
+	if len(encodedReturnData) == 0 {
+		return nil, nil
+	}
+
+	// Decode the return data
+	decodedReturnValues, err := method.Outputs.Unpack(encodedReturnData)
+	if err != nil {
+		return nil, err
+	}
+
+	return decodedReturnValues, nil
+}
+
 // String returns a displayable string representing the CallSequenceElement.
 func (cse *CallSequenceElement) String() string {
 	// Obtain our contract name

fuzzing/coverage/coverage_tracer.go

@@ -98,7 +98,7 @@ func (t *CoverageTracer) NativeTracer() *chain.TestChainTracer {
 	return t.nativeTracer
 }
 
-// CaptureTxStart is called upon the start of transaction execution, as defined by tracers.Tracer.
+// OnTxStart is called upon the start of transaction execution, as defined by tracers.Tracer.
 func (t *CoverageTracer) OnTxStart(vm *tracing.VMContext, tx *coretypes.Transaction, from common.Address) {
 	// Reset our call frame states
 	t.callDepth = 0

fuzzing/fuzzer.go

@@ -575,10 +575,10 @@ func defaultCallSequenceGeneratorConfigFunc(fuzzer *Fuzzer, valueSet *valuegener
 	mutationalGeneratorConfig := &valuegeneration.MutationalValueGeneratorConfig{
 		MinMutationRounds:               0,
 		MaxMutationRounds:               1,
-		GenerateRandomAddressBias:       0.5,
+		GenerateRandomAddressBias:       0.05,
 		GenerateRandomIntegerBias:       0.5,
-		GenerateRandomStringBias:        0.5,
-		GenerateRandomBytesBias:         0.5,
+		GenerateRandomStringBias:        0.05,
+		GenerateRandomBytesBias:         0.05,
 		MutateAddressProbability:        0.1,
 		MutateArrayStructureProbability: 0.1,
 		MutateBoolProbability:           0.1,

fuzzing/fuzzer_worker.go

@@ -257,9 +257,14 @@ func (fw *FuzzerWorker) updateMethods() {
 // deployed in the Chain.
 // Returns the length of the call sequence tested, any requests for call sequence shrinking, or an error if one occurs.
 func (fw *FuzzerWorker) testNextCallSequence() (calls.CallSequence, []ShrinkCallSequenceRequest, error) {
+	// We will make a copy of the worker's base value set so that we can rollback to it at the end of the call sequence
+	originalValueSet := fw.valueSet.Clone()
+
 	// After testing the sequence, we'll want to rollback changes to reset our testing state.
 	var err error
 	defer func() {
+		// Reset the value set back to the original
+		fw.valueSet = originalValueSet
 		if err == nil {
 			err = fw.chain.RevertToBlockNumber(fw.testingBaseBlockNumber)
 		}
@@ -282,11 +287,21 @@ func (fw *FuzzerWorker) testNextCallSequence() (calls.CallSequence, []ShrinkCall
 
 	// Our "post execution check function" method will check coverage and call all testing functions. If one returns a
 	// request for a shrunk call sequence, we exit our call sequence execution immediately to go fulfill the shrink
-	// request.
+	// request. Additionally, the execution check function will also attempt to add any return data to the value set for
+	// this call sequence. Note that the value set is reset after each call sequence (see the defer section above)
 	executionCheckFunc := func(currentlyExecutedSequence calls.CallSequence) (bool, error) {
+		// Get the last call sequence element that was executed
+		latestCallSequenceElement := currentlyExecutedSequence[len(currentlyExecutedSequence)-1]
+		// Get the decoded return values and add it to the base value set
+		// Don't throw an error since we care more about coverage than adding the return values to the base value set
+		decodedReturnValues, err := latestCallSequenceElement.DecodedReturnValues()
+		if decodedReturnValues != nil && err == nil {
+			fw.valueSet.Add(decodedReturnValues)
+		}
+
 		// Check for updates to coverage and corpus.
 		// If we detect coverage changes, add this sequence with weight as 1 + sequences tested (to avoid zero weights)
-		err := fw.fuzzer.corpus.CheckSequenceCoverageAndUpdate(currentlyExecutedSequence, fw.getNewCorpusCallSequenceWeight(), true)
+		err = fw.fuzzer.corpus.CheckSequenceCoverageAndUpdate(currentlyExecutedSequence, fw.getNewCorpusCallSequenceWeight(), true)
 		if err != nil {
 			return true, err
 		}

fuzzing/fuzzer_worker_sequence_generator.go

@@ -287,9 +287,9 @@ func (g *CallSequenceGenerator) generateNewElement() (*calls.CallSequenceElement
 	}
 
 	// Select a random method
-	// There is a 1/100 chance that a pure method will be invoked or if there are only pure functions that are callable
+	// There is a 1/1000 chance that a pure method will be invoked or if there are only pure functions that are callable
 	var selectedMethod *contracts.DeployedContractMethod
-	if (len(g.worker.pureMethods) > 0 && g.worker.randomProvider.Intn(100) == 0) || callOnlyPureFunctions {
+	if (len(g.worker.pureMethods) > 0 && g.worker.randomProvider.Intn(1000) == 0) || callOnlyPureFunctions {
 		selectedMethod = &g.worker.pureMethods[g.worker.randomProvider.Intn(len(g.worker.pureMethods))]
 	} else {
 		selectedMethod = &g.worker.stateChangingMethods[g.worker.randomProvider.Intn(len(g.worker.stateChangingMethods))]

fuzzing/valuegeneration/generator_mutational.go

@@ -43,7 +43,7 @@ type MutationalValueGeneratorConfig struct {
 	// GenerateRandomStringBias defines the probability in which a string generated by the value generator is entirely
 	// random, rather than mutated. Value range is [0.0, 1.0].
 	GenerateRandomStringBias float32
-	// GenerateRandomStringBias defines the probability in which a byte array generated by the value generator is
+	// GenerateRandomBytesBias defines the probability in which a byte array generated by the value generator is
 	// entirely random, rather than mutated. Value range is [0.0, 1.0].
 	GenerateRandomBytesBias float32
 
@@ -244,14 +244,20 @@ var bytesMutationMethods = []func(*MutationalValueGenerator, []byte, ...[]byte)
 	},
 }
 
-// mutateBytesInternal takes a byte array and returns either a random new byte array, or a mutated value based off the
-// input.
+// mutateBytesInternal takes a byte array and a length. This function returns either a fixed length byte array (based on
+// the provided length) or a byte slice. The returned byte array/slice is either randomly generated or mutated using
+// the provided input.
 // If a nil input is provided, this method uses an existing base value set value as the starting point for mutation.
-func (g *MutationalValueGenerator) mutateBytesInternal(b []byte) []byte {
+func (g *MutationalValueGenerator) mutateBytesInternal(b []byte, length int) []byte {
 	// If we have no inputs or our bias directs us to, use the random generator instead
 	inputs := g.valueSet.Bytes()
 	randomGeneratorDecision := g.randomProvider.Float32()
 	if len(inputs) == 0 || randomGeneratorDecision < g.config.GenerateRandomBytesBias {
+		// If the length is non-zero, generate a fixed byte array
+		if length > 0 {
+			return g.RandomValueGenerator.GenerateFixedBytes(length)
+		}
+		// Otherwise, generate a random byte slice
 		return g.RandomValueGenerator.GenerateBytes()
 	}
 
@@ -269,6 +275,19 @@ func (g *MutationalValueGenerator) mutateBytesInternal(b []byte) []byte {
 		input = bytesMutationMethods[g.randomProvider.Intn(len(bytesMutationMethods))](g, input, inputs...)
 	}
 
+	// If we want a fixed-byte array and the mutated input is smaller than the requested length, then generate a random
+	// byte array and append it to the existing input
+	if length > 0 && len(input) < length {
+		randomSlice := g.RandomValueGenerator.GenerateFixedBytes(length - len(input))
+		input = append(input, randomSlice...)
+	}
+
+	// Similarly, if we want a fixed-byte array and the mutated input is larger than the requested length, then truncate
+	// the array
+	if length > 0 && len(input) > length {
+		return input[:length]
+	}
+
 	return input
 }
 
@@ -415,7 +434,7 @@ func (g *MutationalValueGenerator) MutateBool(bl bool) bool {
 
 // GenerateBytes generates bytes and returns them.
 func (g *MutationalValueGenerator) GenerateBytes() []byte {
-	return g.mutateBytesInternal(nil)
+	return g.mutateBytesInternal(nil, 0)
 }
 
 // MutateBytes takes a dynamic-sized byte array input and returns a mutated value based off the input.
@@ -428,20 +447,20 @@ func (g *MutationalValueGenerator) MutateBytes(b []byte) []byte {
 		if randomGeneratorDecision < g.config.MutateBytesGenerateNewBias {
 			return g.GenerateBytes()
 		} else {
-			return g.mutateBytesInternal(b)
+			return g.mutateBytesInternal(b, 0)
 		}
 	}
 	return b
 }
 
 // MutateFixedBytes takes a fixed-sized byte array input and returns a mutated value based off the input.
 func (g *MutationalValueGenerator) MutateFixedBytes(b []byte) []byte {
-	// Determine whether to perform mutations against this input or just return it as-is.
-	randomGeneratorDecision := g.randomProvider.Float32()
-	if randomGeneratorDecision < g.config.MutateFixedBytesProbability {
-		return g.GenerateFixedBytes(len(b))
-	}
-	return b
+	return g.mutateBytesInternal(b, len(b))
+}
+
+// GenerateFixedBytes generates a fixed-sized byte array to use when populating inputs.
+func (g *MutationalValueGenerator) GenerateFixedBytes(length int) []byte {
+	return g.mutateBytesInternal(nil, length)
 }
 
 // GenerateString generates strings and returns them.

fuzzing/valuegeneration/value_set.go

@@ -2,8 +2,10 @@ package valuegeneration
 
 import (
 	"encoding/hex"
+	"github.com/crytic/medusa/utils/reflectionutils"
 	"hash"
 	"math/big"
+	"reflect"
 
 	"github.com/ethereum/go-ethereum/common"
 	"golang.org/x/crypto/sha3"
@@ -172,3 +174,52 @@ func (vs *ValueSet) RemoveBytes(b []byte) {
 
 	delete(vs.bytes, hashStr)
 }
+
+// Add adds one or more values. Note the values must be a primitive type (signed/unsigned integer, address, string,
+// bytes, fixed bytes)
+func (vs *ValueSet) Add(values []any) {
+	// Iterate across each value and assert on its type
+	for _, value := range values {
+		switch v := value.(type) {
+		case uint8:
+			vs.AddInteger(new(big.Int).SetUint64(uint64(v)))
+		case uint16:
+			vs.AddInteger(new(big.Int).SetUint64(uint64(v)))
+		case uint32:
+			vs.AddInteger(new(big.Int).SetUint64(uint64(v)))
+		case uint64:
+			vs.AddInteger(new(big.Int).SetUint64(v))
+		case int8:
+			vs.AddInteger(new(big.Int).SetInt64(int64(v)))
+		case int16:
+			vs.AddInteger(new(big.Int).SetInt64(int64(v)))
+		case int32:
+			vs.AddInteger(new(big.Int).SetInt64(int64(v)))
+		case int64:
+			vs.AddInteger(new(big.Int).SetInt64(v))
+		case *big.Int:
+			vs.AddInteger(v)
+		case common.Address:
+			vs.AddAddress(v)
+		case bool:
+			if value == true {
+				vs.AddInteger(new(big.Int).SetUint64(1))
+			} else {
+				vs.AddInteger(new(big.Int).SetUint64(0))
+			}
+		case string:
+			vs.AddString(v)
+		case []byte:
+			vs.AddBytes(v)
+		default:
+			// We need to be able to capture fixed bytes. Unfortunately, the only way to do so is using reflection
+			r := reflect.TypeOf(value)
+			// If we have a fixed array of uint8 (aka byte), then we will convert it into a slice and add to value set
+			if r.Kind() == reflect.Array && r.Elem().Kind() == reflect.Uint8 {
+				b := reflectionutils.ArrayToSlice(reflect.ValueOf(value)).([]byte)
+				vs.AddBytes(b)
+			}
+			continue
+		}
+	}
+}