Skip to content

Commit 9cd68dc

Browse files
OBorceOBorce
committed
fix standalon private keys encryption and usage
1 parent a6db1a1 commit 9cd68dc

File tree

7 files changed

+188
-13
lines changed

7 files changed

+188
-13
lines changed

‎wallet/src/account/mod.rs

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1833,12 +1833,7 @@ impl<K: AccountKeyChains> Account<K> {
18331833

18341834
/// Return true if this destination can be spent by this account
18351835
fn is_destination_mine(&self, destination: &Destination) -> bool {
1836-
match destination {
1837-
Destination::PublicKeyHash(pkh) => self.key_chain.is_public_key_hash_mine(pkh),
1838-
Destination::PublicKey(pk) => self.key_chain.is_public_key_mine(pk),
1839-
Destination::AnyoneCanSpend => false,
1840-
Destination::ScriptHash(_) | Destination::ClassicMultisig(_) => false,
1841-
}
1836+
self.key_chain.has_private_key_for_destination(destination)
18421837
}
18431838

18441839
/// 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> {
18471842
Destination::PublicKeyHash(pkh) => {
18481843
self.key_chain.is_public_key_hash_mine_or_watched(*pkh)
18491844
}
1850-
Destination::PublicKey(pk) => self.key_chain.is_public_key_mine(pk),
1845+
Destination::PublicKey(pk) => self.key_chain.is_public_key_mine_or_watched(pk.clone()),
18511846
Destination::AnyoneCanSpend => false,
18521847
Destination::ScriptHash(_) => false,
18531848
Destination::ClassicMultisig(_) => {
@@ -1886,7 +1881,7 @@ impl<K: AccountKeyChains> Account<K> {
18861881
}
18871882
Destination::PublicKey(pk) => {
18881883
let found = self.key_chain.mark_public_key_as_used(db_tx, &pk)?;
1889-
if found {
1884+
if found || self.key_chain.is_public_key_watched(pk.clone()) {
18901885
return Ok(true);
18911886
}
18921887
}

‎wallet/src/key_chain/account_key_chain/mod.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,16 @@ impl<V: VrfKeyChain> AccountKeyChains for AccountKeyChainImpl<V> {
471471
self.is_public_key_hash_mine(&pubkey_hash) || self.is_public_key_hash_watched(pubkey_hash)
472472
}
473473

474+
fn is_public_key_watched(&self, public_key: PublicKey) -> bool {
475+
let dest = Destination::PublicKey(public_key);
476+
self.standalone_watch_only_keys.contains_key(&dest)
477+
|| self.standalone_private_keys.contains_key(&dest)
478+
}
479+
480+
fn is_public_key_mine_or_watched(&self, public_key: PublicKey) -> bool {
481+
self.is_public_key_mine(&public_key) || self.is_public_key_watched(public_key)
482+
}
483+
474484
/// Find the corresponding public key for a given public key hash
475485
fn get_public_key_from_public_key_hash(
476486
&self,

‎wallet/src/key_chain/mod.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -193,14 +193,21 @@ where
193193
// Return true if the provided public key belongs to this key chain
194194
fn is_public_key_mine(&self, public_key: &PublicKey) -> bool;
195195

196+
// Return true if the provided public key is one of the standalone added keys
197+
fn is_public_key_watched(&self, public_key: PublicKey) -> bool;
198+
199+
// Return true if the provided public key belongs to this key chain
200+
// or is one of the standalone added keys
201+
fn is_public_key_mine_or_watched(&self, public_key: PublicKey) -> bool;
202+
196203
// Return true if the provided public key hash belongs to this key chain
197204
fn is_public_key_hash_mine(&self, pubkey_hash: &PublicKeyHash) -> bool;
198205

199-
// Return true if the provided public key hash is one the standalone added keys
206+
// Return true if the provided public key hash is one of the standalone added keys
200207
fn is_public_key_hash_watched(&self, pubkey_hash: PublicKeyHash) -> bool;
201208

202209
// Return true if the provided public key hash belongs to this key chain
203-
// or is one the standalone added keys
210+
// or is one of the standalone added keys
204211
fn is_public_key_hash_mine_or_watched(&self, pubkey_hash: PublicKeyHash) -> bool;
205212

206213
/// Find the corresponding public key for a given public key hash

‎wallet/src/wallet/tests.rs

Lines changed: 133 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -298,18 +298,19 @@ fn create_wallet(chain_config: Arc<ChainConfig>) -> DefaultWallet {
298298
}
299299

300300
#[track_caller]
301-
fn create_block<B, P>(
301+
fn create_block_with_address_reward<B, P>(
302302
chain_config: &Arc<ChainConfig>,
303303
wallet: &mut Wallet<B, P>,
304304
transactions: Vec<SignedTransaction>,
305305
reward: Amount,
306306
block_height: u64,
307+
address: Destination,
307308
) -> (Address<Destination>, Block)
308309
where
309310
B: storage::Backend + 'static,
310311
P: SignerProvider,
311312
{
312-
let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1;
313+
let address = Address::new(chain_config, address).unwrap();
313314

314315
let block1 = Block::new(
315316
transactions,
@@ -327,6 +328,29 @@ where
327328
(address, block1)
328329
}
329330

331+
#[track_caller]
332+
fn create_block<B, P>(
333+
chain_config: &Arc<ChainConfig>,
334+
wallet: &mut Wallet<B, P>,
335+
transactions: Vec<SignedTransaction>,
336+
reward: Amount,
337+
block_height: u64,
338+
) -> (Address<Destination>, Block)
339+
where
340+
B: storage::Backend + 'static,
341+
P: SignerProvider,
342+
{
343+
let address = wallet.get_new_address(DEFAULT_ACCOUNT_INDEX).unwrap().1;
344+
create_block_with_address_reward(
345+
chain_config,
346+
wallet,
347+
transactions,
348+
reward,
349+
block_height,
350+
address.into_object(),
351+
)
352+
}
353+
330354
#[track_caller]
331355
fn test_balance_from_genesis(
332356
chain_type: ChainType,
@@ -1251,6 +1275,112 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) {
12511275
.unwrap();
12521276
}
12531277
}
1278+
1279+
#[rstest]
1280+
#[trace]
1281+
#[case(Seed::from_entropy())]
1282+
fn locked_wallet_standalone_keys(#[case] seed: Seed) {
1283+
let mut rng = make_seedable_rng(seed);
1284+
let chain_config = Arc::new(create_mainnet());
1285+
1286+
let mut wallet = create_wallet(chain_config.clone());
1287+
1288+
let coin_balance = get_coin_balance(&wallet);
1289+
assert_eq!(coin_balance, Amount::ZERO);
1290+
1291+
let (standalone_sk, standalone_pk) =
1292+
PrivateKey::new_from_rng(&mut rng, KeyKind::Secp256k1Schnorr);
1293+
wallet
1294+
.add_standalone_private_key(DEFAULT_ACCOUNT_INDEX, standalone_sk, None)
1295+
.unwrap();
1296+
1297+
// Generate a new block which sends reward to the wallet
1298+
let block1_amount = Amount::from_atoms(rng.gen_range(NETWORK_FEE + 1..NETWORK_FEE + 10000));
1299+
let _ = create_block_with_address_reward(
1300+
&chain_config,
1301+
&mut wallet,
1302+
vec![],
1303+
block1_amount,
1304+
0,
1305+
Destination::PublicKey(standalone_pk),
1306+
);
1307+
1308+
let password = Some(gen_random_password(&mut rng));
1309+
wallet.encrypt_wallet(&password).unwrap();
1310+
wallet.lock_wallet().unwrap();
1311+
1312+
let coin_balance = get_coin_balance(&wallet);
1313+
assert_eq!(coin_balance, block1_amount);
1314+
1315+
let new_output = TxOutput::Transfer(
1316+
OutputValue::Coin(Amount::from_atoms(
1317+
rng.gen_range(1..=block1_amount.into_atoms() - NETWORK_FEE),
1318+
)),
1319+
Destination::AnyoneCanSpend,
1320+
);
1321+
1322+
assert_eq!(
1323+
wallet.create_transaction_to_addresses(
1324+
DEFAULT_ACCOUNT_INDEX,
1325+
[new_output.clone()],
1326+
SelectedInputs::Utxos(vec![]),
1327+
BTreeMap::new(),
1328+
FeeRate::from_amount_per_kb(Amount::ZERO),
1329+
FeeRate::from_amount_per_kb(Amount::ZERO),
1330+
TxAdditionalInfo::new(),
1331+
),
1332+
Err(WalletError::DatabaseError(
1333+
wallet_storage::Error::WalletLocked
1334+
))
1335+
);
1336+
1337+
// success after unlock
1338+
wallet.unlock_wallet(&password.unwrap()).unwrap();
1339+
if rng.gen::<bool>() {
1340+
wallet
1341+
.create_transaction_to_addresses(
1342+
DEFAULT_ACCOUNT_INDEX,
1343+
[new_output],
1344+
SelectedInputs::Utxos(vec![]),
1345+
BTreeMap::new(),
1346+
FeeRate::from_amount_per_kb(Amount::ZERO),
1347+
FeeRate::from_amount_per_kb(Amount::ZERO),
1348+
TxAdditionalInfo::new(),
1349+
)
1350+
.unwrap();
1351+
} else {
1352+
// check if we remove the password it should fail to lock
1353+
wallet.encrypt_wallet(&None).unwrap();
1354+
1355+
let err = wallet.lock_wallet().unwrap_err();
1356+
assert_eq!(
1357+
err,
1358+
WalletError::DatabaseError(wallet_storage::Error::WalletLockedWithoutAPassword)
1359+
);
1360+
1361+
// check that the kdf challenge has been deleted
1362+
assert!(wallet
1363+
.db
1364+
.transaction_ro()
1365+
.unwrap()
1366+
.get_encryption_key_kdf_challenge()
1367+
.unwrap()
1368+
.is_none());
1369+
1370+
wallet
1371+
.create_transaction_to_addresses(
1372+
DEFAULT_ACCOUNT_INDEX,
1373+
[new_output],
1374+
SelectedInputs::Utxos(vec![]),
1375+
BTreeMap::new(),
1376+
FeeRate::from_amount_per_kb(Amount::ZERO),
1377+
FeeRate::from_amount_per_kb(Amount::ZERO),
1378+
TxAdditionalInfo::new(),
1379+
)
1380+
.unwrap();
1381+
}
1382+
}
1383+
12541384
#[rstest]
12551385
#[trace]
12561386
#[case(Seed::from_entropy())]
@@ -5363,7 +5493,7 @@ fn test_add_standalone_private_key(#[case] seed: Seed) {
53635493

53645494
// Check amount is still zero
53655495
let coin_balance = get_coin_balance(&wallet);
5366-
assert_eq!(coin_balance, Amount::ZERO);
5496+
assert_eq!(coin_balance, block1_amount);
53675497

53685498
// but the transaction has been added to the wallet
53695499
let tx_data = wallet

‎wallet/storage/src/internal/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ impl<B: storage::Backend> Store<B> {
102102
};
103103
tx.encrypt_root_keys(&sym_key)?;
104104
tx.encrypt_seed_phrase(&sym_key)?;
105+
tx.encrypt_standalone_private_keys(&sym_key)?;
105106
tx.commit()?;
106107

107108
self.encryption_state = EncryptionState::Unlocked(sym_key);

‎wallet/storage/src/internal/store_tx.rs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -690,6 +690,34 @@ impl<B: storage::Backend> WalletStorageEncryptionWrite for StoreTxRwUnlocked<'_,
690690
.into_iter()
691691
.try_for_each(|(k, v)| self.write::<db::DBSeedPhrase, _, _, _>(k, v))
692692
}
693+
694+
fn encrypt_standalone_private_keys(
695+
&mut self,
696+
new_encryption_key: &Option<SymmetricKey>,
697+
) -> crate::Result<()> {
698+
let encrypted_standalone_private_keys: Vec<_> = self
699+
.storage
700+
.get::<db::DBStandalonePrivateKeys, _>()
701+
.prefix_iter_decoded(&())?
702+
.map(|(k, v)| {
703+
let decrypted = v
704+
.private_key
705+
.try_take(self.encryption_key)
706+
.expect("key was checked when unlocked");
707+
(
708+
k,
709+
StandalonePrivateKey {
710+
label: v.label,
711+
private_key: MaybeEncrypted::new(&decrypted, new_encryption_key),
712+
},
713+
)
714+
})
715+
.collect();
716+
717+
encrypted_standalone_private_keys
718+
.into_iter()
719+
.try_for_each(|(k, v)| self.write::<db::DBStandalonePrivateKeys, _, _, _>(k, v))
720+
}
693721
}
694722

695723
/// Wallet data storage transaction

‎wallet/storage/src/lib.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,10 @@ pub trait WalletStorageEncryptionWrite {
222222
fn del_encryption_kdf_challenge(&mut self) -> Result<()>;
223223
fn encrypt_root_keys(&mut self, new_encryption_key: &Option<SymmetricKey>) -> Result<()>;
224224
fn encrypt_seed_phrase(&mut self, new_encryption_key: &Option<SymmetricKey>) -> Result<()>;
225+
fn encrypt_standalone_private_keys(
226+
&mut self,
227+
new_encryption_key: &Option<SymmetricKey>,
228+
) -> Result<()>;
225229
}
226230

227231
/// Marker trait for types where read/write operations are run in a transaction

0 commit comments

Comments
 (0)