diff --git a/lee/state_machine/src/privacy_preserving_transaction/message.rs b/lee/state_machine/src/privacy_preserving_transaction/message.rs index 9fafad9e..1315cf4a 100644 --- a/lee/state_machine/src/privacy_preserving_transaction/message.rs +++ b/lee/state_machine/src/privacy_preserving_transaction/message.rs @@ -84,6 +84,20 @@ impl Message { Sha256::digest(bytes).into() } + + /// Ensure that the commitments, nullifiers, and ciphertexts agree. + pub fn validate_note_lengths(&self) -> Result { + let count = self.new_nullifiers.len(); + if self.new_commitments.len() != count || self.encrypted_private_post_states.len() != count + { + return Err(LeeError::InvalidInput(format!( + "Note vectors disagree in length with {count} nullifiers, {} commitments, and {} ciphertexts", + self.new_commitments.len(), + self.encrypted_private_post_states.len(), + ))); + } + Ok(count) + } } #[cfg(test)] @@ -140,6 +154,20 @@ pub mod tests { } } + #[test] + fn validate_note_lengths_accepts_matching_and_rejects_mismatched() { + assert_eq!(Message::default().validate_note_lengths().unwrap(), 0); + + let mismatched = Message { + new_commitments: vec![Commitment::new( + &AccountId::new([0; 32]), + &Account::default(), + )], + ..Default::default() + }; + assert!(mismatched.validate_note_lengths().is_err()); + } + #[test] fn hash_privacy_pinned() { let msg = Message { diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index 1e4ba380..3f7a900d 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -506,6 +506,12 @@ impl WalletCore { tx: &lee::privacy_preserving_transaction::PrivacyPreservingTransaction, acc_decode_mask: &[AccDecodeData], ) -> Result<()> { + let note_count = tx.message.validate_note_lengths()?; + anyhow::ensure!( + note_count == acc_decode_mask.len(), + "Decode mask has {} entries but the transaction has {note_count} notes", + acc_decode_mask.len(), + ); for (output_index, acc_decode_data) in acc_decode_mask.iter().enumerate() { match acc_decode_data { AccDecodeData::Decode(secret, acc_account_id) => { @@ -731,6 +737,7 @@ impl WalletCore { let LeeTransaction::PrivacyPreserving(pp_tx) = &tx else { continue; }; + pp_tx.message.validate_note_lengths()?; // Eagerly decrypt note updates using expected nullifiers. let handled = self .storage