diff --git a/CHANGELOG.md b/CHANGELOG.md index 34b36d35b..3b5f38b65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Removed ### Fixed +- [\#129](https://github.com/Manta-Network/manta-rs/pull/129) Reduce cost of signer key-search algorithm by adding dynamic pre-computation table ### Security diff --git a/manta-accounting/src/key.rs b/manta-accounting/src/key.rs index 820c65854..06da64345 100644 --- a/manta-accounting/src/key.rs +++ b/manta-accounting/src/key.rs @@ -21,8 +21,6 @@ //! //! [`BIP-0044`]: https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki -// TODO: Build custom iterator types for [`keypairs`] and [`generate_keys`]. - use alloc::vec::Vec; use core::{ cmp, @@ -35,6 +33,7 @@ use manta_crypto::{ key::kdf::KeyDerivationFunction, rand::{RngCore, Sample}, }; +use manta_util::collections::btree_map::{self, BTreeMap}; #[cfg(feature = "serde")] use manta_util::serde::{Deserialize, Serialize}; @@ -524,6 +523,70 @@ where }) } + /// Returns a new [`ViewKeyTable`] for `self`. + #[inline] + pub fn view_key_table(self) -> ViewKeyTable<'h, H> { + ViewKeyTable::new(self) + } +} + +/// View Key Table +pub struct ViewKeyTable<'h, H> +where + H: HierarchicalKeyDerivationScheme + ?Sized, +{ + /// Account Keys + keys: AccountKeysMut<'h, H>, + + /// Pre-computed View Keys + view_keys: BTreeMap, +} + +impl<'h, H> ViewKeyTable<'h, H> +where + H: HierarchicalKeyDerivationScheme + ?Sized, +{ + /// View Key Buffer Maximum Size Limit + pub const VIEW_KEY_BUFFER_LIMIT: usize = 16 * (H::GAP_LIMIT as usize); + + /// Builds a new [`ViewKeyTable`] over the account `keys`. + #[inline] + pub fn new(keys: AccountKeysMut<'h, H>) -> Self { + Self { + keys, + view_keys: Default::default(), + } + } + + /// Returns the account keys associated to `self`. + #[inline] + pub fn into_keys(self) -> AccountKeysMut<'h, H> { + self.keys + } + + /// Returns the view key for this account at `index`, if it does not exceed the maximum index. + /// + /// # Limits + /// + /// This function uses a view key buffer that stores the computed keys to reduce the number of + /// times a re-compute of the view keys is needed while searching. The buffer only grows past + /// the current key bounds with a call to [`find_index_with_gap`](Self::find_index_with_gap) + /// which extends the buffer by at most [`GAP_LIMIT`]-many keys per round. To prevent allocating + /// too much memory, the internal buffer is capped at [`VIEW_KEY_BUFFER_LIMIT`]-many elements. + /// + /// [`GAP_LIMIT`]: HierarchicalKeyDerivationScheme::GAP_LIMIT + /// [`VIEW_KEY_BUFFER_LIMIT`]: Self::VIEW_KEY_BUFFER_LIMIT + #[inline] + pub fn view_key(&mut self, index: KeyIndex) -> Option<&H::SecretKey> { + btree_map::get_or_mutate(&mut self.view_keys, &index, |map| { + let next_key = self.keys.view_key(index)?; + if map.len() == Self::VIEW_KEY_BUFFER_LIMIT { + btree_map::pop_last(map); + } + Some(btree_map::insert_then_get(map, index, next_key)) + }) + } + /// Applies `f` to the view keys generated by `self` returning the first non-`None` result with /// it's key index and key attached, or returns `None` if every application of `f` returned /// `None`. @@ -534,12 +597,12 @@ where { let mut index = KeyIndex::default(); loop { - let view_key = self.view_key(index)?; - if let Some(item) = f(&view_key) { - self.account.last_used_index = cmp::max(self.account.last_used_index, index); + if let Some(item) = f(self.view_key(index)?) { + self.keys.account.last_used_index = + cmp::max(self.keys.account.last_used_index, index); return Some(ViewKeySelection { index, - keypair: SecretKeyPair::new(self.derive_spend(index), view_key), + keypair: self.keys.derive_pair(index), item, }); } @@ -563,15 +626,15 @@ where where F: FnMut(&H::SecretKey) -> Option, { - let previous_maximum = self.account.maximum_index; - self.account.maximum_index.index += H::GAP_LIMIT; + let previous_maximum = self.keys.account.maximum_index; + self.keys.account.maximum_index.index += H::GAP_LIMIT; match self.find_index(f) { Some(result) => { - self.account.maximum_index = cmp::max(previous_maximum, result.index); + self.keys.account.maximum_index = cmp::max(previous_maximum, result.index); Some(result) } _ => { - self.account.maximum_index = previous_maximum; + self.keys.account.maximum_index = previous_maximum; None } } @@ -612,6 +675,25 @@ where pub item: T, } +impl ViewKeySelection +where + H: HierarchicalKeyDerivationScheme + ?Sized, +{ + /// Computes `f` on `self.item` returning a new [`ViewKeySelection`] with the same `index` and + /// `keypair`. + #[inline] + pub fn map(self, f: F) -> ViewKeySelection + where + F: FnOnce(T) -> U, + { + ViewKeySelection { + index: self.index, + keypair: self.keypair, + item: f(self.item), + } + } +} + /// Account #[cfg_attr( feature = "serde", diff --git a/manta-accounting/src/wallet/signer.rs b/manta-accounting/src/wallet/signer.rs index bbd7932b8..16f4e9d10 100644 --- a/manta-accounting/src/wallet/signer.rs +++ b/manta-accounting/src/wallet/signer.rs @@ -28,7 +28,10 @@ use crate::{ asset::{Asset, AssetId, AssetMap, AssetMetadata, AssetValue}, - key::{self, HierarchicalKeyDerivationScheme, KeyIndex, SecretKeyPair, ViewKeySelection}, + key::{ + self, HierarchicalKeyDerivationScheme, KeyIndex, SecretKeyPair, ViewKeySelection, + ViewKeyTable, + }, transfer::{ self, batch::Join, @@ -630,63 +633,69 @@ where ) } - /// Inserts the new `utxo`-`encrypted_note` pair if a known key can decrypt the note and - /// validate the utxo. + /// Finds the next viewing key that can decrypt the `encrypted_note` from the `view_key_table`. #[inline] - fn insert_next_item( - &mut self, + fn find_next_key<'h>( + view_key_table: &mut ViewKeyTable<'h, C::HierarchicalKeyDerivationScheme>, parameters: &Parameters, with_recovery: bool, - utxo: Utxo, encrypted_note: EncryptedNote, + ) -> Option>> { + let mut finder = DecryptedMessage::find(encrypted_note); + view_key_table + .find_index_with_maybe_gap(with_recovery, move |k| { + finder.decrypt(¶meters.note_encryption_scheme, k) + }) + .map(|selection| selection.map(|item| item.plaintext)) + } + + /// Inserts the new `utxo`-`note` pair into the `utxo_accumulator` adding the spendable amount + /// to `assets` if there is no void number to match it. + #[inline] + fn insert_next_item( + utxo_accumulator: &mut C::UtxoAccumulator, + assets: &mut C::AssetMap, + parameters: &Parameters, + utxo: Utxo, + selection: ViewKeySelection>, void_numbers: &mut Vec>, deposit: &mut Vec, - ) -> Result<(), SyncError> { - let mut finder = DecryptedMessage::find(encrypted_note); - if let Some(ViewKeySelection { + ) { + let ViewKeySelection { index, keypair, - item, - }) = self - .accounts - .get_mut_default() - .find_index_with_maybe_gap(with_recovery, |k| { - finder.decrypt(¶meters.note_encryption_scheme, k) - }) - { - let Note { + item: Note { ephemeral_secret_key, asset, - } = item.plaintext; - if let Some(void_number) = - parameters.check_full_asset(&keypair.spend, &ephemeral_secret_key, &asset, &utxo) - { - if let Some(index) = void_numbers.iter().position(move |v| v == &void_number) { - void_numbers.remove(index); - } else { - self.utxo_accumulator.insert(&utxo); - self.assets.insert((index, ephemeral_secret_key), asset); - if !asset.is_zero() { - deposit.push(asset); - } - return Ok(()); + }, + } = selection; + if let Some(void_number) = + parameters.check_full_asset(&keypair.spend, &ephemeral_secret_key, &asset, &utxo) + { + if let Some(index) = void_numbers.iter().position(move |v| v == &void_number) { + void_numbers.remove(index); + } else { + utxo_accumulator.insert(&utxo); + assets.insert((index, ephemeral_secret_key), asset); + if !asset.is_zero() { + deposit.push(asset); } + return; } } - self.utxo_accumulator.insert_nonprovable(&utxo); - Ok(()) + utxo_accumulator.insert_nonprovable(&utxo); } /// Checks if `asset` matches with `void_number`, removing it from the `utxo_accumulator` and /// inserting it into the `withdraw` set if this is the case. #[inline] fn is_asset_unspent( + utxo_accumulator: &mut C::UtxoAccumulator, parameters: &Parameters, secret_spend_key: &SecretKey, ephemeral_secret_key: &SecretKey, asset: Asset, void_numbers: &mut Vec>, - utxo_accumulator: &mut C::UtxoAccumulator, withdraw: &mut Vec, ) -> bool { let utxo = parameters.utxo( @@ -716,33 +725,44 @@ where inserts: I, mut void_numbers: Vec>, is_partial: bool, - ) -> Result, SyncError> + ) -> SyncResponse where I: Iterator, EncryptedNote)>, { let void_number_count = void_numbers.len(); let mut deposit = Vec::new(); let mut withdraw = Vec::new(); + let mut view_key_table = self.accounts.get_mut_default().view_key_table(); for (utxo, encrypted_note) in inserts { - self.insert_next_item( + if let Some(selection) = Self::find_next_key( + &mut view_key_table, parameters, with_recovery, - utxo, encrypted_note, - &mut void_numbers, - &mut deposit, - )?; + ) { + Self::insert_next_item( + &mut self.utxo_accumulator, + &mut self.assets, + parameters, + utxo, + selection, + &mut void_numbers, + &mut deposit, + ); + } else { + self.utxo_accumulator.insert_nonprovable(&utxo); + } } self.assets.retain(|(index, ephemeral_secret_key), assets| { assets.retain( |asset| match self.accounts.get_default().spend_key(*index) { Some(secret_spend_key) => Self::is_asset_unspent( + &mut self.utxo_accumulator, parameters, &secret_spend_key, ephemeral_secret_key, *asset, &mut void_numbers, - &mut self.utxo_accumulator, &mut withdraw, ), _ => true, @@ -753,7 +773,7 @@ where self.checkpoint.update_from_void_numbers(void_number_count); self.checkpoint .update_from_utxo_accumulator(&self.utxo_accumulator); - Ok(SyncResponse { + SyncResponse { checkpoint: self.checkpoint.clone(), balance_update: if is_partial { // TODO: Whenever we are doing a full update, don't even build the `deposit` and @@ -764,7 +784,7 @@ where assets: self.assets.assets().into(), } }, - }) + } } /// Builds the pre-sender associated to `key` and `asset`. @@ -1144,7 +1164,7 @@ where } else { let has_pruned = request.prune(checkpoint); let SyncData { receivers, senders } = request.data; - let result = self.state.sync_with( + let response = self.state.sync_with( &self.parameters.parameters, request.with_recovery, receivers.into_iter(), @@ -1152,7 +1172,7 @@ where !has_pruned, ); self.state.utxo_accumulator.commit(); - result + Ok(response) } } diff --git a/manta-util/src/array.rs b/manta-util/src/array.rs index 7ae72a768..7814c8035 100644 --- a/manta-util/src/array.rs +++ b/manta-util/src/array.rs @@ -30,7 +30,7 @@ use alloc::{boxed::Box, vec::Vec}; #[cfg(feature = "serde-array")] use crate::serde::{Deserialize, Serialize}; -/// Error Message for the [`into_array_unchecked`] and [`into_boxed_array_unchecked`] messages. +/// Error Message for the [`into_array_unchecked`] and [`into_boxed_array_unchecked`] Functions const INTO_UNCHECKED_ERROR_MESSAGE: &str = "Input did not have the correct length to match the output array of length"; diff --git a/manta-util/src/collections.rs b/manta-util/src/collections.rs new file mode 100644 index 000000000..6fc15551d --- /dev/null +++ b/manta-util/src/collections.rs @@ -0,0 +1,91 @@ +// Copyright 2019-2022 Manta Network. +// This file is part of manta-rs. +// +// manta-rs is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// manta-rs is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with manta-rs. If not, see . + +//! Collection Types + +#[doc(inline)] +pub use alloc::collections::*; + +/// An ordered map based on a B-Tree. +pub mod btree_map { + use core::borrow::Borrow; + + #[doc(inline)] + pub use alloc::collections::btree_map::*; + + /// Error Message for the [`pop_last`] Function + const POP_LAST_ERROR_MESSAGE: &str = + "Querying the last element of the map must guarantee that `BTreeMap::remove` always succeeds."; + + /// Pops the last element in the key-ordering of `map`. + /// + /// # Limitation + /// + /// Until is stabilized this is an equivalent + /// way to pop the last key-value pair out of a [`BTreeMap`] but requires cloning the key. + #[inline] + pub fn pop_last(map: &mut BTreeMap) -> Option<(K, V)> + where + K: Clone + Ord, + { + let key = map.keys().last()?.clone(); + match map.remove(&key) { + Some(value) => Some((key, value)), + _ => unreachable!("{}", POP_LAST_ERROR_MESSAGE), + } + } + + /// Returns the value stored at `key` in the `map` or executes `f` on the map if there was no + /// value stored at `key`. + /// + /// # Limitation + /// + /// The current implementation of the borrow-checker is too conservative and forces a + /// long-living borrow of the `map` whenever we want to return borrowed data from the map. In + /// this case, we do an extra query to work around this. + #[inline] + pub fn get_or_mutate<'m, K, V, Q, F>( + map: &'m mut BTreeMap, + key: &Q, + f: F, + ) -> Option<&'m V> + where + K: Borrow + Ord, + Q: Ord + ?Sized, + F: FnOnce(&mut BTreeMap) -> Option<&V>, + { + if map.contains_key(key) { + map.get(key) + } else { + f(map) + } + } + + /// Inserts the `key`-`value` pair into the `map`, returning a reference to the inserted value. + #[inline] + pub fn insert_then_get(map: &mut BTreeMap, key: K, value: V) -> &mut V + where + K: Ord, + { + match map.entry(key) { + Entry::Vacant(entry) => entry.insert(value), + Entry::Occupied(mut entry) => { + entry.insert(value); + entry.into_mut() + } + } + } +} diff --git a/manta-util/src/lib.rs b/manta-util/src/lib.rs index d932eb31e..ea8f19887 100644 --- a/manta-util/src/lib.rs +++ b/manta-util/src/lib.rs @@ -37,6 +37,10 @@ pub mod ops; pub mod persistence; pub mod pointer; +#[cfg(feature = "alloc")] +#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))] +pub mod collections; + #[cfg(feature = "alloc")] #[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))] pub mod vec; diff --git a/manta-util/src/vec.rs b/manta-util/src/vec.rs index f4848a525..70011c3b3 100644 --- a/manta-util/src/vec.rs +++ b/manta-util/src/vec.rs @@ -14,12 +14,12 @@ // You should have received a copy of the GNU General Public License // along with manta-rs. If not, see . -//! Vector Utilities +//! Vectors use crate::create_seal; #[doc(inline)] -pub use alloc::vec::Vec; +pub use alloc::vec::*; create_seal! {}