diff --git a/nssa/src/public_transaction/witness_set.rs b/nssa/src/public_transaction/witness_set.rs index 8222c5ed..7a32c0ea 100644 --- a/nssa/src/public_transaction/witness_set.rs +++ b/nssa/src/public_transaction/witness_set.rs @@ -1,6 +1,6 @@ use borsh::{BorshDeserialize, BorshSerialize}; -use crate::{PrivateKey, PublicKey, Signature, public_transaction::Message}; +use crate::{PrivateKey, PublicKey, Signature, error::NssaError, public_transaction::Message}; #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct WitnessSet { @@ -8,19 +8,35 @@ pub struct WitnessSet { } impl WitnessSet { - #[must_use] - pub fn from_list(signatures: &[Signature], pub_keys: &[PublicKey]) -> Self { - assert_eq!(signatures.len(), pub_keys.len()); + pub fn from_list( + message: &Message, + signatures: &[Signature], + pub_keys: &[PublicKey], + ) -> Result { + if signatures.len() != pub_keys.len() { + return Err(NssaError::InvalidInput( + "`nssa::public_transaction::witness_set::from_list()`: mismatch in signature and public key counts".to_owned(), + )); + } + let message_hash = message.hash_message(); let signatures_and_public_keys = signatures .iter() .zip(pub_keys.iter()) - .map(|(sig, key)| (sig.clone(), key.clone())) - .collect(); + .map(|(sig, key)| { + if sig.is_valid_for(&message_hash, key) { + Ok((sig.clone(), key.clone())) + } else { + Err(NssaError::InvalidInput( + "`nssa::public_transaction::witness_set::from_list()`: signature does not correspond to public key".to_owned(), + )) + } + }) + .collect::>()?; - Self { + Ok(Self { signatures_and_public_keys, - } + }) } #[must_use] @@ -74,6 +90,73 @@ mod tests { use super::*; use crate::AccountId; + #[test] + fn from_list_accepts_valid_pairs() { + let key1 = PrivateKey::try_new([42; 32]).unwrap(); + let key2 = PrivateKey::try_new([13; 32]).unwrap(); + let pubkey1 = PublicKey::new_from_private_key(&key1); + let pubkey2 = PublicKey::new_from_private_key(&key2); + let addr1 = AccountId::from(&pubkey1); + let addr2 = AccountId::from(&pubkey2); + let message = Message::try_new::>( + [1_u32; 8], + vec![addr1, addr2], + vec![1_u128.into(), 2_u128.into()], + vec![], + ) + .unwrap(); + + let WitnessSet { + signatures_and_public_keys, + } = WitnessSet::for_message(&message, &[&key1, &key2]); + let (sigs, keys): (Vec<_>, Vec<_>) = signatures_and_public_keys.into_iter().unzip(); + + assert!(WitnessSet::from_list(&message, &sigs, &keys).is_ok()); + } + + #[test] + fn from_list_rejects_mismatched_pairs() { + let key1 = PrivateKey::try_new([42; 32]).unwrap(); + let key2 = PrivateKey::try_new([13; 32]).unwrap(); + let pubkey1 = PublicKey::new_from_private_key(&key1); + let pubkey2 = PublicKey::new_from_private_key(&key2); + let addr1 = AccountId::from(&pubkey1); + let addr2 = AccountId::from(&pubkey2); + let message = Message::try_new::>( + [1_u32; 8], + vec![addr1, addr2], + vec![1_u128.into(), 2_u128.into()], + vec![], + ) + .unwrap(); + + let WitnessSet { + signatures_and_public_keys, + } = WitnessSet::for_message(&message, &[&key1, &key2]); + let (sigs, keys): (Vec<_>, Vec<_>) = signatures_and_public_keys.into_iter().unzip(); + + // Swapped keys should be rejected. + assert!( + WitnessSet::from_list(&message, &sigs, &[keys[1].clone(), keys[0].clone()]).is_err() + ); + } + + #[test] + fn from_list_rejects_length_mismatch() { + let key1 = PrivateKey::try_new([1_u8; 32]).unwrap(); + let pubkey1 = PublicKey::new_from_private_key(&key1); + let addr1 = AccountId::from(&pubkey1); + let message = + Message::try_new::>([0; 8], vec![addr1], vec![1_u128.into()], vec![]).unwrap(); + + let WitnessSet { + signatures_and_public_keys, + } = WitnessSet::for_message(&message, &[&key1]); + let (sigs, _keys): (Vec<_>, Vec<_>) = signatures_and_public_keys.into_iter().unzip(); + + assert!(WitnessSet::from_list(&message, &sigs, &[]).is_err()); + } + #[test] fn for_message_constructor() { let key1 = PrivateKey::try_new([1; 32]).unwrap(); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index fdb83d94..6e4ac88f 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -414,7 +414,7 @@ impl WalletCore { ) .unwrap(); - let witness_set = Self::sign_privacy_message(&message, &proof, &acc_manager); + let witness_set = Self::sign_privacy_message(&message, proof.clone(), &acc_manager); let tx = PrivacyPreservingTransaction::new(message, witness_set); let shared_secrets: Vec<_> = private_account_keys @@ -572,12 +572,12 @@ impl WalletCore { #[must_use] pub fn sign_privacy_message( message: &nssa::privacy_preserving_transaction::Message, - proof: &Proof, + proof: Proof, acc_manager: &privacy_preserving_tx::AccountManager, ) -> nssa::privacy_preserving_transaction::witness_set::WitnessSet { nssa::privacy_preserving_transaction::witness_set::WitnessSet::for_message( message, - proof.clone(), + proof, &acc_manager.public_account_auth(), ) } diff --git a/wallet/src/program_facades/native_token_transfer/public.rs b/wallet/src/program_facades/native_token_transfer/public.rs index 2ddb0d87..d99015e0 100644 --- a/wallet/src/program_facades/native_token_transfer/public.rs +++ b/wallet/src/program_facades/native_token_transfer/public.rs @@ -42,7 +42,7 @@ impl NativeTokenTransfer<'_> { // Assumes this now silently skips accounts without signing keys let witness_set = WalletCore::sign_public_message(self.0, &message, &sign_ids) - .expect("`WalletCore::sign_public_message() failed to produce a signature for a NativeTokenTransfer."); + .expect("`WalletCore::sign_public_message()` failed to produce a signature for a NativeTokenTransfer."); let tx = PublicTransaction::new(message, witness_set);