diff --git a/Cargo.lock b/Cargo.lock index 560acdde..dd4c4814 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1999,6 +1999,7 @@ version = "0.0.1" dependencies = [ "alloy-primitives", "alloy-sol-types", + "proptest", "revm", "revm-interpreter", "slh-dsa", diff --git a/crates/pq-precompile/Cargo.toml b/crates/pq-precompile/Cargo.toml index 2bf792c4..55b38005 100644 --- a/crates/pq-precompile/Cargo.toml +++ b/crates/pq-precompile/Cargo.toml @@ -16,5 +16,8 @@ revm.workspace = true revm-interpreter.workspace = true slh-dsa.workspace = true +[dev-dependencies] +proptest.workspace = true + [lints] workspace = true diff --git a/crates/pq-precompile/src/lib.rs b/crates/pq-precompile/src/lib.rs index c6d5dfa8..025385a3 100644 --- a/crates/pq-precompile/src/lib.rs +++ b/crates/pq-precompile/src/lib.rs @@ -15,7 +15,7 @@ // limitations under the License. use alloy_primitives::{address, Address, Bytes}; -use alloy_sol_types::{sol, SolValue}; +use alloy_sol_types::{abi, sol, SolCall, SolType, SolValue}; use revm::precompile::{PrecompileError, PrecompileHalt, PrecompileOutput}; use revm_interpreter::gas::KECCAK256WORD; use revm_interpreter::Gas; @@ -80,17 +80,35 @@ fn early_revert( )) } +/// The `(bytes, bytes, bytes)` parameter tuple of [`IPQ::verifySlhDsaSha2128sCall`]. +type CallParams<'a> = ::Parameters<'a>; + +/// Token form of [`CallParams`]: three `PackedSeqToken`s that borrow from the input. +type CallTokens<'a> = ::Token<'a>; + +/// Gas for hashing `msg_len` message bytes, one KECCAK256-rate charge per 32-byte word. +/// +/// `GAS_PER_MSG_WORD` (6) < 32, so `div_ceil(32) * GAS_PER_MSG_WORD` cannot exceed `u64::MAX`. +#[allow(clippy::arithmetic_side_effects)] +fn msg_word_gas(msg_len: usize) -> u64 { + (msg_len as u64).div_ceil(32) * GAS_PER_MSG_WORD +} + /// Executes the PQ (SLH-DSA-SHA2-128s) precompile. /// /// Early-path failures (short input, wrong selector, ABI decode error) charge a 200-gas /// penalty and return OOG if the caller has insufficient gas — preventing free probing. +/// +/// Base and message gas are reserved before the byte arguments are copied out of the +/// input, so the work a call performs is bounded by the gas it pays. Decoding is split to +/// allow this: `abi_decode_raw_validate` is tokenize, then `type_check`, then `detokenize`, +/// and only `detokenize` copies. The first two steps run before the charges; the copy runs +/// after them. pub fn run_pq_precompile( gas: u64, data: &[u8], reservoir: u64, ) -> Result { - use alloy_sol_types::SolCall; - let mut gas_counter = Gas::new(gas); if data.len() < 4 { @@ -102,8 +120,12 @@ pub fn run_pq_precompile( return early_revert(&mut gas_counter, reservoir, "Invalid selector"); } - let args = match IPQ::verifySlhDsaSha2128sCall::abi_decode_raw_validate(&data[4..]) { - Ok(args) => args, + // Tokenize and type-check without copying: reads the offset and length words, bounds-checks + // them against the body, and yields borrowed slices. O(1) regardless of input size. + let tokens = match abi::decode_sequence::>(&data[4..]) + .and_then(|tokens| CallParams::<'_>::type_check(&tokens).map(|()| tokens)) + { + Ok(tokens) => tokens, Err(_) => return early_revert(&mut gas_counter, reservoir, "Execution reverted"), }; @@ -111,13 +133,13 @@ pub fn run_pq_precompile( return Ok(PrecompileOutput::halt(PrecompileHalt::OutOfGas, reservoir)); } - // GAS_PER_MSG_WORD (6) < 32, so the product cannot exceed u64::MAX - #[allow(clippy::arithmetic_side_effects)] - let msg_word_gas = (args.message.len() as u64).div_ceil(32) * GAS_PER_MSG_WORD; - if !gas_counter.record_regular_cost(msg_word_gas) { + if !gas_counter.record_regular_cost(msg_word_gas(tokens.1.as_slice().len())) { return Ok(PrecompileOutput::halt(PrecompileHalt::OutOfGas, reservoir)); } + // Paid for: copy the byte arguments out of the input. + let args = IPQ::verifySlhDsaSha2128sCall::new(CallParams::<'_>::detokenize(tokens)); + if args.vk.len() != VK_LEN { return Ok(PrecompileOutput::revert( gas_counter.used(), @@ -167,15 +189,20 @@ pub fn run_pq_precompile( #[cfg(test)] mod tests { use super::{ - run_pq_precompile, EARLY_REVERT_GAS, GAS_PER_MSG_WORD, IPQ, SIG_LEN, VERIFY_BASE_GAS, - VK_LEN, + early_revert, revert_message_to_bytes, run_pq_precompile, EARLY_REVERT_GAS, + GAS_PER_MSG_WORD, IPQ, SIG_LEN, VERIFY_BASE_GAS, VK_LEN, }; + use alloy_primitives::U256; use alloy_sol_types::{SolCall, SolValue}; - use revm::precompile::PrecompileOutput; + use proptest::prelude::*; + use revm::precompile::{PrecompileError, PrecompileHalt, PrecompileOutput}; + use revm_interpreter::Gas; use slh_dsa::{ - signature::{Keypair, Signer}, - Sha2_128s, SigningKey, + signature::{Keypair, Signer, Verifier}, + Sha2_128s, Signature, SigningKey, VerifyingKey as SlhDsaVerifyingKey, }; + use std::sync::atomic::Ordering; + use std::sync::LazyLock; fn make_keypair() -> SigningKey { SigningKey::::slh_keygen_internal(&[1u8; 16], &[2u8; 16], &[3u8; 16]) @@ -292,4 +319,258 @@ mod tests { let expected = VERIFY_BASE_GAS + 2 * GAS_PER_MSG_WORD; assert_eq!(output.gas_used, expected); } + + // --- Gas is reserved before the byte arguments are copied --- + + /// Calldata whose three dynamic offsets alias one `shared_len`-byte array. The ABI decoder + /// accepts it; `vk` and `sig` then fail their length checks, but only after base and + /// message gas have been charged. + fn aliased_calldata(shared_len: usize) -> Vec { + let mut calldata = IPQ::verifySlhDsaSha2128sCall::SELECTOR.to_vec(); + for _ in 0..3 { + calldata.extend_from_slice(&U256::from(0x60).to_be_bytes::<32>()); + } + calldata.extend_from_slice(&U256::from(shared_len).to_be_bytes::<32>()); + let padded_len = shared_len.div_ceil(32).saturating_mul(32); + calldata.resize(calldata.len().saturating_add(padded_len), 0); + calldata + } + + #[test] + fn zero_gas_aliased_input_halts_out_of_gas() { + let calldata = aliased_calldata(1 << 20); + let output = run_pq_precompile(0, &calldata, 0).expect("Ok(halt)"); + assert!(output.is_halt()); + assert_eq!(output.halt_reason(), Some(&PrecompileHalt::OutOfGas)); + } + + #[test] + fn underfunded_aliased_input_below_base_gas_halts_out_of_gas() { + let calldata = aliased_calldata(1 << 20); + let output = + run_pq_precompile(VERIFY_BASE_GAS.saturating_sub(1), &calldata, 0).expect("Ok(halt)"); + assert!(output.is_halt()); + assert_eq!(output.halt_reason(), Some(&PrecompileHalt::OutOfGas)); + } + + /// An underfunded call must not copy the byte arguments. 64 zero-gas calls with a 1 MiB + /// aliased body would copy 3 MiB each (192 MiB total) if the arguments were copied before + /// gas is reserved; reserving first allocates nothing for the arguments. Other tests run + /// concurrently, so the bound is loose (16 MiB) but still an order of magnitude below the + /// copying figure. + #[test] + fn underfunded_aliased_input_is_not_copied() { + const ITERATIONS: usize = 64; + const BODY_LEN: usize = 1 << 20; + const MAX_ALLOCATED_BYTES: usize = 16 << 20; + + let calldata = aliased_calldata(BODY_LEN); + let before = alloc_counter::ALLOCATED_BYTES.load(Ordering::Relaxed); + for _ in 0..ITERATIONS { + let output = run_pq_precompile(0, &calldata, 0).expect("Ok(halt)"); + assert!(output.is_halt()); + } + let allocated = alloc_counter::ALLOCATED_BYTES + .load(Ordering::Relaxed) + .saturating_sub(before); + assert!( + allocated < MAX_ALLOCATED_BYTES, + "underfunded calls allocated {allocated} bytes; the byte arguments are being copied \ + before gas is reserved" + ); + } + + /// Global allocator that counts requested bytes. Only compiled into this crate's test + /// binary; every operation is delegated to `System` unchanged. + mod alloc_counter { + use std::alloc::{GlobalAlloc, Layout, System}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + pub(super) static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0); + + struct CountingAllocator; + + // SAFETY: every call is forwarded to `System` with the caller's arguments unchanged; + // the counter is a relaxed atomic and never touches the allocation itself. + unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + // SAFETY: same contract as this function's caller. + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: same contract as this function's caller. + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + ALLOCATED_BYTES.fetch_add(new_size, Ordering::Relaxed); + // SAFETY: same contract as this function's caller. + unsafe { System.realloc(ptr, layout, new_size) } + } + } + + #[global_allocator] + static GLOBAL: CountingAllocator = CountingAllocator; + } + + /// Reference implementation that decodes the full input before charging gas. The + /// differential test checks that `run_pq_precompile` returns the same output for every + /// `(gas, input, reservoir)`, so reserving gas first changes no observable result. + #[allow(clippy::arithmetic_side_effects)] + fn legacy_run_pq_precompile( + gas: u64, + data: &[u8], + reservoir: u64, + ) -> Result { + let mut gas_counter = Gas::new(gas); + + if data.len() < 4 { + return early_revert(&mut gas_counter, reservoir, "Input too short"); + } + + let selector = [data[0], data[1], data[2], data[3]]; + if selector != IPQ::verifySlhDsaSha2128sCall::SELECTOR { + return early_revert(&mut gas_counter, reservoir, "Invalid selector"); + } + + let args = match IPQ::verifySlhDsaSha2128sCall::abi_decode_raw_validate(&data[4..]) { + Ok(args) => args, + Err(_) => return early_revert(&mut gas_counter, reservoir, "Execution reverted"), + }; + + if !gas_counter.record_regular_cost(VERIFY_BASE_GAS) { + return Ok(PrecompileOutput::halt(PrecompileHalt::OutOfGas, reservoir)); + } + + let msg_word_gas = (args.message.len() as u64).div_ceil(32) * GAS_PER_MSG_WORD; + if !gas_counter.record_regular_cost(msg_word_gas) { + return Ok(PrecompileOutput::halt(PrecompileHalt::OutOfGas, reservoir)); + } + + if args.vk.len() != VK_LEN { + return Ok(PrecompileOutput::revert( + gas_counter.used(), + revert_message_to_bytes("Invalid verifying key length"), + reservoir, + )); + } + if args.sig.len() != SIG_LEN { + return Ok(PrecompileOutput::revert( + gas_counter.used(), + revert_message_to_bytes("Invalid signature length"), + reservoir, + )); + } + + let verifying_key = match SlhDsaVerifyingKey::::try_from(args.vk.as_ref()) { + Ok(vk) => vk, + Err(_) => { + return Ok(PrecompileOutput::revert( + gas_counter.used(), + revert_message_to_bytes("Failed to parse verifying key"), + reservoir, + )); + } + }; + let signature = match Signature::::try_from(args.sig.as_ref()) { + Ok(sig) => sig, + Err(_) => { + return Ok(PrecompileOutput::revert( + gas_counter.used(), + revert_message_to_bytes("Failed to parse signature"), + reservoir, + )); + } + }; + + let is_valid = verifying_key + .verify(args.message.as_ref(), &signature) + .is_ok(); + Ok(PrecompileOutput::new( + gas_counter.used(), + is_valid.abi_encode().into(), + reservoir, + )) + } + + /// A valid signature over a fixed message, signed once (SLH-DSA signing is slow). + static VALID_CALLDATA: LazyLock> = LazyLock::new(|| { + let sk = make_keypair(); + let msg = b"differential vector"; + let sig = sk.sign(msg); + encode_call(&sk.verifying_key().to_bytes(), msg, &sig.to_bytes()) + }); + + /// Canonical encodings with the exact `vk`/`sig` sizes over-represented so the length + /// checks and the verify path are both reached. + fn arb_canonical_body() -> impl Strategy> { + let vk_len = prop_oneof![3 => Just(VK_LEN), 1 => 0usize..=64]; + let msg_len = prop_oneof![1 => Just(0usize), 3 => 0usize..=4096]; + let sig_len = prop_oneof![3 => Just(SIG_LEN), 1 => 0usize..=8192]; + (vk_len, msg_len, sig_len).prop_map(|(vk, msg, sig)| { + encode_call(&vec![0x11; vk], &vec![0x22; msg], &vec![0x33; sig]) + }) + } + + /// Every input shape the decoder and gas checks branch on: canonical, a real valid + /// signature, aliased offsets, truncated bodies, corrupted offset words, random bodies + /// behind the right selector, and fully random bytes. + fn arb_input() -> impl Strategy> { + let selector = IPQ::verifySlhDsaSha2128sCall::SELECTOR; + prop_oneof![ + 3 => arb_canonical_body(), + 1 => Just(VALID_CALLDATA.clone()), + 2 => (0usize..=4096).prop_map(aliased_calldata), + 2 => (arb_canonical_body(), 0usize..=256).prop_map(|(mut body, cut)| { + body.truncate(body.len().saturating_sub(cut)); + body + }), + 2 => (arb_canonical_body(), 0usize..3, any::()).prop_map( + |(mut body, word, value)| { + let start = 4usize.saturating_add(word.saturating_mul(32)); + if let Some(slot) = body.get_mut(start..start.saturating_add(32)) { + slot.copy_from_slice(&U256::from(value).to_be_bytes::<32>()); + } + body + } + ), + 1 => prop::collection::vec(any::(), 0..=512).prop_map(move |mut body| { + let mut calldata = selector.to_vec(); + calldata.append(&mut body); + calldata + }), + 1 => prop::collection::vec(any::(), 0..=512), + ] + } + + /// Gas values clustered around every threshold the precompile checks. + fn arb_gas() -> impl Strategy { + prop_oneof![ + 1 => Just(0u64), + 2 => 0u64..=1_000, + 3 => VERIFY_BASE_GAS.saturating_sub(64)..=VERIFY_BASE_GAS.saturating_add(8_192), + 1 => any::(), + ] + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(512))] + + /// Reserve-first must be observably identical to the legacy decode-then-charge order + /// for every input, gas, and reservoir: same status, bytes, gas used, and reservoir. + #[test] + fn reserve_first_matches_legacy_ordering( + input in arb_input(), + gas in arb_gas(), + reservoir in any::(), + ) { + let current = run_pq_precompile(gas, &input, reservoir) + .expect("pq precompile never returns Err"); + let legacy = legacy_run_pq_precompile(gas, &input, reservoir) + .expect("pq precompile never returns Err"); + prop_assert_eq!(current, legacy); + } + } }