[pallet-revive] fix prestate tracer current address

prdoc/pr_10239.prdoc

@@ -0,0 +1,7 @@
+title: '[pallet-revive] fix prestate tracer current address'
+doc:
+- audience: Runtime Dev
+  description: Fix prestate tracer not reporting the contract addresses properly.
+crates:
+- name: pallet-revive
+  bump: patch

substrate/frame/revive/src/evm/tracing/prestate_tracing.rs

@@ -22,17 +22,26 @@ use crate::{
 use alloc::{collections::BTreeMap, vec::Vec};
 use sp_core::{H160, U256};
 
+/// A call in the call stack.
+#[derive(Debug, Default, Clone, PartialEq)]
+struct Call {
+	/// The address of the call.
+	addr: H160,
+	/// The init code if this is a contract creation.
+	create_code: Option<Code>,
+}
+
 /// A tracer that traces the prestate.
 #[derive(frame_support::DefaultNoBound, Debug, Clone, PartialEq)]
 pub struct PrestateTracer<T> {
 	/// The tracer configuration.
 	config: PrestateTracerConfig,
 
-	/// The current address of the contract's which storage is being accessed.
-	current_addr: H160,
+	/// Stack of calls.
+	calls: Vec<Call>,
 
-	/// Whether the current call is a contract creation.
-	is_create: Option<Code>,
+	// The code used by create transaction
+	create_code: Option<Code>,
 
 	// pre / post state
 	trace: (BTreeMap<H160, PrestateTraceInfo>, BTreeMap<H160, PrestateTraceInfo>),
@@ -49,6 +58,14 @@ where
 		Self { config, ..Default::default() }
 	}
 
+	fn current_addr(&self) -> H160 {
+		self.calls.last().map(|c| c.addr.clone()).unwrap_or_default()
+	}
+
+	fn current_is_create(&self) -> bool {
+		self.calls.last().map(|c| c.create_code.is_some()).unwrap_or_default()
+	}
+
 	/// Returns an empty trace.
 	pub fn empty_trace(&self) -> PrestateTrace {
 		if self.config.diff_mode {
@@ -172,7 +189,11 @@ where
 	}
 
 	fn instantiate_code(&mut self, code: &crate::Code, _salt: Option<&[u8; 32]>) {
-		self.is_create = Some(code.clone());
+		if let Some(current) = self.calls.last_mut() {
+			current.create_code = Some(code.clone());
+		} else {
+			self.create_code = Some(code.clone());
+		}
 	}
 
 	fn enter_child_span(
@@ -194,7 +215,11 @@ where
 			)
 		});
 
-		if self.is_create.is_none() {
+		if !is_delegate_call {
+			self.calls.push(Call { addr: to, create_code: self.create_code.take() });
+		}
+
+		if !self.current_is_create() {
 			self.trace.0.entry(to).or_insert_with_key(|addr| {
 				Self::prestate_info(
 					addr,
@@ -203,18 +228,14 @@ where
 				)
 			});
 		}
-
-		if !is_delegate_call {
-			self.current_addr = to;
-		}
 	}
 
 	fn exit_child_span_with_error(&mut self, _error: crate::DispatchError, _gas_used: Weight) {
-		self.is_create = None;
+		self.calls.pop();
 	}
 
 	fn exit_child_span(&mut self, output: &ExecReturnValue, _gas_used: Weight) {
-		let create_code = self.is_create.take();
+		let Call { addr: current_addr, create_code } = self.calls.pop().unwrap_or_default();
 		if output.did_revert() {
 			return
 		}
@@ -228,12 +249,12 @@ where
 					PristineCode::<T>::get(&code_hash).map(|code| Bytes::from(code.to_vec())),
 			}
 		} else {
-			Self::bytecode(&self.current_addr)
+			Self::bytecode(&current_addr)
 		};
 
 		Self::update_prestate_info(
-			self.trace.1.entry(self.current_addr).or_default(),
-			&self.current_addr,
+			self.trace.1.entry(current_addr).or_default(),
+			&current_addr,
 			code,
 		);
 	}
@@ -244,7 +265,7 @@ where
 		let old_value = self
 			.trace
 			.0
-			.entry(self.current_addr)
+			.entry(self.current_addr())
 			.or_default()
 			.storage
 			.entry(key.clone())
@@ -257,19 +278,19 @@ where
 		if old_value.as_ref().map(|v| v.0.as_ref()) != new_value {
 			self.trace
 				.1
-				.entry(self.current_addr)
+				.entry(self.current_addr())
 				.or_default()
 				.storage
 				.insert(key, new_value.map(|v| v.to_vec().into()));
 		} else {
-			self.trace.1.entry(self.current_addr).or_default().storage.remove(&key);
+			self.trace.1.entry(self.current_addr()).or_default().storage.remove(&key);
 		}
 	}
 
 	fn storage_read(&mut self, key: &Key, value: Option<&[u8]>) {
 		self.trace
 			.0
-			.entry(self.current_addr)
+			.entry(self.current_addr())
 			.or_default()
 			.storage
 			.entry(key.unhashed().to_vec().into())

substrate/frame/revive/src/tests/pvm.rs

@@ -4240,19 +4240,31 @@ fn prestate_tracing_works() {
 		// redact balance so that tests are resilient to weight changes
 		let alice_redacted_balance = Some(U256::from(1));
 
-		let test_cases: Vec<(Box<dyn FnOnce()>, _, _)> = vec![
-			(
-				Box::new(|| {
-					builder::bare_call(addr)
-						.data((3u32, addr_callee).encode())
-						.build_and_unwrap_result();
+		struct TestCase {
+			description: &'static str,
+			exec_call: Box<dyn FnOnce()>,
+			config: PrestateTracerConfig,
+			expected_trace: PrestateTrace,
+		}
+
+		let test_cases: Vec<TestCase> = vec![
+			TestCase {
+				description: "prestate mode with cross-contract call",
+				exec_call: Box::new({
+					let addr = addr;
+					let addr_callee = addr_callee;
+					move || {
+						builder::bare_call(addr)
+							.data((3u32, addr_callee).encode())
+							.build_and_unwrap_result();
+					}
 				}),
-				PrestateTracerConfig {
+				config: PrestateTracerConfig {
 					diff_mode: false,
 					disable_storage: false,
 					disable_code: false,
 				},
-				PrestateTrace::Prestate(BTreeMap::from([
+				expected_trace: PrestateTrace::Prestate(BTreeMap::from([
 					(
 						ALICE_ADDR,
 						PrestateTraceInfo {
@@ -4284,19 +4296,24 @@ fn prestate_tracing_works() {
 						},
 					),
 				])),
-			),
-			(
-				Box::new(|| {
-					builder::bare_call(addr)
-						.data((3u32, addr_callee).encode())
-						.build_and_unwrap_result();
+			},
+			TestCase {
+				description: "diff mode with cross-contract call",
+				exec_call: Box::new({
+					let addr = addr;
+					let addr_callee = addr_callee;
+					move || {
+						builder::bare_call(addr)
+							.data((3u32, addr_callee).encode())
+							.build_and_unwrap_result();
+					}
 				}),
-				PrestateTracerConfig {
+				config: PrestateTracerConfig {
 					diff_mode: true,
 					disable_storage: false,
 					disable_code: false,
 				},
-				PrestateTrace::DiffMode {
+				expected_trace: PrestateTrace::DiffMode {
 					pre: BTreeMap::from([
 						(
 							BOB_ADDR,
@@ -4332,19 +4349,23 @@ fn prestate_tracing_works() {
 						),
 					]),
 				},
-			),
-			(
-				Box::new(|| {
-					builder::bare_instantiate(Code::Upload(dummy_code.clone()))
-						.salt(None)
-						.build_and_unwrap_result();
+			},
+			TestCase {
+				description: "diff mode with contract instantiation",
+				exec_call: Box::new({
+					let dummy_code = dummy_code.clone();
+					move || {
+						builder::bare_instantiate(Code::Upload(dummy_code))
+							.salt(None)
+							.build_and_unwrap_result();
+					}
 				}),
-				PrestateTracerConfig {
+				config: PrestateTracerConfig {
 					diff_mode: true,
 					disable_storage: false,
 					disable_code: false,
 				},
-				PrestateTrace::DiffMode {
+				expected_trace: PrestateTrace::DiffMode {
 					pre: BTreeMap::from([(
 						ALICE_ADDR,
 						PrestateTraceInfo {
@@ -4373,10 +4394,10 @@ fn prestate_tracing_works() {
 						),
 					]),
 				},
-			),
+			},
 		];
 
-		for (exec_call, config, expected_trace) in test_cases.into_iter() {
+		for TestCase { description, exec_call, config, expected_trace } in test_cases.into_iter() {
 			let mut tracer = PrestateTracer::<Test>::new(config);
 			trace(&mut tracer, || {
 				exec_call();
@@ -4401,7 +4422,7 @@ fn prestate_tracing_works() {
 				},
 			}
 
-			assert_eq!(trace, expected_trace);
+			assert_eq!(trace, expected_trace, "Trace mismatch for: {description}");
 		}
 	});
 }