Skip to content

Commit

Permalink
Use LazyLock for all static values
Browse files Browse the repository at this point in the history
  • Loading branch information
@MartinquaXD
MartinquaXD committed Jan 16, 2025
1 parent c17c509 commit ea1f22a
Showing 1 changed file with 83 additions and 60 deletions.
143 changes: 83 additions & 60 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,10 +2,11 @@ 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, LazyLock},
},
Expand All @@ -21,17 +22,7 @@ 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;

/// Creates a new balance override detector.
pub fn new(simulator: Arc<dyn CodeSimulating>) -> Self {
Self { simulator }
Expand All @@ -41,57 +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> {
// 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::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
// 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 @@ -101,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 Down Expand Up @@ -154,7 +177,7 @@ mod tests {
assert_eq!(
storage,
Strategy::SolidityMapping {
slot: *OPEN_ZEPPELIN_ERC20_UPGRADEABLE
slot: U256::from(OPEN_ZEPPELIN_ERC20_UPGRADEABLE),
}
);

Expand Down

0 comments on commit ea1f22a

Please sign in to comment.