New cheatcodes (draft)

.gitignore

@@ -29,3 +29,5 @@ docs/book
 
 # Build results
 result
+
+.vscode/*

chain/cheat_code_tracer.go

@@ -1,6 +1,7 @@
 package chain
 
 import (
+	"fmt"
 	"math/big"
 
 	"github.com/crytic/medusa-geth/common"
@@ -9,6 +10,7 @@ import (
 	"github.com/crytic/medusa-geth/core/vm"
 	"github.com/crytic/medusa-geth/eth/tracers"
 	"github.com/crytic/medusa/chain/types"
+	"github.com/holiman/uint256"
 )
 
 // cheatCodeTracer represents an EVM.Logger which tracks and patches EVM execution state to enable extended
@@ -232,6 +234,43 @@ func (t *cheatCodeTracer) OnOpcode(pc uint64, op byte, gas, cost uint64, scope t
 	if t.callDepth > 0 {
 		t.callFrames[t.callDepth-1].onNextFrameEnterHooks.Execute(true, true)
 	}
+
+	// Support for expectRevert cheatcode (see standard_cheat_code_contrat.go)
+	// TODO : support dynamic value for expectRevert
+	if _, ok := currentCallFrame.extraData["expectRevert"]; ok {
+		// expectRevert does not affect the calls to the cheatcode VM (ex: prank)
+		// So if the next call is another call to the VM
+		// We forward the extraData to the next Frame
+		if scope.Address() == StandardCheatcodeContractAddress {
+			// TODO refactor the following to be an internal function used in both here and standard_cheat_code_contrat
+			delete(currentCallFrame.extraData, "expectRevert")
+
+			cheatCodeCallerFrame := t.PreviousCallFrame()
+			cheatCodeCallerFrame.onNextFrameEnterHooks.Push(func() {
+				revertFrame := t.PreviousCallFrame()
+				// TODO : support dynamic value for expectRevert instead of a bool
+				revertFrame.extraData["expectRevert"] = true
+			})
+		} else {
+			delete(currentCallFrame.extraData, "expectRevert")
+
+			stack := scope.StackData()
+			index := len(stack) - 1
+			return_value := stack[index]
+			if return_value.Eq(uint256.NewInt(0)) {
+				stack[index] = *uint256.NewInt(1)
+			} else {
+
+				if !return_value.Eq(uint256.NewInt(1)) {
+					// TODO: find a better error handling
+					panic(fmt.Sprintf("expected revert but got return value %v", return_value))
+				}
+
+				stack[index] = *uint256.NewInt(0)
+			}
+
+		}
+	}
 }
 
 // CaptureTxEndSetAdditionalResults can be used to set additional results captured from execution tracing. If this

chain/standard_cheat_code_contract.go

@@ -10,13 +10,12 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/crytic/medusa/chain/types"
-
 	"github.com/crytic/medusa-geth/accounts/abi"
 	"github.com/crytic/medusa-geth/common"
 	"github.com/crytic/medusa-geth/core/tracing"
 	"github.com/crytic/medusa-geth/core/vm"
 	"github.com/crytic/medusa-geth/crypto"
+	"github.com/crytic/medusa/chain/types"
 	"github.com/crytic/medusa/utils"
 	"github.com/holiman/uint256"
 )
@@ -59,6 +58,10 @@ func getStandardCheatCodeContract(tracer *cheatCodeTracer) (*CheatCodeContract,
 	if err != nil {
 		return nil, err
 	}
+	typeBytes4, err := abi.NewType("bytes4", "", nil)
+	if err != nil {
+		return nil, err
+	}
 	typeBytes32, err := abi.NewType("bytes32", "", nil)
 	if err != nil {
 		return nil, err
@@ -810,10 +813,234 @@ func getStandardCheatCodeContract(tracer *cheatCodeTracer) (*CheatCodeContract,
 		},
 	)
 
+	// getCode: Retrieves the runtime bytecode for a contract
+	contract.addMethod("getDeployedCode", abi.Arguments{{Type: typeString}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+
+			contractPath := inputs[0].(string)
+
+			_, contractName, err := parseContractPath(contractPath)
+			if err != nil {
+				fmt.Println("getDeployedCode error: invalid path format: %v", err)
+				return nil, cheatCodeRevertData([]byte(fmt.Sprintf("getCode error: invalid path format: %v", err)))
+			}
+
+			compiledContract, exists := tracer.chain.CompiledContracts[contractName]
+			if !exists {
+				// TODO: this is probably not the best to print given it will be printed in loop
+				// But it should be shown once to the user to avoid mistakes
+				// Same is true for getCode
+				fmt.Println("getDeployedCode error: contract not found: %s (did you forget to deploy the contract with predeployedContracts?)", contractName)
+				return nil, cheatCodeRevertData([]byte(fmt.Sprintf("getCode error: contract not found: %s", contractName)))
+			}
+
+			bytecode := compiledContract.RuntimeBytecode
+			if len(bytecode) == 0 {
+				fmt.Println("getDeployedCode error: contract bytecode is empty: %s", contractName)
+				return nil, cheatCodeRevertData([]byte(fmt.Sprintf("getCode error: contract bytecode is empty: %s", contractName)))
+			}
+
+			fmt.Println("getDeployedCode found")
+			// Return the bytecode
+			return []any{bytecode}, nil
+		},
+	)
+
+	// assertTrue
+	contract.addMethod("assertTrue", abi.Arguments{{Type: typeBool}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			bool_result := inputs[0].(bool)
+
+			if !bool_result {
+				return nil, cheatCodeRevertData([]byte("assertFalse failed"))
+			}
+
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// assertTrue with reason
+	contract.addMethod("assertTrue", abi.Arguments{{Type: typeBool}, {Type: typeString}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			bool_result := inputs[0].(bool)
+
+			if !bool_result {
+				return nil, cheatCodeRevertData([]byte(inputs[1].(string)))
+			}
+
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// assertFalse
+	contract.addMethod("assertFalse", abi.Arguments{{Type: typeBool}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			bool_result := inputs[0].(bool)
+
+			if bool_result {
+				return nil, cheatCodeRevertData([]byte("assertFalse failed"))
+			}
+
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// assertFalse with reason
+	contract.addMethod("assertFalse", abi.Arguments{{Type: typeBool}, {Type: typeString}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			bool_result := inputs[0].(bool)
+
+			if bool_result {
+				return nil, cheatCodeRevertData([]byte(inputs[1].(string)))
+			}
+
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// assume: Revert if the condition if false
+	contract.addMethod("assume", abi.Arguments{{Type: typeBool}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			bool_result := inputs[0].(bool)
+
+			if !bool_result {
+				return nil, cheatCodeRevertData([]byte("assume failed"))
+			}
+
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// expectEmit: NOOP for now (TODO)
+	contract.addMethod("expectEmit", abi.Arguments{}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+			// Return nothing
+			return nil, nil
+		},
+	)
+
+	// expectRevert: Follow the expect revert logic
+	// TODO: merge the different expectRevert to reduce the code dupplicate and handle properly the reason check
+	contract.addMethod("expectRevert", abi.Arguments{}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+
+			// To implement expectRevert we follow this logic:
+			// We add a hook on the next call that happen after the call to the cheatcode's VM
+			// Which is the "next" frame of the "previous frame"
+			// The previous frame being the caller of the cheatcode, and its next frame the next call
+			// The hook just set expectRevert to true
+			// The actual update is done in the OnOpcode hook (cheat_code_tracer.go)
+
+			cheatCodeCallerFrame := tracer.PreviousCallFrame()
+			cheatCodeCallerFrame.onNextFrameEnterHooks.Push(func() {
+				revertFrame := tracer.PreviousCallFrame()
+				// TODO : support dynamic value for expectRevert instead of a bool
+				revertFrame.extraData["expectRevert"] = true
+			})
+			return nil, nil
+		},
+	)
+
+	// expectRevert: Follow the expect revert logic
+	contract.addMethod("expectRevert", abi.Arguments{{Type: typeBytes4}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+
+			// To implement expectRevert we follow this logic:
+			// We add a hook on the next call that happen after the call to the cheatcode's VM
+			// Which is the "next" frame of the "previous frame"
+			// The previous frame being the caller of the cheatcode, and its next frame the next call
+			// The hook just set expectRevert to true
+			// The actual update is done in the OnOpcode hook (cheat_code_tracer.go)
+
+			cheatCodeCallerFrame := tracer.PreviousCallFrame()
+			cheatCodeCallerFrame.onNextFrameEnterHooks.Push(func() {
+				revertFrame := tracer.PreviousCallFrame()
+				// TODO : support dynamic value for expectRevert instead of a bool
+				revertFrame.extraData["expectRevert"] = true
+			})
+			return nil, nil
+		},
+	)
+
+	// expectRevert: Follow the expect revert logic
+	contract.addMethod("expectRevert", abi.Arguments{{Type: typeBytes}}, abi.Arguments{},
+		func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+
+			// To implement expectRevert we follow this logic:
+			// We add a hook on the next call that happen after the call to the cheatcode's VM
+			// Which is the "next" frame of the "previous frame"
+			// The previous frame being the caller of the cheatcode, and its next frame the next call
+			// The hook just set expectRevert to true
+			// The actual update is done in the OnOpcode hook (cheat_code_tracer.go)
+
+			cheatCodeCallerFrame := tracer.PreviousCallFrame()
+			cheatCodeCallerFrame.onNextFrameEnterHooks.Push(func() {
+				revertFrame := tracer.PreviousCallFrame()
+				// TODO : support dynamic value for expectRevert instead of a bool
+				revertFrame.extraData["expectRevert"] = true
+			})
+			return nil, nil
+		},
+	)
+
+	// assertEq: Register assertEq for all supported types
+	assertEqTypes := []abi.Type{typeAddress, typeBytes, typeBytes4, typeBytes32, typeUint8, typeUint64, typeUint256, typeInt256, typeString, typeBool}
+	for _, t := range assertEqTypes {
+		currentType := t // capture range variable
+		contract.addMethod("assertEq", abi.Arguments{{Type: currentType}, {Type: currentType}}, abi.Arguments{},
+			func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+				if !assertEqGenerator(inputs, currentType) {
+					return nil, cheatCodeRevertData([]byte("assertEq failed"))
+				}
+				return nil, nil
+			},
+		)
+
+		contract.addMethod("assertEq", abi.Arguments{{Type: currentType}, {Type: currentType}, {Type: typeString}}, abi.Arguments{},
+			func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) {
+				if !assertEqGenerator(inputs, currentType) {
+					reason := inputs[2].(string)
+					return nil, cheatCodeRevertData([]byte(reason))
+				}
+				return nil, nil
+			},
+		)
+	}
+
 	// Return our precompile contract information.
 	return contract, nil
 }
 
+func assertEqGenerator(inputs []any, t abi.Type) bool {
+	l := inputs[0]
+	r := inputs[1]
+
+	// Use type-specific comparisons based on the ABI type
+	switch t.T {
+	case abi.AddressTy:
+		return l.(common.Address) == r.(common.Address)
+	case abi.BoolTy:
+		return l.(bool) == r.(bool)
+	case abi.IntTy, abi.UintTy:
+		return l.(*big.Int).Cmp(r.(*big.Int)) == 0
+	case abi.StringTy:
+		return l.(string) == r.(string)
+	case abi.BytesTy:
+		return string(l.([]byte)) == string(r.([]byte))
+	case abi.FixedBytesTy:
+		lBytes := l.([32]byte)
+		rBytes := r.([32]byte)
+		return string(lBytes[:]) == string(rBytes[:])
+	default:
+		return l == r
+	}
+}
+
 // parseContractPath parses a contract path in the following formats:
 // - "MyContract.sol:MyContract"
 // - "MyContract"

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"
 

docs/src/project_configuration/fuzzing_config.md

@@ -103,13 +103,34 @@ 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

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

docs/src/static/medusa.json

@@ -13,7 +13,9 @@
     "targetContracts": [],
     "predeployedContracts": {},
     "targetContractsBalances": [],
+    "TargetContractsInitFunctions": [],
     "constructorArgs": {},
+    "initializationArgs": {},
     "deployerAddress": "0x30000",
     "senderAddresses": ["0x10000", "0x20000", "0x30000"],
     "blockNumberDelayMax": 60480,