From 9cd68dce1ea44ae3103a12f3d5f3132099f665d0 Mon Sep 17 00:00:00 2001 From: Boris Oncev Date: Tue, 19 Aug 2025 02:20:16 +0200 Subject: [PATCH 1/2] fix standalon private keys encryption and usage --- wallet/src/account/mod.rs | 11 +- wallet/src/key_chain/account_key_chain/mod.rs | 10 ++ wallet/src/key_chain/mod.rs | 11 +- wallet/src/wallet/tests.rs | 136 +++++++++++++++++- wallet/storage/src/internal/mod.rs | 1 + wallet/storage/src/internal/store_tx.rs | 28 ++++ wallet/storage/src/lib.rs | 4 + 7 files changed, 188 insertions(+), 13 deletions(-) diff --git a/wallet/src/account/mod.rs b/wallet/src/account/mod.rs index e30799afcd..218b572216 100644 --- a/wallet/src/account/mod.rs +++ b/wallet/src/account/mod.rs @@ -1833,12 +1833,7 @@ impl Account { /// 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 Account { 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 Account { } 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); } } diff --git a/wallet/src/key_chain/account_key_chain/mod.rs b/wallet/src/key_chain/account_key_chain/mod.rs index 3cf89a9dbe..38e3cc0599 100644 --- a/wallet/src/key_chain/account_key_chain/mod.rs +++ b/wallet/src/key_chain/account_key_chain/mod.rs @@ -471,6 +471,16 @@ impl AccountKeyChains for AccountKeyChainImpl { 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, diff --git a/wallet/src/key_chain/mod.rs b/wallet/src/key_chain/mod.rs index 087c0b5b17..4684dbd9b9 100644 --- a/wallet/src/key_chain/mod.rs +++ b/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 diff --git a/wallet/src/wallet/tests.rs b/wallet/src/wallet/tests.rs index 12679c060b..8060b9ca7e 100644 --- a/wallet/src/wallet/tests.rs +++ b/wallet/src/wallet/tests.rs @@ -298,18 +298,19 @@ fn create_wallet(chain_config: Arc) -> DefaultWallet { } #[track_caller] -fn create_block( +fn create_block_with_address_reward( chain_config: &Arc, wallet: &mut Wallet, transactions: Vec, reward: Amount, block_height: u64, + address: Destination, ) -> (Address, 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( + chain_config: &Arc, + wallet: &mut Wallet, + transactions: Vec, + reward: Amount, + block_height: u64, +) -> (Address, 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::() { + 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 diff --git a/wallet/storage/src/internal/mod.rs b/wallet/storage/src/internal/mod.rs index e3b045e9c5..727999d168 100644 --- a/wallet/storage/src/internal/mod.rs +++ b/wallet/storage/src/internal/mod.rs @@ -102,6 +102,7 @@ impl Store { }; 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); diff --git a/wallet/storage/src/internal/store_tx.rs b/wallet/storage/src/internal/store_tx.rs index 0f4b123bf4..799ae5def3 100644 --- a/wallet/storage/src/internal/store_tx.rs +++ b/wallet/storage/src/internal/store_tx.rs @@ -690,6 +690,34 @@ impl WalletStorageEncryptionWrite for StoreTxRwUnlocked<'_, .into_iter() .try_for_each(|(k, v)| self.write::(k, v)) } + + fn encrypt_standalone_private_keys( + &mut self, + new_encryption_key: &Option, + ) -> crate::Result<()> { + let encrypted_standalone_private_keys: Vec<_> = self + .storage + .get::() + .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::(k, v)) + } } /// Wallet data storage transaction diff --git a/wallet/storage/src/lib.rs b/wallet/storage/src/lib.rs index b1cf8bb69f..d01d17c16a 100644 --- a/wallet/storage/src/lib.rs +++ b/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) -> Result<()>; fn encrypt_seed_phrase(&mut self, new_encryption_key: &Option) -> Result<()>; + fn encrypt_standalone_private_keys( + &mut self, + new_encryption_key: &Option, + ) -> Result<()>; } /// Marker trait for types where read/write operations are run in a transaction From 7331ff456e30d625a7479657a10bca39e5dd6a6b Mon Sep 17 00:00:00 2001 From: Boris Oncev Date: Wed, 20 Aug 2025 01:27:07 +0200 Subject: [PATCH 2/2] fix review comments --- .../wallet_tokens_change_metadata_uri.py | 1 - wallet/src/wallet/tests.rs | 223 ++++++++---------- 2 files changed, 100 insertions(+), 124 deletions(-) diff --git a/test/functional/wallet_tokens_change_metadata_uri.py b/test/functional/wallet_tokens_change_metadata_uri.py index 25a62452db..44456d6714 100644 --- a/test/functional/wallet_tokens_change_metadata_uri.py +++ b/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 diff --git a/wallet/src/wallet/tests.rs b/wallet/src/wallet/tests.rs index 8060b9ca7e..86e85edf68 100644 --- a/wallet/src/wallet/tests.rs +++ b/wallet/src/wallet/tests.rs @@ -298,34 +298,29 @@ fn create_wallet(chain_config: Arc) -> DefaultWallet { } #[track_caller] -fn create_block_with_address_reward( +fn create_block_with_reward_address( chain_config: &Arc, wallet: &mut Wallet, transactions: Vec, reward: Amount, block_height: u64, address: Destination, -) -> (Address, Block) +) -> Block where B: storage::Backend + 'static, P: SignerProvider, { - let address = Address::new(chain_config, address).unwrap(); - 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] @@ -341,14 +336,15 @@ where P: SignerProvider, { let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1; - create_block_with_address_reward( + let block = create_block_with_reward_address( chain_config, wallet, transactions, reward, block_height, - address.into_object(), - ) + address.clone().into_object(), + ); + (address, block) } #[track_caller] @@ -1279,7 +1275,11 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) { #[rstest] #[trace] #[case(Seed::from_entropy())] -fn locked_wallet_standalone_keys(#[case] seed: Seed) { +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()); @@ -1290,28 +1290,90 @@ fn locked_wallet_standalone_keys(#[case] seed: Seed) { 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(); + 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(); + } - // 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(); + let standalone_destination = if rng.gen::() { + Destination::PublicKey(standalone_pk) + } else { + Destination::PublicKeyHash((&standalone_pk).into()) + }; + + if rng.gen::() { + // 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), @@ -1336,49 +1398,17 @@ fn locked_wallet_standalone_keys(#[case] seed: Seed) { // success after unlock wallet.unlock_wallet(&password.unwrap()).unwrap(); - if rng.gen::() { - 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(); - } + 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] @@ -5450,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, block1_amount); - - // 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())]