Allow Medusa to run init function (ex: setUp) once per fuzz run

chain/test_chain.go

@@ -3,9 +3,10 @@ package chain
 import (
 	"errors"
 	"fmt"
-	compilationTypes "github.com/crytic/medusa/compilation/types"
 	"math/big"
 
+	compilationTypes "github.com/crytic/medusa/compilation/types"
+
 	"github.com/crytic/medusa/chain/state"
 	"golang.org/x/net/context"
 

cmd/fuzz_flags.go

@@ -46,6 +46,12 @@ func addFuzzFlags() error {
 	fuzzCmd.Flags().StringSlice("target-contracts", []string{},
 		fmt.Sprintf("target contracts for fuzz testing (unless a config file is provided, default is %v)", defaultConfig.Fuzzing.TargetContracts))
 
+	// Will call post-deployment initialization function defined by `targetContractsInitFunctions` to be called on all contracts that have the implementation
+	fuzzCmd.Flags().Bool("use-init-fns", false, "runs init functions (`setUp`, `initialize`) on all contracts that have the implementation")
+
+	// Will call setUp() function if implemented
+	fuzzCmd.Flags().Bool("enable-foundry-setup", false, "runs `setUp` function on all contracts that have it implemented")
+
 	// Corpus directory
 	fuzzCmd.Flags().String("corpus-dir", "",
 		fmt.Sprintf("directory path for corpus items and coverage reports (unless a config file is provided, default is %q)", defaultConfig.Fuzzing.CorpusDirectory))
@@ -85,6 +91,7 @@ func addFuzzFlags() error {
 	// Log level
 	fuzzCmd.Flags().String("log-level", "", "set which level of log messages will be displayed (trace, debug, info, warn, error, or panic; default: info)")
 	return nil
+
 }
 
 // updateProjectConfigWithFuzzFlags will update the given projectConfig with any CLI arguments that were provided to the fuzz command
@@ -260,6 +267,30 @@ func updateProjectConfigWithFuzzFlags(cmd *cobra.Command, projectConfig *config.
 		}
 	}
 
+	// Update configuration to run init functions
+	if cmd.Flags().Changed("use-init-fns") {
+		useInitFns, err := cmd.Flags().GetBool("use-init-fns")
+		if err != nil {
+			return err
+		}
+		if useInitFns {
+			// Enable the init functions feature but the actual functions need to be specified in config
+			projectConfig.Fuzzing.UseInitFunctions = true
+		}
+	}
+
+	// Update configuration to run `setUp` function where implemented
+	if cmd.Flags().Changed("enable-foundry-setup") {
+		enableFoundrySetUp, err := cmd.Flags().GetBool("enable-foundry-setup")
+		if err != nil {
+			return err
+		}
+		if enableFoundrySetUp {
+			projectConfig.Fuzzing.UseInitFunctions = true
+			projectConfig.Fuzzing.TargetContractsInitFunctions = []string{"setUp"}
+		}
+	}
+
 	// Update log level
 	if cmd.Flags().Changed("log-level") {
 		levelStr, err := cmd.Flags().GetString("log-level")
@@ -273,6 +304,7 @@ func updateProjectConfigWithFuzzFlags(cmd *cobra.Command, projectConfig *config.
 		}
 
 		projectConfig.Logging.Level = level
+
 	}
 
 	return nil

docs/src/project_configuration/fuzzing_config.md

@@ -111,13 +111,33 @@ The fuzzing configuration defines the parameters for the fuzzing campaign.
   then `A` will have a starting balance of `1,234 wei`, `B` will have `4,660 wei (0x1234 in decimal)`, and `C` will have `1.2 ETH (1.2 × 10^18 wei)`.
 - **Default**: `[]`
 
+### `targetContractsInitFunctions`
+
+- **Type**: [String] (e.g. `["setUp", "initialize", ""]`)
+- **Description**: Specifies post-deployment initialization functions to call for each contract in `targetContracts`. This array has a one-to-one mapping with `targetContracts`, where each element corresponds to the initialization function for the contract at the same index. Empty strings indicate no initialization for that contract.
+- **Default**: `[]`
+
 ### `constructorArgs`
 
 - **Type**: `{"contractName": {"variableName": _value}}`
 - **Description**: If a contract in the `targetContracts` has a `constructor` that takes in variables, these can be specified here.
   An example can be found [here](#using-constructorargs).
 - **Default**: `{}`
 
+### `initializationArgs`
+
+- **Type**: `{"contractName": {"parameterName": _value}}`
+- **Description**: Specifies arguments to pass to initialization functions defined in `targetContractsInitFunctions`. The keys in this map must match the contract names exactly, and the parameter names must match the parameter names in the function signature.
+  For example, if contract `MyContract` has an initialization function `initialize(uint256 _value, address _owner)`, then you would configure:
+  ```json
+  {
+    "MyContract": {
+      "_value": "100",
+      "_owner": "0x1234..."
+    }
+  }
+  ```
+
 ### `deployerAddress`
 
 - **Type**: Address

docs/src/static/function_level_testing_medusa.json

@@ -10,7 +10,9 @@
     "coverageEnabled": true,
     "targetContracts": ["TestDepositContract"],
     "targetContractsBalances": ["21267647932558653966460912964485513215"],
+    "TargetContractsInitFunctions": [],
     "constructorArgs": {},
+    "initializationArgs": {},
     "deployerAddress": "0x30000",
     "senderAddresses": ["0x10000", "0x20000", "0x30000"],
     "blockNumberDelayMax": 60480,

docs/src/static/medusa.json

@@ -14,7 +14,10 @@
     "targetContracts": [],
     "predeployedContracts": {},
     "targetContractsBalances": [],
+    "targetContractsInitFunctions": [],
+    "useInitFunctions": false,
     "constructorArgs": {},
+    "initializationArgs": {},
     "deployerAddress": "0x30000",
     "senderAddresses": ["0x10000", "0x20000", "0x30000"],
     "blockNumberDelayMax": 60480,

fuzzing/config/config.go

@@ -88,6 +88,16 @@ type FuzzingConfig struct {
 	// TargetContracts
 	TargetContractsBalances []*ContractBalance `json:"targetContractsBalances"`
 
+	// Holds the logic whether to run initialization functions supplied by `enable-foundry-setup` or `use-init-fns`
+	UseInitFunctions bool `json:"useInitFunctions"`
+
+	// TargetContractsInitFunctions is the list of functions to users to specify an "init function" (with setUp() as the default)
+	TargetContractsInitFunctions []string `json:"targetContractsInitFunctions"`
+
+	// InitializationArgs holds the arguments for TargetContractsInitFunctions deployments. It is available via the project
+	// configuration
+	InitializationArgs map[string]map[string]any `json:"initializationArgs"`
+
 	// ConstructorArgs holds the constructor arguments for TargetContracts deployments. It is available via the project
 	// configuration
 	ConstructorArgs map[string]map[string]any `json:"constructorArgs"`
@@ -482,3 +492,21 @@ func (p *ProjectConfig) Validate() error {
 
 	return nil
 }
+
+// Helper function to enable init functions with specific functions
+func (p *ProjectConfig) EnableInitFunctions(initFunctions []string) {
+	p.Fuzzing.UseInitFunctions = true
+	p.Fuzzing.TargetContractsInitFunctions = initFunctions
+}
+
+// Helper function to enable Foundry setup
+func (p *ProjectConfig) EnableFoundrySetup() {
+	p.Fuzzing.UseInitFunctions = true
+	p.Fuzzing.TargetContractsInitFunctions = []string{"setUp"}
+}
+
+// Helper function to disable init functions
+func (p *ProjectConfig) DisableInitFunctions() {
+	p.Fuzzing.UseInitFunctions = false
+	p.Fuzzing.TargetContractsInitFunctions = []string{}
+}

fuzzing/config/config_defaults.go

@@ -39,21 +39,23 @@ func GetDefaultProjectConfig(platform string) (*ProjectConfig, error) {
 	// Create a project configuration
 	projectConfig := &ProjectConfig{
 		Fuzzing: FuzzingConfig{
-			Workers:                 10,
-			WorkerResetLimit:        50,
-			Timeout:                 0,
-			TestLimit:               0,
-			ShrinkLimit:             5_000,
-			CallSequenceLength:      100,
-			PruneFrequency:          5,
-			TargetContracts:         []string{},
-			TargetContractsBalances: []*ContractBalance{},
-			PredeployedContracts:    map[string]string{},
-			ConstructorArgs:         map[string]map[string]any{},
-			CorpusDirectory:         "",
-			CoverageEnabled:         true,
-			CoverageFormats:         []string{"html", "lcov"},
-			CoverageExclusions:      []string{},
+			Workers:                      10,
+			WorkerResetLimit:             50,
+			Timeout:                      0,
+			TestLimit:                    0,
+			ShrinkLimit:                  5_000,
+			CallSequenceLength:           100,
+			PruneFrequency:               5,
+			TargetContracts:              []string{},
+			TargetContractsBalances:      []*ContractBalance{},
+			TargetContractsInitFunctions: []string{},
+			UseInitFunctions:             false,
+			PredeployedContracts:         map[string]string{},
+			ConstructorArgs:              map[string]map[string]any{},
+			InitializationArgs:           map[string]map[string]any{},
+			CorpusDirectory:              "",
+			CoverageEnabled:              true,
+			CoverageFormats:              []string{"html", "lcov"},
 			SenderAddresses: []string{
 				"0x10000",
 				"0x20000",