From ba594e678801fc899756b3f114df6119ba350c15 Mon Sep 17 00:00:00 2001 From: moana Date: Wed, 14 Aug 2024 13:44:50 +0200 Subject: [PATCH] wallet-core: Add filter_map function Resolves #2129 --- wallet-core/src/keys.rs | 18 ++++- wallet-core/src/lib.rs | 28 ++++++- wallet-core/tests/balance.rs | 72 ----------------- wallet-core/tests/keys.rs | 26 +++++- wallet-core/tests/notes.rs | 148 +++++++++++++++++++++++++++++++++++ 5 files changed, 217 insertions(+), 75 deletions(-) delete mode 100644 wallet-core/tests/balance.rs create mode 100644 wallet-core/tests/notes.rs diff --git a/wallet-core/src/keys.rs b/wallet-core/src/keys.rs index 74eae1a81a..76535042c7 100644 --- a/wallet-core/src/keys.rs +++ b/wallet-core/src/keys.rs @@ -6,8 +6,12 @@ //! Utilities to derive keys from the seed. +use alloc::vec::Vec; +use core::ops::Range; + use rand_chacha::{rand_core::SeedableRng, ChaCha12Rng}; use sha2::{Digest, Sha256}; +use zeroize::Zeroize; use execution_core::{ signatures::bls::SecretKey as BlsSecretKey, @@ -16,7 +20,6 @@ use execution_core::{ ViewKey as PhoenixViewKey, }, }; -use zeroize::Zeroize; use crate::RNG_SEED; @@ -40,6 +43,19 @@ pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey { PhoenixSecretKey::random(&mut rng_with_index(seed, index, b"PSK")) } +/// Generates multiple [`PhoenixSecretKey`] from a seed and a range of indices. +/// +/// The randomness is generated using [`rng_with_index`]. +#[must_use] +pub fn derive_multiple_phoenix_sk( + seed: &[u8; RNG_SEED], + index_range: Range, +) -> Vec { + index_range + .map(|index| derive_phoenix_sk(seed, index)) + .collect() +} + /// Generates a [`PheonixPublicKey`] from its seed and index. /// /// First the [`PhoenixSecretKey`] is derived with [`derive_phoenix_sk`], then diff --git a/wallet-core/src/lib.rs b/wallet-core/src/lib.rs index 37ed06f247..69612a40dd 100644 --- a/wallet-core/src/lib.rs +++ b/wallet-core/src/lib.rs @@ -22,9 +22,35 @@ pub const RNG_SEED: usize = 64; // phoenix-transaction const MAX_INPUT_NOTES: usize = 4; +use alloc::collections::btree_map::BTreeMap; use alloc::vec::Vec; -use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey}; +use execution_core::{ + transfer::phoenix::{ + Note, SecretKey as PhoenixSecretKey, ViewKey as PhoenixViewKey, + }, + BlsScalar, +}; + +/// Filter all notes that are owned by the given keys, mapped to their +/// nullifiers. +pub fn map_owned( + keys: impl AsRef<[PhoenixSecretKey]>, + notes: impl AsRef<[Note]>, +) -> BTreeMap { + notes + .as_ref() + .iter() + .fold(BTreeMap::new(), |mut notes_map, note| { + for sk in keys.as_ref() { + if sk.owns(note.stealth_address()) { + let nullifier = note.gen_nullifier(sk); + notes_map.insert(nullifier, note.clone()); + } + } + notes_map + }) +} /// Calculate the sum for all the given [`Note`]s that belong to the given /// [`PhoenixViewKey`]. diff --git a/wallet-core/tests/balance.rs b/wallet-core/tests/balance.rs deleted file mode 100644 index 7bbf6a560a..0000000000 --- a/wallet-core/tests/balance.rs +++ /dev/null @@ -1,72 +0,0 @@ -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at http://mozilla.org/MPL/2.0/. -// -// Copyright (c) DUSK NETWORK. All rights reserved. - -use ff::Field; -use rand::rngs::StdRng; -use rand::SeedableRng; - -use execution_core::{ - transfer::phoenix::{ - Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey, - }, - JubJubScalar, -}; - -use wallet_core::{phoenix_balance, BalanceInfo}; - -#[test] -fn test_balance() { - let mut rng = StdRng::seed_from_u64(0xdab); - - let owner_sk = PhoenixSecretKey::random(&mut rng); - let owner_pk = PhoenixPublicKey::from(&owner_sk); - let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng)); - - let mut notes = Vec::new(); - - // create the notes - for value in 0..=10 { - let value_blinder = JubJubScalar::random(&mut rng); - let sender_blinder = [ - JubJubScalar::random(&mut rng), - JubJubScalar::random(&mut rng), - ]; - - // we want to test with a mix of transparent and obfuscated notes so we - // make every 10th note transparent - let note = if value % 10 == 0 { - Note::transparent( - &mut rng, - &sender_pk, - &owner_pk, - value, - sender_blinder, - ) - } else { - Note::obfuscated( - &mut rng, - &sender_pk, - &owner_pk, - value, - value_blinder, - sender_blinder, - ) - }; - notes.push(note); - } - - // the sum of these notes should be 5 * 11 = 55 - // and the spendable notes are 7 + 8 + 9 + 10 = 34 - let expected_balance = BalanceInfo { - value: 55, - spendable: 34, - }; - - assert_eq!( - phoenix_balance(&(&owner_sk).into(), notes), - expected_balance - ); -} diff --git a/wallet-core/tests/keys.rs b/wallet-core/tests/keys.rs index 5eaf1b2442..681592c49f 100644 --- a/wallet-core/tests/keys.rs +++ b/wallet-core/tests/keys.rs @@ -7,7 +7,8 @@ use dusk_bytes::Serializable; use wallet_core::keys::{ - derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk, + derive_bls_sk, derive_multiple_phoenix_sk, derive_phoenix_pk, + derive_phoenix_sk, derive_phoenix_vk, }; const SEED: [u8; 64] = [0; 64]; @@ -26,6 +27,29 @@ fn test_derive_phoenix_sk() { assert_eq!(derive_phoenix_sk(&SEED, INDEX).to_bytes(), sk_bytes); } +#[test] +fn test_derive_multiple_phoenix_sk() { + // it is important that we always derive the same key from a fixed seed + let sk_bytes_0 = [ + 160, 210, 234, 8, 94, 23, 76, 60, 130, 143, 137, 225, 37, 83, 68, 218, + 207, 192, 171, 235, 252, 130, 133, 62, 18, 232, 6, 49, 245, 123, 220, + 12, 250, 111, 39, 88, 24, 41, 156, 174, 241, 14, 118, 173, 11, 53, 192, + 126, 7, 119, 70, 69, 212, 230, 124, 79, 223, 140, 93, 153, 33, 147, + 163, 0, + ]; + let sk_bytes_1 = [ + 0, 229, 97, 222, 152, 25, 163, 173, 84, 216, 211, 69, 205, 122, 63, + 227, 98, 234, 163, 66, 145, 71, 217, 221, 29, 78, 36, 77, 68, 29, 144, + 2, 235, 80, 237, 21, 95, 54, 16, 89, 74, 200, 124, 248, 119, 216, 38, + 167, 19, 129, 205, 138, 218, 57, 198, 4, 4, 202, 115, 62, 55, 213, 141, + 0, + ]; + + let keys = derive_multiple_phoenix_sk(&SEED, INDEX..INDEX + 2); + assert_eq!(keys[0].to_bytes(), sk_bytes_0,); + assert_eq!(keys[1].to_bytes(), sk_bytes_1,); +} + #[test] fn test_derive_phoenix_pk() { // it is important that we always derive the same key from a fixed seed diff --git a/wallet-core/tests/notes.rs b/wallet-core/tests/notes.rs new file mode 100644 index 0000000000..c9ecfce44f --- /dev/null +++ b/wallet-core/tests/notes.rs @@ -0,0 +1,148 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. +// +// Copyright (c) DUSK NETWORK. All rights reserved. + +use ff::Field; +use rand::rngs::StdRng; +use rand::SeedableRng; + +use execution_core::{ + transfer::phoenix::{ + Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey, + }, + JubJubScalar, +}; + +use wallet_core::{ + keys::derive_multiple_phoenix_sk, map_owned, phoenix_balance, BalanceInfo, +}; + +#[test] +fn test_map_owned() { + // Assuming this set of notes where the number used as suffix is the + // "owner": + // notes := [A1, B1, C2, D2, E1, F3] + + let mut rng = StdRng::seed_from_u64(0xdab); + const SEED_1: [u8; 64] = [1; 64]; + const SEED_2: [u8; 64] = [2; 64]; + + let owner_1_sks = derive_multiple_phoenix_sk(&SEED_1, 0..3); + let owner_1_pks = [ + PhoenixPublicKey::from(&owner_1_sks[0]), + PhoenixPublicKey::from(&owner_1_sks[1]), + PhoenixPublicKey::from(&owner_1_sks[2]), + ]; + let owner_2_sks = derive_multiple_phoenix_sk(&SEED_2, 0..2); + let owner_2_pks = [ + PhoenixPublicKey::from(&owner_2_sks[0]), + PhoenixPublicKey::from(&owner_2_sks[1]), + ]; + let owner_3_pk = + PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng)); + + let value = 42; + let notes = vec![ + gen_note(&mut rng, true, &owner_1_pks[0], value), // owner 1 + gen_note(&mut rng, true, &owner_1_pks[1], value), // owner 1 + gen_note(&mut rng, true, &owner_2_pks[0], value), // owner 2 + gen_note(&mut rng, true, &owner_2_pks[1], value), // owner 2 + gen_note(&mut rng, true, &owner_1_pks[2], value), // owner 1 + gen_note(&mut rng, true, &owner_3_pk, value), // owner 3 + ]; + + // notes with idx 0, 1 and 4 are owned by owner_1 + let notes_by_1 = map_owned(&owner_1_sks, ¬es); + assert_eq!(notes_by_1.len(), 3); + let note = ¬es[0]; + let nullifier = note.gen_nullifier(&owner_1_sks[0]); + assert_eq!(¬es_by_1[&nullifier], note); + let note = ¬es[1]; + let nullifier = note.gen_nullifier(&owner_1_sks[1]); + assert_eq!(¬es_by_1[&nullifier], note); + let note = ¬es[4]; + let nullifier = note.gen_nullifier(&owner_1_sks[2]); + assert_eq!(¬es_by_1[&nullifier], note); + + // notes with idx 2 and 3 are owned by owner_2 + let notes_by_2 = map_owned(&owner_2_sks, ¬es); + assert_eq!(notes_by_2.len(), 2); + let note = ¬es[2]; + let nullifier = note.gen_nullifier(&owner_2_sks[0]); + assert_eq!(¬es_by_2[&nullifier], note); + let note = ¬es[3]; + let nullifier = note.gen_nullifier(&owner_2_sks[1]); + assert_eq!(¬es_by_2[&nullifier], note); +} + +#[test] +fn test_balance() { + let mut rng = StdRng::seed_from_u64(0xdab); + + let owner_sk = PhoenixSecretKey::random(&mut rng); + let owner_pk = PhoenixPublicKey::from(&owner_sk); + + let mut notes = Vec::new(); + + // create the notes + for value in 0..=10 { + // we want to test with a mix of transparent and obfuscated notes so we + // make every 10th note transparent + let obfuscated_note = if value % 10 == 0 { false } else { true }; + + notes.push(gen_note(&mut rng, obfuscated_note, &owner_pk, value)); + + // also push some notes that are not owned + if value % 4 == 0 { + let not_owner_pk = + PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng)); + notes.push(gen_note( + &mut rng, + obfuscated_note, + ¬_owner_pk, + value, + )); + } + } + + // the sum of these notes should be 5 * 11 = 55 + // and the spendable notes are 7 + 8 + 9 + 10 = 34 + let expected_balance = BalanceInfo { + value: 55, + spendable: 34, + }; + + assert_eq!( + phoenix_balance(&(&owner_sk).into(), notes), + expected_balance + ); +} + +fn gen_note( + rng: &mut StdRng, + obfuscated_note: bool, + owner_pk: &PhoenixPublicKey, + value: u64, +) -> Note { + let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(rng)); + + let value_blinder = JubJubScalar::random(&mut *rng); + let sender_blinder = [ + JubJubScalar::random(&mut *rng), + JubJubScalar::random(&mut *rng), + ]; + if obfuscated_note { + Note::obfuscated( + rng, + &sender_pk, + &owner_pk, + value, + value_blinder, + sender_blinder, + ) + } else { + Note::transparent(rng, &sender_pk, &owner_pk, value, sender_blinder) + } +}