From 71519c11278aadcb1714b664229b0c0ea7080b59 Mon Sep 17 00:00:00 2001 From: zizou <111426680+flopell@users.noreply.github.com> Date: Wed, 6 Nov 2024 17:13:09 +0700 Subject: [PATCH 1/2] feat(simulation-py): make token bruteforce compatible with vyper --- .../tycho_simulation_py/evm/pool_state.py | 16 ++-- .../python/tycho_simulation_py/evm/token.py | 87 ++++++++++--------- .../python/tycho_simulation_py/evm/utils.py | 18 ++-- 3 files changed, 69 insertions(+), 52 deletions(-) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py b/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py index 59108933..ddc15c7b 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py @@ -19,6 +19,7 @@ from ..exceptions import RecoverableSimulationException from ..models import EVMBlock, Capability, Address, EthereumToken from .utils import ( + ContractCompiler, create_engine, get_contract_bytecode, frac_to_decimal, @@ -97,8 +98,8 @@ def __init__( self.involved_contracts: set[Address] = involved_contracts or set() """A set of all contract addresses involved in the simulation of this pool.""" - self.token_storage_slots: dict[Address, tuple[int, int]] = ( - token_storage_slots or {} + self.token_storage_slots: dict[Address, tuple[tuple[int, int], ContractCompiler]] = ( + token_storage_slots or {} ) """Allows the specification of custom storage slots for token allowances and balances. This is particularly useful for token contracts involved in protocol @@ -177,7 +178,7 @@ def _set_marginal_prices(self): t1, [sell_amount], block=self.block, - overwrites=self.block_lasting_overwrites, + overwrites=self._get_overwrites(t0,t1), )[0] if Capability.ScaledPrices in self.capabilities: self.marginal_prices[(t0, t1)] = frac_to_decimal(frac) @@ -298,9 +299,11 @@ def _get_token_overwrites( max_amount = sell_token.to_onchain_amount( self.get_sell_amount_limit(sell_token, buy_token) ) + slots, compiler = self.token_storage_slots.get(sell_token.address, ((0, 1), ContractCompiler.Solidity)) overwrites = ERC20OverwriteFactory( sell_token, - token_slots=self.token_storage_slots.get(sell_token.address, (0, 1)), + token_slots=slots, + compiler=compiler ) overwrites.set_balance(max_amount, EXTERNAL_ACCOUNT) overwrites.set_allowance( @@ -318,9 +321,10 @@ def _get_balance_overwrites(self) -> dict[Address, dict[int, int]]: address = self.balance_owner or self.id_ for t in self.tokens: slots = (0, 1) + compiler = ContractCompiler.Solidity if t.address in self.involved_contracts: - slots = self.token_storage_slots.get(t.address) - overwrites = ERC20OverwriteFactory(t, token_slots=slots) + slots, compiler = self.token_storage_slots.get(t.address) + overwrites = ERC20OverwriteFactory(t, token_slots=slots, compiler=compiler) overwrites.set_balance( t.to_onchain_amount(self.balances[t.address]), address ) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/token.py b/tycho_simulation_py/python/tycho_simulation_py/evm/token.py index a52ed97e..b6c71f40 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/token.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/token.py @@ -1,5 +1,5 @@ from .adapter_contract import TychoSimulationContract -from .utils import ERC20OverwriteFactory +from .utils import ContractCompiler, ERC20OverwriteFactory from .constants import EXTERNAL_ACCOUNT from . import SimulationEngine from ..models import EVMBlock, EthereumToken @@ -14,7 +14,7 @@ class SlotDetectionFailure(Exception): def brute_force_slots( t: EthereumToken, block: EVMBlock, engine: SimulationEngine -) -> tuple[int, int]: +) -> tuple[tuple[int, int], ContractCompiler]: """Brute-force detection of storage slots for token allowances and balances. This function attempts to determine the storage slots used by the token contract for @@ -37,9 +37,9 @@ def brute_force_slots( Returns ------- - tuple[int, int] - A tuple containing the detected balance storage slot and the allowance - storage slot, respectively. + tuple[tuple[int, int], ContractCompiler] + A tuple containing a tuple containing the detected balance storage slot and the allowance + storage slot, respectively and in what compiler was used for this contract. Raises ------ @@ -49,47 +49,52 @@ def brute_force_slots( """ token_contract = TychoSimulationContract(t.address, "ERC20", engine) balance_slot = None - for i in range(20): - overwrite_factory = ERC20OverwriteFactory(t, (i, 1)) - overwrite_factory.set_balance(_MARKER_VALUE, EXTERNAL_ACCOUNT) - res = token_contract.call( - "balanceOf", - [EXTERNAL_ACCOUNT], - block_number=block.id, - timestamp=int(block.ts.timestamp()), - overrides=overwrite_factory.get_tycho_overwrites(), - caller=EXTERNAL_ACCOUNT, - value=0, - ) - if res.return_value is None: - continue - if res.return_value[0] == _MARKER_VALUE: - balance_slot = i - break + compiler = ContractCompiler.Solidity + for i in range(100): + for compiler_flag in [ContractCompiler.Solidity, ContractCompiler.Vyper]: + overwrite_factory = ERC20OverwriteFactory(t, (i, 1), compiler=compiler_flag) + overwrite_factory.set_balance(_MARKER_VALUE, EXTERNAL_ACCOUNT) + res = token_contract.call( + "balanceOf", + [EXTERNAL_ACCOUNT], + block_number=block.id, + timestamp=int(block.ts.timestamp()), + overrides=overwrite_factory.get_tycho_overwrites(), + caller=EXTERNAL_ACCOUNT, + value=0, + ) - allowance_slot = None - for i in range(20): - overwrite_factory = ERC20OverwriteFactory(t, (0, i)) - overwrite_factory.set_allowance(_MARKER_VALUE, _SPENDER, EXTERNAL_ACCOUNT) - res = token_contract.call( - "allowance", - [EXTERNAL_ACCOUNT, _SPENDER], - block_number=block.id, - timestamp=int(block.ts.timestamp()), - overrides=overwrite_factory.get_tycho_overwrites(), - caller=EXTERNAL_ACCOUNT, - value=0, - ) - if res.return_value is None: - continue - if res.return_value[0] == _MARKER_VALUE: - allowance_slot = i - break + if res.return_value is None: + continue + if res.return_value[0] == _MARKER_VALUE: + balance_slot = i + compiler = compiler_flag + break if balance_slot is None: raise SlotDetectionFailure(f"Failed to infer balance slot for {t.address}") + allowance_slot = None + for i in range(100): + overwrite_factory = ERC20OverwriteFactory(t, (0, i), compiler=compiler) + overwrite_factory.set_allowance(_MARKER_VALUE, _SPENDER, EXTERNAL_ACCOUNT) + res = token_contract.call( + "allowance", + [EXTERNAL_ACCOUNT, _SPENDER], + block_number=block.id, + timestamp=int(block.ts.timestamp()), + overrides=overwrite_factory.get_tycho_overwrites(), + caller=EXTERNAL_ACCOUNT, + value=0, + ) + if res.return_value is None: + continue + if res.return_value[0] == _MARKER_VALUE: + allowance_slot = i + break + + if allowance_slot is None: raise SlotDetectionFailure(f"Failed to infer allowance slot for {t.address}") - return balance_slot, allowance_slot + return ((balance_slot, allowance_slot), compiler) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py b/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py index 8ffc8bac..6156d5f5 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py @@ -1,3 +1,4 @@ +import enum import json import os from decimal import Decimal @@ -61,9 +62,13 @@ def create_engine( return engine +class ContractCompiler(enum.Enum): + Solidity = enum.auto() + Vyper = enum.auto() + class ERC20OverwriteFactory: - def __init__(self, token: EthereumToken, token_slots=(0, 1)): + def __init__(self, token: EthereumToken, token_slots=(0, 1), compiler: ContractCompiler = ContractCompiler.Solidity): """ Initialize the ERC20OverwriteFactory. @@ -72,6 +77,7 @@ def __init__(self, token: EthereumToken, token_slots=(0, 1)): """ self._token = token self._overwrites = dict() + self._contract_compiler = compiler self._balance_slot: int = token_slots[0] self._allowance_slot: int = token_slots[1] self._total_supply_slot: Final[int] = 2 @@ -84,7 +90,7 @@ def set_balance(self, balance: int, owner: Address): balance: The balance value. owner: The owner's address. """ - storage_index = get_storage_slot_at_key(HexStr(owner), self._balance_slot) + storage_index = get_storage_slot_at_key(HexStr(owner), self._balance_slot, self._contract_compiler) self._overwrites[storage_index] = balance log.log( 5, @@ -103,8 +109,8 @@ def set_allowance(self, allowance: int, spender: Address, owner: Address): """ storage_index = get_storage_slot_at_key( HexStr(spender), - get_storage_slot_at_key(HexStr(owner), self._allowance_slot), - ) + get_storage_slot_at_key(HexStr(owner), self._allowance_slot, self._contract_compiler), + self._contract_compiler) self._overwrites[storage_index] = allowance log.log( 5, @@ -153,7 +159,7 @@ def get_geth_overwrites(self) -> dict[Address, dict[int, int]]: return {self._token.address: {"stateDiff": formatted_overwrites, "code": code}} -def get_storage_slot_at_key(key: Address, mapping_slot: int) -> int: +def get_storage_slot_at_key(key: Address, mapping_slot: int, lang = ContractCompiler.Solidity) -> int: """Get storage slot index of a value stored at a certain key in a mapping Parameters @@ -194,6 +200,8 @@ def get_storage_slot_at_key(key: Address, mapping_slot: int) -> int: key_bytes = bytes.fromhex(key[2:]).rjust(32, b"\0") mapping_slot_bytes = int.to_bytes(mapping_slot, 32, "big") slot_bytes = eth_utils.keccak(key_bytes + mapping_slot_bytes) + if lang == ContractCompiler.Vyper: + slot_bytes = eth_utils.keccak(mapping_slot_bytes + key_bytes) return int.from_bytes(slot_bytes, "big") From b1331c612c3fb57c78f99c187de888cd221ea48a Mon Sep 17 00:00:00 2001 From: zizou <111426680+flopell@users.noreply.github.com> Date: Thu, 7 Nov 2024 10:24:02 +0700 Subject: [PATCH 2/2] refactor(simulation_py): miscellaneous improvements Introduce a struct for ERC20 slots, docs, and more small code improvement --- .../tycho_simulation_py/evm/pool_state.py | 9 ++++-- .../python/tycho_simulation_py/evm/token.py | 10 +++--- .../python/tycho_simulation_py/evm/utils.py | 31 ++++++++++++++----- 3 files changed, 34 insertions(+), 16 deletions(-) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py b/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py index ddc15c7b..264a15cd 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/pool_state.py @@ -20,6 +20,7 @@ from ..models import EVMBlock, Capability, Address, EthereumToken from .utils import ( ContractCompiler, + ERC20Slots, create_engine, get_contract_bytecode, frac_to_decimal, @@ -98,12 +99,14 @@ def __init__( self.involved_contracts: set[Address] = involved_contracts or set() """A set of all contract addresses involved in the simulation of this pool.""" - self.token_storage_slots: dict[Address, tuple[tuple[int, int], ContractCompiler]] = ( + self.token_storage_slots: dict[Address, tuple[ERC20Slots, ContractCompiler]] = ( token_storage_slots or {} ) """Allows the specification of custom storage slots for token allowances and balances. This is particularly useful for token contracts involved in protocol logic that extends beyond simple transfer functionality. + Each entry also specify the compiler with which the target contract was compiled. + This is later used to compute storage slot for maps. """ self._engine: Optional[SimulationEngine] = None @@ -299,7 +302,7 @@ def _get_token_overwrites( max_amount = sell_token.to_onchain_amount( self.get_sell_amount_limit(sell_token, buy_token) ) - slots, compiler = self.token_storage_slots.get(sell_token.address, ((0, 1), ContractCompiler.Solidity)) + slots, compiler = self.token_storage_slots.get(sell_token.address, (ERC20Slots(0, 1), ContractCompiler.Solidity)) overwrites = ERC20OverwriteFactory( sell_token, token_slots=slots, @@ -320,7 +323,7 @@ def _get_balance_overwrites(self) -> dict[Address, dict[int, int]]: balance_overwrites = {} address = self.balance_owner or self.id_ for t in self.tokens: - slots = (0, 1) + slots = ERC20Slots(0, 1) compiler = ContractCompiler.Solidity if t.address in self.involved_contracts: slots, compiler = self.token_storage_slots.get(t.address) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/token.py b/tycho_simulation_py/python/tycho_simulation_py/evm/token.py index b6c71f40..42d99393 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/token.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/token.py @@ -1,5 +1,5 @@ from .adapter_contract import TychoSimulationContract -from .utils import ContractCompiler, ERC20OverwriteFactory +from .utils import ContractCompiler, ERC20OverwriteFactory, ERC20Slots from .constants import EXTERNAL_ACCOUNT from . import SimulationEngine from ..models import EVMBlock, EthereumToken @@ -14,7 +14,7 @@ class SlotDetectionFailure(Exception): def brute_force_slots( t: EthereumToken, block: EVMBlock, engine: SimulationEngine -) -> tuple[tuple[int, int], ContractCompiler]: +) -> tuple[ERC20Slots, ContractCompiler]: """Brute-force detection of storage slots for token allowances and balances. This function attempts to determine the storage slots used by the token contract for @@ -52,7 +52,7 @@ def brute_force_slots( compiler = ContractCompiler.Solidity for i in range(100): for compiler_flag in [ContractCompiler.Solidity, ContractCompiler.Vyper]: - overwrite_factory = ERC20OverwriteFactory(t, (i, 1), compiler=compiler_flag) + overwrite_factory = ERC20OverwriteFactory(t, ERC20Slots(i, 1), compiler=compiler_flag) overwrite_factory.set_balance(_MARKER_VALUE, EXTERNAL_ACCOUNT) res = token_contract.call( "balanceOf", @@ -76,7 +76,7 @@ def brute_force_slots( allowance_slot = None for i in range(100): - overwrite_factory = ERC20OverwriteFactory(t, (0, i), compiler=compiler) + overwrite_factory = ERC20OverwriteFactory(t, ERC20Slots(0, i), compiler=compiler) overwrite_factory.set_allowance(_MARKER_VALUE, _SPENDER, EXTERNAL_ACCOUNT) res = token_contract.call( "allowance", @@ -97,4 +97,4 @@ def brute_force_slots( if allowance_slot is None: raise SlotDetectionFailure(f"Failed to infer allowance slot for {t.address}") - return ((balance_slot, allowance_slot), compiler) + return (ERC20Slots(balance_slot, allowance_slot), compiler) diff --git a/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py b/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py index 6156d5f5..78480399 100644 --- a/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py +++ b/tycho_simulation_py/python/tycho_simulation_py/evm/utils.py @@ -5,7 +5,7 @@ from fractions import Fraction from functools import lru_cache from logging import getLogger -from typing import Final, Any +from typing import Final, Any, NamedTuple import eth_abi import eth_utils @@ -65,10 +65,23 @@ def create_engine( class ContractCompiler(enum.Enum): Solidity = enum.auto() Vyper = enum.auto() + + def compute_map_slot(self, map_base_slot: bytes, key: bytes) -> bytes: + if self == ContractCompiler.Solidity: + return eth_utils.keccak(key + map_base_slot) + elif self == ContractCompiler.Vyper: + return eth_utils.keccak(map_base_slot + key) + else: + raise NotImplementedError(f"compute_map_slot not implemented for {self.name}") + + +class ERC20Slots(NamedTuple): + balance_map: int + allowance_map: int class ERC20OverwriteFactory: - def __init__(self, token: EthereumToken, token_slots=(0, 1), compiler: ContractCompiler = ContractCompiler.Solidity): + def __init__(self, token: EthereumToken, token_slots: ERC20Slots = ERC20Slots(0, 1), compiler: ContractCompiler = ContractCompiler.Solidity): """ Initialize the ERC20OverwriteFactory. @@ -78,8 +91,8 @@ def __init__(self, token: EthereumToken, token_slots=(0, 1), compiler: ContractC self._token = token self._overwrites = dict() self._contract_compiler = compiler - self._balance_slot: int = token_slots[0] - self._allowance_slot: int = token_slots[1] + self._balance_slot: int = token_slots.balance_map + self._allowance_slot: int = token_slots.allowance_map self._total_supply_slot: Final[int] = 2 def set_balance(self, balance: int, owner: Address): @@ -159,7 +172,7 @@ def get_geth_overwrites(self) -> dict[Address, dict[int, int]]: return {self._token.address: {"stateDiff": formatted_overwrites, "code": code}} -def get_storage_slot_at_key(key: Address, mapping_slot: int, lang = ContractCompiler.Solidity) -> int: +def get_storage_slot_at_key(key: Address, mapping_slot: int, compiler = ContractCompiler.Solidity) -> int: """Get storage slot index of a value stored at a certain key in a mapping Parameters @@ -170,6 +183,10 @@ def get_storage_slot_at_key(key: Address, mapping_slot: int, lang = ContractComp mapping_slot Storage slot at which the mapping itself is stored. See the examples for more explanation. + + compiler + The compiler with which the target contract was compiled. Solidity and Vyper handle + maps differently. This defaults to Solidity because it's the most used. Returns ------- @@ -199,9 +216,7 @@ def get_storage_slot_at_key(key: Address, mapping_slot: int, lang = ContractComp """ key_bytes = bytes.fromhex(key[2:]).rjust(32, b"\0") mapping_slot_bytes = int.to_bytes(mapping_slot, 32, "big") - slot_bytes = eth_utils.keccak(key_bytes + mapping_slot_bytes) - if lang == ContractCompiler.Vyper: - slot_bytes = eth_utils.keccak(mapping_slot_bytes + key_bytes) + slot_bytes = compiler.compute_map_slot(mapping_slot_bytes, key_bytes) return int.from_bytes(slot_bytes, "big")