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

@@ -7,7 +7,7 @@ use {
     maplit::hashmap,
     std::{
         fmt::{self, Debug, Formatter},
-        sync::Arc,
+        sync::{Arc, LazyLock},
     },
     thiserror::Error,
     web3::{signing::keccak256, types::CallRequest},
@@ -21,6 +21,13 @@ pub struct Detector {
     simulator: Arc<dyn CodeSimulating>,
 }
 
+/// 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: LazyLock<U256> = LazyLock::new(|| {
+    U256::from("52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00")
+});
+
 impl Detector {
     /// Number of different slots to try out.
     const TRIES: u8 = 25;
@@ -34,21 +41,29 @@ 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> {
+        // 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::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)
-        });
+            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
@@ -111,3 +126,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: *OPEN_ZEPPELIN_ERC20_UPGRADEABLE
+            }
+        );
+
+        let storage = detector
+            .detect(addr!("0000000000c5dc95539589fbd24be07c6c14eca4"))
+            .await
+            .unwrap();
+        assert_eq!(storage, Strategy::SoladyMapping);
+    }
+}