fix standalon private keys encryption and usage

Author: OBorce

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,18 +298,19 @@ fn create_wallet(chain_config: Arc<ChainConfig>) -> DefaultWallet {
 }
 
 #[track_caller]
-fn create_block<B, P>(
+fn create_block_with_address_reward<B, P>(
     chain_config: &Arc<ChainConfig>,
     wallet: &mut Wallet<B, P>,
     transactions: Vec<SignedTransaction>,
     reward: Amount,
     block_height: u64,
+    address: Destination,
 ) -> (Address<Destination>, Block)
 where
     B: storage::Backend + 'static,
     P: SignerProvider,
 {
-    let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1;
+    let address = Address::new(chain_config, address).unwrap();
 
     let block1 = Block::new(
         transactions,
@@ -327,6 +328,29 @@ where
     (address, 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;
+    create_block_with_address_reward(
+        chain_config,
+        wallet,
+        transactions,
+        reward,
+        block_height,
+        address.into_object(),
+    )
+}
+
 #[track_caller]
 fn test_balance_from_genesis(
     chain_type: ChainType,
@@ -1251,6 +1275,112 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) {
             .unwrap();
     }
 }
+
+#[rstest]
+#[trace]
+#[case(Seed::from_entropy())]
+fn locked_wallet_standalone_keys(#[case] seed: Seed) {
+    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);
+    wallet
+        .add_standalone_private_key(DEFAULT_ACCOUNT_INDEX, standalone_sk, None)
+        .unwrap();
+
+    // Generate a new block which sends reward to the wallet
+    let block1_amount = Amount::from_atoms(rng.gen_range(NETWORK_FEE + 1..NETWORK_FEE + 10000));
+    let _ = create_block_with_address_reward(
+        &chain_config,
+        &mut wallet,
+        vec![],
+        block1_amount,
+        0,
+        Destination::PublicKey(standalone_pk),
+    );
+
+    let password = Some(gen_random_password(&mut rng));
+    wallet.encrypt_wallet(&password).unwrap();
+    wallet.lock_wallet().unwrap();
+
+    let coin_balance = get_coin_balance(&wallet);
+    assert_eq!(coin_balance, block1_amount);
+
+    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();
+    if rng.gen::<bool>() {
+        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();
+    } else {
+        // check if we remove the password it should fail to lock
+        wallet.encrypt_wallet(&None).unwrap();
+
+        let err = wallet.lock_wallet().unwrap_err();
+        assert_eq!(
+            err,
+            WalletError::DatabaseError(wallet_storage::Error::WalletLockedWithoutAPassword)
+        );
+
+        // check that the kdf challenge has been deleted
+        assert!(wallet
+            .db
+            .transaction_ro()
+            .unwrap()
+            .get_encryption_key_kdf_challenge()
+            .unwrap()
+            .is_none());
+
+        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())]
@@ -5363,7 +5493,7 @@ fn test_add_standalone_private_key(#[case] seed: Seed) {
 
     // Check amount is still zero
     let coin_balance = get_coin_balance(&wallet);
-    assert_eq!(coin_balance, Amount::ZERO);
+    assert_eq!(coin_balance, block1_amount);
 
     // but the transaction has been added to the wallet
     let tx_data = wallet

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

wallet/storage/src/lib.rs

@@ -222,6 +222,10 @@ pub trait WalletStorageEncryptionWrite {
     fn del_encryption_kdf_challenge(&mut self) -> Result<()>;
     fn encrypt_root_keys(&mut self, new_encryption_key: &Option<SymmetricKey>) -> Result<()>;
     fn encrypt_seed_phrase(&mut self, new_encryption_key: &Option<SymmetricKey>) -> Result<()>;
+    fn encrypt_standalone_private_keys(
+        &mut self,
+        new_encryption_key: &Option<SymmetricKey>,
+    ) -> Result<()>;
 }
 
 /// Marker trait for types where read/write operations are run in a transaction