From a88bfbd4cfcb3981188bcd3090cc7f66607ffdcc Mon Sep 17 00:00:00 2001 From: Nicholas Rodrigues Lordello Date: Fri, 6 Dec 2024 10:02:28 +0100 Subject: [PATCH] Add A Balance Override Detector (#3148) # Description This PR adds a balance override detector for quote verification. This allows tokens to automatically support balance overrides without manual configurations. # Changes A new (but as of yet) unused component is added that can detect balance override strategies for tokens. The current implementation is quite simple - just try a bunch of different slots and see which one works. For now, the number of slots to try is hard-coded at 25, but this can trivially be changed to be configurable in the future. The component is marked with `#[allow(dead_code)]` as it isn't being used anywhere, In the interest of keeping the PR size down, I will add all of the glue code to make it used in a follow up (which will also test whether or not the detector works as expected). ## How to Test In the meantime, I did some local testing to check that it works with tokens for which I know the mapping slot on Ethereum Mainnet (test not included in the PR): ```rust async fn temporary_test() { let web3 = Arc::new(ethrpc::Web3::new(ethrpc::create_test_transport( "https://eth.llamarpc.com", ))); let detector = Detector::new(web3); let result = detector .detect(addr!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")) .await; println!("{result:?}"); assert!(matches!(result, Ok(Strategy::Mapping { slot }) if slot == U256::from(9_u8))); } ``` This test successfully passes! --- .../trade_verifier/balance_overrides.rs | 63 ++++++----- .../balance_overrides/detector.rs | 102 ++++++++++++++++++ 2 files changed, 140 insertions(+), 25 deletions(-) create mode 100644 crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs 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 65e04f0665..ae7550f9a2 100644 --- a/crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs +++ b/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); +/// 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]: 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) -> Self { Self(config) @@ -47,20 +80,11 @@ impl ConfigurationBalanceOverrides { impl BalanceOverriding for ConfigurationBalanceOverrides { fn state_override(&self, request: &BalanceOverrideRequest) -> Option { 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 . -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}; 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 new file mode 100644 index 0000000000..0ce48951e3 --- /dev/null +++ b/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, +} + +#[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) -> 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 { + // 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), +}