diff --git a/crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs b/crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs index d1f85c03d3..d12d3c9621 100644 --- a/crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs +++ b/crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs @@ -86,7 +86,8 @@ impl Display for TokenConfiguration { fn fmt(&self, f: &mut Formatter) -> fmt::Result { let format_entry = |f: &mut Formatter, (addr, strategy): (&Address, &Strategy)| match strategy { - Strategy::Mapping { slot } => write!(f, "{addr:?}@{slot}"), + Strategy::SolidityMapping { slot } => write!(f, "{addr:?}@{slot}"), + Strategy::SoladyMapping => write!(f, "SoladyMapping({addr:?})"), }; let mut entries = self.0.iter(); @@ -121,7 +122,7 @@ impl FromStr for TokenConfiguration { .context("expected {addr}@{slot} format")?; Ok(( addr.parse()?, - Strategy::Mapping { + Strategy::SolidityMapping { slot: slot.parse()?, }, )) @@ -151,7 +152,7 @@ pub struct BalanceOverrideRequest { } /// Balance override strategy for a token. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum Strategy { /// Balance override strategy for tokens whose balances are stored in a /// direct Solidity mapping from token holder to balance amount in the @@ -160,29 +161,40 @@ pub enum Strategy { /// The strategy is configured with the storage slot [^1] of the mapping. /// /// [^1]: - Mapping { slot: U256 }, + SolidityMapping { slot: U256 }, + /// Strategy computing storage slot for balances based on the Solady library + /// [^1]. + /// + /// [^1]: + SoladyMapping, } impl Strategy { /// Computes the storage slot and value to override for a particular token /// holder and amount. fn state_override(&self, holder: &Address, amount: &U256) -> (H256, H256) { - match self { - Self::Mapping { slot } => { - let key = { - let mut buf = [0; 64]; - buf[12..32].copy_from_slice(holder.as_fixed_bytes()); - slot.to_big_endian(&mut buf[32..64]); - H256(signing::keccak256(&buf)) - }; - let value = { - let mut buf = [0; 32]; - amount.to_big_endian(&mut buf); - H256(buf) - }; - (key, value) + let key = match self { + Self::SolidityMapping { slot } => { + let mut buf = [0; 64]; + buf[12..32].copy_from_slice(holder.as_fixed_bytes()); + slot.to_big_endian(&mut buf[32..64]); + H256(signing::keccak256(&buf)) } - } + Self::SoladyMapping => { + let mut buf = [0; 32]; + buf[0..20].copy_from_slice(holder.as_fixed_bytes()); + buf[28..32].copy_from_slice(&[0x87, 0xa2, 0x11, 0xa2]); + H256(signing::keccak256(&buf)) + } + }; + + let value = { + let mut buf = [0; 32]; + amount.to_big_endian(&mut buf); + H256(buf) + }; + + (key, value) } } @@ -264,7 +276,7 @@ mod tests { async fn balance_override_computation() { let balance_overrides = BalanceOverrides { hardcoded: hashmap! { - addr!("DEf1CA1fb7FBcDC777520aa7f396b4E015F497aB") => Strategy::Mapping { + addr!("DEf1CA1fb7FBcDC777520aa7f396b4E015F497aB") => Strategy::SolidityMapping { slot: U256::from(0), }, }, @@ -290,7 +302,7 @@ mod tests { // You can verify the state override computation is correct by running: // ``` - // curl -X POST $RPC -H 'Content-Type: application/data' --data '{ + // curl -X POST $RPC -H 'Content-Type: application/json' --data '{ // "jsonrpc": "2.0", // "id": 0, // "method": "eth_call", @@ -327,4 +339,55 @@ mod tests { None, ); } + + #[tokio::test] + async fn balance_override_computation_solady() { + let balance_overrides = BalanceOverrides { + hardcoded: hashmap! { + addr!("0000000000c5dc95539589fbd24be07c6c14eca4") => Strategy::SoladyMapping, + }, + ..Default::default() + }; + + assert_eq!( + balance_overrides + .state_override(BalanceOverrideRequest { + token: addr!("0000000000c5dc95539589fbd24be07c6c14eca4"), + holder: addr!("d8dA6BF26964aF9D7eEd9e03E53415D37aA96045"), + amount: 0x42_u64.into(), + }) + .await, + Some(StateOverride { + state_diff: Some(hashmap! { + H256(hex!("f6a6656ed2d14bad3cdd3e8871db3f535a136a1b6cd5ae2dced8eb813f3d4e4f")) => + H256(hex!("0000000000000000000000000000000000000000000000000000000000000042")), + }), + ..Default::default() + }), + ); + + // You can verify the state override computation is correct by running: + // ``` + // curl -X POST $RPC -H 'Content-Type: application/json' --data '{ + // "jsonrpc": "2.0", + // "id": 0, + // "method": "eth_call", + // "params": [ + // { + // "to": "0x0000000000c5dc95539589fbd24be07c6c14eca4", + // "data": "0x70a08231000000000000000000000000d8dA6BF26964aF9D7eEd9e03E53415D37aA96045" + // }, + // "latest", + // { + // "0x0000000000c5dc95539589fbd24be07c6c14eca4": { + // "stateDiff": { + // "f6a6656ed2d14bad3cdd3e8871db3f535a136a1b6cd5ae2dced8eb813f3d4e4f": + // "0x0000000000000000000000000000000000000000000000000000000000000042" + // } + // } + // } + // ] + // }' + // ``` + } } 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 d8795b7ce4..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,12 +2,13 @@ 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, + sync::{Arc, LazyLock}, }, thiserror::Error, web3::{signing::keccak256, types::CallRequest}, @@ -22,9 +23,6 @@ pub struct Detector { } 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 } @@ -34,49 +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 { - // 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::Mapping { - slot: U256::from(i), - }; - // 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. - let amount = U256::from(u64::from_be_bytes([i; 8])); - - (strategy, amount) - }); - - // 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() }, }; @@ -86,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") @@ -111,3 +149,42 @@ pub enum DetectionError { #[error(transparent)] Simulation(#[from] SimulationError), } + +#[cfg(test)] +mod tests { + use {super::*, ethrpc::create_env_test_transport, web3::Web3}; + + /// Tests that we can detect storage slots by probing the first + /// n slots or by checking hardcoded known slots. + /// Set `NODE_URL` environment to a mainnet RPC URL. + #[ignore] + #[tokio::test] + async fn detects_storage_slots() { + let detector = Detector { + simulator: Arc::new(Web3::new(create_env_test_transport())), + }; + + let storage = detector + .detect(addr!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2")) + .await + .unwrap(); + assert_eq!(storage, Strategy::SolidityMapping { slot: 3.into() }); + + let storage = detector + .detect(addr!("4956b52ae2ff65d74ca2d61207523288e4528f96")) + .await + .unwrap(); + assert_eq!( + storage, + Strategy::SolidityMapping { + slot: U256::from(OPEN_ZEPPELIN_ERC20_UPGRADEABLE), + } + ); + + let storage = detector + .detect(addr!("0000000000c5dc95539589fbd24be07c6c14eca4")) + .await + .unwrap(); + assert_eq!(storage, Strategy::SoladyMapping); + } +}