Support faking balances for more tokens

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]: <https://docs.soliditylang.org/en/latest/internals/layout_in_storage.html#mappings-and-dynamic-arrays>
-    Mapping { slot: U256 },
+    SolidityMapping { slot: U256 },
+    /// Strategy computing storage slot for balances based on the Solady library
+    /// [^1].
+    ///
+    /// [^1]: <https://github.com/Vectorized/solady/blob/6122858a3aed96ee9493b99f70a245237681a95f/src/tokens/ERC20.sol#L75-L81>
+    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),
                 },
             },

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<dyn CodeSimulating>) -> 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<Strategy, DetectionError> {
-        // 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]: <https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable/blob/master/contracts/token/ERC20/ERC20Upgradeable.sol#L43-L44>
+static OPEN_ZEPPELIN_ERC20_UPGRADEABLE: &str =
+    "52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00";
+
+/// Address which we try to override the balances for.
+static HOLDER: LazyLock<Address> = 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<Vec<StrategyHelper>> = 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<HashMap<H256, H256>> = LazyLock::new(|| {
+    TESTED_STRATEGIES
+        .iter()
+        .map(|helper| helper.strategy.state_override(&HOLDER, &helper.balance))
+        .collect::<HashMap<_, _>>()
+});
+
 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);
+    }
+}