From 3098315e842f54da1404c5aaa2b324c65c60493a Mon Sep 17 00:00:00 2001 From: Murisi Tarusenga Date: Thu, 26 Sep 2024 16:30:47 +0200 Subject: [PATCH] Separate the storage of shielded keys from their birthdays. --- crates/apps_lib/src/cli.rs | 3 +- crates/apps_lib/src/cli/client.rs | 5 ++- crates/apps_lib/src/cli/context.rs | 25 ++++++------ crates/benches/native_vps.rs | 6 +-- crates/node/src/bench_utils.rs | 7 +++- crates/wallet/src/lib.rs | 20 ++++++--- crates/wallet/src/store.rs | 43 ++++++++++++-------- genesis/hardware/src/pre-genesis/wallet.toml | 2 + genesis/localnet/src/pre-genesis/wallet.toml | 2 + 9 files changed, 68 insertions(+), 45 deletions(-) diff --git a/crates/apps_lib/src/cli.rs b/crates/apps_lib/src/cli.rs index c309248c5e..facd39f4db 100644 --- a/crates/apps_lib/src/cli.rs +++ b/crates/apps_lib/src/cli.rs @@ -7645,8 +7645,7 @@ pub mod args { find_viewing_key(&mut wallet) } else { find_viewing_key(&mut ctx.borrow_mut_chain_or_exit().wallet) - } - .key; + }; Ok(PayAddressGen:: { alias: self.alias, diff --git a/crates/apps_lib/src/cli/client.rs b/crates/apps_lib/src/cli/client.rs index a11f6b785b..19f63744f0 100644 --- a/crates/apps_lib/src/cli/client.rs +++ b/crates/apps_lib/src/cli/client.rs @@ -4,6 +4,7 @@ use color_eyre::eyre::Result; use namada_sdk::io::{display_line, Io, NamadaIo}; use namada_sdk::masp::ShieldedContext; use namada_sdk::{Namada, NamadaImpl}; +use namada_sdk::wallet::DatedViewingKey; use crate::cli; use crate::cli::api::{CliApi, CliClient}; @@ -349,8 +350,8 @@ impl CliApi { chain_ctx .wallet .get_viewing_keys() - .values() - .copied(), + .into_iter() + .map(|(k, v)| DatedViewingKey::new(v, chain_ctx.wallet.find_birthday(k).copied())), ); crate::client::masp::syncing( diff --git a/crates/apps_lib/src/cli/context.rs b/crates/apps_lib/src/cli/context.rs index 8cc2da6c78..fd86a2b80d 100644 --- a/crates/apps_lib/src/cli/context.rs +++ b/crates/apps_lib/src/cli/context.rs @@ -587,7 +587,6 @@ impl ArgFromMutContext for ExtendedSpendingKey { // Or it is a stored alias of one ctx.wallet .find_spending_key(raw, None) - .map(|k| k.key) .map_err(|_find_err| format!("Unknown spending key {}", raw)) }) } @@ -612,9 +611,7 @@ impl ArgFromMutContext for PseudoExtendedKey { ctx.wallet .find_spending_key(raw, None) .map(|k| { - PseudoExtendedKey::from(MaspExtendedSpendingKey::from( - k.key, - )) + PseudoExtendedKey::from(MaspExtendedSpendingKey::from(k)) }) .map_err(|_find_err| { format!("Unknown spending key {}", raw) @@ -626,9 +623,7 @@ impl ArgFromMutContext for PseudoExtendedKey { .find_viewing_key(raw) .copied() .map(|k| { - PseudoExtendedKey::from(MaspExtendedViewingKey::from( - k.key, - )) + PseudoExtendedKey::from(MaspExtendedViewingKey::from(k)) }) .map_err(|_find_err| format!("Unknown viewing key {}", raw)) }) @@ -644,9 +639,12 @@ impl ArgFromMutContext for DatedSpendingKey { // Either the string is a raw extended spending key FromStr::from_str(raw).or_else(|_parse_err| { // Or it is a stored alias of one - ctx.wallet + let sk = ctx.wallet .find_spending_key(raw, None) - .map_err(|_find_err| format!("Unknown spending key {}", raw)) + .map_err(|_find_err| format!("Unknown spending key {}", raw))?; + let birthday = ctx.wallet + .find_birthday(raw); + Ok(DatedSpendingKey::new(sk, birthday.copied())) }) } } @@ -663,7 +661,6 @@ impl ArgFromMutContext for ExtendedViewingKey { ctx.wallet .find_viewing_key(raw) .copied() - .map(|k| k.key) .map_err(|_find_err| format!("Unknown viewing key {}", raw)) }) } @@ -678,10 +675,12 @@ impl ArgFromMutContext for DatedViewingKey { // Either the string is a raw full viewing key FromStr::from_str(raw).or_else(|_parse_err| { // Or it is a stored alias of one - ctx.wallet + let vk = ctx.wallet .find_viewing_key(raw) - .copied() - .map_err(|_find_err| format!("Unknown viewing key {}", raw)) + .map_err(|_find_err| format!("Unknown viewing key {}", raw))?; + let birthday = ctx.wallet + .find_birthday(raw); + Ok(DatedViewingKey::new(*vk, birthday.copied())) }) } } diff --git a/crates/benches/native_vps.rs b/crates/benches/native_vps.rs index 6e12b24523..474067abcd 100644 --- a/crates/benches/native_vps.rs +++ b/crates/benches/native_vps.rs @@ -412,7 +412,7 @@ fn prepare_ibc_tx_and_ctx(bench_name: &str) -> (BenchShieldedCtx, BatchedTx) { shielded_ctx.generate_shielded_action( Amount::native_whole(10), TransferSource::ExtendedSpendingKey( - ExtendedSpendingKey::from(albert_spending_key.key).into(), + ExtendedSpendingKey::from(albert_spending_key).into(), ), defaults::bertha_address().to_string(), ) @@ -606,14 +606,14 @@ fn setup_storage_for_masp_verification( "unshielding" => shielded_ctx.generate_masp_tx( amount, TransferSource::ExtendedSpendingKey( - ExtendedSpendingKey::from(albert_spending_key.key).into(), + ExtendedSpendingKey::from(albert_spending_key).into(), ), TransferTarget::Address(defaults::albert_address()), ), "shielded" => shielded_ctx.generate_masp_tx( amount, TransferSource::ExtendedSpendingKey( - ExtendedSpendingKey::from(albert_spending_key.key).into(), + ExtendedSpendingKey::from(albert_spending_key).into(), ), TransferTarget::PaymentAddress(bertha_payment_addr), ), diff --git a/crates/node/src/bench_utils.rs b/crates/node/src/bench_utils.rs index d86dee1ad0..1490d600c3 100644 --- a/crates/node/src/bench_utils.rs +++ b/crates/node/src/bench_utils.rs @@ -106,7 +106,7 @@ pub use namada_sdk::tx::{ TX_UPDATE_STEWARD_COMMISSION, TX_VOTE_PROPOSAL as TX_VOTE_PROPOSAL_WASM, TX_WITHDRAW_WASM, VP_USER_WASM, }; -use namada_sdk::wallet::Wallet; +use namada_sdk::wallet::{DatedSpendingKey, Wallet}; use namada_sdk::{ parameters, proof_of_stake, tendermint, Namada, NamadaImpl, PaymentAddress, TransferSource, TransferTarget, @@ -1136,7 +1136,6 @@ impl Default for BenchShieldedCtx { .wallet .find_viewing_key(viewing_alias) .unwrap() - .key .to_string(), ); let viewing_key = ExtendedFullViewingKey::from( @@ -1179,6 +1178,10 @@ impl BenchShieldedCtx { .wallet .find_spending_key(ALBERT_SPENDING_KEY, None) .unwrap(); + let spending_key = DatedSpendingKey::new( + spending_key, + self.wallet.find_birthday(ALBERT_SPENDING_KEY).copied(), + ); self.shielded = async_runtime .block_on(namada_apps_lib::client::masp::syncing( self.shielded, diff --git a/crates/wallet/src/lib.rs b/crates/wallet/src/lib.rs index 68e6fec27b..338ee84003 100644 --- a/crates/wallet/src/lib.rs +++ b/crates/wallet/src/lib.rs @@ -274,7 +274,7 @@ pub struct Wallet { utils: U, store: Store, decrypted_key_cache: HashMap, - decrypted_spendkey_cache: HashMap, + decrypted_spendkey_cache: HashMap, } impl From> for Store { @@ -436,12 +436,20 @@ impl Wallet { pub fn find_viewing_key( &self, alias: impl AsRef, - ) -> Result<&DatedViewingKey, FindKeyError> { + ) -> Result<&ExtendedViewingKey, FindKeyError> { self.store.find_viewing_key(alias.as_ref()).ok_or_else(|| { FindKeyError::KeyNotFound(alias.as_ref().to_string()) }) } + /// Find the birthday of the given alias + pub fn find_birthday( + &self, + alias: impl AsRef, + ) -> Option<&BlockHeight> { + self.store.find_birthday(alias.as_ref()) + } + /// Find the payment address with the given alias in the wallet and return /// it pub fn find_payment_addr( @@ -501,7 +509,7 @@ impl Wallet { } /// Get all known viewing keys by their alias - pub fn get_viewing_keys(&self) -> HashMap { + pub fn get_viewing_keys(&self) -> HashMap { self.store .get_viewing_keys() .iter() @@ -512,7 +520,7 @@ impl Wallet { /// Get all known viewing keys by their alias pub fn get_spending_keys( &self, - ) -> HashMap> { + ) -> HashMap> { self.store .get_spending_keys() .iter() @@ -944,7 +952,7 @@ impl Wallet { &mut self, alias: impl AsRef, password: Option>, - ) -> Result { + ) -> Result { // Try cache first if let Some(cached_key) = self .decrypted_spendkey_cache @@ -1196,7 +1204,7 @@ impl Wallet { // Cache the newly added key self.decrypted_spendkey_cache.insert( alias.clone(), - DatedKeypair::new(spend_key, birthday), + spend_key, ); }) .map(Into::into) diff --git a/crates/wallet/src/store.rs b/crates/wallet/src/store.rs index 8237abbab5..cb15a6c359 100644 --- a/crates/wallet/src/store.rs +++ b/crates/wallet/src/store.rs @@ -22,7 +22,6 @@ use zeroize::Zeroizing; use super::alias::{self, Alias}; use super::derivation_path::DerivationPath; use super::pre_genesis; -use crate::keys::{DatedKeypair, DatedSpendingKey, DatedViewingKey}; use crate::{StoredKeypair, WalletIo}; /// Actions that can be taken when there is an alias conflict @@ -63,10 +62,12 @@ pub struct ValidatorData { /// A Storage area for keys and addresses #[derive(Serialize, Deserialize, Debug, Default)] pub struct Store { + /// Known birthdays + birthdays: BTreeMap, /// Known viewing keys - view_keys: BTreeMap, + view_keys: BTreeMap, /// Known spending keys - spend_keys: BTreeMap>, + spend_keys: BTreeMap>, /// Payment address book payment_addrs: BiBTreeMap, /// Cryptographic keypairs @@ -135,7 +136,7 @@ impl Store { pub fn find_spending_key( &self, alias: impl AsRef, - ) -> Option<&StoredKeypair> { + ) -> Option<&StoredKeypair> { self.spend_keys.get(&alias.into()) } @@ -143,10 +144,18 @@ impl Store { pub fn find_viewing_key( &self, alias: impl AsRef, - ) -> Option<&DatedViewingKey> { + ) -> Option<&ExtendedViewingKey> { self.view_keys.get(&alias.into()) } + /// Find the birthday of the given alias + pub fn find_birthday( + &self, + alias: impl AsRef, + ) -> Option<&BlockHeight> { + self.birthdays.get(&alias.into()) + } + /// Find the payment address with the given alias and return it pub fn find_payment_addr( &self, @@ -200,7 +209,7 @@ impl Store { viewing_key: &ExtendedViewingKey, ) -> Option { for (alias, vk) in &self.view_keys { - if *viewing_key == vk.key { + if *viewing_key == *vk { return self.derivation_paths.get(alias).cloned(); } } @@ -267,14 +276,14 @@ impl Store { } /// Get all known viewing keys by their alias. - pub fn get_viewing_keys(&self) -> &BTreeMap { + pub fn get_viewing_keys(&self) -> &BTreeMap { &self.view_keys } /// Get all known spending keys by their alias. pub fn get_spending_keys( &self, - ) -> &BTreeMap> { + ) -> &BTreeMap> { &self.spend_keys } @@ -412,15 +421,11 @@ impl Store { } self.remove_alias(&alias); - let (spendkey_to_store, _raw_spendkey) = - StoredKeypair::new(DatedKeypair::new(spendkey, birthday), password); + let (spendkey_to_store, _raw_spendkey) = StoredKeypair::new(spendkey, password); self.spend_keys.insert(alias.clone(), spendkey_to_store); // Simultaneously add the derived viewing key to ease balance viewing - let viewkey = DatedKeypair::new( - zip32::ExtendedFullViewingKey::from(&spendkey.into()).into(), - birthday, - ); - self.view_keys.insert(alias.clone(), viewkey); + birthday.map(|x| self.birthdays.insert(alias.clone(), x)); + self.view_keys.insert(alias.clone(), zip32::ExtendedFullViewingKey::from(&spendkey.into()).into()); path.map(|p| self.derivation_paths.insert(alias.clone(), p)); Some(alias) } @@ -456,8 +461,8 @@ impl Store { } } self.remove_alias(&alias); - self.view_keys - .insert(alias.clone(), DatedKeypair::new(viewkey, birthday)); + birthday.map(|x| self.birthdays.insert(alias.clone(), x)); + self.view_keys.insert(alias.clone(), viewkey); path.map(|p| self.derivation_paths.insert(alias.clone(), p)); Some(alias) } @@ -600,6 +605,7 @@ impl Store { || self.pkhs.values().contains(alias) || self.public_keys.contains_key(alias) || self.derivation_paths.contains_key(alias) + || self.birthdays.contains_key(alias) } /// Completely remove the given alias from all maps in the wallet @@ -612,12 +618,14 @@ impl Store { self.pkhs.retain(|_key, val| val != alias); self.public_keys.remove(alias); self.derivation_paths.remove(alias); + self.birthdays.remove(alias); } /// Extend this store from another store (typically pre-genesis). /// Note that this method ignores `validator_data` if any. pub fn extend(&mut self, store: Store) { let Self { + birthdays, view_keys, spend_keys, payment_addrs, @@ -629,6 +637,7 @@ impl Store { validator_data: _, address_vp_types, } = self; + birthdays.extend(store.birthdays); view_keys.extend(store.view_keys); spend_keys.extend(store.spend_keys); payment_addrs.extend(store.payment_addrs); diff --git a/genesis/hardware/src/pre-genesis/wallet.toml b/genesis/hardware/src/pre-genesis/wallet.toml index 1f36ffe432..7866cb9c2f 100644 --- a/genesis/hardware/src/pre-genesis/wallet.toml +++ b/genesis/hardware/src/pre-genesis/wallet.toml @@ -1,3 +1,5 @@ +[birthdays] + [view_keys] [spend_keys] diff --git a/genesis/localnet/src/pre-genesis/wallet.toml b/genesis/localnet/src/pre-genesis/wallet.toml index fbf083e87b..eb231c1e0e 100644 --- a/genesis/localnet/src/pre-genesis/wallet.toml +++ b/genesis/localnet/src/pre-genesis/wallet.toml @@ -1,3 +1,5 @@ +[birthdays] + [view_keys] [spend_keys]