Skip to content

Commit

Permalink
Support faking balances for more tokens (#3238)
Browse files Browse the repository at this point in the history 
# Description
To verify more quotes we introduced some logic to figure out which
storage slots need to be overwritten to fake a token balance for a given
owner. That way our quote verification simulation will not revert when a
user currently doesn't hold the necessary amounts of token which
improves UX a lot.

While looking into some quote we couldn't verify
([simulation](https://dashboard.tenderly.co/cow-protocol/production/simulator/e3456eaf-90e6-4212-9758-deef4b80948f?trace=0.3.0.3.0.0.0.5.0.0.1.0.0.0.0))
it became apparent that our current strategy to detect which storage
slot needs to get overwritten doesn't work for some tokens.
The current logic assumes a "normal" contract where member variables are
laid out using the default solidity logic.
This is not sufficient to detect proxy contracts which store their state
at a hardcoded address since we can't predict where the author decided
to put their member variables in their custom memory layout.

However, most of the time people use libraries to handle that logic for
them. These libraries hardcode the necessary information to find the
member variable storing the balances in the token so we can just add
support by probing for these known hardcoded values too.

Kudos to @fedgiac for helping me debug the problem with the original
proxy contract.

# Changes
This PR adds support for 2 libraries: OpenZeppelin and Solady.
The OpenZeppeline implementation simply puts the `balances` member
variable at a hardcoded position so we can use the existing encoding
logic with the hardcoded address we found the in their source code.

Solady is a bit different since it doesn't hash 64 bytes but instead
only hashes 48 bytes to compute the storage slot of an addresses
balance. For this I introduced another enum variant.

This PR also fixes a bug with the original implementation. The original
implementation works by calling `token.balanceOf(address)` with 25
storage slots. Each slot gets overwritten with a specific value. When
`balanceOf()` returns that value we recover which storage slot we
originally overwrote.
This idea makes a lot of sense but had a problem with the very first
override. The amount we override was computed with
`U256::from(u64::from_be_bytes([i; 8]))` where `i` is an index variable.
That means the first override always tried to set the balance to `0`. If
you call `token.balanceOf(random_address)` it will return 0 because the
address doesn't have any balances. The original logic interpreted this
`0` as the sentinel `0` value it originally overwrote and therefore
returned that the `balances` mapping lives in the `0th` storage slot.
To fix that we now compute the sentinel values with
`U256::from(u64::from_be_bytes([i + 1; 8]))` to ensure that we always
search for non-zero values to detect the correct storage override.
Note that this error only surfaced when none of the storage slots we
tried to override was the correct one. If we were able to override the
storage slot `balanceOf()` would no longer return `0` to cause the false
positive.

## How to test
added an ignored unit test detecting 1 storage slot of each variant:
1. standard solidity mapping + standard solidity layout
2. standard solidity mapping + hardcoded OpenZeppelin storage slot
3. custom Solady mapping

Also there was already 1 e2e test that already tested the original
logic.
  • Loading branch information
@MartinquaXD
MartinquaXD authored Jan 16, 2025
1 parent 50f7c9a commit 395d605
Show file tree
Hide file tree
Showing 2 changed files with 206 additions and 66 deletions.
105 changes: 84 additions & 21 deletions crates/shared/src/price_estimation/trade_verifier/balance_overrides.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -121,7 +122,7 @@ impl FromStr for TokenConfiguration {
.context("expected {addr}@{slot} format")?;
Ok((
addr.parse()?,
Strategy::Mapping {
Strategy::SolidityMapping {
slot: slot.parse()?,
},
))
Expand Down Expand Up @@ -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
Expand All @@ -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)
}
}

Expand Down Expand Up @@ -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),
},
},
Expand All @@ -290,7 +302,7 @@ mod tests {

// You can verify the state override computation is correct by running:
// ```
// curl -X POST $RPC -H 'Content-Type: application/data' --data '{
// curl -X POST $RPC -H 'Content-Type: application/json' --data '{
// "jsonrpc": "2.0",
// "id": 0,
// "method": "eth_call",
Expand Down Expand Up @@ -327,4 +339,55 @@ mod tests {
None,
);
}

#[tokio::test]
async fn balance_override_computation_solady() {
let balance_overrides = BalanceOverrides {
hardcoded: hashmap! {
addr!("0000000000c5dc95539589fbd24be07c6c14eca4") => Strategy::SoladyMapping,
},
..Default::default()
};

assert_eq!(
balance_overrides
.state_override(BalanceOverrideRequest {
token: addr!("0000000000c5dc95539589fbd24be07c6c14eca4"),
holder: addr!("d8dA6BF26964aF9D7eEd9e03E53415D37aA96045"),
amount: 0x42_u64.into(),
})
.await,
Some(StateOverride {
state_diff: Some(hashmap! {
H256(hex!("f6a6656ed2d14bad3cdd3e8871db3f535a136a1b6cd5ae2dced8eb813f3d4e4f")) =>
H256(hex!("0000000000000000000000000000000000000000000000000000000000000042")),
}),
..Default::default()
}),
);

// You can verify the state override computation is correct by running:
// ```
// curl -X POST $RPC -H 'Content-Type: application/json' --data '{
// "jsonrpc": "2.0",
// "id": 0,
// "method": "eth_call",
// "params": [
// {
// "to": "0x0000000000c5dc95539589fbd24be07c6c14eca4",
// "data": "0x70a08231000000000000000000000000d8dA6BF26964aF9D7eEd9e03E53415D37aA96045"
// },
// "latest",
// {
// "0x0000000000c5dc95539589fbd24be07c6c14eca4": {
// "stateDiff": {
// "f6a6656ed2d14bad3cdd3e8871db3f535a136a1b6cd5ae2dced8eb813f3d4e4f":
// "0x0000000000000000000000000000000000000000000000000000000000000042"
// }
// }
// }
// ]
// }'
// ```
}
}
167 changes: 122 additions & 45 deletions crates/shared/src/price_estimation/trade_verifier/balance_overrides/detector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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 }
Expand All @@ -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()
},
};
Expand All @@ -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")
Expand All @@ -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);
}
}

0 comments on commit 395d605

Please sign in to comment.