From 5dc20f3e307193104109dd5768f5e1cd6584491d Mon Sep 17 00:00:00 2001 From: Boris Oncev Date: Thu, 18 Jan 2024 13:25:59 +0100 Subject: [PATCH] fix wallet migration v5 to not alter existing DB tables --- wallet/src/account/mod.rs | 1 - wallet/src/key_chain/account_key_chain/mod.rs | 31 +++---- wallet/src/key_chain/tests.rs | 2 - wallet/src/key_chain/vrf_key_chain/mod.rs | 19 ++-- wallet/src/wallet/mod.rs | 90 ++++++++++++++----- wallet/storage/src/internal/store_tx.rs | 35 ++++---- wallet/storage/src/lib.rs | 21 ++--- wallet/storage/src/schema.rs | 5 +- wallet/types/src/account_info.rs | 13 ++- 9 files changed, 127 insertions(+), 90 deletions(-) diff --git a/wallet/src/account/mod.rs b/wallet/src/account/mod.rs index 59c908e386..b55e8d4e0d 100644 --- a/wallet/src/account/mod.rs +++ b/wallet/src/account/mod.rs @@ -170,7 +170,6 @@ impl Account { &chain_config, key_chain.account_index(), key_chain.account_public_key().clone(), - key_chain.account_vrf_public_key().clone(), key_chain.lookahead_size(), name, ); diff --git a/wallet/src/key_chain/account_key_chain/mod.rs b/wallet/src/key_chain/account_key_chain/mod.rs index 88c29a0cd3..b148344892 100644 --- a/wallet/src/key_chain/account_key_chain/mod.rs +++ b/wallet/src/key_chain/account_key_chain/mod.rs @@ -99,21 +99,27 @@ impl AccountKeyChain { let sub_chains = WithPurpose::new(receiving_key_chain, change_key_chain); let legacy_key_path = make_path_to_vrf_key(&chain_config, account_index); - let legacy_public_key = + let legacy_vrf_key = root_vrf_key.clone().derive_absolute_path(&legacy_key_path)?.to_public_key(); - db_tx.set_legacy_vrf_public_key(&account_id, &legacy_public_key)?; - let account_vrf_pub_key = root_vrf_key .derive_absolute_path(&account_path)? .derive_child(VRF_INDEX)? .to_public_key(); + db_tx.set_account_vrf_public_keys( + &account_id, + &wallet_types::account_info::AccountVrfKeys { + account_vrf_key: account_vrf_pub_key.clone(), + legacy_vrf_key: legacy_vrf_key.clone(), + }, + )?; + let vrf_chain = VrfKeySoftChain::new_empty( chain_config.clone(), account_id, account_vrf_pub_key.clone(), - legacy_public_key, + legacy_vrf_key, ); vrf_chain.save_usage_state(db_tx)?; @@ -169,18 +175,13 @@ impl AccountKeyChain { id, )?; - let vrf_chain = VrfKeySoftChain::load_keys( - chain_config.clone(), - account_info.account_vrf_key().clone(), - db_tx, - id, - )?; + let vrf_chain = VrfKeySoftChain::load_keys(chain_config.clone(), db_tx, id)?; Ok(AccountKeyChain { chain_config, account_index: account_info.account_index(), account_public_key: pubkey_id, - account_vrf_public_key: account_info.account_vrf_key().clone().into(), + account_vrf_public_key: vrf_chain.get_account_vrf_public_key().clone().into(), sub_chains, vrf_chain, lookahead_size: account_info.lookahead_size().into(), @@ -247,12 +248,8 @@ impl AccountKeyChain { &self.get_account_id(), )?; - self.vrf_chain = VrfKeySoftChain::load_keys( - self.chain_config.clone(), - self.account_vrf_public_key.clone().take(), - db_tx, - &self.get_account_id(), - )?; + self.vrf_chain = + VrfKeySoftChain::load_keys(self.chain_config.clone(), db_tx, &self.get_account_id())?; Ok(()) } diff --git a/wallet/src/key_chain/tests.rs b/wallet/src/key_chain/tests.rs index a07ebfcc4c..48bc2318d5 100644 --- a/wallet/src/key_chain/tests.rs +++ b/wallet/src/key_chain/tests.rs @@ -166,7 +166,6 @@ fn key_lookahead(#[case] purpose: KeyPurpose) { &chain_config, key_chain.account_index(), key_chain.account_public_key().clone(), - key_chain.account_vrf_public_key().clone(), key_chain.lookahead_size(), None, ); @@ -235,7 +234,6 @@ fn top_up_and_lookahead(#[case] purpose: KeyPurpose) { &chain_config, key_chain.account_index(), key_chain.account_public_key().clone(), - key_chain.account_vrf_public_key().clone(), key_chain.lookahead_size(), None, ); diff --git a/wallet/src/key_chain/vrf_key_chain/mod.rs b/wallet/src/key_chain/vrf_key_chain/mod.rs index 9eba9afcab..4c6d25c8e9 100644 --- a/wallet/src/key_chain/vrf_key_chain/mod.rs +++ b/wallet/src/key_chain/vrf_key_chain/mod.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use utils::const_value::ConstValue; use utils::ensure; use wallet_storage::{WalletStorageReadLocked, WalletStorageWriteLocked}; +use wallet_types::account_info::AccountVrfKeys; use wallet_types::keys::KeychainUsageState; use wallet_types::AccountId; @@ -97,12 +98,14 @@ impl VrfKeySoftChain { pub fn load_keys( chain_config: Arc, - account_pubkey: ExtendedVRFPublicKey, db_tx: &impl WalletStorageReadLocked, id: &AccountId, ) -> KeyChainResult { - let legacy_public_key = db_tx - .get_legacy_vrf_public_key(id)? + let AccountVrfKeys { + account_vrf_key, + legacy_vrf_key, + } = db_tx + .get_account_vrf_public_keys(id)? .ok_or(KeyChainError::CouldNotLoadKeyChain)?; let usage = db_tx @@ -114,7 +117,7 @@ impl VrfKeySoftChain { let child_number = ChildNumber::from_index_with_hardened_bit(index); Ok(( child_number, - account_pubkey.clone().derive_child(child_number)?, + account_vrf_key.clone().derive_child(child_number)?, )) }) .collect::>()?; @@ -122,13 +125,17 @@ impl VrfKeySoftChain { VrfKeySoftChain::new_from_parts( chain_config.clone(), id.clone(), - account_pubkey, + account_vrf_key, public_keys, usage, - legacy_public_key, + legacy_vrf_key, ) } + pub fn get_account_vrf_public_key(&self) -> &ExtendedVRFPublicKey { + &self.parent_pubkey + } + /// Issue a new key pub fn issue_new( &mut self, diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index b191a1b9f7..521e4f2693 100644 --- a/wallet/src/wallet/mod.rs +++ b/wallet/src/wallet/mod.rs @@ -22,7 +22,10 @@ use crate::account::{ Currency, CurrentFeeRate, DelegationData, PartiallySignedTransaction, PoolData, UnconfirmedTokenInfo, UtxoSelectorError, }; -use crate::key_chain::{KeyChainError, MasterKeyChain, LOOKAHEAD_SIZE}; +use crate::key_chain::{ + make_account_path, make_path_to_vrf_key, KeyChainError, MasterKeyChain, LOOKAHEAD_SIZE, + VRF_INDEX, +}; use crate::send_request::{make_issue_token_outputs, IssueNftArguments, StakePoolDataArguments}; use crate::wallet_events::{WalletEvents, WalletEventsNoOp}; use crate::{Account, SendRequest}; @@ -42,6 +45,7 @@ use common::primitives::{Amount, BlockHeight, Id}; use common::size_estimation::SizeEstimationError; use consensus::PoSGenerateBlockInputData; use crypto::key::hdkd::child_number::ChildNumber; +use crypto::key::hdkd::derivable::Derivable; use crypto::key::hdkd::u31::U31; use crypto::key::PublicKey; use crypto::vrf::VRFPublicKey; @@ -366,7 +370,34 @@ impl Wallet { fn migration_v5(db: &Store, chain_config: Arc) -> WalletResult<()> { let mut db_tx = db.transaction_rw_unlocked(None)?; - Self::reset_wallet_transactions(chain_config.clone(), &mut db_tx)?; + for (id, info) in db_tx.get_accounts_info()? { + let root_vrf_key = MasterKeyChain::load_root_vrf_key(&db_tx)?; + let account_path = make_account_path(&chain_config, info.account_index()); + let legacy_key_path = make_path_to_vrf_key(&chain_config, info.account_index()); + let legacy_vrf_key = root_vrf_key + .clone() + .derive_absolute_path(&legacy_key_path) + .map_err(|err| WalletError::KeyChainError(KeyChainError::Derivation(err)))? + .to_public_key(); + + let account_vrf_pub_key = root_vrf_key + .derive_absolute_path(&account_path) + .map_err(|err| WalletError::KeyChainError(KeyChainError::Derivation(err)))? + .derive_child(VRF_INDEX) + .map_err(|err| WalletError::KeyChainError(KeyChainError::Derivation(err)))? + .to_public_key(); + + db_tx.set_account_vrf_public_keys( + &id, + &wallet_types::account_info::AccountVrfKeys { + account_vrf_key: account_vrf_pub_key.clone(), + legacy_vrf_key: legacy_vrf_key.clone(), + }, + )?; + } + + Self::reset_wallet_transactions_and_load(chain_config.clone(), &mut db_tx)?; + db_tx.set_storage_version(WALLET_VERSION_V5)?; db_tx.commit()?; logging::log::info!( @@ -429,7 +460,8 @@ impl Wallet { /// this will cause the wallet to rescan the blockchain pub fn reset_wallet_to_genesis(&mut self) -> WalletResult<()> { let mut db_tx = self.db.transaction_rw_unlocked(None)?; - let mut accounts = Self::reset_wallet_transactions(self.chain_config.clone(), &mut db_tx)?; + let mut accounts = + Self::reset_wallet_transactions_and_load(self.chain_config.clone(), &mut db_tx)?; self.next_unused_account = accounts.pop_last().expect("not empty accounts"); self.accounts = accounts; db_tx.commit()?; @@ -439,34 +471,45 @@ impl Wallet { fn reset_wallet_transactions( chain_config: Arc, db_tx: &mut impl WalletStorageWriteUnlocked, - ) -> WalletResult> { + ) -> WalletResult<()> { db_tx.clear_transactions()?; db_tx.clear_addresses()?; db_tx.clear_public_keys()?; let lookahead_size = db_tx.get_lookahead_size()?; + // set all accounts best block to genesis + for (id, mut info) in db_tx.get_accounts_info()? { + info.update_best_block(BlockHeight::new(0), chain_config.genesis_block_id()); + info.set_lookahead_size(lookahead_size); + db_tx.set_account(&id, &info)?; + db_tx.set_account_unconfirmed_tx_counter(&id, 0)?; + db_tx.set_keychain_usage_state( + &AccountKeyPurposeId::new(id.clone(), KeyPurpose::Change), + &KeychainUsageState::new(None, None), + )?; + db_tx.set_keychain_usage_state( + &AccountKeyPurposeId::new(id.clone(), KeyPurpose::ReceiveFunds), + &KeychainUsageState::new(None, None), + )?; + db_tx + .set_vrf_keychain_usage_state(&id.clone(), &KeychainUsageState::new(None, None))?; + } + + Ok(()) + } + + fn reset_wallet_transactions_and_load( + chain_config: Arc, + db_tx: &mut impl WalletStorageWriteUnlocked, + ) -> WalletResult> { + Self::reset_wallet_transactions(chain_config.clone(), db_tx)?; + // set all accounts best block to genesis db_tx .get_accounts_info()? - .into_iter() - .map(|(id, mut info)| { - info.update_best_block(BlockHeight::new(0), chain_config.genesis_block_id()); - info.set_lookahead_size(lookahead_size); - db_tx.set_account(&id, &info)?; - db_tx.set_account_unconfirmed_tx_counter(&id, 0)?; - db_tx.set_keychain_usage_state( - &AccountKeyPurposeId::new(id.clone(), KeyPurpose::Change), - &KeychainUsageState::new(None, None), - )?; - db_tx.set_keychain_usage_state( - &AccountKeyPurposeId::new(id.clone(), KeyPurpose::ReceiveFunds), - &KeychainUsageState::new(None, None), - )?; - db_tx.set_vrf_keychain_usage_state( - &id.clone(), - &KeychainUsageState::new(None, None), - )?; + .into_keys() + .map(|id| { let mut account = Account::load_from_database(chain_config.clone(), db_tx, &id)?; account.top_up_addresses(db_tx)?; account.scan_genesis(db_tx, &WalletEventsNoOp)?; @@ -603,7 +646,8 @@ impl Wallet { let mut db_tx = self.db.transaction_rw_unlocked(None)?; db_tx.set_lookahead_size(lookahead_size)?; - let mut accounts = Self::reset_wallet_transactions(self.chain_config.clone(), &mut db_tx)?; + let mut accounts = + Self::reset_wallet_transactions_and_load(self.chain_config.clone(), &mut db_tx)?; self.next_unused_account = accounts.pop_last().expect("not empty accounts"); self.accounts = accounts; db_tx.commit()?; diff --git a/wallet/storage/src/internal/store_tx.rs b/wallet/storage/src/internal/store_tx.rs index b81ecdd738..5847ab501b 100644 --- a/wallet/storage/src/internal/store_tx.rs +++ b/wallet/storage/src/internal/store_tx.rs @@ -15,14 +15,16 @@ use std::collections::BTreeMap; +use crate::{ + schema::{self as db, Schema}, + WalletStorageEncryptionRead, WalletStorageEncryptionWrite, WalletStorageReadLocked, + WalletStorageReadUnlocked, WalletStorageWriteLocked, WalletStorageWriteUnlocked, +}; use common::{ address::Address, chain::{block::timestamp::BlockTimestamp, Destination, SignedTransaction}, }; -use crypto::{ - kdf::KdfChallenge, key::extended::ExtendedPublicKey, symkey::SymmetricKey, - vrf::ExtendedVRFPublicKey, -}; +use crypto::{kdf::KdfChallenge, key::extended::ExtendedPublicKey, symkey::SymmetricKey}; use serialization::{Codec, DecodeAll, Encode, EncodeLike}; use storage::{schema, MakeMapRef}; use utils::{ @@ -30,6 +32,7 @@ use utils::{ maybe_encrypted::{MaybeEncrypted, MaybeEncryptedError}, }; use wallet_types::{ + account_info::AccountVrfKeys, chain_info::ChainInfo, keys::{RootKeyConstant, RootKeys}, seed_phrase::{SeedPhraseConstant, SerializableSeedPhrase}, @@ -37,16 +40,10 @@ use wallet_types::{ AccountWalletTxId, KeychainUsageState, WalletTx, }; -use crate::{ - schema::{self as db, Schema}, - WalletStorageEncryptionRead, WalletStorageEncryptionWrite, WalletStorageReadLocked, - WalletStorageReadUnlocked, WalletStorageWriteLocked, WalletStorageWriteUnlocked, -}; - mod well_known { use common::chain::block::timestamp::BlockTimestamp; - use crypto::{kdf::KdfChallenge, vrf::ExtendedVRFPublicKey}; - use wallet_types::chain_info::ChainInfo; + use crypto::kdf::KdfChallenge; + use wallet_types::{account_info::AccountVrfKeys, chain_info::ChainInfo}; use super::Codec; @@ -73,7 +70,7 @@ mod well_known { declare_entry!(MedianTime: BlockTimestamp); declare_entry!(StoreChainInfo: ChainInfo); declare_entry!(LookaheadSize: u32); - declare_entry!(LegacyVfrPubKey: ExtendedVRFPublicKey); + declare_entry!(LegacyVfrPubKey: AccountVrfKeys); } #[derive(PartialEq, Clone)] @@ -227,11 +224,11 @@ macro_rules! impl_read_ops { self.read::(account_id) } - fn get_legacy_vrf_public_key( + fn get_account_vrf_public_keys( &self, account_id: &AccountId, - ) -> crate::Result> { - self.read::(account_id) + ) -> crate::Result> { + self.read::(account_id) } fn get_keychain_usage_state( @@ -427,12 +424,12 @@ macro_rules! impl_write_ops { self.write::(id, counter) } - fn set_legacy_vrf_public_key( + fn set_account_vrf_public_keys( &mut self, id: &AccountId, - legacy_testnet_public_key: &ExtendedVRFPublicKey, + account_vrf_keys: &AccountVrfKeys, ) -> crate::Result<()> { - self.write::(id, legacy_testnet_public_key) + self.write::(id, account_vrf_keys) } fn set_user_transaction( diff --git a/wallet/storage/src/lib.rs b/wallet/storage/src/lib.rs index b209a36d75..c83b4d0b5c 100644 --- a/wallet/storage/src/lib.rs +++ b/wallet/storage/src/lib.rs @@ -23,17 +23,14 @@ use common::{ address::{Address, AddressError}, chain::{block::timestamp::BlockTimestamp, Destination, SignedTransaction}, }; -use crypto::{ - kdf::KdfChallenge, key::extended::ExtendedPublicKey, symkey::SymmetricKey, - vrf::ExtendedVRFPublicKey, -}; +use crypto::{kdf::KdfChallenge, key::extended::ExtendedPublicKey, symkey::SymmetricKey}; pub use internal::{Store, StoreTxRo, StoreTxRoUnlocked, StoreTxRw, StoreTxRwUnlocked}; use std::collections::BTreeMap; use wallet_types::{ - chain_info::ChainInfo, keys::RootKeys, seed_phrase::SerializableSeedPhrase, - AccountDerivationPathId, AccountId, AccountInfo, AccountKeyPurposeId, AccountWalletCreatedTxId, - AccountWalletTxId, KeychainUsageState, WalletTx, + account_info::AccountVrfKeys, chain_info::ChainInfo, keys::RootKeys, + seed_phrase::SerializableSeedPhrase, AccountDerivationPathId, AccountId, AccountInfo, + AccountKeyPurposeId, AccountWalletCreatedTxId, AccountWalletTxId, KeychainUsageState, WalletTx, }; /// Wallet Errors @@ -74,10 +71,8 @@ pub trait WalletStorageReadLocked { ) -> Result>; fn get_user_transactions(&self) -> Result>; fn get_account_unconfirmed_tx_counter(&self, account_id: &AccountId) -> Result>; - fn get_legacy_vrf_public_key( - &self, - account_id: &AccountId, - ) -> Result>; + fn get_account_vrf_public_keys(&self, account_id: &AccountId) + -> Result>; fn get_accounts_info(&self) -> crate::Result>; fn get_address(&self, id: &AccountDerivationPathId) -> Result>; fn get_addresses( @@ -124,10 +119,10 @@ pub trait WalletStorageWriteLocked: WalletStorageReadLocked { fn del_transaction(&mut self, id: &AccountWalletTxId) -> Result<()>; fn clear_transactions(&mut self) -> Result<()>; fn set_account_unconfirmed_tx_counter(&mut self, id: &AccountId, counter: u64) -> Result<()>; - fn set_legacy_vrf_public_key( + fn set_account_vrf_public_keys( &mut self, id: &AccountId, - legacy_testnet_public_key: &ExtendedVRFPublicKey, + vrf_public_keys: &AccountVrfKeys, ) -> Result<()>; fn set_user_transaction( &mut self, diff --git a/wallet/storage/src/schema.rs b/wallet/storage/src/schema.rs index 86b32628c5..7cd4700288 100644 --- a/wallet/storage/src/schema.rs +++ b/wallet/storage/src/schema.rs @@ -16,9 +16,10 @@ //! Wallet database schema use common::chain::SignedTransaction; -use crypto::{key::extended::ExtendedPublicKey, vrf::ExtendedVRFPublicKey}; +use crypto::key::extended::ExtendedPublicKey; use utils::maybe_encrypted::MaybeEncrypted; use wallet_types::{ + account_info::AccountVrfKeys, keys::{RootKeyConstant, RootKeys}, seed_phrase::{SeedPhraseConstant, SerializableSeedPhrase}, AccountDerivationPathId, AccountId, AccountInfo, AccountKeyPurposeId, AccountWalletCreatedTxId, @@ -51,6 +52,6 @@ storage::decl_schema! { /// Store for each account's unconfirmed transaction order counter pub DBUnconfirmedTxCounters: Map, /// Store for each account's legacy VRF public key - pub DBLegacyVRFPublicKeys: Map, + pub DBVRFPublicKeys: Map, } } diff --git a/wallet/types/src/account_info.rs b/wallet/types/src/account_info.rs index 3aef2effdc..88b728d74b 100644 --- a/wallet/types/src/account_info.rs +++ b/wallet/types/src/account_info.rs @@ -35,7 +35,6 @@ pub const DEFAULT_ACCOUNT_INDEX: U31 = match U31::from_u32(0) { pub struct AccountInfo { account_index: U31, account_key: ExtendedPublicKey, - account_vrf_key: ExtendedVRFPublicKey, lookahead_size: u32, best_block_height: BlockHeight, best_block_id: Id, @@ -47,14 +46,12 @@ impl AccountInfo { chain_config: &ChainConfig, account_index: U31, account_key: ExtendedPublicKey, - account_vrf_key: ExtendedVRFPublicKey, lookahead_size: u32, name: Option, ) -> Self { Self { account_index, account_key, - account_vrf_key, lookahead_size, best_block_height: BlockHeight::zero(), best_block_id: chain_config.genesis_block_id(), @@ -70,10 +67,6 @@ impl AccountInfo { &self.account_key } - pub fn account_vrf_key(&self) -> &ExtendedVRFPublicKey { - &self.account_vrf_key - } - pub fn lookahead_size(&self) -> u32 { self.lookahead_size } @@ -107,3 +100,9 @@ impl AccountInfo { self.best_block_id = best_block_id; } } + +#[derive(Debug, Clone, Encode, Decode)] +pub struct AccountVrfKeys { + pub account_vrf_key: ExtendedVRFPublicKey, + pub legacy_vrf_key: ExtendedVRFPublicKey, +}