diff --git a/chia/_tests/environments/wallet.py b/chia/_tests/environments/wallet.py index eb23d97ec651..eae392243497 100644 --- a/chia/_tests/environments/wallet.py +++ b/chia/_tests/environments/wallet.py @@ -1,9 +1,11 @@ from __future__ import annotations +import contextlib import json import operator +import unittest from dataclasses import asdict, dataclass, field -from typing import TYPE_CHECKING, ClassVar, Optional, Union, cast +from typing import TYPE_CHECKING, Any, ClassVar, Iterator, Union, cast from chia._tests.environments.common import ServiceEnvironment from chia.rpc.full_node_rpc_client import FullNodeRpcClient @@ -15,7 +17,6 @@ from chia.simulator.full_node_simulator import FullNodeSimulator from chia.types.blockchain_format.sized_bytes import bytes32 from chia.util.ints import uint32 -from chia.wallet.derivation_record import DerivationRecord from chia.wallet.transaction_record import TransactionRecord from chia.wallet.util.transaction_type import CLAWBACK_INCOMING_TRANSACTION_TYPES from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, TXConfig @@ -260,6 +261,22 @@ async def wait_for_transactions_to_settle( return pending_txs +class NewPuzzleHashError(Exception): + pass + + +def catch_puzzle_hash_errors(func: Any) -> Any: + @contextlib.asynccontextmanager + async def catching_puzhash_errors(self: WalletStateManager, *args: Any, **kwargs: Any) -> Any: + try: + async with func(self, *args, **kwargs) as action_scope: + yield action_scope + except NewPuzzleHashError: + pass + + return catching_puzhash_errors + + @dataclass class WalletTestFramework: full_node: FullNodeSimulator @@ -268,6 +285,15 @@ class WalletTestFramework: environments: list[WalletEnvironment] tx_config: TXConfig = DEFAULT_TX_CONFIG + @staticmethod + @contextlib.contextmanager + def new_puzzle_hashes_allowed() -> Iterator[None]: + with unittest.mock.patch( + "chia.wallet.wallet_state_manager.WalletStateManager.new_action_scope", + catch_puzzle_hash_errors(WalletStateManager.new_action_scope), + ): + yield + async def process_pending_states( self, state_transitions: list[WalletStateTransition], invalid_transactions: list[bytes32] = [] ) -> None: @@ -284,13 +310,11 @@ async def process_pending_states( """ # Take note of the number of puzzle hashes if we're supposed to be reusing if self.tx_config.reuse_puzhash: - puzzle_hash_indexes: list[dict[uint32, Optional[DerivationRecord]]] = [] + puzzle_hash_indexes: list[dict[uint32, int]] = [] for env in self.environments: - ph_indexes: dict[uint32, Optional[DerivationRecord]] = {} + ph_indexes: dict[uint32, int] = {} for wallet_id in env.wallet_state_manager.wallets: - ph_indexes[ - wallet_id - ] = await env.wallet_state_manager.puzzle_store.get_current_derivation_record_for_wallet(wallet_id) + ph_indexes[wallet_id] = await env.wallet_state_manager.puzzle_store.get_unused_count(wallet_id) puzzle_hash_indexes.append(ph_indexes) pending_txs: list[list[TransactionRecord]] = [] @@ -359,5 +383,5 @@ async def process_pending_states( for env, ph_indexes_before in zip(self.environments, puzzle_hash_indexes): for wallet_id, ph_index in zip(env.wallet_state_manager.wallets, ph_indexes_before): assert ph_indexes_before[wallet_id] == ( - await env.wallet_state_manager.puzzle_store.get_current_derivation_record_for_wallet(wallet_id) + await env.wallet_state_manager.puzzle_store.get_unused_count(wallet_id) ) diff --git a/chia/_tests/wallet/conftest.py b/chia/_tests/wallet/conftest.py index 4a0a117e5639..43e231a731b0 100644 --- a/chia/_tests/wallet/conftest.py +++ b/chia/_tests/wallet/conftest.py @@ -1,5 +1,7 @@ from __future__ import annotations +import contextlib +import unittest from collections.abc import AsyncIterator, Awaitable from contextlib import AsyncExitStack from dataclasses import replace @@ -14,7 +16,7 @@ run_block_generator2, ) -from chia._tests.environments.wallet import WalletEnvironment, WalletState, WalletTestFramework +from chia._tests.environments.wallet import NewPuzzleHashError, WalletEnvironment, WalletState, WalletTestFramework from chia._tests.util.setup_nodes import setup_simulators_and_wallets_service from chia._tests.wallet.wallet_block_tools import WalletBlockTools from chia.consensus.constants import ConsensusConstants @@ -144,6 +146,28 @@ def tx_config(request: Any) -> TXConfig: return replace(DEFAULT_TX_CONFIG, reuse_puzhash=request.param) +def new_action_scope_wrapper(func: Any) -> Any: + @contextlib.asynccontextmanager + async def wrapped_new_action_scope(self: WalletStateManager, *args: Any, **kwargs: Any) -> Any: + # Take note of the number of puzzle hashes if we're supposed to be reusing + ph_indexes: dict[uint32, int] = {} + for wallet_id in self.wallets: + ph_indexes[wallet_id] = await self.puzzle_store.get_unused_count(wallet_id) + + async with func(self, *args, **kwargs) as action_scope: + yield action_scope + + # Finally, check that the number of puzzle hashes did or did not increase by the specified amount + if action_scope.config.tx_config.reuse_puzhash: + for wallet_id, ph_index in zip(self.wallets, ph_indexes): + if not ph_indexes[wallet_id] == (await self.puzzle_store.get_unused_count(wallet_id)): + raise NewPuzzleHashError( + f"wallet ID {wallet_id} generated new puzzle hashes while reuse_puzhash was False" + ) + + return wrapped_new_action_scope + + # This fixture automatically creates 4 parametrized tests trusted/untrusted x reuse/new derivations # These parameterizations can be skipped by manually specifying "trusted" or "reuse puzhash" to the fixture @pytest.fixture(scope="function") @@ -174,77 +198,81 @@ async def wallet_environments( full_node[0]._api.full_node.config = {**full_node[0]._api.full_node.config, **config_overrides} - wallet_rpc_clients: list[WalletRpcClient] = [] - async with AsyncExitStack() as astack: - for service in wallet_services: - service._node.config = { - **service._node.config, - "trusted_peers": ( - {full_node[0]._api.server.node_id.hex(): full_node[0]._api.server.node_id.hex()} - if trusted_full_node - else {} - ), - **config_overrides, - } - service._node.wallet_state_manager.config = service._node.config - # Shorten the 10 seconds default value - service._node.coin_state_retry_seconds = 2 - await service._node.server.start_client( - PeerInfo(bt.config["self_hostname"], full_node[0]._api.full_node.server.get_port()), None - ) - wallet_rpc_clients.append( - await astack.enter_async_context( - WalletRpcClient.create_as_context( - bt.config["self_hostname"], - # Semantics guarantee us a non-None value here - service.rpc_server.listen_port, # type: ignore[union-attr] - service.root_path, - service.config, + new_action_scope_wrapped = new_action_scope_wrapper(WalletStateManager.new_action_scope) + with unittest.mock.patch( + "chia.wallet.wallet_state_manager.WalletStateManager.new_action_scope", new=new_action_scope_wrapped + ): + wallet_rpc_clients: list[WalletRpcClient] = [] + async with AsyncExitStack() as astack: + for service in wallet_services: + service._node.config = { + **service._node.config, + "trusted_peers": ( + {full_node[0]._api.server.node_id.hex(): full_node[0]._api.server.node_id.hex()} + if trusted_full_node + else {} + ), + **config_overrides, + } + service._node.wallet_state_manager.config = service._node.config + # Shorten the 10 seconds default value + service._node.coin_state_retry_seconds = 2 + await service._node.server.start_client( + PeerInfo(bt.config["self_hostname"], full_node[0]._api.full_node.server.get_port()), None + ) + wallet_rpc_clients.append( + await astack.enter_async_context( + WalletRpcClient.create_as_context( + bt.config["self_hostname"], + # Semantics guarantee us a non-None value here + service.rpc_server.listen_port, # type: ignore[union-attr] + service.root_path, + service.config, + ) ) ) - ) - wallet_states: list[WalletState] = [] - for service, blocks_needed in zip(wallet_services, request.param["blocks_needed"]): - if blocks_needed > 0: - await full_node[0]._api.farm_blocks_to_wallet( - count=blocks_needed, wallet=service._node.wallet_state_manager.main_wallet + wallet_states: list[WalletState] = [] + for service, blocks_needed in zip(wallet_services, request.param["blocks_needed"]): + if blocks_needed > 0: + await full_node[0]._api.farm_blocks_to_wallet( + count=blocks_needed, wallet=service._node.wallet_state_manager.main_wallet + ) + await full_node[0]._api.wait_for_wallet_synced(wallet_node=service._node, timeout=20) + wallet_states.append( + WalletState( + Balance( + confirmed_wallet_balance=uint128(2_000_000_000_000 * blocks_needed), + unconfirmed_wallet_balance=uint128(2_000_000_000_000 * blocks_needed), + spendable_balance=uint128(2_000_000_000_000 * blocks_needed), + pending_change=uint64(0), + max_send_amount=uint128(2_000_000_000_000 * blocks_needed), + unspent_coin_count=uint32(2 * blocks_needed), + pending_coin_removal_count=uint32(0), + ), + ) ) - await full_node[0]._api.wait_for_wallet_synced(wallet_node=service._node, timeout=20) - wallet_states.append( - WalletState( - Balance( - confirmed_wallet_balance=uint128(2_000_000_000_000 * blocks_needed), - unconfirmed_wallet_balance=uint128(2_000_000_000_000 * blocks_needed), - spendable_balance=uint128(2_000_000_000_000 * blocks_needed), - pending_change=uint64(0), - max_send_amount=uint128(2_000_000_000_000 * blocks_needed), - unspent_coin_count=uint32(2 * blocks_needed), - pending_coin_removal_count=uint32(0), - ), + + assert full_node[0].rpc_server is not None + client_node = await astack.enter_async_context( + FullNodeRpcClient.create_as_context( + bt.config["self_hostname"], + full_node[0].rpc_server.listen_port, + full_node[0].root_path, + full_node[0].config, ) ) - - assert full_node[0].rpc_server is not None - client_node = await astack.enter_async_context( - FullNodeRpcClient.create_as_context( - bt.config["self_hostname"], - full_node[0].rpc_server.listen_port, - full_node[0].root_path, - full_node[0].config, + yield WalletTestFramework( + full_node[0]._api, + client_node, + trusted_full_node, + [ + WalletEnvironment( + service=service, + rpc_client=rpc_client, + wallet_states={uint32(1): wallet_state}, + ) + for service, rpc_client, wallet_state in zip(wallet_services, wallet_rpc_clients, wallet_states) + ], + tx_config, ) - ) - yield WalletTestFramework( - full_node[0]._api, - client_node, - trusted_full_node, - [ - WalletEnvironment( - service=service, - rpc_client=rpc_client, - wallet_states={uint32(1): wallet_state}, - ) - for service, rpc_client, wallet_state in zip(wallet_services, wallet_rpc_clients, wallet_states) - ], - tx_config, - ) diff --git a/chia/_tests/wallet/did_wallet/test_did.py b/chia/_tests/wallet/did_wallet/test_did.py index a7f89ec0e80c..2197cc7755ce 100644 --- a/chia/_tests/wallet/did_wallet/test_did.py +++ b/chia/_tests/wallet/did_wallet/test_did.py @@ -843,7 +843,7 @@ async def test_did_find_lost_did(self, wallet_environments: WalletTestFramework) # Delete the coin and wallet coin = await did_wallet.get_coin() await wallet_node_0.wallet_state_manager.coin_store.delete_coin_record(coin.name()) - await wallet_node_0.wallet_state_manager.user_store.delete_wallet(did_wallet.wallet_info.id) + await wallet_node_0.wallet_state_manager.delete_wallet(did_wallet.wallet_info.id) wallet_node_0.wallet_state_manager.wallets.pop(did_wallet.wallet_info.id) assert len(wallet_node_0.wallet_state_manager.wallets) == 1 # Find lost DID @@ -857,6 +857,21 @@ async def test_did_find_lost_did(self, wallet_environments: WalletTestFramework) ) ) did_wallet = wallet_node_0.wallet_state_manager.wallets[did_wallets[0].id] + env_0.wallet_aliases["did_found"] = did_wallets[0].id + await env_0.change_balances( + { + "did_found": { + "init": True, + "confirmed_wallet_balance": 101, + "unconfirmed_wallet_balance": 101, + "spendable_balance": 101, + "max_send_amount": 101, + "unspent_coin_count": 1, + } + } + ) + await env_0.check_balances() + # Spend DID recovery_list = [bytes32.fromhex(did_wallet.get_my_DID())] await did_wallet.update_recovery_list(recovery_list, uint64(1)) @@ -866,31 +881,21 @@ async def test_did_find_lost_did(self, wallet_environments: WalletTestFramework) ) as action_scope: await did_wallet.create_update_spend(action_scope) - env_0.wallet_aliases["did_found"] = 3 - await wallet_environments.process_pending_states( [ WalletStateTransition( pre_block_balance_updates={ "did_found": { - "init": True, - "confirmed_wallet_balance": 101, - "unconfirmed_wallet_balance": 202, # Seems strange - "spendable_balance": 101, - "max_send_amount": 101, - "unspent_coin_count": 1, + "spendable_balance": -101, + "max_send_amount": -101, "pending_change": 101, "pending_coin_removal_count": 1, - "set_remainder": True, }, }, post_block_balance_updates={ "did_found": { - "confirmed_wallet_balance": 0, - "unconfirmed_wallet_balance": -101, - "spendable_balance": 0, - "max_send_amount": 0, - "unspent_coin_count": 0, + "spendable_balance": 101, + "max_send_amount": 101, "pending_change": -101, "pending_coin_removal_count": -1, }, @@ -902,7 +907,7 @@ async def test_did_find_lost_did(self, wallet_environments: WalletTestFramework) # Delete the coin and change inner puzzle coin = await did_wallet.get_coin() await wallet_node_0.wallet_state_manager.coin_store.delete_coin_record(coin.name()) - new_inner_puzzle = await did_wallet.get_new_did_innerpuz() + new_inner_puzzle = await did_wallet.get_did_innerpuz(new=True) did_wallet.did_info = dataclasses.replace(did_wallet.did_info, current_inner=new_inner_puzzle) # Recovery the coin assert did_wallet.did_info.origin_coin is not None # mypy @@ -1051,7 +1056,7 @@ async def test_did_attest_after_recovery(self, wallet_environments: WalletTestFr backup_data, ) env_0.wallet_aliases["did_2"] = 3 - new_ph = await did_wallet_3.get_new_did_inner_hash() + new_ph = await did_wallet_3.get_did_inner_hash(new=True) coin = await did_wallet_2.get_coin() pubkey = ( await did_wallet_3.wallet_state_manager.get_unused_derivation_record(did_wallet_3.wallet_info.id) @@ -1138,7 +1143,7 @@ async def test_did_attest_after_recovery(self, wallet_environments: WalletTestFr ) env_1.wallet_aliases["did_2"] = 3 coin = await did_wallet.get_coin() - new_ph = await did_wallet_4.get_new_did_inner_hash() + new_ph = await did_wallet_4.get_did_inner_hash(new=True) pubkey = ( await did_wallet_4.wallet_state_manager.get_unused_derivation_record(did_wallet_4.wallet_info.id) ).pubkey diff --git a/chia/_tests/wallet/rpc/test_wallet_rpc.py b/chia/_tests/wallet/rpc/test_wallet_rpc.py index be30431d2e9f..5b9996d748fa 100644 --- a/chia/_tests/wallet/rpc/test_wallet_rpc.py +++ b/chia/_tests/wallet/rpc/test_wallet_rpc.py @@ -2728,10 +2728,11 @@ async def test_split_coins(wallet_environments: WalletTestFramework) -> None: ) assert response == SplitCoinsResponse([], []) - await env.rpc_client.split_coins( - xch_request, - wallet_environments.tx_config, - ) + with wallet_environments.new_puzzle_hashes_allowed(): + await env.rpc_client.split_coins( + xch_request, + wallet_environments.tx_config, + ) await wallet_environments.process_pending_states( [ @@ -2797,10 +2798,11 @@ async def test_split_coins(wallet_environments: WalletTestFramework) -> None: push=True, ) - await env.rpc_client.split_coins( - cat_request, - wallet_environments.tx_config, - ) + with wallet_environments.new_puzzle_hashes_allowed(): + await env.rpc_client.split_coins( + cat_request, + wallet_environments.tx_config, + ) await wallet_environments.process_pending_states( [ @@ -2966,7 +2968,7 @@ async def test_combine_coins(wallet_environments: WalletTestFramework) -> None: async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope: await cat_wallet.generate_signed_transaction( [BIG_COIN_AMOUNT, SMALL_COIN_AMOUNT, REALLY_SMALL_COIN_AMOUNT], - [await env.xch_wallet.get_puzzle_hash(new=action_scope.config.tx_config.reuse_puzhash)] * 3, + [await env.xch_wallet.get_puzzle_hash(new=not action_scope.config.tx_config.reuse_puzhash)] * 3, action_scope, ) diff --git a/chia/_tests/wallet/vc_wallet/test_vc_wallet.py b/chia/_tests/wallet/vc_wallet/test_vc_wallet.py index 701d44be501f..135576ed605f 100644 --- a/chia/_tests/wallet/vc_wallet/test_vc_wallet.py +++ b/chia/_tests/wallet/vc_wallet/test_vc_wallet.py @@ -528,16 +528,18 @@ async def test_vc_lifecycle(wallet_environments: WalletTestFramework) -> None: ) # Test melting a CRCAT - tx = ( - await client_1.cat_spend( - env_1.dealias_wallet_id("crcat"), - wallet_environments.tx_config, - uint64(20), - wallet_1_addr, - uint64(0), - cat_discrepancy=(-50, Program.to(None), Program.to(None)), - ) - ).transaction + # This is intended to trigger an edge case where the output and change are the same forcing a new puzhash + with wallet_environments.new_puzzle_hashes_allowed(): + tx = ( + await client_1.cat_spend( + env_1.dealias_wallet_id("crcat"), + wallet_environments.tx_config, + uint64(20), + wallet_1_addr, + uint64(0), + cat_discrepancy=(-50, Program.to(None), Program.to(None)), + ) + ).transaction [tx] = await wallet_node_1.wallet_state_manager.add_pending_transactions([tx]) await wallet_environments.process_pending_states( [ diff --git a/chia/rpc/wallet_rpc_api.py b/chia/rpc/wallet_rpc_api.py index 74aaf8d851f8..c9d0229677d6 100644 --- a/chia/rpc/wallet_rpc_api.py +++ b/chia/rpc/wallet_rpc_api.py @@ -1256,7 +1256,7 @@ async def combine_coins( primary_output_amount = uint64(primary_output_amount - request.fee) await wallet.generate_signed_transaction( primary_output_amount, - await wallet.get_puzzle_hash(new=action_scope.config.tx_config.reuse_puzhash), + await wallet.get_puzzle_hash(new=not action_scope.config.tx_config.reuse_puzhash), action_scope, request.fee, set(coins), @@ -1266,7 +1266,7 @@ async def combine_coins( assert isinstance(wallet, CATWallet) await wallet.generate_signed_transaction( [primary_output_amount], - [await wallet.get_puzzle_hash(new=action_scope.config.tx_config.reuse_puzhash)], + [await wallet.get_puzzle_hash(new=not action_scope.config.tx_config.reuse_puzhash)], action_scope, request.fee, coins=set(coins), @@ -3310,14 +3310,18 @@ async def nft_mint_nft( if isinstance(royalty_address, str): royalty_puzhash = decode_puzzle_hash(royalty_address) elif royalty_address is None: - royalty_puzhash = await nft_wallet.standard_wallet.get_new_puzzlehash() + royalty_puzhash = await nft_wallet.standard_wallet.get_puzzle_hash( + new=not action_scope.config.tx_config.reuse_puzhash + ) else: royalty_puzhash = royalty_address target_address = request.get("target_address") if isinstance(target_address, str): target_puzhash = decode_puzzle_hash(target_address) elif target_address is None: - target_puzhash = await nft_wallet.standard_wallet.get_new_puzzlehash() + target_puzhash = await nft_wallet.standard_wallet.get_puzzle_hash( + new=not action_scope.config.tx_config.reuse_puzhash + ) else: target_puzhash = target_address if "uris" not in request: diff --git a/chia/wallet/cat_wallet/cat_wallet.py b/chia/wallet/cat_wallet/cat_wallet.py index c9264e5a9eeb..66227e8201d6 100644 --- a/chia/wallet/cat_wallet/cat_wallet.py +++ b/chia/wallet/cat_wallet/cat_wallet.py @@ -144,10 +144,10 @@ async def create_new_cat_wallet( ) assert self.cat_info.limitations_program_hash != empty_bytes except Exception: - await wallet_state_manager.user_store.delete_wallet(self.id()) + await wallet_state_manager.delete_wallet(self.id()) raise if spend_bundle is None: - await wallet_state_manager.user_store.delete_wallet(self.id()) + await wallet_state_manager.delete_wallet(self.id()) raise ValueError("Failed to create spend.") await self.wallet_state_manager.add_new_wallet(self) @@ -422,6 +422,12 @@ async def puzzle_solution_received(self, coin: Coin, parent_coin_data: Optional[ # We also need to make sure there's no record of the transaction await self.wallet_state_manager.tx_store.delete_transaction_record(record.coin.name()) + async def get_inner_puzzle(self, new: bool) -> Program: + return await self.standard_wallet.get_puzzle(new=new) + + async def get_inner_puzzle_hash(self, new: bool) -> bytes32: + return await self.standard_wallet.get_puzzle_hash(new=new) + async def get_new_inner_hash(self) -> bytes32: puzzle = await self.get_new_inner_puzzle() return puzzle.get_tree_hash() diff --git a/chia/wallet/dao_wallet/dao_wallet.py b/chia/wallet/dao_wallet/dao_wallet.py index 43cd6680068b..3693d5a8c6a9 100644 --- a/chia/wallet/dao_wallet/dao_wallet.py +++ b/chia/wallet/dao_wallet/dao_wallet.py @@ -184,7 +184,7 @@ async def create_new_dao_and_wallet( fee_for_cat=fee_for_cat, ) except Exception as e_info: # pragma: no cover - await wallet_state_manager.user_store.delete_wallet(self.id()) + await wallet_state_manager.delete_wallet(self.id()) self.log.exception(f"Failed to create dao wallet: {e_info}") raise diff --git a/chia/wallet/did_wallet/did_wallet.py b/chia/wallet/did_wallet/did_wallet.py index f77b0deba339..37b01c94ebc0 100644 --- a/chia/wallet/did_wallet/did_wallet.py +++ b/chia/wallet/did_wallet/did_wallet.py @@ -142,7 +142,7 @@ async def create_new_did_wallet( try: await self.generate_new_decentralised_id(amount, action_scope, fee, extra_conditions) except Exception: - await wallet_state_manager.user_store.delete_wallet(self.id()) + await wallet_state_manager.delete_wallet(self.id()) raise await self.wallet_state_manager.add_new_wallet(self) @@ -208,7 +208,6 @@ async def create_new_did_wallet_from_coin_spend( :param name: Wallet name :return: DID wallet """ - self = DIDWallet() self.wallet_state_manager = wallet_state_manager if name is None: @@ -317,6 +316,7 @@ async def get_pending_change_balance(self) -> uint64: for record in unconfirmed_tx: our_spend = False + # Need to check belonging with hint_dict for coin in record.removals: if await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id()): our_spend = True @@ -326,7 +326,12 @@ async def get_pending_change_balance(self) -> uint64: continue for coin in record.additions: - if (await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id())) and ( + hint_dict = { + coin_id: bytes32(memos[0]) + for coin_id, memos in record.memos + if len(memos) > 0 and len(memos[0]) == 32 + } + if (await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id(), hint_dict)) and ( coin not in record.removals ): addition_amount += coin.amount @@ -576,7 +581,7 @@ async def create_update_spend( assert self.did_info.current_inner is not None assert self.did_info.origin_coin is not None coin = await self.get_coin() - new_inner_puzzle = await self.get_new_did_innerpuz() + new_inner_puzzle = await self.get_did_innerpuz(new=not action_scope.config.tx_config.reuse_puzhash) uncurried = did_wallet_puzzles.uncurry_innerpuz(new_inner_puzzle) assert uncurried is not None p2_puzzle = uncurried[0] @@ -1097,6 +1102,12 @@ async def recovery_spend( ) await self.save_info(new_did_info) + async def get_p2_inner_hash(self, new: bool) -> bytes32: + return await self.standard_wallet.get_puzzle_hash(new=new) + + async def get_p2_inner_puzzle(self, new: bool) -> Program: + return await self.standard_wallet.get_puzzle(new=new) + async def get_new_p2_inner_hash(self) -> bytes32: puzzle = await self.get_new_p2_inner_puzzle() return puzzle.get_tree_hash() @@ -1104,7 +1115,7 @@ async def get_new_p2_inner_hash(self) -> bytes32: async def get_new_p2_inner_puzzle(self) -> Program: return await self.standard_wallet.get_new_puzzle() - async def get_new_did_innerpuz(self, origin_id: Optional[bytes32] = None) -> Program: + async def get_did_innerpuz(self, new: bool, origin_id: Optional[bytes32] = None) -> Program: if self.did_info.origin_coin is not None: launcher_id = self.did_info.origin_coin.name() elif origin_id is not None: @@ -1112,15 +1123,15 @@ async def get_new_did_innerpuz(self, origin_id: Optional[bytes32] = None) -> Pro else: raise ValueError("must have origin coin") return did_wallet_puzzles.create_innerpuz( - p2_puzzle_or_hash=await self.get_new_p2_inner_puzzle(), + p2_puzzle_or_hash=await self.get_p2_inner_puzzle(new=new), recovery_list=self.did_info.backup_ids, num_of_backup_ids_needed=self.did_info.num_of_backup_ids_needed, launcher_id=launcher_id, metadata=did_wallet_puzzles.metadata_to_program(json.loads(self.did_info.metadata)), ) - async def get_new_did_inner_hash(self) -> bytes32: - innerpuz = await self.get_new_did_innerpuz() + async def get_did_inner_hash(self, new: bool) -> bytes32: + innerpuz = await self.get_did_innerpuz(new=new) return innerpuz.get_tree_hash() async def get_innerpuz_for_new_innerhash(self, pubkey: G1Element): @@ -1215,7 +1226,9 @@ async def generate_new_decentralised_id( genesis_launcher_puz = SINGLETON_LAUNCHER_PUZZLE launcher_coin = Coin(origin.name(), genesis_launcher_puz.get_tree_hash(), amount) - did_inner: Program = await self.get_new_did_innerpuz(launcher_coin.name()) + did_inner: Program = await self.get_did_innerpuz( + new=not action_scope.config.tx_config.reuse_puzhash, origin_id=launcher_coin.name() + ) did_inner_hash = did_inner.get_tree_hash() did_full_puz = create_singleton_puzzle(did_inner, launcher_coin.name()) did_puzzle_hash = did_full_puz.get_tree_hash() diff --git a/chia/wallet/nft_wallet/nft_wallet.py b/chia/wallet/nft_wallet/nft_wallet.py index 9c95ac25c2f1..71f109afec6c 100644 --- a/chia/wallet/nft_wallet/nft_wallet.py +++ b/chia/wallet/nft_wallet/nft_wallet.py @@ -292,7 +292,7 @@ async def remove_coin(self, coin: Coin, height: uint32) -> None: if did_wallet_info.origin_coin.name() == self.did_id: return self.log.info(f"No NFT, deleting wallet {self.wallet_info.name} ...") - await self.wallet_state_manager.user_store.delete_wallet(self.wallet_info.id) + await self.wallet_state_manager.delete_wallet(self.wallet_info.id) self.wallet_state_manager.wallets.pop(self.wallet_info.id) else: self.log.info("Tried removing NFT coin that doesn't exist: %s", coin.name()) diff --git a/chia/wallet/puzzles/tails.py b/chia/wallet/puzzles/tails.py index efbdfc24c034..93b749525a8f 100644 --- a/chia/wallet/puzzles/tails.py +++ b/chia/wallet/puzzles/tails.py @@ -97,7 +97,7 @@ async def generate_issuance_bundle( origin = coins.copy().pop() origin_id = origin.name() - cat_inner: Program = await wallet.get_new_inner_puzzle() + cat_inner: Program = await wallet.get_inner_puzzle(new=not action_scope.config.tx_config.reuse_puzhash) tail: Program = cls.construct([Program.to(origin_id)]) wallet.lineage_store = await CATLineageStore.create( diff --git a/chia/wallet/wallet_puzzle_store.py b/chia/wallet/wallet_puzzle_store.py index f4e3d3cbdad1..ebf16d267db1 100644 --- a/chia/wallet/wallet_puzzle_store.py +++ b/chia/wallet/wallet_puzzle_store.py @@ -343,6 +343,22 @@ async def get_unused_derivation_path(self) -> Optional[uint32]: return None + async def get_unused_derivation_path_for_wallet(self, wallet_id: uint32) -> Optional[uint32]: + """ + Returns the first unused derivation path by derivation_index. + """ + async with self.db_wrapper.reader_no_transaction() as conn: + row = await execute_fetchone( + conn, + "SELECT MIN(derivation_index) FROM derivation_paths WHERE wallet_id=? AND used=0 AND hardened=0;", + (wallet_id,), + ) + + if row is not None and row[0] is not None: + return uint32(row[0]) + + return None + async def delete_wallet(self, wallet_id: uint32) -> None: async with self.db_wrapper.writer_maybe_transaction() as conn: # First fetch all puzzle hashes since we need them to drop them from the cache @@ -363,3 +379,18 @@ async def delete_wallet(self, wallet_id: uint32) -> None: except KeyError: pass self.last_derivation_index = None + + async def get_unused_count(self, wallet_id: uint32) -> int: + """ + Returns a count of unused derivation indexes + """ + + async with self.db_wrapper.reader_no_transaction() as conn: + row = await execute_fetchone( + conn, + "SELECT COUNT(*) FROM derivation_paths WHERE wallet_id=? AND used=1", + (wallet_id,), + ) + row_count = 0 if row is None else row[0] + + return row_count diff --git a/chia/wallet/wallet_state_manager.py b/chia/wallet/wallet_state_manager.py index 420aecc35610..832a1c559975 100644 --- a/chia/wallet/wallet_state_manager.py +++ b/chia/wallet/wallet_state_manager.py @@ -519,7 +519,7 @@ async def update_wallet_puzzle_hashes(self, wallet_id: uint32) -> None: derivation_paths: list[DerivationRecord] = [] target_wallet = self.wallets[wallet_id] last: Optional[uint32] = await self.puzzle_store.get_last_derivation_path_for_wallet(wallet_id) - unused: Optional[uint32] = await self.puzzle_store.get_unused_derivation_path() + unused: Optional[uint32] = await self.puzzle_store.get_unused_derivation_path_for_wallet(wallet_id) if unused is None: # This handles the case where the database has entries but they have all been used unused = await self.puzzle_store.get_last_derivation_path() @@ -1268,7 +1268,7 @@ async def handle_did( and launch_id == wallet.did_info.origin_coin.name() and not wallet.did_info.sent_recovery_transaction ): - await self.user_store.delete_wallet(wallet.id()) + await self.delete_wallet(wallet.id()) removed_wallet_ids.append(wallet.id()) for remove_id in removed_wallet_ids: self.wallets.pop(remove_id) @@ -1543,7 +1543,7 @@ async def handle_nft( break if is_empty and nft_wallet.did_id is not None and not has_did: self.log.info(f"No NFT, deleting wallet {nft_wallet.did_id.hex()} ...") - await self.user_store.delete_wallet(nft_wallet.wallet_info.id) + await self.delete_wallet(nft_wallet.wallet_info.id) self.wallets.pop(nft_wallet.wallet_info.id) if nft_wallet.nft_wallet_info.did_id == new_did_id and new_derivation_record is not None: self.log.info( @@ -2485,7 +2485,7 @@ async def reorg_rollback(self, height: int) -> list[uint32]: if remove: remove_ids.append(wallet_id) for wallet_id in remove_ids: - await self.user_store.delete_wallet(wallet_id) + await self.delete_wallet(wallet_id) self.state_changed("wallet_removed", wallet_id) return remove_ids @@ -2817,3 +2817,7 @@ async def new_action_scope( extra_spends=extra_spends, ) as action_scope: yield action_scope + + async def delete_wallet(self, wallet_id: uint32) -> None: + await self.user_store.delete_wallet(wallet_id) + await self.puzzle_store.delete_wallet(wallet_id) diff --git a/poetry.lock b/poetry.lock index 13636b1e3f42..7c5ece9a2081 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2803,13 +2803,13 @@ files = [ [[package]] name = "types-aiofiles" -version = "23.2.0.20240311" +version = "24.1.0.20240626" description = "Typing stubs for aiofiles" optional = true python-versions = ">=3.8" files = [ - {file = "types-aiofiles-23.2.0.20240311.tar.gz", hash = "sha256:208e6b090de732739ef74ab8f133c954479c8e77e614f276f9e475a0cc986430"}, - {file = "types_aiofiles-23.2.0.20240311-py3-none-any.whl", hash = "sha256:ed10a8002d88c94220597b77304cf1a1d8cf489c7143fc3ffa2c96488b20fec7"}, + {file = "types-aiofiles-24.1.0.20240626.tar.gz", hash = "sha256:48604663e24bc2d5038eac05ccc33e75799b0779e93e13d6a8f711ddc306ac08"}, + {file = "types_aiofiles-24.1.0.20240626-py3-none-any.whl", hash = "sha256:7939eca4a8b4f9c6491b6e8ef160caee9a21d32e18534a57d5ed90aee47c66b4"}, ] [[package]] @@ -2836,13 +2836,13 @@ files = [ [[package]] name = "types-setuptools" -version = "70.0.0.20240524" +version = "75.5.0.20241122" description = "Typing stubs for setuptools" optional = true python-versions = ">=3.8" files = [ - {file = "types-setuptools-70.0.0.20240524.tar.gz", hash = "sha256:e31fee7b9d15ef53980526579ac6089b3ae51a005a281acf97178e90ac71aff6"}, - {file = "types_setuptools-70.0.0.20240524-py3-none-any.whl", hash = "sha256:8f5379b9948682d72a9ab531fbe52932e84c4f38deda570255f9bae3edd766bc"}, + {file = "types_setuptools-75.5.0.20241122-py3-none-any.whl", hash = "sha256:d69c445f7bdd5e49d1b2441aadcee1388febcc9ad9d9d5fd33648b555e0b1c31"}, + {file = "types_setuptools-75.5.0.20241122.tar.gz", hash = "sha256:196aaf1811cbc1c77ac1d4c4879d5308b6fdf426e56b73baadbca2a1827dadef"}, ] [[package]] @@ -3228,4 +3228,4 @@ upnp = ["miniupnpc"] [metadata] lock-version = "2.0" python-versions = ">=3.9, <3.13" -content-hash = "7d32f4c04dde72a555390500b070f6c482453ef71e187b6103bfe182bf95e728" +content-hash = "7a9d62216865e7f468b021d9d4e93e79c97e2361591a081cf6ee6b2a1bafdc81" diff --git a/pyproject.toml b/pyproject.toml index b913135249f8..a41c0fce19f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,10 +90,10 @@ pytest-cov = { version = "5.0.0", optional = true } pytest-mock = { version = "3.14.0", optional = true } pytest-monitor = { version = "1.6.6", platform = "linux", optional = true } pytest-xdist = { version = "3.6.1", optional = true } -types-aiofiles = { version = "23.2.0.20240311", optional = true } +types-aiofiles = { version = "24.1.0.20240626", optional = true } types-cryptography = { version = "3.3.23.2", optional = true } types-pyyaml = { version = "6.0.12.20240311", optional = true } -types-setuptools = { version = "70.0.0.20240524", optional = true } +types-setuptools = { version = "75.5.0.20241122", optional = true } lxml = { version = "5.2.2", optional = true } miniupnpc = { version = "2.2.2", source = "chia", optional = true } # big-o = {version = "0.11.0", optional = true}