From 3841f5a430288a63c8207853eca11560bf7a5712 Mon Sep 17 00:00:00 2001 From: Alex Ostrovski Date: Tue, 21 Jan 2025 13:03:20 +0200 Subject: [PATCH] feat: Allow customizing precompiles logic (#79) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # What ❔ Allows customizing the precompiles logic by extending the `World` trait. ## Why ❔ - Allows caching `ecrecover` calls (= signature verification) with the default account abstraction. - Precompiles is the only entry point for crypto backends, so this allows theoretically customizing a backend (not implemented in this PR). --- crates/vm2-interface/src/state_interface.rs | 4 +- crates/vm2-interface/src/tracer_interface.rs | 4 +- .../src/instruction_handlers/precompiles.rs | 225 ++++----- crates/vm2/src/lib.rs | 7 + crates/vm2/src/precompiles/legacy.rs | 470 ++++++++++++++++++ crates/vm2/src/precompiles/mod.rs | 111 +++++ 6 files changed, 699 insertions(+), 122 deletions(-) create mode 100644 crates/vm2/src/precompiles/legacy.rs create mode 100644 crates/vm2/src/precompiles/mod.rs diff --git a/crates/vm2-interface/src/state_interface.rs b/crates/vm2-interface/src/state_interface.rs index 3f75597..046630d 100644 --- a/crates/vm2-interface/src/state_interface.rs +++ b/crates/vm2-interface/src/state_interface.rs @@ -216,7 +216,7 @@ pub struct L2ToL1Log { } #[cfg(test)] -pub mod tests { +pub(crate) mod testonly { use primitive_types::{H160, U256}; use super::{ @@ -224,7 +224,7 @@ pub mod tests { }; #[derive(Debug)] - pub struct DummyState; + pub(crate) struct DummyState; impl StateInterface for DummyState { fn read_register(&self, _: u8) -> (U256, bool) { diff --git a/crates/vm2-interface/src/tracer_interface.rs b/crates/vm2-interface/src/tracer_interface.rs index 29539b9..d31aaff 100644 --- a/crates/vm2-interface/src/tracer_interface.rs +++ b/crates/vm2-interface/src/tracer_interface.rs @@ -303,7 +303,7 @@ impl ShouldStop { } /// Cycle statistics emitted by the VM and supplied to [`Tracer::on_extra_prover_cycles()`]. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq)] pub enum CycleStats { /// Call to the `keccak256` precompile with the specified number of hash cycles. Keccak256(u32), @@ -349,7 +349,7 @@ impl Tracer for (A, B) { #[cfg(test)] mod tests { use super::{CallingMode, OpcodeType}; - use crate::{opcodes, tests::DummyState, GlobalStateInterface, Tracer}; + use crate::{opcodes, testonly::DummyState, GlobalStateInterface, Tracer}; struct FarCallCounter(usize); diff --git a/crates/vm2/src/instruction_handlers/precompiles.rs b/crates/vm2/src/instruction_handlers/precompiles.rs index 7e6c4bd..0a0be1d 100644 --- a/crates/vm2/src/instruction_handlers/precompiles.rs +++ b/crates/vm2/src/instruction_handlers/precompiles.rs @@ -1,141 +1,130 @@ -use primitive_types::{H160, U256}; -use zk_evm_abstractions::{ - aux::Timestamp, - precompiles::{ - ecrecover::ecrecover_function, keccak256::keccak256_rounds_function, - secp256r1_verify::secp256r1_verify_function, sha256::sha256_rounds_function, - }, - queries::LogQuery, - vm::Memory, -}; -use zkevm_opcode_defs::{ - system_params::{ - ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, - SECP256R1_VERIFY_PRECOMPILE_ADDRESS, SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, - }, - PrecompileAuxData, PrecompileCallABI, -}; -use zksync_vm2_interface::{opcodes, CycleStats, HeapId, Tracer}; +use primitive_types::U256; +use zksync_vm2_interface::{opcodes, HeapId, Tracer}; use super::{common::boilerplate_ext, ret::spontaneous_panic}; use crate::{ addressing_modes::{Arguments, Destination, Register1, Register2, Source}, - heap::Heaps, instruction::ExecutionStatus, + precompiles::{PrecompileMemoryReader, Precompiles}, Instruction, VirtualMachine, World, }; -fn precompile_call>( - vm: &mut VirtualMachine, - world: &mut W, - tracer: &mut T, -) -> ExecutionStatus { - boilerplate_ext::(vm, world, tracer, |vm, args, _, tracer| { - // The user gets to decide how much gas to burn - // This is safe because system contracts are trusted - let aux_data = PrecompileAuxData::from_u256(Register2::get(args, &mut vm.state)); - let Ok(()) = vm.state.use_gas(aux_data.extra_ergs_cost) else { - vm.state.current_frame.pc = spontaneous_panic(); - return; - }; +#[derive(Debug)] +struct PrecompileAuxData { + extra_ergs_cost: u32, + extra_pubdata_cost: u32, +} - #[allow(clippy::cast_possible_wrap)] - { - vm.world_diff.pubdata.0 += aux_data.extra_pubdata_cost as i32; - } +impl PrecompileAuxData { + #[allow(clippy::cast_possible_truncation)] + fn from_u256(raw_value: U256) -> Self { + let raw = raw_value.0; + let extra_ergs_cost = raw[0] as u32; + let extra_pubdata_cost = (raw[0] >> 32) as u32; - let mut abi = PrecompileCallABI::from_u256(Register1::get(args, &mut vm.state)); - if abi.memory_page_to_read == 0 { - abi.memory_page_to_read = vm.state.current_frame.heap.as_u32(); - } - if abi.memory_page_to_write == 0 { - abi.memory_page_to_write = vm.state.current_frame.heap.as_u32(); + Self { + extra_ergs_cost, + extra_pubdata_cost, } + } +} - let query = LogQuery { - timestamp: Timestamp(0), - key: abi.to_u256(), - // only two first fields are read by the precompile - tx_number_in_block: Default::default(), - aux_byte: Default::default(), - shard_id: Default::default(), - address: H160::default(), - read_value: U256::default(), - written_value: U256::default(), - rw_flag: Default::default(), - rollback: Default::default(), - is_service: Default::default(), - }; +#[derive(Debug)] +struct PrecompileCallAbi { + input_memory_offset: u32, + input_memory_length: u32, + output_memory_offset: u32, + output_memory_length: u32, + memory_page_to_read: HeapId, + memory_page_to_write: HeapId, + precompile_interpreted_data: u64, +} - let address_bytes = vm.state.current_frame.address.0; - let address_low = u16::from_le_bytes([address_bytes[19], address_bytes[18]]); - let heaps = &mut vm.state.heaps; +impl PrecompileCallAbi { + #[allow(clippy::cast_possible_truncation)] + fn from_u256(raw_value: U256) -> Self { + let raw = raw_value.0; + let input_memory_offset = raw[0] as u32; + let input_memory_length = (raw[0] >> 32) as u32; + let output_memory_offset = raw[1] as u32; + let output_memory_length = (raw[1] >> 32) as u32; + let memory_page_to_read = HeapId::from_u32_unchecked(raw[2] as u32); + let memory_page_to_write = HeapId::from_u32_unchecked((raw[2] >> 32) as u32); + let precompile_interpreted_data = raw[3]; - #[allow(clippy::cast_possible_truncation)] - // if we're having `> u32::MAX` cycles, we've got larger issues - match address_low { - KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => { - tracer.on_extra_prover_cycles(CycleStats::Keccak256( - keccak256_rounds_function::<_, false>(0, query, heaps).0 as u32, - )); - } - SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => { - tracer.on_extra_prover_cycles(CycleStats::Sha256( - sha256_rounds_function::<_, false>(0, query, heaps).0 as u32, - )); - } - ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS => { - tracer.on_extra_prover_cycles(CycleStats::EcRecover( - ecrecover_function::<_, false>(0, query, heaps).0 as u32, - )); - } - SECP256R1_VERIFY_PRECOMPILE_ADDRESS => { - tracer.on_extra_prover_cycles(CycleStats::Secp256r1Verify( - secp256r1_verify_function::<_, false>(0, query, heaps).0 as u32, - )); - } - _ => { - // A precompile call may be used just to burn gas - } + Self { + input_memory_offset, + input_memory_length, + output_memory_offset, + output_memory_length, + memory_page_to_read, + memory_page_to_write, + precompile_interpreted_data, } - - Register1::set(args, &mut vm.state, 1.into()); - }) + } } -impl Memory for Heaps { - fn execute_partial_query( - &mut self, - _monotonic_cycle_counter: u32, - mut query: zk_evm_abstractions::queries::MemoryQuery, - ) -> zk_evm_abstractions::queries::MemoryQuery { - let page = HeapId::from_u32_unchecked(query.location.page.0); +fn precompile_call>( + vm: &mut VirtualMachine, + world: &mut W, + tracer: &mut T, +) -> ExecutionStatus { + boilerplate_ext::( + vm, + world, + tracer, + |vm, args, world, tracer| { + // The user gets to decide how much gas to burn + // This is safe because system contracts are trusted + let aux_data = PrecompileAuxData::from_u256(Register2::get(args, &mut vm.state)); + let Ok(()) = vm.state.use_gas(aux_data.extra_ergs_cost) else { + vm.state.current_frame.pc = spontaneous_panic(); + return; + }; - let start = query.location.index.0 * 32; - if query.rw_flag { - self.write_u256(page, start, query.value); - } else { - query.value = self[page].read_u256(start); - query.value_is_pointer = false; - } - query - } + #[allow(clippy::cast_possible_wrap)] + { + vm.world_diff.pubdata.0 += aux_data.extra_pubdata_cost as i32; + } - fn specialized_code_query( - &mut self, - _monotonic_cycle_counter: u32, - _query: zk_evm_abstractions::queries::MemoryQuery, - ) -> zk_evm_abstractions::queries::MemoryQuery { - todo!() - } + let mut abi = PrecompileCallAbi::from_u256(Register1::get(args, &mut vm.state)); + if abi.memory_page_to_read.as_u32() == 0 { + abi.memory_page_to_read = vm.state.current_frame.heap; + } + if abi.memory_page_to_write.as_u32() == 0 { + abi.memory_page_to_write = vm.state.current_frame.heap; + } - fn read_code_query( - &self, - _monotonic_cycle_counter: u32, - _query: zk_evm_abstractions::queries::MemoryQuery, - ) -> zk_evm_abstractions::queries::MemoryQuery { - todo!() - } + let address_bytes = vm.state.current_frame.address.0; + let address_low = u16::from_le_bytes([address_bytes[19], address_bytes[18]]); + let heap_to_read = &vm.state.heaps[abi.memory_page_to_read]; + let memory = PrecompileMemoryReader::new( + heap_to_read, + abi.input_memory_offset, + abi.input_memory_length, + ); + let output = world.precompiles().call_precompile( + address_low, + memory, + abi.precompile_interpreted_data, + ); + + if let Some(cycle_stats) = output.cycle_stats { + tracer.on_extra_prover_cycles(cycle_stats); + } + + let mut write_offset = abi.output_memory_offset * 32; + for i in 0..output.len.min(abi.output_memory_length) { + vm.state.heaps.write_u256( + abi.memory_page_to_write, + write_offset, + output.buffer[i as usize], + ); + write_offset += 32; + } + Register1::set(args, &mut vm.state, 1.into()); + }, + ) } impl> Instruction { diff --git a/crates/vm2/src/lib.rs b/crates/vm2/src/lib.rs index c49514e..1422943 100644 --- a/crates/vm2/src/lib.rs +++ b/crates/vm2/src/lib.rs @@ -20,6 +20,7 @@ pub use self::{ vm::{Settings, VirtualMachine}, world_diff::{Snapshot, StorageChange, WorldDiff}, }; +use crate::precompiles::{LegacyPrecompiles, Precompiles}; pub mod addressing_modes; #[cfg(not(feature = "single_instruction_test"))] @@ -33,6 +34,7 @@ mod heap; mod instruction; mod instruction_handlers; mod mode_requirements; +pub mod precompiles; mod predication; #[cfg(not(feature = "single_instruction_test"))] mod program; @@ -98,6 +100,11 @@ pub trait World: StorageInterface + Sized { /// Loads bytecode bytes for the `decommit` opcode. fn decommit_code(&mut self, hash: U256) -> Vec; + + /// Returns precompiles to be used. + fn precompiles(&self) -> &impl Precompiles { + &LegacyPrecompiles + } } /// Deterministic (across program runs and machines) hash that can be used for `Debug` implementations diff --git a/crates/vm2/src/precompiles/legacy.rs b/crates/vm2/src/precompiles/legacy.rs new file mode 100644 index 0000000..14db53e --- /dev/null +++ b/crates/vm2/src/precompiles/legacy.rs @@ -0,0 +1,470 @@ +use primitive_types::{H160, U256}; +use zk_evm_abstractions::{ + aux::Timestamp, + precompiles::{ + ecrecover::ecrecover_function, keccak256::keccak256_rounds_function, + secp256r1_verify::secp256r1_verify_function, sha256::sha256_rounds_function, + }, + queries::{LogQuery, MemoryQuery}, + vm::Memory, +}; +use zkevm_opcode_defs::{ + PrecompileCallABI, ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, + KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, SECP256R1_VERIFY_PRECOMPILE_ADDRESS, + SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, +}; +use zksync_vm2_interface::CycleStats; + +use super::{PrecompileMemoryReader, PrecompileOutput, Precompiles}; + +fn create_query(input_offset: u32, input_len: u32, aux_data: u64) -> LogQuery { + let abi = PrecompileCallABI { + input_memory_offset: input_offset, + input_memory_length: input_len, + output_memory_offset: 0, + output_memory_length: 2, // not read by implementations + // Pages are fake; we assume that precompiles are implemented correctly and don't read / write anywhere but the specified pages + memory_page_to_read: 1, + memory_page_to_write: 2, + precompile_interpreted_data: aux_data, + }; + LogQuery { + timestamp: Timestamp(0), + key: abi.to_u256(), + // only two first fields are read by the precompile + tx_number_in_block: Default::default(), + aux_byte: Default::default(), + shard_id: Default::default(), + address: H160::default(), + read_value: U256::default(), + written_value: U256::default(), + rw_flag: Default::default(), + rollback: Default::default(), + is_service: Default::default(), + } +} + +#[derive(Debug)] +struct LegacyIo<'a> { + input: PrecompileMemoryReader<'a>, + output: PrecompileOutput, +} + +impl<'a> LegacyIo<'a> { + fn new(input: PrecompileMemoryReader<'a>) -> Self { + Self { + input, + output: PrecompileOutput::default(), + } + } +} + +impl Memory for LegacyIo<'_> { + fn execute_partial_query( + &mut self, + _monotonic_cycle_counter: u32, + mut query: MemoryQuery, + ) -> MemoryQuery { + let start_word = query.location.index.0; + if query.rw_flag { + assert!(start_word < 2, "standard precompiles never write >2 words"); + self.output.buffer[start_word as usize] = query.value; + self.output.len = self.output.len.max(start_word + 1); + } else { + // Access `Heap` directly for a speed-up + query.value = self.input.heap.read_u256(start_word * 32); + query.value_is_pointer = false; + } + query + } + + fn specialized_code_query( + &mut self, + _monotonic_cycle_counter: u32, + _query: MemoryQuery, + ) -> MemoryQuery { + unimplemented!("should not be called") + } + + fn read_code_query(&self, _monotonic_cycle_counter: u32, _query: MemoryQuery) -> MemoryQuery { + unimplemented!("should not be called") + } +} + +/// Precompiles implementation using legacy VM code. +#[derive(Debug)] +pub struct LegacyPrecompiles; + +impl Precompiles for LegacyPrecompiles { + #[allow(clippy::cast_possible_truncation)] + fn call_precompile( + &self, + address_low: u16, + memory: PrecompileMemoryReader<'_>, + aux_input: u64, + ) -> PrecompileOutput { + let query = create_query(memory.offset, memory.len, aux_input); + let mut io = LegacyIo::new(memory); + match address_low { + KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => { + let cycles = keccak256_rounds_function::<_, false>(0, query, &mut io).0; + io.output + .with_cycle_stats(CycleStats::Keccak256(cycles as u32)) + } + SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => { + let cycles = sha256_rounds_function::<_, false>(0, query, &mut io).0; + io.output + .with_cycle_stats(CycleStats::Sha256(cycles as u32)) + } + ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS => { + let cycles = ecrecover_function::<_, false>(0, query, &mut io).0; + io.output + .with_cycle_stats(CycleStats::EcRecover(cycles as u32)) + } + SECP256R1_VERIFY_PRECOMPILE_ADDRESS => { + let cycles = secp256r1_verify_function::<_, false>(0, query, &mut io).0; + io.output + .with_cycle_stats(CycleStats::Secp256r1Verify(cycles as u32)) + } + _ => PrecompileOutput::default(), + } + } +} + +#[allow(clippy::cast_possible_truncation)] // OK for tests +#[cfg(test)] +mod tests { + use proptest::{array, collection, num, option, prelude::*}; + use zkevm_opcode_defs::{ + k256::ecdsa::{SigningKey as K256SigningKey, VerifyingKey as K256VerifyingKey}, + p256::ecdsa::SigningKey as P256SigningKey, + sha3::{self, Digest}, + }; + use zksync_vm2_interface::HeapId; + + use super::*; + use crate::heap::Heaps; + + const MAX_LEN: usize = 2_048; + + fn arbitrary_aligned_bytes(alignment: usize) -> impl Strategy> { + (0..=(MAX_LEN / alignment)).prop_flat_map(move |len_in_words| { + collection::vec(num::u8::ANY, len_in_words * alignment) + }) + } + + fn key_to_address(key: &K256VerifyingKey) -> U256 { + let encoded_key = key.to_encoded_point(false); + let encoded_key = &encoded_key.as_bytes()[1..]; + debug_assert_eq!(encoded_key.len(), 64); + let address_digest = sha3::Keccak256::digest(encoded_key); + let address_u256 = U256::from_big_endian(&address_digest); + // Mask out upper bytes of the hash. + address_u256 & U256::MAX >> (256 - 160) + } + + fn test_keccak_precompile(input: &[u8], initial_offset: u32) -> Result<(), TestCaseError> { + let input_len = input.len() as u32; + assert_eq!(input_len % 32, 0); + + let mut heaps = Heaps::new(&[]); + for (i, u256_chunk) in input.chunks(32).enumerate() { + let offset = i as u32 * 32 + initial_offset; + heaps.write_u256(HeapId::FIRST, offset, U256::from_big_endian(u256_chunk)); + } + + let memory = PrecompileMemoryReader::new(&heaps[HeapId::FIRST], initial_offset, input_len); + let output = LegacyPrecompiles.call_precompile( + KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, + memory, + 0, + ); + + prop_assert_eq!(output.len, 1); + let expected_hash = sha3::Keccak256::digest(input); + let expected_hash = U256::from_big_endian(&expected_hash); + prop_assert_eq!(output.buffer[0], expected_hash); + prop_assert!(matches!(output.cycle_stats, Some(CycleStats::Keccak256(_)))); + Ok(()) + } + + fn test_sha256_precompile( + input: &[u8], + initial_offset_in_words: u32, + ) -> Result<(), TestCaseError> { + assert_eq!(input.len() % 64, 0); + let mut heaps = Heaps::new(&[]); + + for (i, u256_chunk) in input.chunks(32).enumerate() { + let offset = i as u32 * 32 + initial_offset_in_words * 32; + heaps.write_u256(HeapId::FIRST, offset, U256::from_big_endian(u256_chunk)); + } + + let max_round_count = input.len() as u32 / 64; + for round_count in 0..=max_round_count { + let memory = PrecompileMemoryReader::new( + &heaps[HeapId::FIRST], + initial_offset_in_words, + round_count * 2, + ); + let output = LegacyPrecompiles.call_precompile( + SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, + memory, + round_count.into(), + ); + + if round_count == 0 { + prop_assert_eq!(output.len, 0); + } else { + prop_assert_eq!(output.len, 1); + prop_assert_ne!(output.buffer[0], U256::zero()); + } + prop_assert!(matches!(output.cycle_stats, Some(CycleStats::Sha256(_)))); + } + Ok(()) + } + + #[derive(Debug, Clone, Copy)] + enum EcRecoverMutation { + RecoveryId, + Digest(usize), + R(usize), + S(usize), + } + + impl EcRecoverMutation { + fn gen() -> impl Strategy { + (0..4).prop_flat_map(|raw| match raw { + 0 => Just(Self::RecoveryId).boxed(), + 1 => (0_usize..32).prop_map(Self::Digest).boxed(), + 2 => (0_usize..32).prop_map(Self::R).boxed(), + 3 => (0_usize..32).prop_map(Self::S).boxed(), + _ => unreachable!(), + }) + } + } + + fn test_ecrecover_precompile( + signing_key: &K256SigningKey, + mutation: Option, + initial_offset_in_words: u32, + ) -> Result<(), TestCaseError> { + let mut heaps = Heaps::new(&[]); + let initial_offset = initial_offset_in_words * 32; + + let message = "test message!"; + let mut message_digest = sha3::Keccak256::digest(message); + + let (signature, recovery_id) = signing_key + .sign_prehash_recoverable(&message_digest) + .unwrap(); + if recovery_id.is_x_reduced() { + return Ok(()); + } + let mut recovery_id = recovery_id.to_byte(); + + println!( + "testing key {:?} with mutation {mutation:?}", + signing_key.verifying_key().to_encoded_point(true) + ); + let mut signature_bytes = signature.to_bytes(); + + match mutation { + Some(EcRecoverMutation::Digest(byte)) => { + message_digest[byte] ^= 1; + } + Some(EcRecoverMutation::RecoveryId) => { + recovery_id = 1 - recovery_id; + } + Some(EcRecoverMutation::R(byte)) => { + signature_bytes[byte] ^= 1; + } + Some(EcRecoverMutation::S(byte)) => { + signature_bytes[byte + 32] ^= 1; + } + None => { /* Do nothing */ } + } + + heaps.write_u256( + HeapId::FIRST, + initial_offset, + U256::from_big_endian(&message_digest), + ); + heaps.write_u256(HeapId::FIRST, initial_offset + 32, recovery_id.into()); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 64, + U256::from_big_endian(&signature_bytes[..32]), + ); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 96, + U256::from_big_endian(&signature_bytes[32..]), + ); + + let memory = PrecompileMemoryReader::new(&heaps[HeapId::FIRST], initial_offset_in_words, 4); + let output = LegacyPrecompiles.call_precompile( + ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, + memory, + 0, + ); + + prop_assert_eq!(output.len, 2); + let expected_address = key_to_address(signing_key.verifying_key()); + let [is_success, address] = output.buffer; + if mutation.is_some() { + prop_assert_ne!(address, expected_address); + } else { + prop_assert_eq!(is_success, U256::one()); + prop_assert_eq!(address, expected_address); + } + prop_assert!(matches!(output.cycle_stats, Some(CycleStats::EcRecover(1)))); + Ok(()) + } + + #[derive(Debug, Clone, Copy)] + enum P256Mutation { + Digest(usize), + R(usize), + S(usize), + Key(usize), + } + + impl P256Mutation { + fn gen() -> impl Strategy { + (0..4).prop_flat_map(|raw| match raw { + 0 => (0_usize..32).prop_map(Self::Digest).boxed(), + 1 => (0_usize..32).prop_map(Self::R).boxed(), + 2 => (0_usize..32).prop_map(Self::S).boxed(), + 3 => (0_usize..64).prop_map(Self::Key).boxed(), + _ => unreachable!(), + }) + } + } + + fn test_secp256r1_precompile( + signing_key: &P256SigningKey, + mutation: Option, + initial_offset_in_words: u32, + ) -> Result<(), TestCaseError> { + use zkevm_opcode_defs::p256::ecdsa::{signature::hazmat::PrehashSigner, Signature}; + + let mut heaps = Heaps::new(&[]); + let initial_offset = initial_offset_in_words * 32; + + let message = "test message!"; + let mut message_digest = sha3::Keccak256::digest(message); + + let signature: Signature = signing_key.sign_prehash(&message_digest).unwrap(); + + println!( + "testing key {:?} with mutation {mutation:?}", + signing_key.verifying_key().to_encoded_point(true) + ); + let mut signature_bytes = signature.to_bytes(); + let mut key_bytes = signing_key + .verifying_key() + .to_encoded_point(false) + .as_bytes()[1..] + .to_vec(); + assert_eq!(key_bytes.len(), 64); + + match mutation { + Some(P256Mutation::Digest(byte)) => { + message_digest[byte] ^= 1; + } + Some(P256Mutation::R(byte)) => { + signature_bytes[byte] ^= 1; + } + Some(P256Mutation::S(byte)) => { + signature_bytes[byte + 32] ^= 1; + } + Some(P256Mutation::Key(byte)) => { + key_bytes[byte] ^= 1; + } + None => { /* Do nothing */ } + } + + heaps.write_u256( + HeapId::FIRST, + initial_offset, + U256::from_big_endian(&message_digest), + ); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 32, + U256::from_big_endian(&signature_bytes[..32]), + ); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 64, + U256::from_big_endian(&signature_bytes[32..]), + ); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 96, + U256::from_big_endian(&key_bytes[..32]), + ); + heaps.write_u256( + HeapId::FIRST, + initial_offset + 128, + U256::from_big_endian(&key_bytes[32..]), + ); + + let memory = PrecompileMemoryReader::new(&heaps[HeapId::FIRST], initial_offset_in_words, 5); + let output = + LegacyPrecompiles.call_precompile(SECP256R1_VERIFY_PRECOMPILE_ADDRESS, memory, 0); + + prop_assert_eq!(output.len, 2); + let [is_ok, is_verified] = output.buffer; + if mutation.is_none() { + prop_assert_eq!(is_ok, U256::one()); + prop_assert_eq!(is_verified, U256::one()); + } else { + prop_assert!(is_ok.is_zero() || is_verified.is_zero()); + } + prop_assert!(matches!( + output.cycle_stats, + Some(CycleStats::Secp256r1Verify(1)) + )); + Ok(()) + } + + proptest! { + #[test] + fn keccak_precompile_works( + bytes in arbitrary_aligned_bytes(32), + initial_offset in 0..u32::MAX / 2, + ) { + test_keccak_precompile(&bytes, initial_offset)?; + } + + #[test] + fn sha256_precompile_works( + bytes in arbitrary_aligned_bytes(64), + initial_offset_in_words in 0..u32::MAX / 64, + ) { + test_sha256_precompile(&bytes, initial_offset_in_words)?; + } + + #[test] + fn ecrecover_precompile_works( + signing_key in array::uniform32(num::u8::ANY) + .prop_filter_map("not a key", |bytes| K256SigningKey::from_bytes(&bytes.into()).ok()), + mutation in option::of(EcRecoverMutation::gen()), + initial_offset_in_words in 0..u32::MAX / 64, + ) { + test_ecrecover_precompile(&signing_key, mutation, initial_offset_in_words)?; + } + + #[test] + fn secp256r1_precompile_works( + signing_key in array::uniform32(num::u8::ANY) + .prop_filter_map("not a key", |bytes| P256SigningKey::from_bytes(&bytes.into()).ok()), + mutation in option::of(P256Mutation::gen()), + initial_offset_in_words in 0..u32::MAX / 64, + ) { + test_secp256r1_precompile(&signing_key, mutation, initial_offset_in_words)?; + } + } +} diff --git a/crates/vm2/src/precompiles/mod.rs b/crates/vm2/src/precompiles/mod.rs new file mode 100644 index 0000000..b7fa26a --- /dev/null +++ b/crates/vm2/src/precompiles/mod.rs @@ -0,0 +1,111 @@ +//! Precompiles support. + +use primitive_types::U256; +pub use zkevm_opcode_defs::system_params::{ + ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, + SECP256R1_VERIFY_PRECOMPILE_ADDRESS, SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS, +}; +use zksync_vm2_interface::CycleStats; + +pub use self::legacy::LegacyPrecompiles; +use crate::heap::Heap; + +mod legacy; + +/// Provides access to the input memory for a precompile call. +#[derive(Debug, Clone)] +pub struct PrecompileMemoryReader<'a, const IN_WORDS: bool = false> { + heap: &'a Heap, + offset: u32, + len: u32, +} + +impl<'a> PrecompileMemoryReader<'a> { + pub(crate) fn new(heap: &'a Heap, offset: u32, len: u32) -> Self { + Self { heap, offset, len } + } + + /// Assumes that the input offset and length passed via ABI are measured in 32-byte words, rather than bytes. + pub fn assume_offset_in_words(self) -> PrecompileMemoryReader<'a, true> { + PrecompileMemoryReader { + heap: self.heap, + offset: self.offset * 32, + len: self.len * 32, + } + } +} + +/// Iterates over input bytes. +impl Iterator for PrecompileMemoryReader<'_, IN_WORDS> { + type Item = u8; + + fn next(&mut self) -> Option { + if self.len == 0 { + return None; + } + + // This assumes the offset never overflows + let output = self.heap.read_byte(self.offset); + self.offset += 1; + self.len -= 1; + Some(output) + } + + fn size_hint(&self) -> (usize, Option) { + (self.len as usize, Some(self.len as usize)) + } +} + +impl ExactSizeIterator for PrecompileMemoryReader<'_> { + fn len(&self) -> usize { + self.len as usize + } +} + +/// Output of a precompile call returned from [`Precompiles::call_precompile()`]. +#[derive(Debug, Default)] +pub struct PrecompileOutput { + pub(crate) buffer: [U256; 2], + pub(crate) len: u32, + pub(crate) cycle_stats: Option, +} + +impl PrecompileOutput { + /// Assigns cycle stats for this output. + #[must_use] + pub fn with_cycle_stats(mut self, stats: CycleStats) -> Self { + self.cycle_stats = Some(stats); + self + } +} + +impl From for PrecompileOutput { + fn from(value: U256) -> Self { + Self { + buffer: [value, U256::zero()], + len: 1, + cycle_stats: None, + } + } +} + +impl From<[U256; 2]> for PrecompileOutput { + fn from(value: [U256; 2]) -> Self { + Self { + buffer: value, + len: 2, + cycle_stats: None, + } + } +} + +/// Encapsulates precompiles used during VM execution. +pub trait Precompiles { + /// Calls to a precompile. + fn call_precompile( + &self, + address_low: u16, + memory: PrecompileMemoryReader<'_>, + aux_input: u64, + ) -> PrecompileOutput; +}