diff --git a/chain/test_chain.go b/chain/test_chain.go index 05321c4a..098b8e99 100644 --- a/chain/test_chain.go +++ b/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" diff --git a/cmd/fuzz_flags.go b/cmd/fuzz_flags.go index b08ca7c4..7a9e160e 100644 --- a/cmd/fuzz_flags.go +++ b/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 diff --git a/docs/src/project_configuration/fuzzing_config.md b/docs/src/project_configuration/fuzzing_config.md index 452f7d79..f85df4ee 100644 --- a/docs/src/project_configuration/fuzzing_config.md +++ b/docs/src/project_configuration/fuzzing_config.md @@ -111,6 +111,12 @@ 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}}` @@ -118,6 +124,20 @@ The fuzzing configuration defines the parameters for the fuzzing campaign. 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 diff --git a/docs/src/static/function_level_testing_medusa.json b/docs/src/static/function_level_testing_medusa.json index 5e284eef..6e400baf 100644 --- a/docs/src/static/function_level_testing_medusa.json +++ b/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, diff --git a/docs/src/static/medusa.json b/docs/src/static/medusa.json index 4326aa56..5d4ce40f 100644 --- a/docs/src/static/medusa.json +++ b/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, diff --git a/fuzzing/config/config.go b/fuzzing/config/config.go index c1d42bf6..12c8f1ba 100644 --- a/fuzzing/config/config.go +++ b/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{} +} diff --git a/fuzzing/config/config_defaults.go b/fuzzing/config/config_defaults.go index 2c11e38e..559cff3d 100644 --- a/fuzzing/config/config_defaults.go +++ b/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", diff --git a/fuzzing/fuzzer.go b/fuzzing/fuzzer.go index d04d9ef0..81dc881e 100644 --- a/fuzzing/fuzzer.go +++ b/fuzzing/fuzzer.go @@ -527,11 +527,14 @@ func chainSetupFromCompilations(fuzzer *Fuzzer, testChain *chain.TestChain) (*ex // 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 @@ func chainSetupFromCompilations(fuzzer *Fuzzer, testChain *chain.TestChain) (*ex 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 @@ func chainSetupFromCompilations(fuzzer *Fuzzer, testChain *chain.TestChain) (*ex } // 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 @@ func chainSetupFromCompilations(fuzzer *Fuzzer, testChain *chain.TestChain) (*ex return nil, fmt.Errorf("%v was specified in the target contracts but was not found in the compilation artifacts", contractName) } } + return nil, nil } diff --git a/fuzzing/fuzzer_hooks.go b/fuzzing/fuzzer_hooks.go index f6be4324..ede734bf 100644 --- a/fuzzing/fuzzer_hooks.go +++ b/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" diff --git a/fuzzing/fuzzer_test.go b/fuzzing/fuzzer_test.go index 017090de..ec73cf2a 100644 --- a/fuzzing/fuzzer_test.go +++ b/fuzzing/fuzzer_test.go @@ -349,6 +349,7 @@ func TestConsoleLog(t *testing.T) { filePath: filePath, configUpdates: func(config *config.ProjectConfig) { config.Fuzzing.TargetContracts = []string{"TestContract"} + config.Fuzzing.TargetContractsInitFunctions = []string{"testConsoleLog"} config.Fuzzing.TestLimit = 10000 config.Fuzzing.Testing.PropertyTesting.Enabled = false config.Fuzzing.Testing.OptimizationTesting.Enabled = false @@ -514,6 +515,52 @@ func TestDeploymentsWithPayableConstructors(t *testing.T) { }) } +// TestInitializationFunctions runs a test to ensure initialization functions work both with and without arguments +func TestInitializationWithParam(t *testing.T) { + runFuzzerTest(t, &fuzzerSolcFileTest{ + filePath: "testdata/contracts/deployments/deploy_with_init_fns.sol", + configUpdates: func(pkgConfig *config.ProjectConfig) { + pkgConfig.EnableInitFunctions([]string{"initWithParam"}) + // Just a single contract + pkgConfig.Fuzzing.TargetContracts = []string{"SimpleInitParamTest"} + + // With zero balance + pkgConfig.Fuzzing.TargetContractsBalances = []*config.ContractBalance{ + {Int: *big.NewInt(0)}, + } + + // Initialization function with a parameter + pkgConfig.Fuzzing.TargetContractsInitFunctions = []string{"initWithParam"} + + // Create the initialization args map if it doesn't exist + if pkgConfig.Fuzzing.InitializationArgs == nil { + pkgConfig.Fuzzing.InitializationArgs = make(map[string]map[string]any) + } + + // Specify the parameter value - must match the exact parameter name + pkgConfig.Fuzzing.InitializationArgs["SimpleInitParamTest"] = map[string]any{ + "_value": "42", + } + + // Enable property testing + pkgConfig.Fuzzing.Testing.PropertyTesting.Enabled = false + pkgConfig.Fuzzing.TestLimit = 10 + pkgConfig.Fuzzing.Testing.AssertionTesting.Enabled = true + pkgConfig.Fuzzing.Testing.OptimizationTesting.Enabled = false + pkgConfig.Slither.UseSlither = false + + }, + method: func(f *fuzzerTestContext) { + // Start the fuzzer + err := f.fuzzer.Start() + assert.NoError(t, err) + + assertFailedTestsExpected(f, false) + + }, + }) +} + // TestDeploymentsSelfDestruct runs a test to ensure dynamically deployed contracts are detected by the Fuzzer and // their properties are tested appropriately. func TestDeploymentsSelfDestruct(t *testing.T) { @@ -697,7 +744,7 @@ func TestTestingScope(t *testing.T) { // TestDeploymentsWithArgs runs tests to ensure contracts deployed with config provided constructor arguments are // deployed as expected. It expects all properties should fail (indicating values provided were set accordingly). func TestDeploymentsWithArgs(t *testing.T) { - // This contract deploys a contract with specific constructor arguments. Property tests will fail if they are + // This contract deploys a contract with specific constructor arguments as well as init functions with arguments. Property tests will fail if they are // set correctly. runFuzzerTest(t, &fuzzerSolcFileTest{ filePath: "testdata/contracts/deployments/deployment_with_args.sol", @@ -716,6 +763,16 @@ func TestDeploymentsWithArgs(t *testing.T) { "_deployed": "DeployedContract:DeploymentWithArgs", }, } + config.Fuzzing.TargetContractsInitFunctions = []string{"dummyFunction", "dummyFunction"} // this should execute predefined functions in the respective contracts + config.Fuzzing.InitializationArgs = map[string]map[string]any{ + "DeploymentWithArgs": { + "a": "100", // argument for DeploymentWithArgs.dummyFunction + }, + "Dependent": { + "a": "200", // argument for Dependent.dummyFunction + }, + } + config.Fuzzing.Testing.StopOnFailedTest = false config.Fuzzing.TestLimit = 500 // this test should expose a failure quickly. config.Fuzzing.Testing.AssertionTesting.Enabled = false diff --git a/fuzzing/testdata/contracts/deployments/deploy_with_init_fns.sol b/fuzzing/testdata/contracts/deployments/deploy_with_init_fns.sol new file mode 100644 index 00000000..6556de80 --- /dev/null +++ b/fuzzing/testdata/contracts/deployments/deploy_with_init_fns.sol @@ -0,0 +1,19 @@ +// Ultra-simple test for initialization functions with parameters +contract SimpleInitParamTest { + // Track if functions were called and parameter values + bool public initCalled; + uint public initValue; + + // Empty constructor + constructor() {} + + // Initialization function with a parameter + function initWithParam(uint _value) public { + initCalled = true; + initValue = _value; + emit InitCalled("initWithParam", _value); + } + + // Event for tracking + event InitCalled(string functionName, uint value); +} \ No newline at end of file