Merge pull request #1474 from mintlayer/fix/wallet-migration-v5

Fix wallet migration v5 to not alter existing DB tables

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,
         );

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(())
     }

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,
     );

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<ChainConfig>,
-        account_pubkey: ExtendedVRFPublicKey,
         db_tx: &impl WalletStorageReadLocked,
         id: &AccountId,
     ) -> KeyChainResult<VrfKeySoftChain> {
-        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,21 +117,25 @@ 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::<Result<_, DerivationError>>()?;
 
         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,

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<B: storage::Backend> Wallet<B> {
     fn migration_v5(db: &Store<B>, chain_config: Arc<ChainConfig>) -> 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<B: storage::Backend> Wallet<B> {
     /// 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<B: storage::Backend> Wallet<B> {
     fn reset_wallet_transactions(
         chain_config: Arc<ChainConfig>,
         db_tx: &mut impl WalletStorageWriteUnlocked,
-    ) -> WalletResult<BTreeMap<U31, Account>> {
+    ) -> 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<ChainConfig>,
+        db_tx: &mut impl WalletStorageWriteUnlocked,
+    ) -> WalletResult<BTreeMap<U31, Account>> {
+        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<B: storage::Backend> Wallet<B> {
 
         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()?;

wallet/storage/src/internal/store_tx.rs

@@ -15,38 +15,35 @@
 
 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::{
     ensure,
     maybe_encrypted::{MaybeEncrypted, MaybeEncryptedError},
 };
 use wallet_types::{
+    account_info::AccountVrfKeys,
     chain_info::ChainInfo,
     keys::{RootKeyConstant, RootKeys},
     seed_phrase::{SeedPhraseConstant, SerializableSeedPhrase},
     AccountDerivationPathId, AccountId, AccountInfo, AccountKeyPurposeId, AccountWalletCreatedTxId,
     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::<db::DBUnconfirmedTxCounters, _, _>(account_id)
             }
 
-            fn get_legacy_vrf_public_key(
+            fn get_account_vrf_public_keys(
                 &self,
                 account_id: &AccountId,
-            ) -> crate::Result<Option<ExtendedVRFPublicKey>> {
-                self.read::<db::DBLegacyVRFPublicKeys, _, _>(account_id)
+            ) -> crate::Result<Option<AccountVrfKeys>> {
+                self.read::<db::DBVRFPublicKeys, _, _>(account_id)
             }
 
             fn get_keychain_usage_state(
@@ -427,12 +424,12 @@ macro_rules! impl_write_ops {
                 self.write::<db::DBUnconfirmedTxCounters, _, _, _>(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::<db::DBLegacyVRFPublicKeys, _, _, _>(id, legacy_testnet_public_key)
+                self.write::<db::DBVRFPublicKeys, _, _, _>(id, account_vrf_keys)
             }
 
             fn set_user_transaction(