diff --git a/pallets/wormhole/src/lib.rs b/pallets/wormhole/src/lib.rs index 9b4ff8fe..65290cb3 100644 --- a/pallets/wormhole/src/lib.rs +++ b/pallets/wormhole/src/lib.rs @@ -347,14 +347,14 @@ pub mod pallet { }, }; - // Mark nullifiers as used (validate_proof only checks existence) - let mut nullifier_list = Vec::<[u8; 32]>::new(); + // Collect nullifier bytes (no state writes yet - defer until after all checks) + let mut nullifier_list = + Vec::<[u8; 32]>::with_capacity(aggregated_inputs.nullifiers.len()); for nullifier in &aggregated_inputs.nullifiers { let nullifier_bytes: [u8; 32] = (*nullifier) .as_ref() .try_into() .map_err(|_| Error::::InvalidAggregatedPublicInputs)?; - UsedNullifiers::::insert(nullifier_bytes, true); nullifier_list.push(nullifier_bytes); } @@ -405,6 +405,12 @@ pub mod pallet { Error::::TransferAmountBelowMinimum ); + // Cheap checks passed; #[pallet::call] auto-wraps this dispatchable in + // with_storage_layer so any later failure rolls these writes back. + for nullifier_bytes in &nullifier_list { + UsedNullifiers::::insert(*nullifier_bytes, true); + } + // Emit event for each exit account Self::deposit_event(Event::ProofVerified { exit_amount: total_exit_amount, diff --git a/pallets/wormhole/src/mock.rs b/pallets/wormhole/src/mock.rs index dab7c7f1..858e2fca 100644 --- a/pallets/wormhole/src/mock.rs +++ b/pallets/wormhole/src/mock.rs @@ -118,8 +118,8 @@ parameter_types! { /// The "from" account used when recording transfer proofs for minted tokens. /// Uses the shared MINTING_ACCOUNT constant from qp_wormhole. pub const MintingAccount: AccountId = MINTING_ACCOUNT; - /// Minimum transfer amount (10 QUAN) - pub const MinimumTransferAmount: Balance = 10 * UNIT; + /// Minimum transfer amount (10 QUAN). Storage-backed so rollback tests can override it. + pub storage MinimumTransferAmount: Balance = 10 * UNIT; /// Volume fee rate in basis points (10 bps = 0.1%) pub const VolumeFeeRateBps: u32 = 10; /// Proportion of volume fees to burn (50% burned, 50% to miner) diff --git a/pallets/wormhole/src/tests.rs b/pallets/wormhole/src/tests.rs index 1ea8c487..87791075 100644 --- a/pallets/wormhole/src/tests.rs +++ b/pallets/wormhole/src/tests.rs @@ -281,10 +281,14 @@ mod aggregated_proof_tests { mock::*, pallet::{Error, UsedNullifiers}, }; - use frame_support::{assert_noop, assert_ok}; + use codec::Decode; + use frame_support::{assert_noop, assert_ok, traits::fungible::Unbalanced}; use frame_system::RawOrigin; - use qp_wormhole_verifier::{parse_aggregated_public_inputs, ProofWithPublicInputs, C, F}; + use qp_wormhole_verifier::{ + parse_aggregated_public_inputs, AggregatedPublicCircuitInputs, ProofWithPublicInputs, C, F, + }; use sp_core::H256; + use sp_runtime::{ArithmeticError, DispatchError}; /// The D const parameter for plonky2 proofs (extension degree = 2) const D: usize = 2; @@ -306,6 +310,46 @@ mod aggregated_proof_tests { .expect("Proof should deserialize") } + /// Install the proof's expected block hash and advance past it so `validate_proof` passes. + fn setup_valid_proof_state() -> (Vec, AggregatedPublicCircuitInputs) { + let proof = deserialize_test_proof(); + let inputs = parse_aggregated_public_inputs(&proof).expect("Should parse"); + let block_number = inputs.block_data.block_number as u64; + let block_hash_bytes: [u8; 32] = inputs.block_data.block_hash.as_ref().try_into().unwrap(); + frame_system::BlockHash::::insert(block_number, H256::from(block_hash_bytes)); + System::set_block_number(block_number + 10); + (get_test_proof_bytes(), inputs) + } + + fn nullifier_bytes(inputs: &AggregatedPublicCircuitInputs) -> Vec<[u8; 32]> { + inputs.nullifiers.iter().map(|n| n.as_ref().try_into().unwrap()).collect() + } + + /// Decode the real (non-dummy) exit accounts from a parsed proof, matching the + /// filter logic the dispatchable uses when iterating `account_data`. + fn exit_accounts(inputs: &AggregatedPublicCircuitInputs) -> Vec { + inputs + .account_data + .iter() + .filter_map(|a| { + let bytes: [u8; 32] = (*a.exit_account).as_ref().try_into().ok()?; + if bytes == [0u8; 32] || a.summed_output_amount == 0 { + return None; + } + AccountId::decode(&mut &bytes[..]).ok() + }) + .collect() + } + + fn assert_nullifiers_unused(inputs: &AggregatedPublicCircuitInputs, ctx: &str) { + for bytes in nullifier_bytes(inputs) { + assert!( + !UsedNullifiers::::contains_key(bytes), + "Nullifier must not be consumed ({ctx})" + ); + } + } + #[test] fn test_proof_deserialization_succeeds() { // Just test that the proof deserializes correctly @@ -379,29 +423,13 @@ mod aggregated_proof_tests { #[test] fn test_verify_aggregated_proof_fails_with_nullifier_already_used() { new_test_ext().execute_with(|| { - let proof = deserialize_test_proof(); - let inputs = parse_aggregated_public_inputs(&proof).expect("Should parse"); - - // Set up block hash to match the proof - let block_number = inputs.block_data.block_number as u64; - let block_hash_bytes: [u8; 32] = - inputs.block_data.block_hash.as_ref().try_into().unwrap(); - let block_hash = H256::from(block_hash_bytes); - - // Insert a matching block hash - frame_system::BlockHash::::insert(block_number, block_hash); + let (proof_bytes, inputs) = setup_valid_proof_state(); - // Mark one of the nullifiers as already used - if let Some(nullifier) = inputs.nullifiers.first() { - let nullifier_bytes: [u8; 32] = nullifier.as_ref().try_into().unwrap(); - UsedNullifiers::::insert(nullifier_bytes, true); - } - - let proof_bytes = get_test_proof_bytes(); + let pre_used = nullifier_bytes(&inputs).into_iter().next().expect("nullifier"); + UsedNullifiers::::insert(pre_used, true); - let result = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes); - assert!(result.is_err()); - let err = result.unwrap_err(); + let err = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes) + .expect_err("Expected NullifierAlreadyUsed"); assert_eq!(err.error, Error::::NullifierAlreadyUsed.into()); }); } @@ -430,86 +458,83 @@ mod aggregated_proof_tests { #[test] fn test_verify_aggregated_proof_succeeds_with_valid_state() { new_test_ext().execute_with(|| { - let proof = deserialize_test_proof(); - let inputs = parse_aggregated_public_inputs(&proof).expect("Should parse"); - - // Set up block hash to match the proof - let block_number = inputs.block_data.block_number as u64; - let block_hash_bytes: [u8; 32] = - inputs.block_data.block_hash.as_ref().try_into().unwrap(); - let block_hash = H256::from(block_hash_bytes); - - frame_system::BlockHash::::insert(block_number, block_hash); + let (proof_bytes, inputs) = setup_valid_proof_state(); - // Set current block number higher than the proof's block - System::set_block_number(block_number + 10); - - let proof_bytes = get_test_proof_bytes(); - - // This should succeed - proof is valid and state matches assert_ok!(Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes)); - // Verify nullifiers are now marked as used - for nullifier in &inputs.nullifiers { - let nullifier_bytes: [u8; 32] = nullifier.as_ref().try_into().unwrap(); + for bytes in nullifier_bytes(&inputs) { assert!( - UsedNullifiers::::contains_key(nullifier_bytes), + UsedNullifiers::::contains_key(bytes), "Nullifier should be marked as used" ); } - // Verify event was emitted + let exit_amount: u128 = inputs + .account_data + .iter() + .filter(|a| a.summed_output_amount > 0) + .map(|a| (a.summed_output_amount as u128) * crate::SCALE_DOWN_FACTOR) + .sum(); System::assert_has_event( crate::Event::::ProofVerified { - exit_amount: { - // Calculate expected exit amount from public inputs - let mut total = 0u128; - for account_data in &inputs.account_data { - if account_data.summed_output_amount > 0 { - total += (account_data.summed_output_amount as u128) * - crate::SCALE_DOWN_FACTOR; - } - } - total - }, - nullifiers: inputs - .nullifiers - .iter() - .map(|n| n.as_ref().try_into().unwrap()) - .collect(), + exit_amount, + nullifiers: nullifier_bytes(&inputs), } .into(), ); }); } + /// Regression: cheap settlement check (TransferAmountBelowMinimum) trips after ZK + /// verification. The reorder in the dispatchable means we never even reach the + /// nullifier write — so a legitimate user whose proof became under-minimum can retry. #[test] - fn test_verify_aggregated_proof_cannot_replay() { + fn test_verify_aggregated_proof_below_minimum_does_not_consume_nullifiers() { new_test_ext().execute_with(|| { - let proof = deserialize_test_proof(); - let inputs = parse_aggregated_public_inputs(&proof).expect("Should parse"); + let (proof_bytes, inputs) = setup_valid_proof_state(); + MinimumTransferAmount::set(&Balance::MAX); - // Set up block hash to match the proof - let block_number = inputs.block_data.block_number as u64; - let block_hash_bytes: [u8; 32] = - inputs.block_data.block_hash.as_ref().try_into().unwrap(); - let block_hash = H256::from(block_hash_bytes); + let err = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes) + .expect_err("Expected TransferAmountBelowMinimum"); + assert_eq!(err.error, Error::::TransferAmountBelowMinimum.into()); - frame_system::BlockHash::::insert(block_number, block_hash); - System::set_block_number(block_number + 10); + assert_nullifiers_unused(&inputs, "below-minimum failure"); + }); + } - let proof_bytes = get_test_proof_bytes(); + /// Regression: an attacker (or chain state) leaves an exit account at `Balance::MAX` + /// so the dispatchable's settlement mint overflows AFTER nullifiers are written. + /// Without `#[transactional]` the nullifiers would persist, permanently locking the + /// legitimate withdrawal. With the wrapper, storage is fully rolled back. + #[test] + fn test_verify_aggregated_proof_settlement_overflow_rolls_back_nullifiers() { + new_test_ext().execute_with(|| { + let (proof_bytes, inputs) = setup_valid_proof_state(); + + let target = exit_accounts(&inputs).into_iter().next().expect("exit account"); + >::write_balance(&target, Balance::MAX) + .expect("write_balance to MAX should succeed"); + + let err = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes) + .expect_err("Settlement mint must overflow"); + assert_eq!(err.error, DispatchError::Arithmetic(ArithmeticError::Overflow)); + + assert_nullifiers_unused(&inputs, "settlement-mint overflow"); + }); + } + + #[test] + fn test_verify_aggregated_proof_cannot_replay() { + new_test_ext().execute_with(|| { + let (proof_bytes, _inputs) = setup_valid_proof_state(); - // First submission should succeed assert_ok!(Wormhole::verify_aggregated_proof( RawOrigin::None.into(), proof_bytes.clone() )); - // Second submission with same proof should fail (nullifiers already used) - let result = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes); - assert!(result.is_err()); - let err = result.unwrap_err(); + let err = Wormhole::verify_aggregated_proof(RawOrigin::None.into(), proof_bytes) + .expect_err("Replay must fail"); assert_eq!(err.error, Error::::NullifierAlreadyUsed.into()); }); }