anvil/eth/backend/
cheats.rs

1//! Support for "cheat codes" / bypass functions
2
3use alloy_evm::precompiles::{Precompile, PrecompileInput};
4use alloy_primitives::{
5    Address, Bytes,
6    map::{AddressHashSet, foldhash::HashMap},
7};
8use parking_lot::RwLock;
9use revm::precompile::{
10    PrecompileError, PrecompileId, PrecompileOutput, PrecompileResult, secp256k1::ec_recover_run,
11    utilities::right_pad,
12};
13use std::{borrow::Cow, sync::Arc};
14
15/// ID for the [`CheatEcrecover::precompile_id`] precompile.
16static PRECOMPILE_ID_CHEAT_ECRECOVER: PrecompileId =
17    PrecompileId::Custom(Cow::Borrowed("cheat_ecrecover"));
18
19/// Manages user modifications that may affect the node's behavior
20///
21/// Contains the state of executed, non-eth standard cheat code RPC
22#[derive(Clone, Debug, Default)]
23pub struct CheatsManager {
24    /// shareable state
25    state: Arc<RwLock<CheatsState>>,
26}
27
28impl CheatsManager {
29    /// Sets the account to impersonate
30    ///
31    /// Returns `true` if the account is already impersonated
32    pub fn impersonate(&self, addr: Address) -> bool {
33        trace!(target: "cheats", "Start impersonating {:?}", addr);
34        let mut state = self.state.write();
35        // When somebody **explicitly** impersonates an account we need to store it so we are able
36        // to return it from `eth_accounts`. That's why we do not simply call `is_impersonated()`
37        // which does not check that list when auto impersonation is enabled.
38        if state.impersonated_accounts.contains(&addr) {
39            // need to check if already impersonated, so we don't overwrite the code
40            return true;
41        }
42        state.impersonated_accounts.insert(addr)
43    }
44
45    /// Removes the account that from the impersonated set
46    pub fn stop_impersonating(&self, addr: &Address) {
47        trace!(target: "cheats", "Stop impersonating {:?}", addr);
48        self.state.write().impersonated_accounts.remove(addr);
49    }
50
51    /// Returns true if the `addr` is currently impersonated
52    pub fn is_impersonated(&self, addr: Address) -> bool {
53        if self.auto_impersonate_accounts() {
54            true
55        } else {
56            self.state.read().impersonated_accounts.contains(&addr)
57        }
58    }
59
60    /// Returns true is auto impersonation is enabled
61    pub fn auto_impersonate_accounts(&self) -> bool {
62        self.state.read().auto_impersonate_accounts
63    }
64
65    /// Sets the auto impersonation flag which if set to true will make the `is_impersonated`
66    /// function always return true
67    pub fn set_auto_impersonate_account(&self, enabled: bool) {
68        trace!(target: "cheats", "Auto impersonation set to {:?}", enabled);
69        self.state.write().auto_impersonate_accounts = enabled
70    }
71
72    /// Returns all accounts that are currently being impersonated.
73    pub fn impersonated_accounts(&self) -> AddressHashSet {
74        self.state.read().impersonated_accounts.clone()
75    }
76
77    /// Registers an override so that `ecrecover(signature)` returns `addr`.
78    pub fn add_recover_override(&self, sig: Bytes, addr: Address) {
79        self.state.write().signature_overrides.insert(sig, addr);
80    }
81
82    /// If an override exists for `sig`, returns the address; otherwise `None`.
83    pub fn get_recover_override(&self, sig: &Bytes) -> Option<Address> {
84        self.state.read().signature_overrides.get(sig).copied()
85    }
86
87    /// Returns true if any ecrecover overrides have been registered.
88    pub fn has_recover_overrides(&self) -> bool {
89        !self.state.read().signature_overrides.is_empty()
90    }
91}
92
93/// Container type for all the state variables
94#[derive(Clone, Debug, Default)]
95pub struct CheatsState {
96    /// All accounts that are currently impersonated
97    pub impersonated_accounts: AddressHashSet,
98    /// If set to true will make the `is_impersonated` function always return true
99    pub auto_impersonate_accounts: bool,
100    /// Overrides for ecrecover: Signature => Address
101    pub signature_overrides: HashMap<Bytes, Address>,
102}
103
104impl CheatEcrecover {
105    pub fn new(cheats: Arc<CheatsManager>) -> Self {
106        Self { cheats }
107    }
108}
109
110impl Precompile for CheatEcrecover {
111    fn call(&self, input: PrecompileInput<'_>) -> PrecompileResult {
112        if !self.cheats.has_recover_overrides() {
113            return ec_recover_run(input.data, input.gas);
114        }
115
116        const ECRECOVER_BASE: u64 = 3_000;
117        if input.gas < ECRECOVER_BASE {
118            return Err(PrecompileError::OutOfGas);
119        }
120        let padded = right_pad::<128>(input.data);
121        let v = padded[63];
122        let mut sig_bytes = [0u8; 65];
123        sig_bytes[..64].copy_from_slice(&padded[64..128]);
124        sig_bytes[64] = v;
125        let sig_bytes_wrapped = Bytes::copy_from_slice(&sig_bytes);
126        if let Some(addr) = self.cheats.get_recover_override(&sig_bytes_wrapped) {
127            let mut out = [0u8; 32];
128            out[12..].copy_from_slice(addr.as_slice());
129            return Ok(PrecompileOutput::new(ECRECOVER_BASE, Bytes::copy_from_slice(&out)));
130        }
131        ec_recover_run(input.data, input.gas)
132    }
133
134    fn precompile_id(&self) -> &PrecompileId {
135        &PRECOMPILE_ID_CHEAT_ECRECOVER
136    }
137
138    fn is_pure(&self) -> bool {
139        false
140    }
141}
142
143/// A custom ecrecover precompile that supports cheat-based signature overrides.
144#[derive(Clone, Debug)]
145pub struct CheatEcrecover {
146    cheats: Arc<CheatsManager>,
147}
anvil/eth/backend/cheats.rs

anvil/eth/backend/
cheats.rs
