moved logic to contract creation + cleanup + add tests

crytic · Jun 18, 2024 · 1793913 · 1793913
1 parent c1432f5
commit 1793913
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 76 deletions.
diff --git a/fuzzing/contracts/contract.go b/fuzzing/contracts/contract.go
@@ -41,15 +41,21 @@ type Contract struct {
 }
 
 // NewContract returns a new Contract instance with the provided information.
-// TODO investigate whether you can just compute the available functions and add them here (probably)
-// TODO implement temporary whitelist/blacklist logic here (if blacklist mode on remove all by default, whitelist enable all by default)
 func NewContract(name string, sourcePath string, compiledContract *types.CompiledContract, compilation *types.Compilation) *Contract {
+	abi := compiledContract.Abi
+	callableMethods := make(map[string]bool)
+	for _, method := range abi.Methods {
+		// Whitelist all functions by default
+		callableMethods[method.Sig] = true
+
+	}
+
 	return &Contract{
 		name:             name,
 		sourcePath:       sourcePath,
 		compiledContract: compiledContract,
 		compilation:      compilation,
-		callableMethods:  make(map[string]bool),
+		callableMethods:  callableMethods,
 	}
 }
 

diff --git a/fuzzing/fuzzer.go b/fuzzing/fuzzer.go
@@ -290,44 +290,41 @@ func (f *Fuzzer) AddCompilationTargets(compilations []compilationTypes.Compilati
 			f.baseValueSet.SeedFromAst(source.Ast)
 
 			// Loop for every contract and register it in our contract definitions
-			//TODO grab the abi and populate whitelisted array here
 			for contractName := range source.Contracts {
 				contract := source.Contracts[contractName]
 				contractDefinition := fuzzerTypes.NewContract(contractName, sourcePath, &contract, compilation)
 
 				// Populate black/white list
-				if len(f.config.Fuzzing.Testing.FilterBlacklist) > 0 {
-					//TODO abstract into own function
-					for _, filterFunction := range f.config.Fuzzing.Testing.FilterBlacklist {
-						contractNameAndMethodSig := strings.Split(filterFunction, ".")
-						name := contractNameAndMethodSig[0]
-						methodSig := contractNameAndMethodSig[1]
-						if contractName == name {
-							for _, method := range contract.Abi.Methods {
-								if method.Name == methodSig {
-									if method.IsConstant() {
-										// Warn user non state-changing functions should not be in blacklist
-										f.logger.Warn("non state-changing functions should not be in the blacklist")
-									} else {
-										contractDefinition.BlackListFunction(method.Sig)
-									}
+				for _, filterFunction := range f.config.Fuzzing.Testing.FilterBlacklist {
+					contractNameAndMethodSig := strings.Split(filterFunction, ".")
+					name := contractNameAndMethodSig[0]
+					methodSig := contractNameAndMethodSig[1]
+					if contractName == name {
+						for _, method := range contract.Abi.Methods {
+							if method.Name == methodSig {
+								if method.IsConstant() {
+									// Warn user non state-changing functions should not be in blacklist
+									f.logger.Warn("non state-changing functions should not be in the blacklist")
+								} else {
+									// Add a function to the blacklist
+									contractDefinition.BlackListFunction(method.Sig)
 								}
 							}
 						}
 					}
-				} else if len(f.config.Fuzzing.Testing.FilterWhitelist) > 0 {
 					for _, filterFunction := range f.config.Fuzzing.Testing.FilterWhitelist {
 						contractNameAndMethodSig := strings.Split(filterFunction, ".")
 						name := contractNameAndMethodSig[0]
 						methodSig := contractNameAndMethodSig[1]
 						if contractName == name {
 							for _, method := range contract.Abi.Methods {
-								if method.Name == methodSig {
+								if method.Name != methodSig {
 									if method.IsConstant() {
 										// Warn user non state-changing functions should not be in whitelist
 										f.logger.Warn("non state-changing functions should not be in the whitelist")
 									} else {
-										contractDefinition.WhiteListFunction(method.Sig)
+										// Remove a function from the whitelist
+										contractDefinition.BlackListFunction(method.Sig)
 									}
 								}
 							}

diff --git a/fuzzing/fuzzer_test.go b/fuzzing/fuzzer_test.go
@@ -147,18 +147,16 @@ func TestAssertionsAndProperties(t *testing.T) {
 	})
 }
 
-// TestFilterBlacklistTrue runs a test to ensure that FilterBlacklist works as expected when set to true
-func TestFilterBlacklistTrue(t *testing.T) {
+// TestFilterBlacklist runs a test to ensure that FilterBlacklist works as expected when a blacklist is set
+func TestFilterBlacklist(t *testing.T) {
 	runFuzzerTest(t, &fuzzerSolcFileTest{
 		filePath: "testdata/contracts/assertions/assert_filter_functions.sol",
 		configUpdates: func(config *config.ProjectConfig) {
-			config.Fuzzing.DeploymentOrder = []string{"TestContract"}
+			config.Fuzzing.TargetContracts = []string{"TestContract"}
 			config.Fuzzing.TestLimit = 500
 			config.Fuzzing.Testing.StopOnFailedTest = true
 			config.Fuzzing.Testing.PropertyTesting.Enabled = true
-
-			config.Fuzzing.Testing.FilterBlacklist = true
-			config.Fuzzing.Testing.FilterFunctions = []string{"TestContract.reset1()", "TestContract.reset2()"}
+			config.Fuzzing.Testing.FilterBlacklist = []string{"TestContract.reset1()", "TestContract.reset2()"}
 		},
 		method: func(f *fuzzerTestContext) {
 			// Start the fuzzer
@@ -171,18 +169,16 @@ func TestFilterBlacklistTrue(t *testing.T) {
 	})
 }
 
-// TestFilterBlacklistFalse runs a test to ensure that FilterBlacklist works as expected when set to false
-func TestFilterBlacklistFalse(t *testing.T) {
+// TestFilterWhitelist runs a test to ensure that FilterWhitelist works as expected when a whitelist is set
+func TestFilterWhitelist(t *testing.T) {
 	runFuzzerTest(t, &fuzzerSolcFileTest{
 		filePath: "testdata/contracts/assertions/assert_filter_functions.sol",
 		configUpdates: func(config *config.ProjectConfig) {
-			config.Fuzzing.DeploymentOrder = []string{"TestContract"}
+			config.Fuzzing.TargetContracts = []string{"TestContract"}
 			config.Fuzzing.TestLimit = 500
 			config.Fuzzing.Testing.StopOnFailedTest = true
 			config.Fuzzing.Testing.PropertyTesting.Enabled = true
-
-			config.Fuzzing.Testing.FilterBlacklist = false
-			config.Fuzzing.Testing.FilterFunctions = []string{"TestContract.f(uint256)", "TestContract.g(uint256)", "TestContract.h(uint256)", "TestContract.i()"}
+			config.Fuzzing.Testing.FilterWhitelist = []string{"TestContract.f(uint256)", "TestContract.g(uint256)", "TestContract.h(uint256)", "TestContract.i()"}
 		},
 		method: func(f *fuzzerTestContext) {
 			// Start the fuzzer

diff --git a/fuzzing/fuzzer_worker.go b/fuzzing/fuzzer_worker.go
@@ -2,19 +2,17 @@ package fuzzing
 
 import (
 	"fmt"
-	"github.com/ethereum/go-ethereum/accounts/abi"
-	"math/big"
-	"math/rand"
-	"strings"
-
 	"github.com/crytic/medusa/chain"
 	"github.com/crytic/medusa/fuzzing/calls"
 	fuzzerTypes "github.com/crytic/medusa/fuzzing/contracts"
 	"github.com/crytic/medusa/fuzzing/coverage"
 	"github.com/crytic/medusa/fuzzing/valuegeneration"
 	"github.com/crytic/medusa/utils"
+	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"golang.org/x/exp/maps"
+	"math/big"
+	"math/rand"
 )
 
 // FuzzerWorker describes a single thread worker utilizing its own go-ethereum test node to run property tests against
@@ -232,30 +230,10 @@ func (fw *FuzzerWorker) updateStateChangingMethods() {
 	// Loop through each deployed contract
 	for contractAddress, contractDefinition := range fw.deployedContracts {
 		// Check if functions should be filtered based on the configuration
-		if len(fw.fuzzer.config.Fuzzing.Testing.FilterWhitelist) > 0 {
-
-			for _, method := range contractDefinition.CompiledContract().Abi.Methods {
-				//Todo can prob get rid of the isConstant check here
-				if contractDefinition.CallableMethods()[method.Sig] && !method.IsConstant() {
-					// Any whitelisted non-constant method should be tracked as a state changing method.
-					fw.appendStateChangingMethod(contractAddress, contractDefinition, method)
-				}
-			}
-		} else if len(fw.fuzzer.config.Fuzzing.Testing.FilterBlacklist) > 0 {
-			for _, method := range contractDefinition.CompiledContract().Abi.Methods {
-				if !contractDefinition.CallableMethods()[method.Sig] && !method.IsConstant() {
-					// Only add method to list of state changing methods if not present in blacklist
-					fw.appendStateChangingMethod(contractAddress, contractDefinition, method)
-				}
+		for _, method := range contractDefinition.CompiledContract().Abi.Methods {
+			if !contractDefinition.CallableMethods()[method.Sig] {
+				fw.appendStateChangingMethod(contractAddress, contractDefinition, method)
 			}
-		} else {
-			for _, method := range contractDefinition.CompiledContract().Abi.Methods {
-				if !method.IsConstant() {
-					// If white/blacklist functionality isn't present just add non-constant methods
-					fw.appendStateChangingMethod(contractAddress, contractDefinition, method)
-				}
-			}
-
 		}
 	}
 }
@@ -273,20 +251,6 @@ func (fw *FuzzerWorker) appendStateChangingMethod(contractAddress common.Address
 	)
 }
 
-// methodNotInBlacklist checks if a method is not present in the blacklist.
-// Todo refactor
-func (fw *FuzzerWorker) methodNotInBlacklist(contractName, methodSignature string) bool {
-	if len(fw.fuzzer.config.Fuzzing.Testing.FilterFunctions) > 0 {
-		for _, filterFunction := range fw.fuzzer.config.Fuzzing.Testing.FilterFunctions {
-			contractNameAndMethodSignature := strings.Split(filterFunction, ".")
-			if contractNameAndMethodSignature[0] == contractName && contractNameAndMethodSignature[1] == methodSignature {
-				return false
-			}
-		}
-	}
-	return true
-}
-
 // testNextCallSequence tests a call message sequence against the underlying FuzzerWorker's Chain and calls every
 // CallSequenceTestFunc registered with the parent Fuzzer to update any test results. If any call message in the
 // sequence is nil, a call message will be created in its place, targeting a state changing method of a contract