Merge pull request #1959 from mintlayer/fix/wallet-standalone-keys-encryption

fix standalone private keys encryption and usage

Author: OBorce

test/functional/wallet_tokens_change_metadata_uri.py

@@ -133,7 +133,6 @@ async def async_test(self):
             assert_in("Success", await wallet.sync())
 
             token_info = node.chainstate_token_info(token_id)
-            print(token_info)
             assert_equal(new_metadata_uri, token_info['content']['metadata_uri']['hex']);
             assert token_info['content']['metadata_uri']['text'] is not metadata_uri
 

wallet/src/account/mod.rs

@@ -1833,12 +1833,7 @@ impl<K: AccountKeyChains> Account<K> {
 
     /// Return true if this destination can be spent by this account
     fn is_destination_mine(&self, destination: &Destination) -> bool {
-        match destination {
-            Destination::PublicKeyHash(pkh) => self.key_chain.is_public_key_hash_mine(pkh),
-            Destination::PublicKey(pk) => self.key_chain.is_public_key_mine(pk),
-            Destination::AnyoneCanSpend => false,
-            Destination::ScriptHash(_) | Destination::ClassicMultisig(_) => false,
-        }
+        self.key_chain.has_private_key_for_destination(destination)
     }
 
     /// Return true if this destination can be spent by this account or if it is being watched.
@@ -1847,7 +1842,7 @@ impl<K: AccountKeyChains> Account<K> {
             Destination::PublicKeyHash(pkh) => {
                 self.key_chain.is_public_key_hash_mine_or_watched(*pkh)
             }
-            Destination::PublicKey(pk) => self.key_chain.is_public_key_mine(pk),
+            Destination::PublicKey(pk) => self.key_chain.is_public_key_mine_or_watched(pk.clone()),
             Destination::AnyoneCanSpend => false,
             Destination::ScriptHash(_) => false,
             Destination::ClassicMultisig(_) => {
@@ -1886,7 +1881,7 @@ impl<K: AccountKeyChains> Account<K> {
                 }
                 Destination::PublicKey(pk) => {
                     let found = self.key_chain.mark_public_key_as_used(db_tx, &pk)?;
-                    if found {
+                    if found || self.key_chain.is_public_key_watched(pk.clone()) {
                         return Ok(true);
                     }
                 }

wallet/src/key_chain/account_key_chain/mod.rs

@@ -471,6 +471,16 @@ impl<V: VrfKeyChain> AccountKeyChains for AccountKeyChainImpl<V> {
         self.is_public_key_hash_mine(&pubkey_hash) || self.is_public_key_hash_watched(pubkey_hash)
     }
 
+    fn is_public_key_watched(&self, public_key: PublicKey) -> bool {
+        let dest = Destination::PublicKey(public_key);
+        self.standalone_watch_only_keys.contains_key(&dest)
+            || self.standalone_private_keys.contains_key(&dest)
+    }
+
+    fn is_public_key_mine_or_watched(&self, public_key: PublicKey) -> bool {
+        self.is_public_key_mine(&public_key) || self.is_public_key_watched(public_key)
+    }
+
     /// Find the corresponding public key for a given public key hash
     fn get_public_key_from_public_key_hash(
         &self,

wallet/src/key_chain/mod.rs

@@ -193,14 +193,21 @@ where
     // Return true if the provided public key belongs to this key chain
     fn is_public_key_mine(&self, public_key: &PublicKey) -> bool;
 
+    // Return true if the provided public key is one of the standalone added keys
+    fn is_public_key_watched(&self, public_key: PublicKey) -> bool;
+
+    // Return true if the provided public key belongs to this key chain
+    // or is one of the standalone added keys
+    fn is_public_key_mine_or_watched(&self, public_key: PublicKey) -> bool;
+
     // Return true if the provided public key hash belongs to this key chain
     fn is_public_key_hash_mine(&self, pubkey_hash: &PublicKeyHash) -> bool;
 
-    // Return true if the provided public key hash is one the standalone added keys
+    // Return true if the provided public key hash is one of the standalone added keys
     fn is_public_key_hash_watched(&self, pubkey_hash: PublicKeyHash) -> bool;
 
     // Return true if the provided public key hash belongs to this key chain
-    // or is one the standalone added keys
+    // or is one of the standalone added keys
     fn is_public_key_hash_mine_or_watched(&self, pubkey_hash: PublicKeyHash) -> bool;
 
     /// Find the corresponding public key for a given public key hash

wallet/src/wallet/tests.rs

@@ -298,33 +298,53 @@ fn create_wallet(chain_config: Arc<ChainConfig>) -> DefaultWallet {
 }
 
 #[track_caller]
-fn create_block<B, P>(
+fn create_block_with_reward_address<B, P>(
     chain_config: &Arc<ChainConfig>,
     wallet: &mut Wallet<B, P>,
     transactions: Vec<SignedTransaction>,
     reward: Amount,
     block_height: u64,
-) -> (Address<Destination>, Block)
+    address: Destination,
+) -> Block
 where
     B: storage::Backend + 'static,
     P: SignerProvider,
 {
-    let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1;
-
     let block1 = Block::new(
         transactions,
         chain_config.genesis_block_id(),
         chain_config.genesis_block().timestamp(),
         ConsensusData::None,
-        BlockReward::new(vec![make_address_output(
-            address.clone().into_object(),
-            reward,
-        )]),
+        BlockReward::new(vec![make_address_output(address, reward)]),
     )
     .unwrap();
 
     scan_wallet(wallet, BlockHeight::new(block_height), vec![block1.clone()]);
-    (address, block1)
+    block1
+}
+
+#[track_caller]
+fn create_block<B, P>(
+    chain_config: &Arc<ChainConfig>,
+    wallet: &mut Wallet<B, P>,
+    transactions: Vec<SignedTransaction>,
+    reward: Amount,
+    block_height: u64,
+) -> (Address<Destination>, Block)
+where
+    B: storage::Backend + 'static,
+    P: SignerProvider,
+{
+    let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1;
+    let block = create_block_with_reward_address(
+        chain_config,
+        wallet,
+        transactions,
+        reward,
+        block_height,
+        address.clone().into_object(),
+    );
+    (address, block)
 }
 
 #[track_caller]
@@ -1251,6 +1271,146 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) {
             .unwrap();
     }
 }
+
+#[rstest]
+#[trace]
+#[case(Seed::from_entropy())]
+fn locked_wallet_standalone_keys(
+    #[case] seed: Seed,
+    #[values(true, false)] insert_before_encrypt: bool,
+    #[values(true, false)] change_password: bool,
+) {
+    let mut rng = make_seedable_rng(seed);
+    let chain_config = Arc::new(create_mainnet());
+
+    let mut wallet = create_wallet(chain_config.clone());
+
+    let coin_balance = get_coin_balance(&wallet);
+    assert_eq!(coin_balance, Amount::ZERO);
+
+    let (standalone_sk, standalone_pk) =
+        PrivateKey::new_from_rng(&mut rng, KeyKind::Secp256k1Schnorr);
+    let mut password = Some(gen_random_password(&mut rng));
+
+    if insert_before_encrypt {
+        wallet
+            .add_standalone_private_key(DEFAULT_ACCOUNT_INDEX, standalone_sk, None)
+            .unwrap();
+        wallet.encrypt_wallet(&password).unwrap();
+    } else {
+        wallet.encrypt_wallet(&password).unwrap();
+        wallet
+            .add_standalone_private_key(DEFAULT_ACCOUNT_INDEX, standalone_sk, None)
+            .unwrap();
+    }
+
+    if change_password {
+        password = Some(gen_random_password(&mut rng));
+        wallet.encrypt_wallet(&password).unwrap();
+    }
+
+    let block1_amount = Amount::from_atoms(rng.gen_range(NETWORK_FEE + 1..NETWORK_FEE + 10000));
+
+    let standalone_destination = if rng.gen::<bool>() {
+        Destination::PublicKey(standalone_pk)
+    } else {
+        Destination::PublicKeyHash((&standalone_pk).into())
+    };
+
+    if rng.gen::<bool>() {
+        // test that wallet will recognise a destination belonging to a standalone key in a block
+        // reward
+        let _ = create_block_with_reward_address(
+            &chain_config,
+            &mut wallet,
+            vec![],
+            block1_amount,
+            0,
+            standalone_destination,
+        );
+    } else {
+        // test that wallet will recognise a destination belonging to a standalone key in a
+        // transaction
+        let output = make_address_output(standalone_destination, block1_amount);
+
+        let tx = SignedTransaction::new(Transaction::new(0, vec![], vec![output]).unwrap(), vec![])
+            .unwrap();
+
+        let block1 = Block::new(
+            vec![tx.clone()],
+            chain_config.genesis_block_id(),
+            chain_config.genesis_block().timestamp(),
+            ConsensusData::None,
+            BlockReward::new(vec![]),
+        )
+        .unwrap();
+
+        scan_wallet(&mut wallet, BlockHeight::new(0), vec![block1.clone()]);
+
+        // check the transaction has been added to the wallet
+        let tx_data = wallet
+            .get_transaction(DEFAULT_ACCOUNT_INDEX, tx.transaction().get_id())
+            .unwrap();
+
+        assert_eq!(tx_data.get_transaction(), tx.transaction());
+    }
+
+    wallet.lock_wallet().unwrap();
+
+    // check balance is recognising spendable UTXOs belonging to the standalone private key
+    let coin_balance = get_coin_balance(&wallet);
+    assert_eq!(coin_balance, block1_amount);
+
+    // also check utxos
+    let utxos = wallet
+        .get_utxos(
+            DEFAULT_ACCOUNT_INDEX,
+            UtxoType::Transfer | UtxoType::LockThenTransfer | UtxoType::IssueNft,
+            UtxoState::Confirmed | UtxoState::Inactive,
+            WithLocked::Unlocked,
+        )
+        .unwrap();
+    assert_eq!(utxos.len(), 1);
+
+    // try to spend the UTXO belonging to the standalone key
+
+    let new_output = TxOutput::Transfer(
+        OutputValue::Coin(Amount::from_atoms(
+            rng.gen_range(1..=block1_amount.into_atoms() - NETWORK_FEE),
+        )),
+        Destination::AnyoneCanSpend,
+    );
+
+    assert_eq!(
+        wallet.create_transaction_to_addresses(
+            DEFAULT_ACCOUNT_INDEX,
+            [new_output.clone()],
+            SelectedInputs::Utxos(vec![]),
+            BTreeMap::new(),
+            FeeRate::from_amount_per_kb(Amount::ZERO),
+            FeeRate::from_amount_per_kb(Amount::ZERO),
+            TxAdditionalInfo::new(),
+        ),
+        Err(WalletError::DatabaseError(
+            wallet_storage::Error::WalletLocked
+        ))
+    );
+
+    // success after unlock
+    wallet.unlock_wallet(&password.unwrap()).unwrap();
+    wallet
+        .create_transaction_to_addresses(
+            DEFAULT_ACCOUNT_INDEX,
+            [new_output],
+            SelectedInputs::Utxos(vec![]),
+            BTreeMap::new(),
+            FeeRate::from_amount_per_kb(Amount::ZERO),
+            FeeRate::from_amount_per_kb(Amount::ZERO),
+            TxAdditionalInfo::new(),
+        )
+        .unwrap();
+}
+
 #[rstest]
 #[trace]
 #[case(Seed::from_entropy())]
@@ -5320,59 +5480,6 @@ fn test_not_exhaustion_of_keys(#[case] seed: Seed) {
     }
 }
 
-#[rstest]
-#[trace]
-#[case(Seed::from_entropy())]
-fn test_add_standalone_private_key(#[case] seed: Seed) {
-    let mut rng = make_seedable_rng(seed);
-    let chain_config = Arc::new(create_regtest());
-
-    let mut wallet = create_wallet(chain_config.clone());
-
-    let coin_balance = get_coin_balance(&wallet);
-    assert_eq!(coin_balance, Amount::ZERO);
-    // generate a random private key unrelated to the wallet and add it
-    let (private_key, pub_key) =
-        crypto::key::PrivateKey::new_from_rng(&mut rng, KeyKind::Secp256k1Schnorr);
-
-    wallet
-        .add_standalone_private_key(DEFAULT_ACCOUNT_INDEX, private_key, None)
-        .unwrap();
-
-    // get the destination address from the new private key and send some coins to it
-    let address =
-        Address::new(&chain_config, Destination::PublicKeyHash((&pub_key).into())).unwrap();
-
-    // Generate a new block which sends reward to the new address
-    let block1_amount = Amount::from_atoms(rng.gen_range(NETWORK_FEE + 100..NETWORK_FEE + 10000));
-    let output = make_address_output(address.clone().into_object(), block1_amount);
-
-    let tx =
-        SignedTransaction::new(Transaction::new(0, vec![], vec![output]).unwrap(), vec![]).unwrap();
-
-    let block1 = Block::new(
-        vec![tx.clone()],
-        chain_config.genesis_block_id(),
-        chain_config.genesis_block().timestamp(),
-        ConsensusData::None,
-        BlockReward::new(vec![]),
-    )
-    .unwrap();
-
-    scan_wallet(&mut wallet, BlockHeight::new(0), vec![block1.clone()]);
-
-    // Check amount is still zero
-    let coin_balance = get_coin_balance(&wallet);
-    assert_eq!(coin_balance, Amount::ZERO);
-
-    // but the transaction has been added to the wallet
-    let tx_data = wallet
-        .get_transaction(DEFAULT_ACCOUNT_INDEX, tx.transaction().get_id())
-        .unwrap();
-
-    assert_eq!(tx_data.get_transaction(), tx.transaction());
-}
-
 #[rstest]
 #[trace]
 #[case(Seed::from_entropy())]

wallet/storage/src/internal/mod.rs

@@ -102,6 +102,7 @@ impl<B: storage::Backend> Store<B> {
         };
         tx.encrypt_root_keys(&sym_key)?;
         tx.encrypt_seed_phrase(&sym_key)?;
+        tx.encrypt_standalone_private_keys(&sym_key)?;
         tx.commit()?;
 
         self.encryption_state = EncryptionState::Unlocked(sym_key);

wallet/storage/src/internal/store_tx.rs

@@ -690,6 +690,34 @@ impl<B: storage::Backend> WalletStorageEncryptionWrite for StoreTxRwUnlocked<'_,
             .into_iter()
             .try_for_each(|(k, v)| self.write::<db::DBSeedPhrase, _, _, _>(k, v))
     }
+
+    fn encrypt_standalone_private_keys(
+        &mut self,
+        new_encryption_key: &Option<SymmetricKey>,
+    ) -> crate::Result<()> {
+        let encrypted_standalone_private_keys: Vec<_> = self
+            .storage
+            .get::<db::DBStandalonePrivateKeys, _>()
+            .prefix_iter_decoded(&())?
+            .map(|(k, v)| {
+                let decrypted = v
+                    .private_key
+                    .try_take(self.encryption_key)
+                    .expect("key was checked when unlocked");
+                (
+                    k,
+                    StandalonePrivateKey {
+                        label: v.label,
+                        private_key: MaybeEncrypted::new(&decrypted, new_encryption_key),
+                    },
+                )
+            })
+            .collect();
+
+        encrypted_standalone_private_keys
+            .into_iter()
+            .try_for_each(|(k, v)| self.write::<db::DBStandalonePrivateKeys, _, _, _>(k, v))
+    }
 }
 
 /// Wallet data storage transaction