From 7616579c04ebc39a1d8a95da35686ef41066d3d2 Mon Sep 17 00:00:00 2001 From: zer0 Date: Tue, 10 Sep 2024 00:28:10 +0200 Subject: [PATCH 1/2] wallet-core: ensure the pick notes code will works on FFI too - Change custom vectors with `owned::NoteList` instead - Change the module structure: avoid clippy pedantic while keeping the same lib's root symbols Resolves #2324 --- wallet-core/src/ffi.rs | 2 +- wallet-core/src/lib.rs | 9 +- wallet-core/src/notes.rs | 157 ++------------------ wallet-core/src/notes/balance.rs | 73 +++++++++ wallet-core/src/notes/owned.rs | 110 ++++++++++++++ wallet-core/src/{input.rs => notes/pick.rs} | 88 ++++------- wallet-core/tests/notes.rs | 94 +++++++----- 7 files changed, 291 insertions(+), 242 deletions(-) create mode 100644 wallet-core/src/notes/balance.rs create mode 100644 wallet-core/src/notes/owned.rs rename wallet-core/src/{input.rs => notes/pick.rs} (58%) diff --git a/wallet-core/src/ffi.rs b/wallet-core/src/ffi.rs index c09459d4c2..2539d89868 100644 --- a/wallet-core/src/ffi.rs +++ b/wallet-core/src/ffi.rs @@ -126,7 +126,7 @@ pub unsafe fn map_owned( let notes: Vec = from_bytes::>(¬es) .or(Err(ErrorCode::UnarchivingError))?; - let owned = notes::map_owned(&keys, notes); + let owned = notes::owned::map(&keys, notes); keys.into_iter().for_each(|mut sk| sk.zeroize()); diff --git a/wallet-core/src/lib.rs b/wallet-core/src/lib.rs index 0b6e80bac7..cf2bb291e7 100644 --- a/wallet-core/src/lib.rs +++ b/wallet-core/src/lib.rs @@ -22,7 +22,6 @@ extern crate alloc; #[macro_use] mod ffi; -pub mod input; pub mod keys; pub mod notes; pub mod transaction; @@ -32,8 +31,12 @@ pub type Seed = [u8; 64]; pub mod prelude { //! Re-export of the most commonly used types and traits. - pub use crate::input::MAX_INPUT_NOTES; pub use crate::keys; + pub use crate::notes::MAX_INPUT_NOTES; } -pub use notes::{map_owned, phoenix_balance, BalanceInfo}; +pub use notes::balance::{ + calculate as phoenix_balance, TotalAmount as BalanceInfo, +}; +pub use notes::owned::map as map_owned; +pub use notes::pick::notes as pick_notes; diff --git a/wallet-core/src/notes.rs b/wallet-core/src/notes.rs index 54c7fcb3fb..62d6fad0c9 100644 --- a/wallet-core/src/notes.rs +++ b/wallet-core/src/notes.rs @@ -6,150 +6,13 @@ //! Provides functions and types for interacting with notes. -use alloc::vec::Vec; -use core::ops::Index; -use dusk_bytes::{DeserializableSlice, Serializable, Write}; -use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey}; -use execution_core::{ - transfer::phoenix::{NoteLeaf, SecretKey as PhoenixSecretKey}, - BlsScalar, -}; - -use rkyv::{Archive, Deserialize, Serialize}; - -// The maximum amount of input notes that can be spend in one -// phoenix-transaction -const MAX_INPUT_NOTES: usize = 4; - -/// A collection of notes stored as key-value pairs. -/// The key is a `BlsScalar` and the value is a `NoteLeaf`. -/// Duplicates are allowed. -#[derive(Default, Archive, Serialize, Deserialize, Debug)] -pub struct OwnedList { - /// The underlying storage of key-value pairs where - /// `BlsScalar` is the key and `NoteLeaf` is the value. - entries: Vec<(BlsScalar, NoteLeaf)>, -} - -impl OwnedList { - /// Inserts a new key-value pair into the collection. - pub fn insert(&mut self, key: BlsScalar, value: NoteLeaf) { - self.entries.push((key, value)); - } - - /// Returns the number of entries (key-value pairs) in the collection. - #[must_use] - pub fn len(&self) -> usize { - self.entries.len() - } - - /// Checks if the collection is empty. - #[must_use] - pub fn is_empty(&self) -> bool { - self.entries.is_empty() - } - - /// Retrieves the value (`NoteLeaf`) associated with a given key - #[must_use] - pub fn get(&self, key: &BlsScalar) -> Option<&NoteLeaf> { - self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v) - } - - /// Retrieves all keys in the collection. - #[must_use] - pub fn keys(&self) -> Vec { - self.entries.iter().map(|(k, _)| *k).collect() - } -} - -impl Index<&BlsScalar> for OwnedList { - type Output = NoteLeaf; - - /// Retrieves the value (`NoteLeaf`) associated with a given key - /// (`BlsScalar`). - /// - /// Panics if the key is not found in the collection. - fn index(&self, index: &BlsScalar) -> &Self::Output { - self.get(index).expect("key not found") - } -} - -/// Filter all notes and their block height that are owned by the given keys, -/// mapped to their nullifiers. -pub fn map_owned( - keys: impl AsRef<[PhoenixSecretKey]>, - notes: impl AsRef<[NoteLeaf]>, -) -> OwnedList { - notes.as_ref().iter().fold( - OwnedList::default(), - |mut notes_map, note_leaf| { - for sk in keys.as_ref() { - if sk.owns(note_leaf.note.stealth_address()) { - let nullifier = note_leaf.note.gen_nullifier(sk); - notes_map.insert(nullifier, note_leaf.clone()); - break; - } - } - notes_map - }, - ) -} - -/// Calculate the sum for all the given [`Note`]s that belong to the given -/// [`PhoenixViewKey`]. -pub fn phoenix_balance( - phoenix_vk: &PhoenixViewKey, - notes: impl Iterator, -) -> BalanceInfo -where - T: AsRef, -{ - let mut values: Vec = notes - .filter_map(|note| note.as_ref().value(Some(phoenix_vk)).ok()) - .collect(); - - values.sort_by(|a, b| b.cmp(a)); - - let spendable = values.iter().take(MAX_INPUT_NOTES).sum(); - let value = spendable + values.iter().skip(MAX_INPUT_NOTES).sum::(); - - BalanceInfo { value, spendable } -} - -/// Information about the balance of a particular key. -#[derive(Debug, Default, Hash, Clone, Copy, PartialEq, Eq)] -pub struct BalanceInfo { - /// The total value of the balance. - pub value: u64, - /// The maximum _spendable_ value in a single transaction. This is - /// different from `value` since there is a maximum number of notes one can - /// spend. - pub spendable: u64, -} - -impl Serializable<{ 2 * u64::SIZE }> for BalanceInfo { - type Error = dusk_bytes::Error; - - fn from_bytes(buf: &[u8; Self::SIZE]) -> Result - where - Self: Sized, - { - let mut reader = &buf[..]; - - let value = u64::from_reader(&mut reader)?; - let spendable = u64::from_reader(&mut reader)?; - - Ok(Self { value, spendable }) - } - - #[allow(unused_must_use)] - fn to_bytes(&self) -> [u8; Self::SIZE] { - let mut buf = [0u8; Self::SIZE]; - let mut writer = &mut buf[..]; - - writer.write(&self.value.to_bytes()); - writer.write(&self.spendable.to_bytes()); - - buf - } -} +/// Module for balance information. +pub mod balance; +/// Module for owned notes. +pub mod owned; +/// Module for picking notes. +pub mod pick; + +/// The maximum amount of input notes that can be spend in one +/// phoenix-transaction +pub const MAX_INPUT_NOTES: usize = 4; diff --git a/wallet-core/src/notes/balance.rs b/wallet-core/src/notes/balance.rs new file mode 100644 index 0000000000..2e05111004 --- /dev/null +++ b/wallet-core/src/notes/balance.rs @@ -0,0 +1,73 @@ +// 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. + +//! Provides functions and types for calculate notes' balance. + +use alloc::vec::Vec; + +use dusk_bytes::{DeserializableSlice, Serializable, Write}; +use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey}; + +use crate::notes::MAX_INPUT_NOTES; + +/// Calculate the sum for all the given [`Note`]s that belong to the given +/// [`PhoenixViewKey`]. +pub fn calculate( + vk: &PhoenixViewKey, + notes: impl Iterator, +) -> TotalAmount +where + T: AsRef, +{ + let mut values: Vec = notes + .filter_map(|note| note.as_ref().value(Some(vk)).ok()) + .collect(); + + values.sort_by(|a, b| b.cmp(a)); + + let spendable = values.iter().take(MAX_INPUT_NOTES).sum(); + let value = spendable + values.iter().skip(MAX_INPUT_NOTES).sum::(); + + TotalAmount { value, spendable } +} + +/// Information about the balance of a particular key. +#[derive(Debug, Default, Hash, Clone, Copy, PartialEq, Eq)] +pub struct TotalAmount { + /// The total value of the balance. + pub value: u64, + /// The maximum _spendable_ value in a single transaction. This is + /// different from `value` since there is a maximum number of notes one can + /// spend. + pub spendable: u64, +} + +impl Serializable<{ 2 * u64::SIZE }> for TotalAmount { + type Error = dusk_bytes::Error; + + fn from_bytes(buf: &[u8; Self::SIZE]) -> Result + where + Self: Sized, + { + let mut reader = &buf[..]; + + let value = u64::from_reader(&mut reader)?; + let spendable = u64::from_reader(&mut reader)?; + + Ok(Self { value, spendable }) + } + + #[allow(unused_must_use)] + fn to_bytes(&self) -> [u8; Self::SIZE] { + let mut buf = [0u8; Self::SIZE]; + let mut writer = &mut buf[..]; + + writer.write(&self.value.to_bytes()); + writer.write(&self.spendable.to_bytes()); + + buf + } +} diff --git a/wallet-core/src/notes/owned.rs b/wallet-core/src/notes/owned.rs new file mode 100644 index 0000000000..02cc625b42 --- /dev/null +++ b/wallet-core/src/notes/owned.rs @@ -0,0 +1,110 @@ +// 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. + +//! Provides functions and types to handle notes' ownership. + +use alloc::vec::Vec; + +use core::ops::Index; +use core::slice::Iter; +use execution_core::{ + transfer::phoenix::{NoteLeaf, SecretKey as PhoenixSecretKey}, + BlsScalar, +}; +use rkyv::{Archive, Deserialize, Serialize}; + +/// A collection of notes stored as key-value pairs. +/// The key is a `BlsScalar` and the value is a `NoteLeaf`. +/// Duplicates are allowed. +#[derive(Default, Archive, Serialize, Deserialize, Debug, PartialEq, Clone)] +pub struct NoteList { + /// The underlying storage of key-value pairs where + /// `BlsScalar` is the key and `NoteLeaf` is the value. + entries: Vec<(BlsScalar, NoteLeaf)>, +} + +impl NoteList { + /// Inserts a new key-value pair into the collection. + pub fn insert(&mut self, key: BlsScalar, value: NoteLeaf) { + self.entries.push((key, value)); + } + + /// Returns the number of entries (key-value pairs) in the collection. + #[must_use] + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Checks if the collection is empty. + #[must_use] + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Retrieves the value (`NoteLeaf`) associated with a given key + #[must_use] + pub fn get(&self, key: &BlsScalar) -> Option<&NoteLeaf> { + self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v) + } + + /// Retrieves all keys in the collection. + #[must_use] + pub fn keys(&self) -> Vec { + self.entries.iter().map(|(k, _)| *k).collect() + } + + /// Returns an iterator over the key-value pairs. + pub fn iter(&self) -> Iter<'_, (BlsScalar, NoteLeaf)> { + self.entries.iter() + } +} + +impl Index<&BlsScalar> for NoteList { + type Output = NoteLeaf; + + /// Retrieves the value (`NoteLeaf`) associated with a given key + /// (`BlsScalar`). + /// + /// Panics if the key is not found in the collection. + fn index(&self, index: &BlsScalar) -> &Self::Output { + self.get(index).expect("key not found") + } +} + +impl<'a> IntoIterator for &'a NoteList { + type IntoIter = core::slice::Iter<'a, (BlsScalar, NoteLeaf)>; + type Item = &'a (BlsScalar, NoteLeaf); + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl From> for NoteList { + fn from(entries: Vec<(BlsScalar, NoteLeaf)>) -> Self { + NoteList { entries } + } +} + +/// Filter all notes and their block height that are owned by the given keys, +/// mapped to their nullifiers. +pub fn map( + keys: impl AsRef<[PhoenixSecretKey]>, + notes: impl AsRef<[NoteLeaf]>, +) -> NoteList { + notes.as_ref().iter().fold( + NoteList::default(), + |mut notes_map, note_leaf| { + for sk in keys.as_ref() { + if sk.owns(note_leaf.note.stealth_address()) { + let nullifier = note_leaf.note.gen_nullifier(sk); + notes_map.insert(nullifier, note_leaf.clone()); + break; + } + } + notes_map + }, + ) +} diff --git a/wallet-core/src/input.rs b/wallet-core/src/notes/pick.rs similarity index 58% rename from wallet-core/src/input.rs rename to wallet-core/src/notes/pick.rs index d9355a13b9..62d5c67f3f 100644 --- a/wallet-core/src/input.rs +++ b/wallet-core/src/notes/pick.rs @@ -8,48 +8,11 @@ use alloc::vec::Vec; -use execution_core::transfer::phoenix::Note; +use crate::notes::owned::NoteList; +use crate::notes::MAX_INPUT_NOTES; +use execution_core::transfer::phoenix::{NoteLeaf, ViewKey as PhoenixViewKey}; use execution_core::BlsScalar; -/// The maximum amount of input notes that can be spend in one -/// phoenix-transaction -pub const MAX_INPUT_NOTES: usize = 4; - -/// Pick the notes to be used in a transaction from a vector of notes. -/// -/// The resulting array is only 4 notes long, the argument of this function can -/// be arbitary amount of notes. -/// -/// # Errors -/// -/// If the target sum is greater than the sum of the notes then an error is -/// returned. If the notes vector is empty then an error is returned. -/// -/// See `InputNotesError` type for possible errors -/// this function can yield. -#[must_use] -pub fn try_input_notes( - nodes: Vec<(Note, u64, BlsScalar)>, - target_sum: u64, -) -> Vec<(Note, BlsScalar)> { - if nodes.is_empty() { - return Vec::new(); - } - - let mut i = 0; - let mut sum = 0; - while sum < target_sum && i < nodes.len() { - sum = sum.saturating_add(nodes[i].1); - i += 1; - } - - if sum < target_sum { - return Vec::new(); - } - - pick_notes(target_sum, nodes) -} - /// Pick the notes to be used in a transaction from a vector of notes. /// /// The notes are picked in a way to maximize the number of notes used, @@ -59,20 +22,34 @@ pub fn try_input_notes( /// larger or equal to the given value. If such a slice is not found, an /// empty vector is returned. /// -/// Note: it is presupposed that the input notes contain enough balance to -/// cover the given `value`. -fn pick_notes( - value: u64, - notes_and_values: Vec<(Note, u64, BlsScalar)>, -) -> Vec<(Note, BlsScalar)> { - let mut notes_and_values = notes_and_values; - let len = notes_and_values.len(); - - if len <= MAX_INPUT_NOTES { - return notes_and_values - .into_iter() - .map(|(note, _, b)| (note, b)) - .collect(); +/// If the target sum is greater than the sum of the notes then an +/// empty vector is returned. +#[must_use] +pub fn notes(vk: &PhoenixViewKey, notes: NoteList, value: u64) -> NoteList { + if notes.is_empty() { + return NoteList::default(); + } + + let mut notes_and_values: Vec<(NoteLeaf, u64, BlsScalar)> = notes + .iter() + .filter_map(|(nullifier, leaf)| { + leaf.as_ref() + .value(Some(vk)) + .ok() + .map(|value| (leaf.clone(), value, *nullifier)) + }) + .collect(); + + let sum: u64 = notes_and_values + .iter() + .fold(0, |sum, &(_, value, _)| sum.saturating_add(value)); + + if sum < value { + return NoteList::default(); + } + + if notes.len() <= MAX_INPUT_NOTES { + return notes; } notes_and_values.sort_by(|(_, aval, _), (_, bval, _)| aval.cmp(bval)); @@ -84,8 +61,9 @@ fn pick_notes( >= value }) .map(|index| notes_and_values[index].clone()) - .map(|(n, _, b)| (n, b)) + .map(|(n, _, b)| (b, n)) .to_vec() + .into() } fn pick_lexicographic bool>( diff --git a/wallet-core/tests/notes.rs b/wallet-core/tests/notes.rs index 3db5bec144..6c7a973a8f 100644 --- a/wallet-core/tests/notes.rs +++ b/wallet-core/tests/notes.rs @@ -10,14 +10,14 @@ use rand::{rngs::StdRng, CryptoRng, RngCore, SeedableRng}; use execution_core::{ transfer::phoenix::{ Note, NoteLeaf, PublicKey as PhoenixPublicKey, - SecretKey as PhoenixSecretKey, + SecretKey as PhoenixSecretKey, ViewKey as PhoenixViewKey, }, - BlsScalar, JubJubScalar, + JubJubScalar, }; use wallet_core::{ - input::try_input_notes, keys::derive_multiple_phoenix_sk, - keys::derive_phoenix_sk, map_owned, phoenix_balance, BalanceInfo, Seed, + keys::derive_multiple_phoenix_sk, keys::derive_phoenix_sk, map_owned, + notes::owned::NoteList, phoenix_balance, pick_notes, BalanceInfo, Seed, }; /// Generate a note, useful for testing purposes @@ -48,6 +48,19 @@ pub fn gen_note( } } +/// Generate a note leaf, useful for testing purposes +pub fn gen_note_leaf( + rng: &mut T, + obfuscated_note: bool, + owner_pk: &PhoenixPublicKey, + value: u64, +) -> NoteLeaf { + let block_height = 1; + let note = gen_note(rng, obfuscated_note, owner_pk, value); + + NoteLeaf { note, block_height } +} + #[test] fn test_map_owned() { let mut rng = StdRng::seed_from_u64(0xdab); @@ -164,53 +177,62 @@ fn knapsack_works() { let mut rng = rand_chacha::ChaCha12Rng::seed_from_u64(0xbeef); - // sanity check - assert_eq!(try_input_notes(vec![], 70), Vec::new(),); - let sk = PhoenixSecretKey::random(&mut rng); + let vk = PhoenixViewKey::from(&sk); let pk = PhoenixPublicKey::from(&sk); + // sanity check + assert!(pick_notes(&vk, NoteList::default(), 70).is_empty()); + // basic check - let note = gen_note(&mut rng, true, &pk, 100); - let n = note.gen_nullifier(&sk); - let available = vec![(note, 100, n)]; - let inputs_notes: Vec<(Note, BlsScalar)> = available - .clone() - .into_iter() - .map(|(a, _, b)| (a, b)) - .collect(); - let input = try_input_notes(available, 70); - assert_eq!(input, inputs_notes); + let leaf = gen_note_leaf(&mut rng, true, &pk, 100); + let n = leaf.note.gen_nullifier(&sk); + let available = NoteList::from(vec![(n, leaf)]); + + let input = pick_notes(&vk, available.clone(), 70); + assert_eq!(input, available); // out of balance basic check - let note = gen_note(&mut rng, true, &pk, 100); - let available = vec![(note, 100, n)]; - assert_eq!(try_input_notes(available, 101), Vec::new()); + let leaf = gen_note_leaf(&mut rng, true, &pk, 100); + let available = NoteList::from(vec![(n, leaf)]); + assert!(pick_notes(&vk, available, 101).is_empty()); // multiple inputs check // note: this test is checking a naive, simple order-based output - let note1 = gen_note(&mut rng, true, &pk, 100); - let note2 = gen_note(&mut rng, true, &pk, 500); - let note3 = gen_note(&mut rng, true, &pk, 300); - - let available = vec![(note1, 100, n), (note2, 500, n), (note3, 300, n)]; + let leaf = [ + gen_note_leaf(&mut rng, true, &pk, 100), + gen_note_leaf(&mut rng, true, &pk, 500), + gen_note_leaf(&mut rng, true, &pk, 300), + ]; - let result: Vec<(Note, BlsScalar)> = available - .clone() - .into_iter() - .map(|(a, _, b)| (a, b)) + let available: Vec<(_, _)> = leaf + .iter() + .map(|l| { + let n = l.note.gen_nullifier(&sk); + (n, l.clone()) + }) .collect(); - assert_eq!(try_input_notes(available.clone(), 600), result); + let available = NoteList::from(available); + + assert_eq!(pick_notes(&vk, available.clone(), 600), available); - let note1 = gen_note(&mut rng, true, &pk, 100); - let note2 = gen_note(&mut rng, true, &pk, 500); - let note3 = gen_note(&mut rng, true, &pk, 300); + let leaf = [ + gen_note_leaf(&mut rng, true, &pk, 100), + gen_note_leaf(&mut rng, true, &pk, 500), + gen_note_leaf(&mut rng, true, &pk, 300), + ]; - let n = note1.gen_nullifier(&sk); + let available: Vec<(_, _)> = leaf + .iter() + .map(|l| { + let n = l.note.gen_nullifier(&sk); + (n, l.clone()) + }) + .collect(); - let available = vec![(note1, 100, n), (note2, 500, n), (note3, 300, n)]; + let available = NoteList::from(available); - assert_eq!(try_input_notes(available, 901), Vec::new()); + assert_eq!(pick_notes(&vk, available, 901), NoteList::default()); } From ceab2beb4cf8cf49b4a0bd3ac75d9854800baa79 Mon Sep 17 00:00:00 2001 From: zer0 Date: Tue, 10 Sep 2024 01:05:42 +0200 Subject: [PATCH 2/2] rusk-wallet: change `try_input_notes` with `pick_notes` --- rusk-wallet/src/clients.rs | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/rusk-wallet/src/clients.rs b/rusk-wallet/src/clients.rs index 4c1a7ff9f9..290389e703 100644 --- a/rusk-wallet/src/clients.rs +++ b/rusk-wallet/src/clients.rs @@ -11,11 +11,14 @@ use execution_core::{ transfer::{phoenix::Prove, Transaction}, Error as ExecutionCoreError, }; + +use execution_core::transfer::phoenix::{Note, NoteLeaf}; + use flume::Receiver; use tokio::time::{sleep, Duration}; use wallet_core::{ - input::try_input_notes, keys::{derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk}, + pick_notes, }; use zeroize::Zeroize; @@ -194,19 +197,19 @@ impl State { .into_iter() .map(|data| { let note = data.note; - let nullifiers = note.gen_nullifier(&sk); - let value = note.value(Some(&vk)).unwrap(); - - Ok((note, value, nullifiers)) + let block_height = data.height; + let nullifier = note.gen_nullifier(&sk); + let leaf = NoteLeaf { note, block_height }; + Ok((nullifier, leaf)) }) .collect(); - let inputs = try_input_notes(inputs?, target) + let inputs = pick_notes(&vk, inputs?.into(), target) .into_iter() - .map(|(note, scalar)| { - let opening = self.fetch_opening(¬e)?; + .map(|(scalar, note)| { + let opening = self.fetch_opening(note.as_ref())?; - Ok((note, opening, scalar)) + Ok((note.note.clone(), opening, *scalar)) }) .collect();