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

@@ -84,6 +84,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"`
@@ -478,3 +488,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,20 +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"},
+			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",

fuzzing/fuzzer.go

@@ -527,11 +527,14 @@
 	// while still being able to use the contract address overrides
 	contractsToDeploy := make([]string, 0)
 	balances := make([]*config.ContractBalance, 0)
+	initFunctions := make([]string, 0)
 
 	for contractName := range fuzzer.config.Fuzzing.PredeployedContracts {
 		contractsToDeploy = append(contractsToDeploy, contractName)
 		// Preserve index of target contract balances
 		balances = append(balances, &config.ContractBalance{Int: *big.NewInt(0)})
+		// Set default empty init function for predeployed contracts
+		initFunctions = append(initFunctions, "")
 	}
 
 	if len(fuzzer.deploymentOrder) > 0 {
@@ -567,6 +570,26 @@
 		balances = append(balances, fuzzer.config.Fuzzing.TargetContractsBalances...)
 	}
 
+	// Process target contracts init functions
+	targetContractsCount := len(fuzzer.config.Fuzzing.TargetContracts)
+	initConfigCount := len(fuzzer.config.Fuzzing.TargetContractsInitFunctions)
+
+	// Add initialization functions for target contracts
+	for i := 0; i < targetContractsCount; i++ {
+		initFunction := "" // Default: no initialization
+
+		if fuzzer.config.Fuzzing.UseInitFunctions {
+			if i < initConfigCount && fuzzer.config.Fuzzing.TargetContractsInitFunctions[i] != "" {
+				// Use explicit per-contract config
+				initFunction = fuzzer.config.Fuzzing.TargetContractsInitFunctions[i]
+			} else if len(fuzzer.config.Fuzzing.TargetContractsInitFunctions) == 1 {
+				// If only one init function specified (like "setUp") apply it to all contracts
+				initFunction = fuzzer.config.Fuzzing.TargetContractsInitFunctions[0]
+			}
+		}
+		initFunctions = append(initFunctions, initFunction)
+	}
+
 	deployedContractAddr := make(map[string]common.Address)
 	// Loop for all contracts to deploy
 	for i, contractName := range contractsToDeploy {
@@ -613,10 +636,116 @@
 				}
 				// Record our deployed contract so the next config-specified constructor args can reference this
 				// contract by name.
+
 				deployedContractAddr[contractName] = result.(common.Address)
+				contractAddr := deployedContractAddr[contractName]
+
+				// Get the initialization function name if exists and feature is enabled
+				if fuzzer.config.Fuzzing.UseInitFunctions && i < len(initFunctions) && initFunctions[i] != "" {
+					initFunction := initFunctions[i]
+					fuzzer.logger.Info(fmt.Sprintf("Checking if init function %s on %s exists", initFunction, contractName))
+
+					// Check if the initialization function exists
+					contractABI := contract.CompiledContract().Abi
+					if method, exists := contractABI.Methods[initFunction]; !exists {
+						fuzzer.logger.Info(fmt.Sprintf("Init function %s not found on %s, skipping", initFunction, contractName))
+					} else {
+						// Initialization function exists, proceed with calling it
+						fuzzer.logger.Info(fmt.Sprintf("Found init function %s with %d inputs", initFunction, len(method.Inputs)))
+
+						// Check if the init function accepts parameters and process them if needed
+						var args []any
+						if len(method.Inputs) > 0 {
+							// Verify InitializationArgs map exists
+							if fuzzer.config.Fuzzing.InitializationArgs == nil {
+								fuzzer.logger.Error(fmt.Errorf("initialization args map is nil but function requires args"))
+								continue
+							}
+
+							// Look for initialization arguments in the config
+							jsonArgs, ok := fuzzer.config.Fuzzing.InitializationArgs[contractName]
+							if !ok {
+								fuzzer.logger.Error(fmt.Errorf("initialization arguments for contract %s not provided", contractName))
+								continue
+							}
+
+							// Debug what args we found
+							fuzzer.logger.Info(fmt.Sprintf("Found args for %s: %+v", contractName, jsonArgs))
+
+							// Decode the arguments
+							decoded, err := valuegeneration.DecodeJSONArgumentsFromMap(method.Inputs,
+								jsonArgs, deployedContractAddr)
+							if err != nil {
+								fuzzer.logger.Error(fmt.Errorf("decoding failed for initialization arguments for contract %s: %v",
+									contractName, err))
+								continue
+							}
+
+							args = decoded
+							fuzzer.logger.Info(fmt.Sprintf("Decoded %d args for %s function %s",
+								len(args), contractName, initFunction))
+						}
+
+						// Log before packing
+						fuzzer.logger.Info(fmt.Sprintf("About to call initialization function %s on contract %s with %d args",
+							initFunction, contractName, len(args)))
+
+						// Pack the function call data with arguments
+						callData, err := contractABI.Pack(initFunction, args...)
+						if err != nil {
+							fuzzer.logger.Error(fmt.Errorf("failed to encode init call to %s: %v", initFunction, err))
+							continue
+						}
+
+						// Create and send the transaction
+						destAddr := contractAddr
+						msg := calls.NewCallMessage(fuzzer.deployer, &destAddr, 0, big.NewInt(0),
+							fuzzer.config.Fuzzing.BlockGasLimit, nil, nil, nil, callData)
+						msg.FillFromTestChainProperties(testChain)
+
+						// Debug log after creating the message
+						fuzzer.logger.Info(fmt.Sprintf("Created message for init function call to %s", initFunction))
+
+						// Create and commit a block with the transaction
+						block, err := testChain.PendingBlockCreate()
+						if err != nil {
+							fuzzer.logger.Error(fmt.Errorf("failed to create pending block for init call: %v", err))
+							continue
+						}
+
+						if err = testChain.PendingBlockAddTx(msg.ToCoreMessage()); err != nil {
+							fuzzer.logger.Error(fmt.Errorf("failed to add initialization transaction for function %s on contract %s to pending block: %v",
+								initFunction, contractName, err))
+							continue
+						}
+
+						if err = testChain.PendingBlockCommit(); err != nil {
+							fuzzer.logger.Error(fmt.Errorf("failed to commit block containing initialization call to function %s on contract %s: %v",
+								initFunction, contractName, err))
+							continue
+						}
+
+						// Check if the call succeeded
+						if block.MessageResults[0].Receipt.Status != types.ReceiptStatusSuccessful {
+							// Create a call sequence element for the trace
+							cse := calls.NewCallSequenceElement(nil, msg, 0, 0)
+							cse.ChainReference = &calls.CallSequenceElementChainReference{
+								Block:            block,
+								TransactionIndex: len(block.Messages) - 1,
+							}
+
+							fuzzer.logger.Error(fmt.Errorf("init function %s call failed on %s: %v",
+								initFunction, contractName,
+								block.MessageResults[0].ExecutionResult.Err))
+						} else {
+							fuzzer.logger.Info(fmt.Sprintf("Successfully called %s on %s with %d args",
+								initFunction, contractName, len(args)))
+						}
+					}
+				}
 
-				// Flag that we found a matching compiled contract definition and deployed it, then exit out of this
-				// inner loop to process the next contract to deploy in the outer loop.
+				// Flag that we found a matching compiled contract definition, deployed it and called available init functions if any,
+				// then exit out of this inner loop to process the next contract to deploy in the outer loop.
 				found = true
 				break
 			}
@@ -627,6 +756,7 @@
 			return nil, fmt.Errorf("%v was specified in the target contracts but was not found in the compilation artifacts", contractName)
 		}
 	}
+
 	return nil, nil
 }
 

fuzzing/fuzzer_hooks.go

@@ -1,9 +1,10 @@
 package fuzzing
 
 import (
-	"github.com/crytic/medusa/fuzzing/config"
 	"math/rand"
 
+	"github.com/crytic/medusa/fuzzing/config"
+
 	"github.com/crytic/medusa/fuzzing/executiontracer"
 
 	"github.com/crytic/medusa/chain"