precompile_utils/testing/

execution.rs

// This file is part of Frontier.

// Copyright (c) Moonsong Labs.
// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{
	solidity::codec::Codec,
	testing::{decode_revert_message, MockHandle, PrettyLog, SubcallHandle, SubcallTrait},
};
use alloc::boxed::Box;
use fp_evm::{
	Context, ExitError, ExitSucceed, Log, PrecompileFailure, PrecompileOutput, PrecompileResult,
	PrecompileSet,
};
use sp_core::{H160, U256};

#[must_use]
pub struct PrecompilesTester<'p, P> {
	precompiles: &'p P,
	handle: MockHandle,

	target_gas: Option<u64>,
	subcall_handle: Option<SubcallHandle>,

	expected_cost: Option<u64>,
	expected_logs: Option<Vec<PrettyLog>>,
	static_call: bool,
}

impl<'p, P: PrecompileSet> PrecompilesTester<'p, P> {
	pub fn new(
		precompiles: &'p P,
		from: impl Into<H160>,
		to: impl Into<H160>,
		data: Vec<u8>,
	) -> Self {
		let to = to.into();
		let mut handle = MockHandle::new(
			to,
			Context {
				address: to,
				caller: from.into(),
				apparent_value: U256::zero(),
			},
		);

		handle.input = data;

		Self {
			precompiles,
			handle,

			target_gas: None,
			subcall_handle: None,

			expected_cost: None,
			expected_logs: None,
			static_call: false,
		}
	}

	pub fn with_value(mut self, value: impl Into<U256>) -> Self {
		self.handle.context.apparent_value = value.into();
		self
	}

	pub fn with_subcall_handle(mut self, subcall_handle: impl SubcallTrait) -> Self {
		self.subcall_handle = Some(Box::new(subcall_handle));
		self
	}

	pub fn with_target_gas(mut self, target_gas: Option<u64>) -> Self {
		self.target_gas = target_gas;
		self
	}

	pub fn with_static_call(mut self, static_call: bool) -> Self {
		self.static_call = static_call;
		self
	}

	pub fn expect_cost(mut self, cost: u64) -> Self {
		self.expected_cost = Some(cost);
		self
	}

	pub fn expect_no_logs(mut self) -> Self {
		self.expected_logs = Some(vec![]);
		self
	}

	pub fn expect_log(mut self, log: Log) -> Self {
		self.expected_logs = Some({
			let mut logs = self.expected_logs.unwrap_or_default();
			logs.push(PrettyLog(log));
			logs
		});
		self
	}

	fn assert_optionals(&self) {
		if let Some(cost) = &self.expected_cost {
			assert_eq!(&self.handle.gas_used, cost);
		}

		if let Some(logs) = &self.expected_logs {
			similar_asserts::assert_eq!(&self.handle.logs, logs);
		}
	}

	fn execute(&mut self) -> Option<PrecompileResult> {
		let handle = &mut self.handle;
		handle.subcall_handle = self.subcall_handle.take();
		handle.is_static = self.static_call;

		if let Some(gas_limit) = self.target_gas {
			handle.gas_limit = gas_limit;
		}

		let res = self.precompiles.execute(handle);

		self.subcall_handle = handle.subcall_handle.take();

		res
	}

	/// Execute the precompile set and expect some precompile to have been executed, regardless of the
	/// result.
	pub fn execute_some(mut self) {
		let res = self.execute();
		assert!(res.is_some());
		self.assert_optionals();
	}

	/// Execute the precompile set and expect no precompile to have been executed.
	pub fn execute_none(mut self) {
		let res = self.execute();
		assert!(res.is_none());
		self.assert_optionals();
	}

	/// Execute the precompile set and check it returns provided output.
	pub fn execute_returns_raw(mut self, output: Vec<u8>) {
		let res = self.execute();

		match res {
			Some(Err(PrecompileFailure::Revert { output, .. })) => {
				let decoded = decode_revert_message(&output);
				eprintln!(
					"Revert message (bytes): {:?}",
					sp_core::hexdisplay::HexDisplay::from(&decoded)
				);
				eprintln!(
					"Revert message (string): {:?}",
					core::str::from_utf8(decoded).ok()
				);
				panic!("Shouldn't have reverted");
			}
			Some(Ok(PrecompileOutput {
				exit_status: ExitSucceed::Returned,
				output: execution_output,
			})) => {
				if execution_output != output {
					eprintln!(
						"Output (bytes): {:?}",
						sp_core::hexdisplay::HexDisplay::from(&execution_output)
					);
					eprintln!(
						"Output (string): {:?}",
						core::str::from_utf8(&execution_output).ok()
					);
					panic!("Output doesn't match");
				}
			}
			other => panic!("Unexpected result: {other:?}"),
		}

		self.assert_optionals();
	}

	/// Execute the precompile set and check it returns provided Solidity encoded output.
	pub fn execute_returns(self, output: impl Codec) {
		self.execute_returns_raw(crate::solidity::encode_return_value(output))
	}

	/// Execute the precompile set and check if it reverts.
	/// Take a closure allowing to perform custom matching on the output.
	pub fn execute_reverts(mut self, check: impl Fn(&[u8]) -> bool) {
		let res = self.execute();

		match res {
			Some(Err(PrecompileFailure::Revert { output, .. })) => {
				let decoded = decode_revert_message(&output);
				if !check(decoded) {
					eprintln!(
						"Revert message (bytes): {:?}",
						sp_core::hexdisplay::HexDisplay::from(&decoded)
					);
					eprintln!(
						"Revert message (string): {:?}",
						core::str::from_utf8(decoded).ok()
					);
					panic!("Revert reason doesn't match !");
				}
			}
			other => panic!("Didn't revert, instead returned {other:?}"),
		}

		self.assert_optionals();
	}

	/// Execute the precompile set and check it returns provided output.
	pub fn execute_error(mut self, error: ExitError) {
		let res = self.execute();
		assert_eq!(
			res,
			Some(Err(PrecompileFailure::Error { exit_status: error }))
		);
		self.assert_optionals();
	}
}

pub trait PrecompileTesterExt: PrecompileSet + Sized {
	fn prepare_test(
		&self,
		from: impl Into<H160>,
		to: impl Into<H160>,
		data: impl Into<Vec<u8>>,
	) -> PrecompilesTester<Self>;
}

impl<T: PrecompileSet> PrecompileTesterExt for T {
	fn prepare_test(
		&self,
		from: impl Into<H160>,
		to: impl Into<H160>,
		data: impl Into<Vec<u8>>,
	) -> PrecompilesTester<Self> {
		PrecompilesTester::new(self, from, to, data.into())
	}
}