diff --git a/evm/src/cpu/kernel/asm/core/create.asm b/evm/src/cpu/kernel/asm/core/create.asm index 0bcddecf..2edd9bbe 100644 --- a/evm/src/cpu/kernel/asm/core/create.asm +++ b/evm/src/cpu/kernel/asm/core/create.asm @@ -159,6 +159,7 @@ after_constructor: %stack (size, ctx) -> (ctx, @SEGMENT_RETURNDATA, 0, size) // context, segment, offset, len KECCAK_GENERAL // stack: codehash, leftover_gas, success, address, kexit_info + %observe_new_contract DUP4 // stack: address, codehash, leftover_gas, success, address, kexit_info %set_codehash @@ -189,7 +190,7 @@ after_constructor: // Pre stack: addr, codehash, redest // Post stack: (empty) // TODO: Should it be copy-on-write (with make_account_copy) instead of mutating the trie? -set_codehash: +global set_codehash: // stack: addr, codehash, retdest %mpt_read_state_trie // stack: account_ptr, codehash, retdest @@ -208,3 +209,20 @@ set_codehash: %num_bytes_to_num_words %mul_const(@INITCODE_WORD_COST) %charge_gas %endmacro + + +// This should be called whenever a new contract is created. +// It does nothing, but just provides a single hook where code can react to newly created contracts. +// When called, the code corresponding to `codehash` should be stored in the return data. +// Pre stack: codehash, retdest +// Post stack: codehash +global observe_new_contract: + // stack codehash, retdest + SWAP1 JUMP + +%macro observe_new_contract + %stack (codehash) -> (codehash, %%after) + %jump(observe_new_contract) +%%after: + // stack: codehash +%endmacro diff --git a/evm/src/cpu/kernel/asm/core/process_txn.asm b/evm/src/cpu/kernel/asm/core/process_txn.asm index 3e16bbd4..a7ccf864 100644 --- a/evm/src/cpu/kernel/asm/core/process_txn.asm +++ b/evm/src/cpu/kernel/asm/core/process_txn.asm @@ -158,6 +158,19 @@ global process_contract_creation_txn_after_constructor: DUP2 DUP2 LT %jumpi(panic) // TODO: need to revert changes here. // stack: leftover_gas, codedeposit_cost, new_ctx, address, retdest SUB + + // Store the code hash of the new contract. + // stack: leftover_gas, new_ctx, address, retdest + GET_CONTEXT + %returndatasize + %stack (size, ctx) -> (ctx, @SEGMENT_RETURNDATA, 0, size) // context, segment, offset, len + KECCAK_GENERAL + // stack: codehash, leftover_gas, new_ctx, address, retdest + %observe_new_contract + DUP4 + // stack: address, codehash, leftover_gas, new_ctx, address, retdest + %set_codehash + // stack: leftover_gas, new_ctx, address, retdest %pay_coinbase_and_refund_sender // TODO: Delete accounts in self-destruct list and empty touched addresses. diff --git a/evm/src/generation/state.rs b/evm/src/generation/state.rs index 49c11611..c0c03e28 100644 --- a/evm/src/generation/state.rs +++ b/evm/src/generation/state.rs @@ -5,10 +5,12 @@ use keccak_hash::keccak; use plonky2::field::types::Field; use crate::cpu::kernel::aggregator::KERNEL; +use crate::cpu::kernel::constants::context_metadata::ContextMetadata; use crate::generation::mpt::all_mpt_prover_inputs_reversed; use crate::generation::rlp::all_rlp_prover_inputs_reversed; use crate::generation::GenerationInputs; -use crate::witness::memory::MemoryState; +use crate::memory::segments::Segment; +use crate::witness::memory::{MemoryAddress, MemoryState}; use crate::witness::state::RegistersState; use crate::witness::traces::{TraceCheckpoint, Traces}; use crate::witness::util::stack_peek; @@ -83,6 +85,10 @@ impl GenerationState { let tip_h256 = H256::from_uint(&tip_u256); let tip_h160 = H160::from(tip_h256); self.observe_address(tip_h160); + } else if dst == KERNEL.global_labels["observe_new_contract"] { + let tip_u256 = stack_peek(self, 0).expect("Empty stack"); + let tip_h256 = H256::from_uint(&tip_u256); + self.observe_contract(tip_h256); } } @@ -93,6 +99,30 @@ impl GenerationState { self.state_key_to_address.insert(state_key, address); } + /// Observe the given code hash and store the associated code. + /// When called, the code corresponding to `codehash` should be stored in the return data. + pub fn observe_contract(&mut self, codehash: H256) { + if self.inputs.contract_code.contains_key(&codehash) { + return; // Return early if the code hash has already been observed. + } + + let ctx = self.registers.context; + let returndata_size_addr = MemoryAddress::new( + ctx, + Segment::ContextMetadata, + ContextMetadata::ReturndataSize as usize, + ); + let returndata_size = self.memory.get(returndata_size_addr).as_usize(); + let code = self.memory.contexts[ctx].segments[Segment::Returndata as usize].content + [..returndata_size] + .iter() + .map(|x| x.as_u32() as u8) + .collect::>(); + debug_assert_eq!(keccak(&code), codehash); + + self.inputs.contract_code.insert(codehash, code); + } + pub fn checkpoint(&self) -> GenerationStateCheckpoint { GenerationStateCheckpoint { registers: self.registers,