Skip to content

Commit

Permalink
Add A Balance Override Detector (#3148)
Browse files Browse the repository at this point in the history
# 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!
  • Loading branch information
nlordell authored Dec 6, 2024
1 parent e75a277 commit a88bfbd
Show file tree
Hide file tree
Showing 2 changed files with 140 additions and 25 deletions.
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
mod detector;

use {
anyhow::Context as _,
ethcontract::{Address, H256, U256},
Expand Down Expand Up @@ -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)
Expand All @@ -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()
})
}
}

Expand Down Expand Up @@ -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};
Expand Down
Original file line number Diff line number Diff line change
@@ -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),
}

0 comments on commit a88bfbd

Please sign in to comment.