From 9f8d153b29513f1f040c2df280659c06811ca6e1 Mon Sep 17 00:00:00 2001 From: Juliya Smith Date: Wed, 11 Dec 2024 15:50:17 -0600 Subject: [PATCH] feat: include proxy ABIs --- src/ape/managers/chain.py | 43 ++++++++++++++++++++++++----- src/ape_ethereum/ecosystem.py | 10 +++---- src/ape_ethereum/proxies.py | 11 +++++++- tests/functional/geth/test_proxy.py | 32 +++++++++++++++++---- 4 files changed, 77 insertions(+), 19 deletions(-) diff --git a/src/ape/managers/chain.py b/src/ape/managers/chain.py index be87fb5ed4..e0b383ca52 100644 --- a/src/ape/managers/chain.py +++ b/src/ape/managers/chain.py @@ -798,26 +798,48 @@ def cache_deployment(self, contract_instance: ContractInstance): """ address = contract_instance.address - contract_type = contract_instance.contract_type + contract_type = contract_instance.contract_type # may be a proxy # Cache contract type in memory before proxy check, # in case it is needed somewhere. It may get overridden. self._local_contract_types[address] = contract_type - proxy_info = self.provider.network.ecosystem.get_proxy_info(address) - if proxy_info: + if proxy_info := self.provider.network.ecosystem.get_proxy_info(address): + # The user is caching a deployment of a proxy with the target already set. self.cache_proxy_info(address, proxy_info) - contract_type = self.get(proxy_info.target) or contract_type - if contract_type: + if implementation_contract := self.get(proxy_info.target): + # Include proxy ABIs but ignore fallback/ctor etc. + proxy_abis = [ + abi for abi in contract_type.abi if abi.type in ("error", "event", "function") + ] + implementation_contract.abi.extend(proxy_abis) + + # Include "hidden" ABIs, such as Safe's `masterCopy()`. + if proxy_info.abi and proxy_info.abi.signature not in [ + abi.signature for abi in contract_type.abi + ]: + implementation_contract.abi.append(proxy_info.abi) + + self._cache_contract_type(address, implementation_contract) + + # Use this contract type in the user's contract instance. + contract_instance.contract_type = implementation_contract + + else: + # No implementation yet. Just cache proxy. self._cache_contract_type(address, contract_type) - return + else: + # Regular contract. Cache normally. + self._cache_contract_type(address, contract_type) + # Cache the deployment now. txn_hash = contract_instance.txn_hash - self._cache_contract_type(address, contract_type) if contract_type.name: self._cache_deployment(address, contract_type, txn_hash) + return contract_type + def cache_proxy_info(self, address: AddressType, proxy_info: ProxyInfoAPI): """ Cache proxy info for a particular address, useful for plugins adding already @@ -1261,6 +1283,12 @@ def instance_at( txn_hash = deployment["transaction_hash"] break + # Include proxy-related ABIs. + if info := self.get_proxy_info(contract_address): + if proxy_abi := info.abi: + if proxy_abi.signature not in [x.signature for x in contract_type.abi]: + contract_type.abi.append(proxy_abi) + return ContractInstance(contract_address, contract_type, txn_hash=txn_hash) def instance_from_receipt( @@ -1271,6 +1299,7 @@ def instance_from_receipt( Args: receipt (:class:`~ape.api.transactions.ReceiptAPI`): The receipt. + contract_type (ContractType): The deployed contract type. Returns: :class:`~ape.contracts.base.ContractInstance` diff --git a/src/ape_ethereum/ecosystem.py b/src/ape_ethereum/ecosystem.py index 50bf9cee65..e0b20c05ac 100644 --- a/src/ape_ethereum/ecosystem.py +++ b/src/ape_ethereum/ecosystem.py @@ -462,8 +462,7 @@ def get_proxy_info(self, address: AddressType) -> Optional[ProxyInfo]: if isinstance(contract_code, bytes): contract_code = to_hex(contract_code) - code = contract_code[2:] - if not code: + if not (code := contract_code[2:]): return None patterns = { @@ -515,7 +514,7 @@ def str_to_slot(text): if _type == ProxyType.Beacon: target = ContractCall(IMPLEMENTATION_ABI, target)(skip_trace=True) - return ProxyInfo(type=_type, target=target) + return ProxyInfo(type=_type, target=target, abi=MASTER_COPY_ABI) # safe >=1.1.0 provides `masterCopy()`, which is also stored in slot 0 # call it and check that target matches @@ -525,7 +524,8 @@ def str_to_slot(text): target = self.conversion_manager.convert(slot_0[-20:], AddressType) # NOTE: `target` is set in initialized proxies if target != ZERO_ADDRESS and target == singleton: - return ProxyInfo(type=ProxyType.GnosisSafe, target=target) + return ProxyInfo(type=ProxyType.GnosisSafe, target=target, abi=MASTER_COPY_ABI) + except ApeException: pass @@ -541,7 +541,7 @@ def str_to_slot(text): target = ContractCall(IMPLEMENTATION_ABI, address)(skip_trace=True) # avoid recursion if target != ZERO_ADDRESS: - return ProxyInfo(type=ProxyType.Delegate, target=target) + return ProxyInfo(type=ProxyType.Delegate, target=target, abi=IMPLEMENTATION_ABI) except (ApeException, ValueError): pass diff --git a/src/ape_ethereum/proxies.py b/src/ape_ethereum/proxies.py index 387b7828db..d3205e05a8 100644 --- a/src/ape_ethereum/proxies.py +++ b/src/ape_ethereum/proxies.py @@ -1,5 +1,5 @@ from enum import IntEnum, auto -from typing import cast +from typing import Optional, cast from eth_pydantic_types.hex import HexStr from ethpm_types import ContractType, MethodABI @@ -69,6 +69,15 @@ class ProxyType(IntEnum): class ProxyInfo(ProxyInfoAPI): type: ProxyType + def __init__(self, **kwargs): + abi = kwargs.pop("abi", None) + super().__init__(**kwargs) + self._abi = abi + + @property + def abi(self) -> Optional[MethodABI]: + return self._abi + MASTER_COPY_ABI = MethodABI( type="function", diff --git a/tests/functional/geth/test_proxy.py b/tests/functional/geth/test_proxy.py index 4d39b186ae..532342365b 100644 --- a/tests/functional/geth/test_proxy.py +++ b/tests/functional/geth/test_proxy.py @@ -57,20 +57,40 @@ def test_uups_proxy(get_contract_type, geth_contract, owner, ethereum): @geth_process_test -def test_gnosis_safe(get_contract_type, geth_contract, owner, ethereum): +def test_gnosis_safe(get_contract_type, geth_contract, owner, ethereum, chain): + # Setup a proxy contract. _type = get_contract_type("SafeProxy") contract = ContractContainer(_type) - target = geth_contract.address + proxy_instance = owner.deploy(contract, target) - contract_instance = owner.deploy(contract, target) - - actual = ethereum.get_proxy_info(contract_instance.address) - + # (test) + actual = ethereum.get_proxy_info(proxy_instance.address) assert actual is not None assert actual.type == ProxyType.GnosisSafe assert actual.target == target + # Ensure we can call the proxy-method. + assert proxy_instance.masterCopy() + + # Ensure we can target methods. + assert isinstance(proxy_instance.myNumber(), int) + + # Ensure this works with new instances. + proxy_instance_ref_2 = chain.contracts.instance_at(proxy_instance.address) + assert proxy_instance_ref_2.masterCopy() + assert isinstance(proxy_instance_ref_2.myNumber(), int) + + # Same - but clear the proxy ABI from the cached type. + chain.contracts._local_contract_types[proxy_instance.address].abi = [ + x + for x in chain.contracts._local_contract_types[proxy_instance.address].abi + if x.type == "function" and x.name != "masterCopy" + ] + proxy_instance_ref_3 = chain.contracts.instance_at(proxy_instance.address) + assert proxy_instance_ref_3.masterCopy() + assert isinstance(proxy_instance_ref_3.myNumber(), int) + @geth_process_test def test_openzeppelin(get_contract_type, geth_contract, owner, ethereum, sender):