revive_strategy/tracing/
storage_tracer.rs

1use alloy_primitives::{Bytes, U256 as RU256};
2use foundry_cheatcodes::Vm::{AccountAccessKind, StorageAccess};
3use polkadot_sdk::{
4    pallet_revive::{self, AccountInfo, Code, tracing::Tracing},
5    sp_core::{H160, U256},
6};
7use revive_env::Runtime;
8
9#[derive(Debug, Default)]
10pub(crate) struct StorageTracer {
11    /// Whether the current call is a contract creation.
12    is_create: Option<Code>,
13
14    records: Vec<AccountAccess>,
15    pending: Vec<AccountAccess>,
16    records_inner: Vec<AccountAccess>,
17    index: usize,
18    calls: Vec<H160>,
19}
20
21/// Represents the account access during vm execution.
22#[derive(Debug, Clone)]
23pub struct AccountAccess {
24    /// Call depth.
25    pub depth: u64,
26    /// Call type.
27    pub kind: AccountAccessKind,
28    /// Account that was accessed.
29    pub account: H160,
30    /// Accessor account.
31    pub accessor: H160,
32    /// Call data.
33    pub data: Bytes,
34    /// Call value.
35    pub value: U256,
36    /// Previous balance of the accessed account.
37    pub old_balance: U256,
38    /// New balance of the accessed account.
39    pub new_balance: U256,
40    /// Storage slots that were accessed.
41    pub storage_accesses: Vec<StorageAccess>,
42    /// is reverted
43    pub reverted: bool,
44    pub index: usize,
45    pub initialized: bool,
46}
47
48impl StorageTracer {
49    pub fn get_records(&self) -> Vec<AccountAccess> {
50        assert!(
51            self.pending.is_empty(),
52            "pending call stack is not empty; found calls without matching returns: {:?}",
53            self.pending
54        );
55        assert!(
56            self.records_inner.is_empty(),
57            "inner stack is not empty; found calls without matching returns: {:?}",
58            self.records_inner
59        );
60        let mut accounts = self.records.clone();
61        accounts.sort_by_key(|f| f.index);
62        accounts
63    }
64
65    fn current_addr(&self) -> H160 {
66        self.calls.last().copied().unwrap_or_default()
67    }
68}
69
70impl Tracing for StorageTracer {
71    fn instantiate_code(&mut self, code: &Code, _salt: Option<&[u8; 32]>) {
72        self.is_create = Some(code.clone());
73    }
74
75    fn enter_child_span(
76        &mut self,
77        from: H160,
78        to: H160,
79        is_delegate_call: Option<H160>,
80        is_read_only: bool,
81        value: U256,
82        input: &[u8],
83        _gas: U256,
84    ) {
85        let code = self.is_create.take();
86
87        if is_delegate_call.is_some() {
88            self.calls.push(self.current_addr());
89        } else {
90            self.calls.push(to);
91        }
92
93        let kind = if code.is_some() {
94            AccountAccessKind::Create
95        } else if is_read_only {
96            AccountAccessKind::StaticCall
97        } else if is_delegate_call.is_some() {
98            AccountAccessKind::DelegateCall
99        } else {
100            AccountAccessKind::Call
101        };
102
103        let last_depth = if !self.pending.is_empty() {
104            self.pending.last().map(|record| record.depth).expect("must have at least one record")
105        } else {
106            self.records.last().map(|record| record.depth).unwrap_or_default()
107        };
108        let new_depth = last_depth.checked_add(1).expect("overflow in recording call depth");
109
110        let mut record = AccountAccess {
111            depth: new_depth,
112            kind,
113            account: to,
114            accessor: from,
115            data: Bytes::from(input.to_vec()),
116            value,
117            reverted: false,
118            old_balance: pallet_revive::Pallet::<Runtime>::evm_balance(&to),
119            new_balance: U256::zero(),
120            storage_accesses: Default::default(),
121            index: self.index,
122            initialized: true,
123        };
124        if let Some(code) = code {
125            match code {
126                Code::Upload(items) => {
127                    record.data = Bytes::from(items);
128                }
129                Code::Existing(_) => (),
130            }
131        }
132        self.pending.push(record);
133        self.index += 1;
134    }
135
136    fn terminate(
137        &mut self,
138        contract_address: H160,
139        beneficiary_address: H160,
140        _gas_left: U256,
141        value: U256,
142    ) {
143        let last_depth = if !self.pending.is_empty() {
144            self.pending.last().map(|record| record.depth).expect("must have at least one record")
145        } else {
146            self.records.last().map(|record| record.depth).unwrap_or_default()
147        };
148        let new_depth = last_depth.checked_add(1).expect("overflow in recording call depth");
149        let account = AccountInfo::<Runtime>::is_contract(&beneficiary_address);
150        let record = AccountAccess {
151            depth: new_depth,
152            kind: AccountAccessKind::SelfDestruct,
153            account: beneficiary_address,
154            accessor: contract_address,
155            data: Bytes::new(),
156            value,
157            reverted: false,
158            old_balance: pallet_revive::Pallet::<Runtime>::evm_balance(&beneficiary_address),
159            new_balance: U256::zero(),
160            storage_accesses: Default::default(),
161            index: self.index,
162            initialized: account,
163        };
164        self.index += 1;
165
166        self.records_inner.push(record);
167    }
168
169    fn exit_child_span_with_error(
170        &mut self,
171        _error: polkadot_sdk::sp_runtime::DispatchError,
172        _gas_left: U256,
173    ) {
174        self.calls.pop();
175
176        let is_create = self.is_create.take();
177        let mut record = self.pending.pop().expect("unexpected return while recording call");
178        record.new_balance = pallet_revive::Pallet::<Runtime>::evm_balance(&self.current_addr());
179        record.reverted = true;
180        record.storage_accesses.iter_mut().for_each(|x| x.reverted = true);
181        self.records_inner.iter_mut().for_each(|x| {
182            if record.reverted {
183                x.reverted = true;
184                x.storage_accesses.iter_mut().for_each(|x| x.reverted = true);
185            }
186        });
187
188        if let Some(code) = is_create {
189            record.kind = AccountAccessKind::Create;
190            match code {
191                Code::Upload(items) => {
192                    record.data = Bytes::from(items);
193                }
194                Code::Existing(_) => (),
195            }
196        }
197
198        if self.pending.is_empty() {
199            // no more pending records, append everything recorded so far.
200            self.records.push(record);
201            // append the inner records.
202            if !self.records_inner.is_empty() {
203                self.records.extend(std::mem::take(&mut self.records_inner));
204            }
205        } else {
206            // we have pending records, so record to inner.
207            self.records_inner.push(record);
208        }
209    }
210
211    fn exit_child_span(
212        &mut self,
213        output: &polkadot_sdk::pallet_revive::ExecReturnValue,
214        _gas_left: U256,
215    ) {
216        self.calls.pop();
217
218        let is_create = self.is_create.take();
219
220        let mut record = self.pending.pop().expect("unexpected return while recording call");
221        record.new_balance = pallet_revive::Pallet::<Runtime>::evm_balance(&self.current_addr());
222        if output.did_revert() {
223            record.reverted = true;
224            record.storage_accesses.iter_mut().for_each(|x| x.reverted = true);
225            self.records_inner.iter_mut().for_each(|x| {
226                if record.reverted {
227                    x.reverted = true;
228                    x.storage_accesses.iter_mut().for_each(|x| x.reverted = true);
229                }
230            });
231        }
232
233        if let Some(code) = is_create {
234            record.kind = AccountAccessKind::Create;
235            match code {
236                Code::Upload(items) => {
237                    record.data = Bytes::from(items);
238                }
239                Code::Existing(_) => (),
240            }
241        }
242
243        if self.pending.is_empty() {
244            // no more pending records, append everything recorded so far.
245            self.records.push(record);
246
247            // append the inner records.
248            if !self.records_inner.is_empty() {
249                self.records.extend(std::mem::take(&mut self.records_inner));
250            }
251        } else {
252            // we have pending records, so record to inner.
253            self.records_inner.push(record);
254        }
255    }
256
257    fn storage_read(&mut self, key: &polkadot_sdk::pallet_revive::Key, value: Option<&[u8]>) {
258        let account = self.current_addr().0.into();
259        let record = self.pending.last_mut().expect("expected at least one record");
260        record.storage_accesses.push(StorageAccess {
261            account,
262            slot: RU256::from_be_slice(key.unhashed()).into(),
263            isWrite: false,
264            previousValue: RU256::from_be_slice(value.unwrap_or_default()).into(),
265            newValue: RU256::from_be_slice(value.unwrap_or_default()).into(),
266            reverted: false,
267        });
268    }
269    fn storage_write(
270        &mut self,
271        key: &polkadot_sdk::pallet_revive::Key,
272        old_value: Option<Vec<u8>>,
273        new_value: Option<&[u8]>,
274    ) {
275        let account = self.current_addr().0.into();
276        let record = self.pending.last_mut().expect("expected at least one record");
277        record.storage_accesses.push(StorageAccess {
278            account,
279            slot: RU256::from_be_slice(key.unhashed()).into(),
280            isWrite: true,
281            previousValue: RU256::from_be_slice(old_value.unwrap_or_default().as_slice()).into(),
282            newValue: RU256::from_be_slice(new_value.unwrap_or_default()).into(),
283            reverted: false,
284        });
285    }
286}
revive_strategy/tracing/storage_tracer.rs

revive_strategy/tracing/
storage_tracer.rs
