Skip to content

Commit

Permalink
Merge bitcoin/bitcoin#28617: test: Add Wallet Unlock Context Manager
Browse files Browse the repository at this point in the history 
004903e test: Add Wallet Unlock Context Manager (Brandon Odiwuor)

Pull request description:

  Fixes #28601, see bitcoin/bitcoin#28403 (comment)

  Add Context Manager to manage the locking and unlocking of locked wallets with a passphrase during testing.

ACKs for top commit:
  kevkevinpal:
    lgtm ACK [004903e](bitcoin/bitcoin@004903e)
  maflcko:
    lgtm ACK 004903e

Tree-SHA512: ab234c167e71531df0d974ff9a31d444f7ce2a1d05aba5ea868cc9452f139845eeb24ca058d88f058bc02482b762adf2d99e63a6640b872cc71a57a0068abfe8
  • Loading branch information
@fanquake
fanquake committed Oct 19, 2023
2 parents 5eb82d5 + 004903e commit 091d29c
Show file tree
Hide file tree
Showing 7 changed files with 143 additions and 128 deletions.
19 changes: 19 additions & 0 deletions test/functional/test_framework/wallet_util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,3 +122,22 @@ def generate_keypair(compressed=True, wif=False):
if wif:
privkey = bytes_to_wif(privkey.get_bytes(), compressed)
return privkey, pubkey

class WalletUnlock():
"""
A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.
"""

MAXIMUM_TIMEOUT = 999000

def __init__(self, wallet, passphrase, timeout=MAXIMUM_TIMEOUT):
self.wallet = wallet
self.passphrase = passphrase
self.timeout = timeout

def __enter__(self):
self.wallet.walletpassphrase(self.passphrase, self.timeout)

def __exit__(self, *args):
_ = args
self.wallet.walletlock()
60 changes: 30 additions & 30 deletions test/functional/wallet_createwallet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
assert_equal,
assert_raises_rpc_error,
)
from test_framework.wallet_util import generate_keypair
from test_framework.wallet_util import generate_keypair, WalletUnlock


EMPTY_PASSPHRASE_MSG = "Empty string given as passphrase, wallet will not be encrypted."
Expand Down Expand Up @@ -108,24 +108,24 @@ def run_test(self):
w4.encryptwallet('pass')
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", w4.getnewaddress)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", w4.getrawchangeaddress)
# Now set a seed and it should work. Wallet should also be encrypted
w4.walletpassphrase("pass", 999000)
if self.options.descriptors:
w4.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPcwuZGKp8TeWppSuLMiLe2d9PupB14QpPeQsqoj3LneJLhGHH13xESfvASyd4EFLJvLrG8b7DrLxEuV7hpF9uUc6XruKA1Wq/0h/*)'),
'timestamp': 'now',
'active': True
},
{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPcwuZGKp8TeWppSuLMiLe2d9PupB14QpPeQsqoj3LneJLhGHH13xESfvASyd4EFLJvLrG8b7DrLxEuV7hpF9uUc6XruKA1Wq/1h/*)'),
'timestamp': 'now',
'active': True,
'internal': True
}])
else:
w4.sethdseed()
w4.getnewaddress()
w4.getrawchangeaddress()
with WalletUnlock(w4, "pass"):
# Now set a seed and it should work. Wallet should also be encrypted
if self.options.descriptors:
w4.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPcwuZGKp8TeWppSuLMiLe2d9PupB14QpPeQsqoj3LneJLhGHH13xESfvASyd4EFLJvLrG8b7DrLxEuV7hpF9uUc6XruKA1Wq/0h/*)'),
'timestamp': 'now',
'active': True
},
{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPcwuZGKp8TeWppSuLMiLe2d9PupB14QpPeQsqoj3LneJLhGHH13xESfvASyd4EFLJvLrG8b7DrLxEuV7hpF9uUc6XruKA1Wq/1h/*)'),
'timestamp': 'now',
'active': True,
'internal': True
}])
else:
w4.sethdseed()
w4.getnewaddress()
w4.getrawchangeaddress()

self.log.info("Test blank creation with privkeys disabled and then encryption")
self.nodes[0].createwallet(wallet_name='w5', disable_private_keys=True, blank=True)
Expand All @@ -142,23 +142,23 @@ def run_test(self):
self.nodes[0].createwallet(wallet_name='wblank', disable_private_keys=False, blank=True, passphrase='thisisapassphrase')
wblank = node.get_wallet_rpc('wblank')
assert_raises_rpc_error(-13, "Error: Please enter the wallet passphrase with walletpassphrase first.", wblank.signmessage, "needanargument", "test")
wblank.walletpassphrase("thisisapassphrase", 999000)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getnewaddress)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getrawchangeaddress)
with WalletUnlock(wblank, "thisisapassphrase"):
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getnewaddress)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getrawchangeaddress)

self.log.info('Test creating a new encrypted wallet.')
# Born encrypted wallet is created (has keys)
self.nodes[0].createwallet(wallet_name='w6', disable_private_keys=False, blank=False, passphrase='thisisapassphrase')
w6 = node.get_wallet_rpc('w6')
assert_raises_rpc_error(-13, "Error: Please enter the wallet passphrase with walletpassphrase first.", w6.signmessage, "needanargument", "test")
w6.walletpassphrase("thisisapassphrase", 999000)
w6.signmessage(w6.getnewaddress('', 'legacy'), "test")
w6.keypoolrefill(1)
# There should only be 1 key for legacy, 3 for descriptors
walletinfo = w6.getwalletinfo()
keys = 4 if self.options.descriptors else 1
assert_equal(walletinfo['keypoolsize'], keys)
assert_equal(walletinfo['keypoolsize_hd_internal'], keys)
with WalletUnlock(w6, "thisisapassphrase"):
w6.signmessage(w6.getnewaddress('', 'legacy'), "test")
w6.keypoolrefill(1)
# There should only be 1 key for legacy, 3 for descriptors
walletinfo = w6.getwalletinfo()
keys = 4 if self.options.descriptors else 1
assert_equal(walletinfo['keypoolsize'], keys)
assert_equal(walletinfo['keypoolsize_hd_internal'], keys)
# Allow empty passphrase, but there should be a warning
resp = self.nodes[0].createwallet(wallet_name='w7', disable_private_keys=False, blank=False, passphrase='')
assert_equal(resp["warnings"], [EMPTY_PASSPHRASE_MSG] if self.options.descriptors else [EMPTY_PASSPHRASE_MSG, LEGACY_WALLET_MSG])
Expand Down
25 changes: 12 additions & 13 deletions test/functional/wallet_descriptor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
assert_equal,
assert_raises_rpc_error
)
from test_framework.wallet_util import WalletUnlock


class WalletDescriptorTest(BitcoinTestFramework):
Expand Down Expand Up @@ -128,11 +129,10 @@ def run_test(self):

# Encrypt wallet 0
send_wrpc.encryptwallet('pass')
send_wrpc.walletpassphrase("pass", 999000)
addr = send_wrpc.getnewaddress()
info2 = send_wrpc.getaddressinfo(addr)
assert info1['hdmasterfingerprint'] != info2['hdmasterfingerprint']
send_wrpc.walletlock()
with WalletUnlock(send_wrpc, "pass"):
addr = send_wrpc.getnewaddress()
info2 = send_wrpc.getaddressinfo(addr)
assert info1['hdmasterfingerprint'] != info2['hdmasterfingerprint']
assert 'hdmasterfingerprint' in send_wrpc.getaddressinfo(send_wrpc.getnewaddress())
info3 = send_wrpc.getaddressinfo(addr)
assert_equal(info2['desc'], info3['desc'])
Expand All @@ -142,14 +142,13 @@ def run_test(self):
send_wrpc.getnewaddress()

self.log.info("Test that unlock is needed when deriving only hardened keys in an encrypted wallet")
send_wrpc.walletpassphrase("pass", 999000)
send_wrpc.importdescriptors([{
"desc": "wpkh(tprv8ZgxMBicQKsPd7Uf69XL1XwhmjHopUGep8GuEiJDZmbQz6o58LninorQAfcKZWARbtRtfnLcJ5MQ2AtHcQJCCRUcMRvmDUjyEmNUWwx8UbK/0h/*h)#y4dfsj7n",
"timestamp": "now",
"range": [0,10],
"active": True
}])
send_wrpc.walletlock()
with WalletUnlock(send_wrpc, "pass"):
send_wrpc.importdescriptors([{
"desc": "wpkh(tprv8ZgxMBicQKsPd7Uf69XL1XwhmjHopUGep8GuEiJDZmbQz6o58LninorQAfcKZWARbtRtfnLcJ5MQ2AtHcQJCCRUcMRvmDUjyEmNUWwx8UbK/0h/*h)#y4dfsj7n",
"timestamp": "now",
"range": [0,10],
"active": True
}])
# Exhaust keypool of 100
for _ in range(100):
send_wrpc.getnewaddress(address_type='bech32')
Expand Down
41 changes: 21 additions & 20 deletions test/functional/wallet_dump.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
assert_equal,
assert_raises_rpc_error,
)
from test_framework.wallet_util import WalletUnlock


def read_dump(file_name, addrs, script_addrs, hd_master_addr_old):
Expand Down Expand Up @@ -172,26 +173,26 @@ def run_test(self):

# encrypt wallet, restart, unlock and dump
self.nodes[0].encryptwallet('test')
self.nodes[0].walletpassphrase("test", 999000)
# Should be a no-op:
self.nodes[0].keypoolrefill()
self.nodes[0].dumpwallet(wallet_enc_dump)

found_comments, found_legacy_addr, found_p2sh_segwit_addr, found_bech32_addr, found_script_addr, found_addr_chg, found_addr_rsv, _ = \
read_dump(wallet_enc_dump, addrs, [multisig_addr], hd_master_addr_unenc)
assert '# End of dump' in found_comments # Check that file is not corrupt
assert_equal(dump_time_str, next(c for c in found_comments if c.startswith('# * Created on')))
assert_equal(dump_best_block_1, next(c for c in found_comments if c.startswith('# * Best block')))
assert_equal(dump_best_block_2, next(c for c in found_comments if c.startswith('# mined on')))
assert_equal(found_legacy_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_p2sh_segwit_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_bech32_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_script_addr, 1)
assert_equal(found_addr_chg, 90 * 2) # old reserve keys are marked as change now
assert_equal(found_addr_rsv, 90 * 2)

# Overwriting should fail
assert_raises_rpc_error(-8, "already exists", lambda: self.nodes[0].dumpwallet(wallet_enc_dump))
with WalletUnlock(self.nodes[0], "test"):
# Should be a no-op:
self.nodes[0].keypoolrefill()
self.nodes[0].dumpwallet(wallet_enc_dump)

found_comments, found_legacy_addr, found_p2sh_segwit_addr, found_bech32_addr, found_script_addr, found_addr_chg, found_addr_rsv, _ = \
read_dump(wallet_enc_dump, addrs, [multisig_addr], hd_master_addr_unenc)
assert '# End of dump' in found_comments # Check that file is not corrupt
assert_equal(dump_time_str, next(c for c in found_comments if c.startswith('# * Created on')))
assert_equal(dump_best_block_1, next(c for c in found_comments if c.startswith('# * Best block')))
assert_equal(dump_best_block_2, next(c for c in found_comments if c.startswith('# mined on')))
assert_equal(found_legacy_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_p2sh_segwit_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_bech32_addr, test_addr_count) # all keys must be in the dump
assert_equal(found_script_addr, 1)
assert_equal(found_addr_chg, 90 * 2) # old reserve keys are marked as change now
assert_equal(found_addr_rsv, 90 * 2)

# Overwriting should fail
assert_raises_rpc_error(-8, "already exists", lambda: self.nodes[0].dumpwallet(wallet_enc_dump))

# Restart node with new wallet, and test importwallet
self.restart_node(0)
Expand Down
22 changes: 10 additions & 12 deletions test/functional/wallet_encryption.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
assert_raises_rpc_error,
assert_equal,
)
from test_framework.wallet_util import WalletUnlock


class WalletEncryptionTest(BitcoinTestFramework):
Expand Down Expand Up @@ -59,19 +60,17 @@ def run_test(self):
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase + "wrong", 10)

# Test walletlock
self.nodes[0].walletpassphrase(passphrase, 999000)
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
with WalletUnlock(self.nodes[0], passphrase):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
assert_raises_rpc_error(-13, "Please enter the wallet passphrase with walletpassphrase first", self.nodes[0].signmessage, address, msg)

# Test passphrase changes
self.nodes[0].walletpassphrasechange(passphrase, passphrase2)
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase, 10)
self.nodes[0].walletpassphrase(passphrase2, 999000)
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
with WalletUnlock(self.nodes[0], passphrase2):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)

# Test timeout bounds
assert_raises_rpc_error(-8, "Timeout cannot be negative.", self.nodes[0].walletpassphrase, passphrase2, -10)
Expand All @@ -97,10 +96,9 @@ def run_test(self):
self.nodes[0].walletpassphrasechange(passphrase2, passphrase_with_nulls)
# walletpassphrasechange should not stop at null characters
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase_with_nulls.partition("\0")[0], 10)
self.nodes[0].walletpassphrase(passphrase_with_nulls, 999000)
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
with WalletUnlock(self.nodes[0], passphrase_with_nulls):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)


if __name__ == '__main__':
Expand Down
48 changes: 23 additions & 25 deletions test/functional/wallet_fundrawtransaction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
find_vout_for_address,
get_fee,
)
from test_framework.wallet_util import generate_keypair
from test_framework.wallet_util import generate_keypair, WalletUnlock

ERR_NOT_ENOUGH_PRESET_INPUTS = "The preselected coins total amount does not cover the transaction target. " \
"Please allow other inputs to be automatically selected or include more coins manually"
Expand Down Expand Up @@ -581,19 +581,18 @@ def test_locked_wallet(self):
wallet.encryptwallet("test")

if self.options.descriptors:
wallet.walletpassphrase("test", 999000)
wallet.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPdYeeZbPSKd2KYLmeVKtcFA7kqCxDvDR13MQ6us8HopUR2wLcS2ZKPhLyKsqpDL2FtL73LMHcgoCL7DXsciA8eX8nbjCR2eG/0h/*h)'),
'timestamp': 'now',
'active': True
},
{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPdYeeZbPSKd2KYLmeVKtcFA7kqCxDvDR13MQ6us8HopUR2wLcS2ZKPhLyKsqpDL2FtL73LMHcgoCL7DXsciA8eX8nbjCR2eG/1h/*h)'),
'timestamp': 'now',
'active': True,
'internal': True
}])
wallet.walletlock()
with WalletUnlock(wallet, "test"):
wallet.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPdYeeZbPSKd2KYLmeVKtcFA7kqCxDvDR13MQ6us8HopUR2wLcS2ZKPhLyKsqpDL2FtL73LMHcgoCL7DXsciA8eX8nbjCR2eG/0h/*h)'),
'timestamp': 'now',
'active': True
},
{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPdYeeZbPSKd2KYLmeVKtcFA7kqCxDvDR13MQ6us8HopUR2wLcS2ZKPhLyKsqpDL2FtL73LMHcgoCL7DXsciA8eX8nbjCR2eG/1h/*h)'),
'timestamp': 'now',
'active': True,
'internal': True
}])

# Drain the keypool.
wallet.getnewaddress()
Expand All @@ -619,9 +618,8 @@ def test_locked_wallet(self):
assert_raises_rpc_error(-4, "Transaction needs a change address, but we can't generate it.", wallet.fundrawtransaction, rawtx)

# Refill the keypool.
wallet.walletpassphrase("test", 999000)
wallet.keypoolrefill(8) #need to refill the keypool to get an internal change address
wallet.walletlock()
with WalletUnlock(wallet, "test"):
wallet.keypoolrefill(8) #need to refill the keypool to get an internal change address

assert_raises_rpc_error(-13, "walletpassphrase", wallet.sendtoaddress, self.nodes[0].getnewaddress(), 1.2)

Expand All @@ -634,16 +632,16 @@ def test_locked_wallet(self):
assert fundedTx["changepos"] != -1

# Now we need to unlock.
wallet.walletpassphrase("test", 999000)
signedTx = wallet.signrawtransactionwithwallet(fundedTx['hex'])
wallet.sendrawtransaction(signedTx['hex'])
self.generate(self.nodes[1], 1)
with WalletUnlock(wallet, "test"):
signedTx = wallet.signrawtransactionwithwallet(fundedTx['hex'])
wallet.sendrawtransaction(signedTx['hex'])
self.generate(self.nodes[1], 1)

# Make sure funds are received at node1.
assert_equal(oldBalance+Decimal('51.10000000'), self.nodes[0].getbalance())
# Make sure funds are received at node1.
assert_equal(oldBalance+Decimal('51.10000000'), self.nodes[0].getbalance())

# Restore pre-test wallet state
wallet.sendall(recipients=[df_wallet.getnewaddress(), df_wallet.getnewaddress(), df_wallet.getnewaddress()])
# Restore pre-test wallet state
wallet.sendall(recipients=[df_wallet.getnewaddress(), df_wallet.getnewaddress(), df_wallet.getnewaddress()])
wallet.unloadwallet()
self.generate(self.nodes[1], 1)

Expand Down
Loading

0 comments on commit 091d29c

Please sign in to comment.