pallet-revive: add interface to implement mocks and pranks (#9909)

Needed for: paritytech/foundry-polkadot#334.

In foundry-polkadot we need the ability to be able to manipulate the
`msg.sender` and the `tx.origin` that a solidity contract sees cheatcode
documentation, plus the ability to mock calls and functions.

Currently all create/call methods use the `bare_instantiate`/`bare_call`
to run things in pallet-revive, the caller then normally gets set
automatically, based on what is the call stack, but for `forge test` we
need to be able to manipulate, so that we can set it to custom values.

Additionally, for delegate_call, bare_call is used, so there is no way
to specify we are dealing with a delegate call, so the call is not
working correcly.
 
For both this paths, we need a way to inject this information into the
execution environment, hence I added an optional hooks interface that we
implement from foundry cheatcodes for prank and mock functionality.

## TODO
- [x] Add tests to make sure the hooks functionality does not regress.

---------

Signed-off-by: Alexandru Gheorghe <[email protected]>
Co-authored-by: cmd[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: alexggh

prdoc/pr_9909.prdoc

@@ -0,0 +1,16 @@
+title: 'pallet-revive: add interface to implement mocks and pranks'
+doc:
+- audience: Node Dev
+  description: |-
+    Needed for: https://github.com/paritytech/foundry-polkadot/pull/334.
+
+    In foundry-polkadot we need the ability to be able to manipulate the `msg.sender` and the `tx.origin` that a solidity contract sees cheatcode documentation, plus the ability to mock calls and functions.
+
+     Currently all create/call methods use the `bare_instantiate`/`bare_call` to run things in pallet-revive,  the caller then normally gets set automatically, based on what is the call stack, but for `forge test` we need to be able to manipulate, so that we can set it to custom values.
+
+    Additionally, for delegate_call, bare_call is used, so there is no way to specify we are dealing with a delegate call, so the call is not working correcly.
+
+    For both this paths, we need a way to inject this information into the execution environment, hence I added an optional hooks interface that we implement from foundry cheatcodes for prank and mock functionality.
+crates:
+- name: pallet-revive
+  bump: major

substrate/frame/revive/src/call_builder.rs

@@ -56,7 +56,7 @@ pub struct CallSetup<T: Config> {
 	value: BalanceOf<T>,
 	data: Vec<u8>,
 	transient_storage_size: u32,
-	exec_config: ExecConfig,
+	exec_config: ExecConfig<T>,
 }
 
 impl<T> Default for CallSetup<T>

substrate/frame/revive/src/exec.rs

@@ -562,7 +562,7 @@ pub struct Stack<'a, T: Config, E> {
 	/// Transient storage used to store data, which is kept for the duration of a transaction.
 	transient_storage: TransientStorage<T>,
 	/// Global behavior determined by the creater of this stack.
-	exec_config: &'a ExecConfig,
+	exec_config: &'a ExecConfig<T>,
 	/// The set of contracts that were created during this call stack.
 	contracts_created: BTreeSet<T::AccountId>,
 	/// The set of contracts that are registered for destruction at the end of this call stack.
@@ -602,7 +602,8 @@ struct Frame<T: Config> {
 
 /// This structure is used to represent the arguments in a delegate call frame in order to
 /// distinguish who delegated the call and where it was delegated to.
-struct DelegateInfo<T: Config> {
+#[derive(Clone)]
+pub struct DelegateInfo<T: Config> {
 	/// The caller of the contract.
 	pub caller: Origin<T>,
 	/// The address of the contract the call was delegated to.
@@ -796,7 +797,7 @@ where
 		storage_meter: &mut storage::meter::Meter<T>,
 		value: U256,
 		input_data: Vec<u8>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> ExecResult {
 		let dest = T::AddressMapper::to_account_id(&dest);
 		if let Some((mut stack, executable)) = Stack::<'_, T, E>::new(
@@ -806,6 +807,7 @@ where
 			storage_meter,
 			value,
 			exec_config,
+			&input_data,
 		)? {
 			stack.run(executable, input_data).map(|_| stack.first_frame.last_frame_output)
 		} else {
@@ -821,14 +823,21 @@ where
 				);
 			});
 
-			let result = Self::transfer_from_origin(
-				&origin,
-				&origin,
-				&dest,
-				value,
-				storage_meter,
-				exec_config,
-			);
+			let result = if let Some(mock_answer) =
+				exec_config.mock_handler.as_ref().and_then(|handler| {
+					handler.mock_call(T::AddressMapper::to_address(&dest), &input_data, value)
+				}) {
+				Ok(mock_answer)
+			} else {
+				Self::transfer_from_origin(
+					&origin,
+					&origin,
+					&dest,
+					value,
+					storage_meter,
+					exec_config,
+				)
+			};
 
 			if_tracing(|t| match result {
 				Ok(ref output) => t.exit_child_span(&output, Weight::zero()),
@@ -854,7 +863,7 @@ where
 		value: U256,
 		input_data: Vec<u8>,
 		salt: Option<&[u8; 32]>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> Result<(H160, ExecReturnValue), ExecError> {
 		let deployer = T::AddressMapper::to_address(&origin);
 		let (mut stack, executable) = Stack::<'_, T, E>::new(
@@ -869,6 +878,7 @@ where
 			storage_meter,
 			value,
 			exec_config,
+			&input_data,
 		)?
 		.expect(FRAME_ALWAYS_EXISTS_ON_INSTANTIATE);
 		let address = T::AddressMapper::to_address(&stack.top_frame().account_id);
@@ -891,7 +901,7 @@ where
 		gas_meter: &'a mut GasMeter<T>,
 		storage_meter: &'a mut storage::meter::Meter<T>,
 		value: BalanceOf<T>,
-		exec_config: &'a ExecConfig,
+		exec_config: &'a ExecConfig<T>,
 	) -> (Self, E) {
 		let call = Self::new(
 			FrameArgs::Call {
@@ -904,6 +914,7 @@ where
 			storage_meter,
 			value.into(),
 			exec_config,
+			&Default::default(),
 		)
 		.unwrap()
 		.unwrap();
@@ -920,7 +931,8 @@ where
 		gas_meter: &'a mut GasMeter<T>,
 		storage_meter: &'a mut storage::meter::Meter<T>,
 		value: U256,
-		exec_config: &'a ExecConfig,
+		exec_config: &'a ExecConfig<T>,
+		input_data: &Vec<u8>,
 	) -> Result<Option<(Self, ExecutableOrPrecompile<T, E, Self>)>, ExecError> {
 		origin.ensure_mapped()?;
 		let Some((first_frame, executable)) = Self::new_frame(
@@ -932,6 +944,8 @@ where
 			BalanceOf::<T>::max_value(),
 			false,
 			true,
+			input_data,
+			exec_config,
 		)?
 		else {
 			return Ok(None);
@@ -968,6 +982,8 @@ where
 		deposit_limit: BalanceOf<T>,
 		read_only: bool,
 		origin_is_caller: bool,
+		input_data: &[u8],
+		exec_config: &ExecConfig<T>,
 	) -> Result<Option<(Frame<T>, ExecutableOrPrecompile<T, E, Self>)>, ExecError> {
 		let (account_id, contract_info, executable, delegate, entry_point) = match frame_args {
 			FrameArgs::Call { dest, cached_info, delegated_call } => {
@@ -996,6 +1012,11 @@ where
 					(None, Some(_)) => CachedContract::None,
 				};
 
+				let delegated_call = delegated_call.or_else(|| {
+					exec_config.mock_handler.as_ref().and_then(|mock_handler| {
+						mock_handler.mock_delegated_caller(address, input_data)
+					})
+				});
 				// in case of delegate the executable is not the one at `address`
 				let executable = if let Some(delegated_call) = &delegated_call {
 					if let Some(precompile) =
@@ -1090,6 +1111,7 @@ where
 		gas_limit: Weight,
 		deposit_limit: BalanceOf<T>,
 		read_only: bool,
+		input_data: &[u8],
 	) -> Result<Option<ExecutableOrPrecompile<T, E, Self>>, ExecError> {
 		if self.frames.len() as u32 == limits::CALL_STACK_DEPTH {
 			return Err(Error::<T>::MaxCallDepthReached.into());
@@ -1121,6 +1143,8 @@ where
 			deposit_limit,
 			read_only,
 			false,
+			input_data,
+			self.exec_config,
 		)? {
 			self.frames.try_push(frame).map_err(|_| Error::<T>::MaxCallDepthReached)?;
 			Ok(Some(executable))
@@ -1152,7 +1176,17 @@ where
 				frame.nested_gas.gas_left(),
 			);
 		});
-
+		let mock_answer = self.exec_config.mock_handler.as_ref().and_then(|handler| {
+			handler.mock_call(
+				frame
+					.delegate
+					.as_ref()
+					.map(|delegate| delegate.callee)
+					.unwrap_or(T::AddressMapper::to_address(&frame.account_id)),
+				&input_data,
+				frame.value_transferred,
+			)
+		});
 		// The output of the caller frame will be replaced by the output of this run.
 		// It is also not accessible from nested frames.
 		// Hence we drop it early to save the memory.
@@ -1335,7 +1369,11 @@ where
 		// transactional storage depth.
 		let transaction_outcome =
 			with_transaction(|| -> TransactionOutcome<Result<_, DispatchError>> {
-				let output = do_transaction();
+				let output = if let Some(mock_answer) = mock_answer {
+					Ok(mock_answer)
+				} else {
+					do_transaction()
+				};
 				match &output {
 					Ok(result) if !result.did_revert() =>
 						TransactionOutcome::Commit(Ok((true, output))),
@@ -1481,7 +1519,7 @@ where
 		to: &T::AccountId,
 		value: U256,
 		storage_meter: &mut storage::meter::GenericMeter<T, S>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> DispatchResult {
 		fn transfer_with_dust<T: Config>(
 			from: &AccountIdOf<T>,
@@ -1592,7 +1630,7 @@ where
 		to: &T::AccountId,
 		value: U256,
 		storage_meter: &mut storage::meter::GenericMeter<T, S>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> ExecResult {
 		// If the from address is root there is no account to transfer from, and therefore we can't
 		// take any `value` other than 0.
@@ -1742,6 +1780,7 @@ where
 			gas_limit,
 			deposit_limit.saturated_into::<BalanceOf<T>>(),
 			self.is_read_only(),
+			&input_data,
 		)? {
 			self.run(executable, input_data)
 		} else {
@@ -1883,6 +1922,7 @@ where
 				gas_limit,
 				deposit_limit.saturated_into::<BalanceOf<T>>(),
 				self.is_read_only(),
+				&input_data,
 			)?
 		};
 		let executable = executable.expect(FRAME_ALWAYS_EXISTS_ON_INSTANTIATE);
@@ -1954,6 +1994,7 @@ where
 				gas_limit,
 				deposit_limit.saturated_into::<BalanceOf<T>>(),
 				is_read_only,
+				&input_data,
 			)? {
 				self.run(executable, input_data)
 			} else {
@@ -1968,8 +2009,13 @@ where
 						Weight::zero(),
 					);
 				});
-
-				let result = if is_read_only && value.is_zero() {
+				let result = if let Some(mock_answer) =
+					self.exec_config.mock_handler.as_ref().and_then(|handler| {
+						handler.mock_call(T::AddressMapper::to_address(&dest), &input_data, value)
+					}) {
+					*self.last_frame_output_mut() = mock_answer.clone();
+					Ok(mock_answer)
+				} else if is_read_only && value.is_zero() {
 					Ok(Default::default())
 				} else if is_read_only {
 					Err(Error::<T>::StateChangeDenied.into())
@@ -2029,6 +2075,16 @@ where
 	}
 
 	fn caller(&self) -> Origin<T> {
+		if let Some(Ok(mock_caller)) = self
+			.exec_config
+			.mock_handler
+			.as_ref()
+			.and_then(|mock_handler| mock_handler.mock_caller(self.frames.len()))
+			.map(|mock_caller| Origin::<T>::from_runtime_origin(mock_caller))
+		{
+			return mock_caller;
+		}
+
 		if let Some(DelegateInfo { caller, .. }) = &self.top_frame().delegate {
 			caller.clone()
 		} else {

substrate/frame/revive/src/lib.rs

@@ -40,6 +40,7 @@ mod weightinfo_extension;
 
 pub mod evm;
 pub mod migrations;
+pub mod mock;
 pub mod precompiles;
 pub mod test_utils;
 pub mod tracing;
@@ -53,11 +54,11 @@ use crate::{
 		runtime::SetWeightLimit,
 		CallTracer, GenericTransaction, PrestateTracer, Trace, Tracer, TracerType, TYPE_EIP1559,
 	},
-	exec::{AccountIdOf, ExecError, Executable, Stack as ExecStack},
+	exec::{AccountIdOf, ExecError, Stack as ExecStack},
 	gas::GasMeter,
 	storage::{meter::Meter as StorageMeter, AccountType, DeletionQueueManager},
 	tracing::if_tracing,
-	vm::{pvm::extract_code_and_data, CodeInfo, ContractBlob, RuntimeCosts},
+	vm::{pvm::extract_code_and_data, CodeInfo, RuntimeCosts},
 	weightinfo_extension::OnFinalizeBlockParts,
 };
 use alloc::{boxed::Box, format, vec};
@@ -95,9 +96,10 @@ pub use crate::{
 	},
 	debug::DebugSettings,
 	evm::{block_hash::ReceiptGasInfo, Address as EthAddress, Block as EthBlock, ReceiptInfo},
-	exec::{Key, MomentOf, Origin as ExecOrigin},
+	exec::{DelegateInfo, Executable, Key, MomentOf, Origin as ExecOrigin},
 	pallet::{genesis, *},
 	storage::{AccountInfo, ContractInfo},
+	vm::ContractBlob,
 };
 pub use codec;
 pub use frame_support::{self, dispatch::DispatchInfo, weights::Weight};
@@ -1491,7 +1493,7 @@ impl<T: Config> Pallet<T> {
 		gas_limit: Weight,
 		storage_deposit_limit: BalanceOf<T>,
 		data: Vec<u8>,
-		exec_config: ExecConfig,
+		exec_config: ExecConfig<T>,
 	) -> ContractResult<ExecReturnValue, BalanceOf<T>> {
 		let mut gas_meter = GasMeter::new(gas_limit);
 		let mut storage_deposit = Default::default();
@@ -1547,7 +1549,7 @@ impl<T: Config> Pallet<T> {
 		code: Code,
 		data: Vec<u8>,
 		salt: Option<[u8; 32]>,
-		exec_config: ExecConfig,
+		exec_config: ExecConfig<T>,
 	) -> ContractResult<InstantiateReturnValue, BalanceOf<T>> {
 		let mut gas_meter = GasMeter::new(gas_limit);
 		let mut storage_deposit = Default::default();
@@ -2129,11 +2131,11 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Uploads new code and returns the Vm binary contract blob and deposit amount collected.
-	fn try_upload_pvm_code(
+	pub fn try_upload_pvm_code(
 		origin: T::AccountId,
 		code: Vec<u8>,
 		storage_deposit_limit: BalanceOf<T>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> Result<(ContractBlob<T>, BalanceOf<T>), DispatchError> {
 		let mut module = ContractBlob::from_pvm_code(code, origin)?;
 		let deposit = module.store_code(exec_config, None)?;
@@ -2161,7 +2163,7 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Convert a weight to a gas value.
-	fn evm_gas_from_weight(weight: Weight) -> U256 {
+	pub fn evm_gas_from_weight(weight: Weight) -> U256 {
 		T::FeeInfo::weight_to_fee(&weight, Combinator::Max).into()
 	}
 
@@ -2174,7 +2176,7 @@ impl<T: Config> Pallet<T> {
 		from: &T::AccountId,
 		to: &T::AccountId,
 		amount: BalanceOf<T>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> DispatchResult {
 		use frame_support::traits::tokens::{Fortitude, Precision, Preservation};
 		match (exec_config.collect_deposit_from_hold.is_some(), hold_reason) {
@@ -2219,7 +2221,7 @@ impl<T: Config> Pallet<T> {
 		from: &T::AccountId,
 		to: &T::AccountId,
 		amount: BalanceOf<T>,
-		exec_config: &ExecConfig,
+		exec_config: &ExecConfig<T>,
 	) -> Result<BalanceOf<T>, DispatchError> {
 		use frame_support::traits::{
 			tokens::{Fortitude, Precision, Preservation, Restriction},