From a5b8545bde5ca0420cf1188bc1261aac4b3964f2 Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 13:03:35 +0200 Subject: [PATCH 1/7] add set_code, get_storage, set_storage to boa.env Fixes #204 --- boa/environment.py | 15 ++++++++++++++- boa/vm/py_evm.py | 11 +++++++++++ tests/unitary/test_fork_utils.py | 32 ++++++++++++++++++++++++++++++++ 3 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/unitary/test_fork_utils.py diff --git a/boa/environment.py b/boa/environment.py index d0656274..7802db08 100644 --- a/boa/environment.py +++ b/boa/environment.py @@ -311,9 +311,22 @@ def _hook_trace_computation(self, computation, contract=None): child_contract = self._lookup_contract_fast(child.msg.code_address) self._hook_trace_computation(child, child_contract) - def get_code(self, address): + def get_code(self, address: _AddressType) -> bytes: return self.evm.get_code(Address(address)) + def set_code(self, address: _AddressType, code: bytes) -> None: + assert self.evm.is_forked, "The EVM is not forked, cannot set code" + self.evm.set_code(Address(address), code) + + def get_storage( + self, address: _AddressType, slot: int, from_journal: bool = True + ) -> int: + return self.evm.get_storage(Address(address), slot, from_journal) + + def set_storage(self, address: _AddressType, slot: int, value: int) -> None: + assert self.evm.is_forked, "The EVM is not forked, cannot set storage" + self.evm.set_storage(Address(address), slot, value) + # function to time travel def time_travel( self, diff --git a/boa/vm/py_evm.py b/boa/vm/py_evm.py index a3f1bee9..0dd07f87 100644 --- a/boa/vm/py_evm.py +++ b/boa/vm/py_evm.py @@ -420,6 +420,17 @@ def set_balance(self, address: Address, value): def get_code(self, address: Address) -> bytes: return self.vm.state.get_code(address.canonical_address) + def set_code(self, address: Address, code: bytes) -> None: + self.vm.state.set_code(address.canonical_address, code) + + def get_storage( + self, address: Address, slot: int, from_journal: bool = True + ) -> int: + return self.vm.state.get_storage(address.canonical_address, slot, from_journal) + + def set_storage(self, address: Address, slot: int, value: int) -> None: + self.vm.state.set_storage(address.canonical_address, slot, value) + def get_gas_limit(self): return self.vm.state.gas_limit diff --git a/tests/unitary/test_fork_utils.py b/tests/unitary/test_fork_utils.py new file mode 100644 index 00000000..e63002cf --- /dev/null +++ b/tests/unitary/test_fork_utils.py @@ -0,0 +1,32 @@ +import pytest + +import boa +from boa.vm.fork import AccountDBFork + +address = "0x0000000000000000000000000000000000000065" + + +def test_code_non_fork(): + with pytest.raises(AssertionError): + boa.env.set_code(address, b"") + + +def test_storage_non_fork(): + with pytest.raises(AssertionError): + boa.env.set_storage(address, 0, 0) + + +def test_code(): + boa.env.evm._set_account_db_class(AccountDBFork) + assert boa.env.get_code(address) == b"" + code = b"some test code" + boa.env.set_code(address, code) + assert boa.env.get_code(address) == code + + +def test_storage(monkeypatch): + storage = 12381920371289 + boa.env.evm._set_account_db_class(AccountDBFork) + assert boa.env.get_storage(address, 0) == 0 + boa.env.set_storage(address, 0, storage) + assert boa.env.get_storage(address, 0) == storage From 49a5411c8f450ecf6628ca279cabd9035396b320 Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 13:14:03 +0200 Subject: [PATCH 2/7] Remove `from_journal` --- boa/environment.py | 6 ++---- boa/vm/py_evm.py | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/boa/environment.py b/boa/environment.py index 7802db08..14bf058d 100644 --- a/boa/environment.py +++ b/boa/environment.py @@ -318,10 +318,8 @@ def set_code(self, address: _AddressType, code: bytes) -> None: assert self.evm.is_forked, "The EVM is not forked, cannot set code" self.evm.set_code(Address(address), code) - def get_storage( - self, address: _AddressType, slot: int, from_journal: bool = True - ) -> int: - return self.evm.get_storage(Address(address), slot, from_journal) + def get_storage(self, address: _AddressType, slot: int) -> int: + return self.evm.get_storage(Address(address), slot) def set_storage(self, address: _AddressType, slot: int, value: int) -> None: assert self.evm.is_forked, "The EVM is not forked, cannot set storage" diff --git a/boa/vm/py_evm.py b/boa/vm/py_evm.py index 0dd07f87..dc7d55fd 100644 --- a/boa/vm/py_evm.py +++ b/boa/vm/py_evm.py @@ -423,10 +423,8 @@ def get_code(self, address: Address) -> bytes: def set_code(self, address: Address, code: bytes) -> None: self.vm.state.set_code(address.canonical_address, code) - def get_storage( - self, address: Address, slot: int, from_journal: bool = True - ) -> int: - return self.vm.state.get_storage(address.canonical_address, slot, from_journal) + def get_storage(self, address: Address, slot: int) -> int: + return self.vm.state.get_storage(address.canonical_address, slot) def set_storage(self, address: Address, slot: int, value: int) -> None: self.vm.state.set_storage(address.canonical_address, slot, value) From 79ef0a62a168d76766006083ba8320a5b721ac1d Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 13:36:44 +0200 Subject: [PATCH 3/7] Move set checks to NetworkEnv --- boa/environment.py | 2 -- boa/network.py | 10 +++++++++- boa/vm/fork.py | 5 ++++- tests/unitary/test_fork_utils.py | 8 +++----- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/boa/environment.py b/boa/environment.py index 14bf058d..2802f55f 100644 --- a/boa/environment.py +++ b/boa/environment.py @@ -315,14 +315,12 @@ def get_code(self, address: _AddressType) -> bytes: return self.evm.get_code(Address(address)) def set_code(self, address: _AddressType, code: bytes) -> None: - assert self.evm.is_forked, "The EVM is not forked, cannot set code" self.evm.set_code(Address(address), code) def get_storage(self, address: _AddressType, slot: int) -> int: return self.evm.get_storage(Address(address), slot) def set_storage(self, address: _AddressType, slot: int, value: int) -> None: - assert self.evm.is_forked, "The EVM is not forked, cannot set storage" self.evm.set_storage(Address(address), slot, value) # function to time travel diff --git a/boa/network.py b/boa/network.py index 27c5b58b..e12e0c73 100644 --- a/boa/network.py +++ b/boa/network.py @@ -8,7 +8,7 @@ from eth_account import Account from requests.exceptions import HTTPError -from boa.environment import Env +from boa.environment import Env, _AddressType from boa.rpc import ( RPC, EthereumRPC, @@ -495,3 +495,11 @@ def _send_txn(self, from_, to=None, gas=None, value=None, data=None): t_obj = TraceObject(trace) if trace is not None else None return receipt, t_obj + + def set_code(self, address: _AddressType, code: bytes) -> None: + assert self.evm.is_forked, "The EVM is not forked, cannot set code" + return super().set_code(address, code) + + def set_storage(self, address: _AddressType, slot: int, value: int) -> None: + assert self.evm.is_forked, "The EVM is not forked, cannot set storage" + super().set_storage(address, slot, value) diff --git a/boa/vm/fork.py b/boa/vm/fork.py index 4cb5e6c2..a9d54ad3 100644 --- a/boa/vm/fork.py +++ b/boa/vm/fork.py @@ -120,7 +120,10 @@ def fetch_multi(self, payload): # fetch_multi is called only with the missing payloads # map the results back to the original indices for result_ix, rpc_result in enumerate(self._rpc.fetch_multi(batch)): - key, item_ix = keys[result_ix] + try: + key, item_ix = keys[result_ix] + except IndexError: + continue ret[item_ix] = rpc_result self._db[key] = json.dumps(rpc_result).encode("utf-8") diff --git a/tests/unitary/test_fork_utils.py b/tests/unitary/test_fork_utils.py index e63002cf..85f46ed0 100644 --- a/tests/unitary/test_fork_utils.py +++ b/tests/unitary/test_fork_utils.py @@ -1,23 +1,22 @@ import pytest import boa -from boa.vm.fork import AccountDBFork +from boa import NetworkEnv address = "0x0000000000000000000000000000000000000065" def test_code_non_fork(): with pytest.raises(AssertionError): - boa.env.set_code(address, b"") + NetworkEnv.set_code(boa.env, address, b"") def test_storage_non_fork(): with pytest.raises(AssertionError): - boa.env.set_storage(address, 0, 0) + NetworkEnv.set_storage(boa.env, address, 0, 0) def test_code(): - boa.env.evm._set_account_db_class(AccountDBFork) assert boa.env.get_code(address) == b"" code = b"some test code" boa.env.set_code(address, code) @@ -26,7 +25,6 @@ def test_code(): def test_storage(monkeypatch): storage = 12381920371289 - boa.env.evm._set_account_db_class(AccountDBFork) assert boa.env.get_storage(address, 0) == 0 boa.env.set_storage(address, 0, storage) assert boa.env.get_storage(address, 0) == storage From 837c037fa56640fa36955487a4f37605b4a964c3 Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 13:42:15 +0200 Subject: [PATCH 4/7] Unused fixture --- tests/unitary/test_fork_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unitary/test_fork_utils.py b/tests/unitary/test_fork_utils.py index 85f46ed0..9085db51 100644 --- a/tests/unitary/test_fork_utils.py +++ b/tests/unitary/test_fork_utils.py @@ -23,7 +23,7 @@ def test_code(): assert boa.env.get_code(address) == code -def test_storage(monkeypatch): +def test_storage(): storage = 12381920371289 assert boa.env.get_storage(address, 0) == 0 boa.env.set_storage(address, 0, storage) From eeb490bf0c28e2c5daf651019283d39c31f592e2 Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 14:42:21 +0200 Subject: [PATCH 5/7] Block set_balance in network mode --- boa/network.py | 9 +++++---- tests/unitary/test_fork_utils.py | 16 ++++++++++++++-- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/boa/network.py b/boa/network.py index e12e0c73..340fd22b 100644 --- a/boa/network.py +++ b/boa/network.py @@ -496,10 +496,11 @@ def _send_txn(self, from_, to=None, gas=None, value=None, data=None): t_obj = TraceObject(trace) if trace is not None else None return receipt, t_obj + def set_balance(self, address, value): + raise NotImplementedError("Cannot use set_balance in network mode") + def set_code(self, address: _AddressType, code: bytes) -> None: - assert self.evm.is_forked, "The EVM is not forked, cannot set code" - return super().set_code(address, code) + raise NotImplementedError("Cannot use set_code in network mode") def set_storage(self, address: _AddressType, slot: int, value: int) -> None: - assert self.evm.is_forked, "The EVM is not forked, cannot set storage" - super().set_storage(address, slot, value) + raise NotImplementedError("Cannot use set_storage in network mode") diff --git a/tests/unitary/test_fork_utils.py b/tests/unitary/test_fork_utils.py index 9085db51..fd0a2cfb 100644 --- a/tests/unitary/test_fork_utils.py +++ b/tests/unitary/test_fork_utils.py @@ -6,16 +6,28 @@ address = "0x0000000000000000000000000000000000000065" +def test_balance_non_fork(): + with pytest.raises(NotImplementedError): + NetworkEnv.set_balance(boa.env, address, 0) + + def test_code_non_fork(): - with pytest.raises(AssertionError): + with pytest.raises(NotImplementedError): NetworkEnv.set_code(boa.env, address, b"") def test_storage_non_fork(): - with pytest.raises(AssertionError): + with pytest.raises(NotImplementedError): NetworkEnv.set_storage(boa.env, address, 0, 0) +def test_balance(): + assert boa.env.get_balance(address) == 0 + balance = 1000 + boa.env.set_balance(address, balance) + assert boa.env.get_balance(address) == balance + + def test_code(): assert boa.env.get_code(address) == b"" code = b"some test code" From 570bf5d1d2e20a91616b76af3bac10053fce8b6c Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 14:43:27 +0200 Subject: [PATCH 6/7] Unintended change --- boa/vm/fork.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/boa/vm/fork.py b/boa/vm/fork.py index a9d54ad3..4cb5e6c2 100644 --- a/boa/vm/fork.py +++ b/boa/vm/fork.py @@ -120,10 +120,7 @@ def fetch_multi(self, payload): # fetch_multi is called only with the missing payloads # map the results back to the original indices for result_ix, rpc_result in enumerate(self._rpc.fetch_multi(batch)): - try: - key, item_ix = keys[result_ix] - except IndexError: - continue + key, item_ix = keys[result_ix] ret[item_ix] = rpc_result self._db[key] = json.dumps(rpc_result).encode("utf-8") From 1102f0768664667f84aaaee7d4a061a9c128a6b3 Mon Sep 17 00:00:00 2001 From: Daniel Schiavini Date: Wed, 8 May 2024 17:00:41 +0200 Subject: [PATCH 7/7] test names --- tests/unitary/test_fork_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unitary/test_fork_utils.py b/tests/unitary/test_fork_utils.py index fd0a2cfb..4b8f1955 100644 --- a/tests/unitary/test_fork_utils.py +++ b/tests/unitary/test_fork_utils.py @@ -6,17 +6,17 @@ address = "0x0000000000000000000000000000000000000065" -def test_balance_non_fork(): +def test_set_balance_network(): with pytest.raises(NotImplementedError): NetworkEnv.set_balance(boa.env, address, 0) -def test_code_non_fork(): +def test_set_code_network(): with pytest.raises(NotImplementedError): NetworkEnv.set_code(boa.env, address, b"") -def test_storage_non_fork(): +def test_set_storage_network(): with pytest.raises(NotImplementedError): NetworkEnv.set_storage(boa.env, address, 0, 0)