diff --git a/crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs b/crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs index ae9d806bd5..257d3ef4d5 100644 --- a/crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs +++ b/crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs @@ -2,10 +2,11 @@ use { super::Strategy, crate::code_simulation::{CodeSimulating, SimulationError}, contracts::{dummy_contract, ERC20}, - ethcontract::{Address, U256}, + ethcontract::{Address, H256, U256}, ethrpc::extensions::StateOverride, maplit::hashmap, std::{ + collections::HashMap, fmt::{self, Debug, Formatter}, sync::{Arc, LazyLock}, }, @@ -21,17 +22,7 @@ pub struct Detector { simulator: Arc, } -/// Storage slot based on OpenZeppelin's ERC20Upgradeable contract [^1]. -/// -/// [^1]: -static OPEN_ZEPPELIN_ERC20_UPGRADEABLE: LazyLock = LazyLock::new(|| { - U256::from("52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00") -}); - impl Detector { - /// Number of different slots to try out. - const TRIES: u8 = 25; - /// Creates a new balance override detector. pub fn new(simulator: Arc) -> Self { Self { simulator } @@ -41,57 +32,15 @@ impl Detector { /// Returns an `Err` if it cannot detect the strategy or an internal /// simulation fails. pub async fn detect(&self, token: Address) -> Result { - // Use an exact value which isn't too large or too small. This helps - // not have false positives for cases where the token balances in - // some other denomination from the actual token balance (such as - // stETH for example) and not run into issues with overflows. - // Also don't use 0 to avoid false postitive when trying to overwrite - // a balance value of 0 which should always succeed. - let marker_amount_for_index = |i| U256::from(u64::from_be_bytes([i + 1; 8])); - - // This is a pretty unsophisticated strategy where we basically try a - // bunch of different slots and see which one sticks. We try balance - // mappings for the first `TRIES` slots; each with a unique value. - let mut tries = (0..Self::TRIES).map(|i| { - let strategy = Strategy::SolidityMapping { slot: U256::from(i) }; - (strategy, marker_amount_for_index(i)) - }) - // Afterwards we try hardcoded storage slots based on popular utility - // libraries like OpenZeppelin. - .chain((Self::TRIES..).zip([ - Strategy::SolidityMapping{ slot: *OPEN_ZEPPELIN_ERC20_UPGRADEABLE }, - Strategy::SoladyMapping, - ]).map(|(index, strategy)| { - (strategy, marker_amount_for_index(index)) - })); - - // On a technical note, Ethereum public addresses are, for the most - // part, generated by taking the 20 last bytes of a Keccak-256 hash (for - // things like contract creation, public address derivation from a - // Secp256k1 public key, etc.), so we use one for our heuristics from a - // 32-byte digest with no know pre-image, to prevent any weird - // interactions with the weird tokens of the world. - let holder = { - let mut address = Address::default(); - address.0.copy_from_slice(&keccak256(b"Moo!")[12..]); - address.0[19] = address.0[19].wrapping_sub(1); - address - }; - let token = dummy_contract!(ERC20, token); let call = CallRequest { to: Some(token.address()), - data: token.methods().balance_of(holder).m.tx.data, + data: token.methods().balance_of(*HOLDER).m.tx.data, ..Default::default() }; let overrides = hashmap! { token.address() => StateOverride { - state_diff: Some( - tries - .clone() - .map(|(strategy, amount)| strategy.state_override(&holder, &amount)) - .collect(), - ), + state_diff: Some(STORAGE_OVERRIDES.clone()), ..Default::default() }, }; @@ -101,13 +50,87 @@ impl Detector { .then(|| U256::from_big_endian(&output)) .ok_or(DetectionError::Decode)?; - let strategy = tries - .find_map(|(strategy, amount)| (amount == balance).then_some(strategy)) - .ok_or(DetectionError::NotFound)?; - Ok(strategy) + TESTED_STRATEGIES + .iter() + .find_map(|helper| (helper.balance == balance).then_some(helper.strategy.clone())) + .ok_or(DetectionError::NotFound) + } +} + +/// Contains all the information we need to determine which state override +/// was successful. +struct StrategyHelper { + /// strategy that was used to compute the state override + strategy: Strategy, + /// balance amount the strategy wrote into the storage + balance: U256, +} + +impl StrategyHelper { + fn new(strategy: Strategy, index: u8) -> Self { + Self { + strategy, + // Use an exact value which isn't too large or too small. This helps + // not have false positives for cases where the token balances in + // some other denomination from the actual token balance (such as + // stETH for example) and not run into issues with overflows. + // We also make sure that we avoid 0 because `balanceOf()` returns + // 0 by default so we can't use it to detect successful state overrides. + balance: U256::from(u64::from_be_bytes([index + 1; 8])), + } } } +/// Storage slot based on OpenZeppelin's ERC20Upgradeable contract [^1]. +/// +/// [^1]: +static OPEN_ZEPPELIN_ERC20_UPGRADEABLE: &str = + "52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00"; + +/// Address which we try to override the balances for. +static HOLDER: LazyLock
= LazyLock::new(|| { + // On a technical note, Ethereum public addresses are, for the most + // part, generated by taking the 20 last bytes of a Keccak-256 hash (for + // things like contract creation, public address derivation from a + // Secp256k1 public key, etc.), so we use one for our heuristics from a + // 32-byte digest with no know pre-image, to prevent any weird + // interactions with the weird tokens of the world. + let mut address = Address::default(); + address.0.copy_from_slice(&keccak256(b"Moo!")[12..]); + address.0[19] = address.0[19].wrapping_sub(1); + address +}); + +/// All the strategies we use to detect where a token stores the balances. +static TESTED_STRATEGIES: LazyLock> = LazyLock::new(|| { + const FIRST_N_SLOTS: u8 = 25; + + // This is a pretty unsophisticated strategy where we basically try a + // bunch of different slots and see which one sticks. We try balance + // mappings for the first `TRIES` slots; each with a unique value. + (0..FIRST_N_SLOTS).map(|i| { + let strategy = Strategy::SolidityMapping { slot: U256::from(i) }; + StrategyHelper::new(strategy, i) + }) + // Afterwards we try hardcoded storage slots based on popular utility + // libraries like OpenZeppelin. + .chain((FIRST_N_SLOTS..).zip([ + Strategy::SolidityMapping{ slot: U256::from(OPEN_ZEPPELIN_ERC20_UPGRADEABLE) }, + Strategy::SoladyMapping, + ]).map(|(index, strategy)| { + StrategyHelper::new(strategy, index) + })) + .collect() +}); + +/// Storage overrides (storage_slot, value) for all tested strategies. +static STORAGE_OVERRIDES: LazyLock> = LazyLock::new(|| { + TESTED_STRATEGIES + .iter() + .map(|helper| helper.strategy.state_override(&HOLDER, &helper.balance)) + .collect::>() +}); + impl Debug for Detector { fn fmt(&self, f: &mut Formatter) -> fmt::Result { f.debug_struct("Detector") @@ -154,7 +177,7 @@ mod tests { assert_eq!( storage, Strategy::SolidityMapping { - slot: *OPEN_ZEPPELIN_ERC20_UPGRADEABLE + slot: U256::from(OPEN_ZEPPELIN_ERC20_UPGRADEABLE), } );