New contract hook (#1002)

* New contract hook

* Minor

* PR feedback

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

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.

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<F: Field> GenerationState<F> {
             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<F: Field> GenerationState<F> {
         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::<Vec<_>>();
+        debug_assert_eq!(keccak(&code), codehash);
+
+        self.inputs.contract_code.insert(codehash, code);
+    }
+
     pub fn checkpoint(&self) -> GenerationStateCheckpoint {
         GenerationStateCheckpoint {
             registers: self.registers,