anvil/eth/backend/
cheats.rs1use alloy_evm::precompiles::{Precompile, PrecompileInput};
4use alloy_primitives::{
5 Address, B256, Bytes,
6 map::{AddressHashSet, foldhash::HashMap},
7};
8use parking_lot::RwLock;
9use revm::precompile::{
10 PrecompileHalt, PrecompileId, PrecompileOutput, PrecompileResult, call_eth_precompile,
11 secp256k1::ec_recover_run, utilities::right_pad,
12};
13use std::{borrow::Cow, sync::Arc};
14
15static PRECOMPILE_ID_CHEAT_ECRECOVER: PrecompileId =
17 PrecompileId::Custom(Cow::Borrowed("cheat_ecrecover"));
18
19#[derive(Clone, Debug, Default)]
23pub struct CheatsManager {
24 state: Arc<RwLock<CheatsState>>,
26}
27
28impl CheatsManager {
29 pub fn impersonate(&self, addr: Address) -> bool {
33 trace!(target: "cheats", %addr, "start impersonating");
34 !self.state.write().impersonated_accounts.insert(addr)
38 }
39
40 pub fn stop_impersonating(&self, addr: &Address) {
42 trace!(target: "cheats", %addr, "stop impersonating");
43 self.state.write().impersonated_accounts.remove(addr);
44 }
45
46 pub fn is_impersonated(&self, addr: Address) -> bool {
48 if self.auto_impersonate_accounts() {
49 true
50 } else {
51 self.state.read().impersonated_accounts.contains(&addr)
52 }
53 }
54
55 pub fn auto_impersonate_accounts(&self) -> bool {
57 self.state.read().auto_impersonate_accounts
58 }
59
60 pub fn set_auto_impersonate_account(&self, enabled: bool) {
63 trace!(target: "cheats", "Auto impersonation set to {:?}", enabled);
64 self.state.write().auto_impersonate_accounts = enabled
65 }
66
67 pub fn impersonated_accounts(&self) -> AddressHashSet {
69 self.state.read().impersonated_accounts.clone()
70 }
71
72 pub fn add_recover_override(&self, sig: Bytes, addr: Address) {
74 self.state.write().signature_overrides.insert(sig, addr);
75 }
76
77 pub fn get_recover_override(&self, sig: &Bytes) -> Option<Address> {
79 self.state.read().signature_overrides.get(sig).copied()
80 }
81
82 pub fn has_recover_overrides(&self) -> bool {
84 !self.state.read().signature_overrides.is_empty()
85 }
86
87 pub fn set_next_block_prevrandao(&self, prevrandao: B256) {
92 trace!(target: "cheats", %prevrandao, "set next block prevrandao");
93 self.state.write().next_block_prevrandao.replace(prevrandao);
94 }
95
96 pub fn take_next_block_prevrandao(&self) -> Option<B256> {
100 self.state.write().next_block_prevrandao.take()
101 }
102
103 pub fn clear_next_block_prevrandao(&self) {
108 self.state.write().next_block_prevrandao.take();
109 }
110}
111
112#[derive(Clone, Debug, Default)]
114pub struct CheatsState {
115 pub impersonated_accounts: AddressHashSet,
117 pub auto_impersonate_accounts: bool,
119 pub signature_overrides: HashMap<Bytes, Address>,
121 pub next_block_prevrandao: Option<B256>,
124}
125
126impl CheatEcrecover {
127 pub const fn new(cheats: Arc<CheatsManager>) -> Self {
128 Self { cheats }
129 }
130}
131
132impl Precompile for CheatEcrecover {
133 fn call(&self, input: PrecompileInput<'_>) -> PrecompileResult {
134 if !self.cheats.has_recover_overrides() {
135 return Ok(call_eth_precompile(ec_recover_run, input.data, input.gas, input.reservoir));
136 }
137
138 const ECRECOVER_BASE: u64 = 3_000;
139 if input.gas < ECRECOVER_BASE {
140 return Ok(PrecompileOutput::halt(PrecompileHalt::OutOfGas, input.reservoir));
141 }
142 let padded = right_pad::<128>(input.data);
143 let v = padded[63];
144 let mut sig_bytes = [0u8; 65];
145 sig_bytes[..64].copy_from_slice(&padded[64..128]);
146 sig_bytes[64] = v;
147 let sig_bytes_wrapped = Bytes::copy_from_slice(&sig_bytes);
148 if let Some(addr) = self.cheats.get_recover_override(&sig_bytes_wrapped) {
149 let mut out = [0u8; 32];
150 out[12..].copy_from_slice(addr.as_slice());
151 return Ok(PrecompileOutput::new(
152 ECRECOVER_BASE,
153 Bytes::copy_from_slice(&out),
154 input.reservoir,
155 ));
156 }
157 Ok(call_eth_precompile(ec_recover_run, input.data, input.gas, input.reservoir))
158 }
159
160 fn precompile_id(&self) -> &PrecompileId {
161 &PRECOMPILE_ID_CHEAT_ECRECOVER
162 }
163
164 fn supports_caching(&self) -> bool {
165 false
166 }
167}
168
169#[derive(Clone, Debug)]
171pub struct CheatEcrecover {
172 cheats: Arc<CheatsManager>,
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 #[test]
180 fn impersonate_returns_false_then_true() {
181 let mgr = CheatsManager::default();
182 let addr = Address::from([1u8; 20]);
183 assert!(!mgr.impersonate(addr));
184 assert!(mgr.impersonate(addr));
185 }
186}