Add A Balance Override Detector

crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs

@@ -1,3 +1,5 @@
+mod detector;
+
 use {
     anyhow::Context as _,
     ethcontract::{Address, H256, U256},
@@ -33,11 +35,42 @@ pub struct BalanceOverrideRequest {
 #[derive(Clone, Debug, Default)]
 pub struct ConfigurationBalanceOverrides(HashMap<Address, Strategy>);
 
+/// Balance override strategy for a token.
 #[derive(Clone, Debug)]
 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
+    /// form `mapping(address holder => uint256 amount)`.
+    ///
+    /// 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 },
 }
 
+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)
+            }
+        }
+    }
+}
+
 impl ConfigurationBalanceOverrides {
     pub fn new(config: HashMap<Address, Strategy>) -> Self {
         Self(config)
@@ -47,20 +80,11 @@ impl ConfigurationBalanceOverrides {
 impl BalanceOverriding for ConfigurationBalanceOverrides {
     fn state_override(&self, request: &BalanceOverrideRequest) -> Option<StateOverride> {
         let strategy = self.0.get(&request.token)?;
-        match strategy {
-            Strategy::Mapping { slot } => {
-                let slot = address_mapping_storage_slot(slot, &request.holder);
-                let value = {
-                    let mut value = H256::default();
-                    request.amount.to_big_endian(&mut value.0);
-                    value
-                };
-                Some(StateOverride {
-                    state_diff: Some(hashmap! { slot => value }),
-                    ..Default::default()
-                })
-            }
-        }
+        let (key, value) = strategy.state_override(&request.holder, &request.amount);
+        Some(StateOverride {
+            state_diff: Some(hashmap! { key => value }),
+            ..Default::default()
+        })
     }
 }
 
@@ -113,17 +137,6 @@ impl FromStr for ConfigurationBalanceOverrides {
     }
 }
 
-/// Computes the storage slot where the value is stored for Solidity mappings
-/// of the form `mapping(address => ...)`.
-///
-/// See <https://docs.soliditylang.org/en/latest/internals/layout_in_storage.html#mappings-and-dynamic-arrays>.
-fn address_mapping_storage_slot(slot: &U256, address: &Address) -> H256 {
-    let mut buf = [0; 64];
-    buf[12..32].copy_from_slice(address.as_fixed_bytes());
-    slot.to_big_endian(&mut buf[32..64]);
-    H256(signing::keccak256(&buf))
-}
-
 #[cfg(test)]
 mod tests {
     use {super::*, hex_literal::hex};

crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs

@@ -0,0 +1,102 @@
+use {
+    super::Strategy,
+    crate::code_simulation::{CodeSimulating, SimulationError},
+    contracts::{dummy_contract, ERC20},
+    ethcontract::{Address, U256},
+    ethrpc::extensions::StateOverride,
+    maplit::hashmap,
+    std::sync::Arc,
+    thiserror::Error,
+    web3::{signing::keccak256, types::CallRequest},
+};
+
+/// A heuristic balance override detector based on `eth_call` simulations.
+///
+/// This has the exact same node requirements as trade verification.
+pub struct Detector {
+    simulator: Arc<dyn CodeSimulating>,
+}
+
+#[allow(dead_code)]
+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 }
+    }
+
+    /// Tries to detect the balance override strategy for the specified token.
+    /// 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,
+            ..Default::default()
+        };
+        let overrides = hashmap! {
+            token.address() => StateOverride {
+                state_diff: Some(
+                    tries
+                        .clone()
+                        .map(|(strategy, amount)| strategy.state_override(&holder, &amount))
+                        .collect(),
+                ),
+                ..Default::default()
+            },
+        };
+
+        let output = self.simulator.simulate(call, overrides, None).await?;
+        let balance = (output.len() == 32)
+            .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)
+    }
+}
+
+/// An error detecting the balance override strategy for a token.
+#[derive(Debug, Error)]
+pub enum DetectionError {
+    #[error("could not detect a balance override strategy")]
+    NotFound,
+    #[error("unable to decode simulation return data")]
+    Decode,
+    #[error(transparent)]
+    Simulation(#[from] SimulationError),
+}