Fix stack after precompiles (#1061)

* Fix precompiles stack

* Fix EXPMOD bugs (#1063)

* Fix expmod gas

* Overflow checks

* Many fixes

* Minor

evm/src/cpu/kernel/asm/core/call.asm

@@ -38,11 +38,13 @@ global sys_call:
     DUP5 DUP5 %address %transfer_eth %jumpi(call_insufficient_balance)
     DUP5 DUP5 %address %journal_add_balance_transfer
     DUP3 %set_new_ctx_gas_limit
-    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
-    DUP9 DUP9 DUP4 DUP4 DUP8 // Duplicate address, new_ctx, kexit_info, ret_offset, and ret_size.
+    DUP4
-    // stack: address, new_ctx, kexit_info, ret_offset, ret_size, ...
+    // stack: address, new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     %handle_precompiles
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
+    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
 
     // Each line in the block below does not change the stack.
     %set_static
@@ -87,11 +89,12 @@ global sys_callcode:
     DUP5 %address %address %transfer_eth %jumpi(call_insufficient_balance)
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     DUP3 %set_new_ctx_gas_limit
-    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
-    DUP9 DUP9 DUP4 DUP4 DUP8 // Duplicate address, new_ctx, kexit_info, ret_offset, and ret_size.
+    DUP4
-    // stack: address, new_ctx, kexit_info, ret_offset, ret_size, ...
+    // stack: address, new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     %handle_precompiles
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
+    %set_new_ctx_parent_pc(after_call_instruction)
 
     // Each line in the block below does not change the stack.
     %set_static
@@ -140,11 +143,13 @@ global sys_staticcall:
     %copy_mem_to_calldata
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     DUP3 %set_new_ctx_gas_limit
-    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
-    DUP9 DUP9 DUP4 DUP4 DUP8 // Duplicate address, new_ctx, kexit_info, ret_offset, and ret_size.
+    DUP4
-    // stack: address, new_ctx, kexit_info, ret_offset, ret_size, ...
+    // stack: address, new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     %handle_precompiles
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
+    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
 
     // Each line in the block below does not change the stack.
     %set_static_true
@@ -190,11 +195,13 @@ global sys_delegatecall:
     %copy_mem_to_calldata
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     DUP3 %set_new_ctx_gas_limit
-    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
-    DUP9 DUP9 DUP4 DUP4 DUP8 // Duplicate address, new_ctx, kexit_info, ret_offset, and ret_size.
+    DUP4
-    // stack: address, new_ctx, kexit_info, ret_offset, ret_size, ...
+    // stack: address, new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
     %handle_precompiles
     // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
+    %set_new_ctx_parent_pc(after_call_instruction)
+    // stack: new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size
 
     // Each line in the block below does not change the stack.
     %set_static

evm/src/cpu/kernel/asm/core/precompiles/blake2_f.asm

@@ -2,6 +2,8 @@ global precompile_blake2_f:
     // stack: retdest, new_ctx, (old stack)
     POP
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/bn_add.asm

@@ -2,6 +2,8 @@ global precompile_bn_add:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/bn_mul.asm

@@ -2,6 +2,8 @@ global precompile_bn_mul:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/ecrec.asm

@@ -2,6 +2,8 @@ global precompile_ecrec:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/expmod.asm

@@ -1,3 +1,12 @@
+// Mod 16 to the range [1, 16].
+%macro mod_16
+    // stack: x
+    %mod_const(16)
+    DUP1 %jumpi(%%after)
+    POP PUSH 16
+%%after:
+%endmacro
+
 // Load bytes, packing 16 bytes into each limb, and store limbs on the stack.
 // We pass around total_num_limbs and len for conveience, because we can't access them from the stack
 // if they're hidden behind the variable number of limbs.
@@ -5,7 +14,7 @@ mload_bytes_as_limbs:
     // stack: ctx, segment, offset, num_bytes, retdest, total_num_limbs, len, ..limbs
     DUP4
     // stack: num_bytes, ctx, segment, offset, num_bytes, retdest, total_num_limbs, len, ..limbs
-    %min_const(16)
+    %mod_16
     // stack: min(16, num_bytes), ctx, segment, offset, num_bytes, retdest, total_num_limbs, len, ..limbs
     %stack (len, addr: 3) -> (addr, len, addr)
     // stack: ctx, segment, offset, min(16, num_bytes), ctx, segment, offset, num_bytes, retdest, total_num_limbs, len, ..limbs
@@ -14,23 +23,26 @@ mload_bytes_as_limbs:
     %stack (new, addr: 3, numb, ret, tot, len) -> (numb, addr, ret, tot, len, new)
     // stack: num_bytes, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
     DUP1
-    %min_const(16)
+    %mod_16
-    SWAP1
+    // stack: num_bytes%16, num_bytes, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
+    DUP1 SWAP2
     SUB
-    // stack: num_bytes_new = num_bytes - min(16, num_bytes), ctx, segment, offset, retdest, total_num_limbs, len, ..limbs
+    // stack:num_bytes_new, num_bytes%16, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
     DUP1
     ISZERO
     %jumpi(mload_bytes_return)
-    // stack: num_bytes_new, ctx, segment, offset, retdest, total_num_limbs, len, ..limbs
+    SWAP1
-    SWAP3
+    // stack: num_bytes%16, num_bytes_new, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
-    %add_const(16)
+    DUP5 // offset
-    SWAP3
+    ADD
-    // stack: num_bytes_new, ctx, segment, offset + 16, retdest, total_num_limbs, len, ..limbs
+    // stack: offset_new, num_bytes_new, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
+    SWAP4 POP
+    // stack: num_bytes_new, ctx, segment, offset_new, retdest, total_num_limbs, len, new_limb, ..limbs
     %stack (num, addr: 3) -> (addr, num)
     %jump(mload_bytes_as_limbs)
 mload_bytes_return:
-    // stack: num_bytes_new, ctx, segment, offset, retdest, total_num_limbs, len, ..limbs
+    // stack: num_bytes_new, num_bytes%16, ctx, segment, offset, retdest, total_num_limbs, len, new_limb, ..limbs
-    %pop4
+    %pop5
     // stack: retdest, total_num_limbs, len, ..limbs
     JUMP
 
@@ -71,6 +83,9 @@ store_limbs_return:
 %endmacro
 
 %macro expmod_gas_f
+    // stack: x
+    // Overflow check
+    DUP1 %ge_const(0x800000000000000000000000000000007) %jumpi(fault_exception)
     // stack: x
     %ceil_div_const(8)
     // stack: ceil(x/8)
@@ -80,9 +95,12 @@ store_limbs_return:
 
 calculate_l_E_prime:
     // stack: l_E, l_B, retdest
-    DUP1
+    // Throw a fault early if the lengths are too large.
-    // stack: l_E, l_E, l_B, retdest
+    DUP2 %gt_const(0x100000000000000000000000000000000) %jumpi(fault_exception)
-    %le_const(32)
+    DUP1 %gt_const(0x100000000000000000000000000000000) %jumpi(fault_exception)
+    DUP1 ISZERO %jumpi(case_le_zero)
+    // stack: l_E, l_B, retdest
+    DUP1 %le_const(32)
     // stack: l_E <= 32, l_E, l_B, retdest
     %jumpi(case_le_32)
     // stack: l_E, l_B, retdest
@@ -101,16 +119,21 @@ calculate_l_E_prime:
     SWAP1
     // stack: l_E, log2(i[96 + l_B..128 + l_B]), l_B, retdest
     %sub_const(32)
+    // Overflow check
+    DUP1 %ge_const(0x2000000000000000000000000000000000000000000000000000000000000000) %jumpi(fault_exception)
     %mul_const(8)
     // stack: 8 * (l_E - 32), log2(i[96 + l_B..128 + l_B]), l_B, retdest
     ADD
     // stack: 8 * (l_E - 32) + log2(i[96 + l_B..128 + l_B]), l_B, retdest
-    SWAP2
+    SWAP1
-    %pop2
+    POP
     // stack: 8 * (l_E - 32) + log2(i[96 + l_B..128 + l_B]), retdest
     SWAP1
     // stack: retdest, 8 * (l_E - 32) + log2(i[96 + l_B..128 + l_B])
     JUMP
+case_le_zero:
+    %stack (l_E, l_B, retdest) -> (retdest, 0)
+    JUMP
 case_le_32:
     // stack: l_E, l_B, retdest
     SWAP1
@@ -131,6 +154,8 @@ global precompile_expmod:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint
@@ -157,7 +182,10 @@ global precompile_expmod:
     GET_CONTEXT
     %mload_packing
     // stack: l_M, l_E, l_B, kexit_info
-
+    DUP3 ISZERO DUP2 ISZERO
+    MUL // AND
+    // stack: l_M==0 && l_B==0, l_M, l_E, l_B, kexit_info
+    %jumpi(zero_base_zero_mod)
     %stack (l: 3) -> (l, l)
     // stack: l_M, l_E, l_B, l_M, l_E, l_B, kexit_info
     %max_3
@@ -365,58 +393,69 @@ expmod_contd:
     // Store return data size: l_M (number of bytes).
     SWAP1
     // stack: l_M, len, kexit_info
-    %mstore_parent_context_metadata(@CTX_METADATA_RETURNDATA_SIZE)
+    DUP1 %mstore_parent_context_metadata(@CTX_METADATA_RETURNDATA_SIZE)
-    // stack: len, kexit_info
+    // stack: l_M, len, kexit_info
-    DUP1
+    DUP1 ISZERO %jumpi(zero_modulus)
-    // stack: len, len, kexit_info
+    // stack: l_M, len, kexit_info
+    DUP1 %ceil_div_const(16)
+    // stack: l_M_128, l_M, len, kexit_info
+    SWAP1 %mod_16
+    // stack: l_M%16, l_M_128, len, kexit_info
+    SWAP2
+    // stack: len, l_M_128, l_M%16, kexit_info
     %mul_const(3)
-    // stack: out=3*len, len, kexit_info
+    // stack: out=3*len, l_M_128, l_M%16, kexit_info
+    %decrement
     DUP2
     DUP2
-    // stack: out, len, out, len, kexit_info
     ADD
-    %decrement
+    // stack: cur_address=out+l_M_128-1, end_address=out-1, l_M_128, l_M%16, kexit_info
+    DUP1 %mload_kernel_general
+    %stack (cur_limb, cur_address, end_address, l_M_128, l_M_mod16, kexit_info) ->
+        (@SEGMENT_RETURNDATA, 0, cur_limb, l_M_mod16, cur_address, end_address, l_M_128, kexit_info)
+    %mload_context_metadata(@CTX_METADATA_PARENT_CONTEXT)
+    %mstore_unpacking
+    // stack: offset, cur_address, end_address, l_M_128, kexit_info
     SWAP1
     %decrement
-    SWAP1
+    // stack: cur_address, offset, end_address, l_M_128, kexit_info
-    // stack: cur_address=out+len-1, end_address=out-1, len, kexit_info
-    PUSH 0
-    // stack: i=0, cur_address, end_address, len, kexit_info
-
     // Store in big-endian format.
 expmod_store_loop:
-    // stack: i, cur_address, end_address, len, kexit_info
+    // stack: cur_address, offset, end_address, l_M_128, kexit_info
-    DUP2
+    DUP3 DUP2 EQ %jumpi(expmod_store_end)
-    // stack: cur_address, i, cur_address, end_address, len, kexit_info
+    // stack: cur_address, offset, end_address, l_M_128, kexit_info
-    %mload_kernel_general
+    DUP1 %mload_kernel_general
-    // stack: cur_limb, i, cur_address, end_address, len, kexit_info
+    %stack (cur_limb, cur_address, offset, end_address, l_M_128, kexit_info) ->
-    DUP2
+         (offset, cur_limb, cur_address, end_address, l_M_128, kexit_info)
-    // stack: i, cur_limb, i, cur_address, end_address, len, kexit_info
-    %mul_const(16)
-    // stack: offset=16*i, cur_limb, i, cur_address, end_address, len, kexit_info
     %stack (offset, cur_limb) -> (@SEGMENT_RETURNDATA, offset, cur_limb, 16)
-    // stack: @SEGMENT_RETURNDATA, offset, cur_limb, 16, i, cur_address, end_address, len, kexit_info
     %mload_context_metadata(@CTX_METADATA_PARENT_CONTEXT)
-    // stack: parent_ctx, @SEGMENT_RETURNDATA, offset, cur_limb, 16, i, cur_address, end_address, len, kexit_info
     %mstore_unpacking
-    // stack: offset', i, cur_address, end_address, len, kexit_info
+    // stack: offset', cur_address, end_address, l_M_128, kexit_info)
-    POP
+    SWAP1 %decrement
-    // stack: i, cur_address, end_address, len, kexit_info
+    // stack: cur_address-1, offset', end_address, l_M_128, kexit_info)
-    %increment
+    %jump(expmod_store_loop)
-    SWAP1
-    %decrement
-    SWAP1
-    // stack: i+1, cur_address-1, end_address, len, kexit_info
-    DUP3
-    DUP2
-    EQ
-    ISZERO
-    %jumpi(expmod_store_loop)
 expmod_store_end:
-    // stack: i, cur_address, end_address, len, kexit_info
+    // stack: cur_address, offset, end_address, l_M_128, kexit_info
     %pop4
+the_end:
     // stack: kexit_info
     %leftover_gas
     // stack: leftover_gas
     PUSH 1 // success
     %jump(terminate_common)
+
+zero_modulus:
+    // stack: l_M, len, kexit_info
+    %pop2
+    %jump(the_end)
+
+zero_base_zero_mod:
+    // stack: l_M, l_E, l_B, kexit_info
+    %mstore_parent_context_metadata(@CTX_METADATA_RETURNDATA_SIZE)
+    // stack: l_E, l_B, kexit_info
+    %pop2
+    // stack: kexit_info
+    PUSH 200
+    %charge_gas
+    // stack: kexit_info
+    %jump(the_end)

evm/src/cpu/kernel/asm/core/precompiles/id.asm

@@ -2,6 +2,8 @@ global precompile_id:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/main.asm

@@ -6,7 +6,6 @@
     %jump(handle_precompiles)
 %%after:
     // stack: new_ctx, (old stack)
-    %pop4
 %endmacro
 
 global handle_precompiles:
@@ -31,6 +30,12 @@ global pop_and_return_success:
     PUSH 1 // success
     %jump(terminate_common)
 
+global after_precompile:
+    %mload_global_metadata(@GLOBAL_METADATA_IS_PRECOMPILE_FROM_EOA) %jumpi(process_message_txn_after_call)
+    %stack (success, leftover_gas, new_ctx, kexit_info, callgas, address, value, args_offset, args_size, ret_offset, ret_size) ->
+        (success, leftover_gas, new_ctx, kexit_info, ret_offset, ret_size)
+    %jump(after_call_instruction)
+
 %macro handle_precompiles_from_eoa
     // stack: retdest
     %mload_txn_field(@TXN_FIELD_TO)
@@ -42,10 +47,10 @@ global pop_and_return_success:
 %endmacro
 
 global handle_precompiles_from_eoa:
+    PUSH 1 %mstore_global_metadata(@GLOBAL_METADATA_IS_PRECOMPILE_FROM_EOA)
     // stack: addr, retdest
     %create_context
     // stack: new_ctx, addr, retdest
-    %set_new_ctx_parent_pc(process_message_txn_after_call)
     %non_intrinisic_gas %set_new_ctx_gas_limit
     // stack: new_ctx, addr, retdest
 

evm/src/cpu/kernel/asm/core/precompiles/rip160.asm

@@ -2,6 +2,8 @@ global precompile_rip160:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/sha256.asm

@@ -2,6 +2,8 @@ global precompile_sha256:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/precompiles/snarkv.asm

@@ -2,6 +2,8 @@ global precompile_snarkv:
     // stack: address, retdest, new_ctx, (old stack)
     %pop2
     // stack: new_ctx, (old stack)
+    %set_new_ctx_parent_pc(after_precompile)
+    // stack: new_ctx, (old stack)
     DUP1
     SET_CONTEXT
     %checkpoint // Checkpoint

evm/src/cpu/kernel/asm/core/process_txn.asm

@@ -302,6 +302,7 @@ process_message_txn_after_call_contd:
 process_message_txn_fail:
     // stack: leftover_gas, new_ctx, retdest
     // Transfer value back to the caller.
+    %mload_txn_field(@TXN_FIELD_VALUE) ISZERO %jumpi(process_message_txn_after_call_contd)
     %mload_txn_field(@TXN_FIELD_VALUE)
     %mload_txn_field(@TXN_FIELD_ORIGIN)
     %mload_txn_field(@TXN_FIELD_TO)

evm/src/cpu/kernel/constants/global_metadata.rs

@@ -67,10 +67,11 @@ pub(crate) enum GlobalMetadata {
     AccessListRlpLen = 32,
     // Boolean flag indicating if the txn is a contract creation txn.
     ContractCreation = 33,
+    IsPrecompileFromEoa = 34,
 }
 
 impl GlobalMetadata {
-    pub(crate) const COUNT: usize = 33;
+    pub(crate) const COUNT: usize = 34;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
         [
@@ -107,6 +108,7 @@ impl GlobalMetadata {
             Self::AccessListRlpStart,
             Self::AccessListRlpLen,
             Self::ContractCreation,
+            Self::IsPrecompileFromEoa,
         ]
     }
 
@@ -146,6 +148,7 @@ impl GlobalMetadata {
             Self::AccessListRlpStart => "GLOBAL_METADATA_ACCESS_LIST_RLP_START",
             Self::AccessListRlpLen => "GLOBAL_METADATA_ACCESS_LIST_RLP_LEN",
             Self::ContractCreation => "GLOBAL_METADATA_CONTRACT_CREATION",
+            Self::IsPrecompileFromEoa => "GLOBAL_METADATA_IS_PRECOMPILE_FROM_EOA",
         }
     }
 }