Fix DID unnecessary wallet deletion (#13925)

chia/rpc/wallet_rpc_api.py

@@ -41,8 +41,15 @@ from chia.wallet.derive_keys import (
     master_sk_to_singleton_owner_sk,
     match_address_to_sk,
 )
+from chia.wallet.did_wallet import did_wallet_puzzles
+from chia.wallet.did_wallet.did_info import DIDInfo
 from chia.wallet.did_wallet.did_wallet import DIDWallet
-from chia.wallet.did_wallet.did_wallet_puzzles import match_did_puzzle, program_to_metadata
+from chia.wallet.did_wallet.did_wallet_puzzles import (
+    match_did_puzzle,
+    program_to_metadata,
+    DID_INNERPUZ_MOD,
+    create_fullpuz,
+)
 from chia.wallet.nft_wallet import nft_puzzles
 from chia.wallet.nft_wallet.nft_info import NFTInfo, NFTCoinInfo
 from chia.wallet.nft_wallet.nft_puzzles import get_metadata_and_phs
@@ -57,6 +64,7 @@ from chia.wallet.trading.offer import Offer
 from chia.wallet.transaction_record import TransactionRecord
 from chia.wallet.uncurried_puzzle import uncurry_puzzle
 from chia.wallet.util.address_type import AddressType, is_valid_address
+from chia.wallet.util.compute_hints import compute_coin_hints
 from chia.wallet.util.compute_memos import compute_memos
 from chia.wallet.util.transaction_type import TransactionType
 from chia.wallet.util.wallet_types import AmountWithPuzzlehash, WalletType
@@ -161,6 +169,7 @@ class WalletRpcApi:
             "/did_transfer_did": self.did_transfer_did,
             "/did_message_spend": self.did_message_spend,
             "/did_get_info": self.did_get_info,
+            "/did_find_lost_did": self.did_find_lost_did,
             # NFT Wallet
             "/nft_mint_nft": self.nft_mint_nft,
             "/nft_get_nfts": self.nft_get_nfts,
@@ -1761,6 +1770,153 @@ class WalletRpcApi:
             "hints": hints,
         }
 
+    async def did_find_lost_did(self, request) -> EndpointResult:
+        """
+        Recover a missing or unspendable DID wallet by a coin id of the DID
+        :param coin_id: It can be DID ID, launcher coin ID or any coin ID of the DID you want to find.
+        The latest coin ID will take less time.
+        :return:
+        """
+        if "coin_id" not in request:
+            return {"success": False, "error": "DID coin ID is required."}
+        coin_id = request["coin_id"]
+        # Check if we have a DID wallet for this
+        if coin_id.startswith(AddressType.DID.hrp(self.service.config)):
+            coin_id = decode_puzzle_hash(coin_id)
+        else:
+            coin_id = bytes32.from_hexstr(coin_id)
+        # Get coin state
+        peer: Optional[WSChiaConnection] = self.service.get_full_node_peer()
+        assert peer is not None
+        coin_spend, coin_state = await self.get_latest_coin_spend(peer, coin_id)
+        full_puzzle: Program = Program.from_bytes(bytes(coin_spend.puzzle_reveal))
+        uncurried = uncurry_puzzle(full_puzzle)
+        curried_args = match_did_puzzle(uncurried.mod, uncurried.args)
+        if curried_args is None:
+            return {"success": False, "error": "The coin is not a DID."}
+        p2_puzzle, recovery_list_hash, num_verification, singleton_struct, metadata = curried_args
+
+        hint_list = compute_coin_hints(coin_spend)
+        old_inner_puzhash = DID_INNERPUZ_MOD.curry(
+            p2_puzzle, recovery_list_hash, num_verification, singleton_struct, metadata
+        ).get_tree_hash()
+        derivation_record = None
+        # Hint is required, if it doesn't have any hint then it should be invalid
+        is_invalid = len(hint_list) == 0
+        for hint in hint_list:
+            derivation_record = (
+                await self.service.wallet_state_manager.puzzle_store.get_derivation_record_for_puzzle_hash(
+                    bytes32(hint)
+                )
+            )
+            if derivation_record is not None:
+                break
+            # Check if the mismatch is because of the memo bug
+            if hint == old_inner_puzhash:
+                is_invalid = True
+                break
+        if is_invalid:
+            # This is an invalid DID, check if we are owner
+            derivation_record = (
+                await self.service.wallet_state_manager.puzzle_store.get_derivation_record_for_puzzle_hash(
+                    p2_puzzle.get_tree_hash()
+                )
+            )
+        launcher_id = singleton_struct.rest().first().as_python()
+        if derivation_record is None:
+            return {"success": False, "error": f"This DID {launcher_id.hex()} is not belong to the connected wallet"}
+        else:
+            our_inner_puzzle: Program = self.service.wallet_state_manager.main_wallet.puzzle_for_pk(
+                derivation_record.pubkey
+            )
+            did_puzzle = DID_INNERPUZ_MOD.curry(
+                our_inner_puzzle, recovery_list_hash, num_verification, singleton_struct, metadata
+            )
+            full_puzzle = create_fullpuz(did_puzzle, launcher_id)
+            did_puzzle_empty_recovery = DID_INNERPUZ_MOD.curry(
+                our_inner_puzzle, Program.to([]).get_tree_hash(), uint64(0), singleton_struct, metadata
+            )
+            full_puzzle_empty_recovery = create_fullpuz(did_puzzle_empty_recovery, launcher_id)
+            if full_puzzle.get_tree_hash() != coin_state.coin.puzzle_hash:
+                if full_puzzle_empty_recovery.get_tree_hash() == coin_state.coin.puzzle_hash:
+                    did_puzzle = did_puzzle_empty_recovery
+                else:
+                    return {
+                        "success": False,
+                        "error": f"Cannot recover DID {launcher_id.hex()} because the last spend is metadata update.",
+                    }
+            # Check if we have the DID wallet
+            did_wallet: Optional[DIDWallet] = None
+            for wallet in self.service.wallet_state_manager.wallets.values():
+                if isinstance(wallet, DIDWallet):
+                    assert wallet.did_info.origin_coin is not None
+                    if wallet.did_info.origin_coin.name() == launcher_id:
+                        did_wallet = wallet
+                        break
+
+            if did_wallet is None:
+                # Create DID wallet
+                response: List[CoinState] = await self.service.get_coin_state([launcher_id], peer=peer)
+                if len(response) == 0:
+                    return {"success": False, "error": f"Could not find the launch coin with ID: {launcher_id.hex()}"}
+                launcher_coin: CoinState = response[0]
+                did_wallet = await DIDWallet.create_new_did_wallet_from_coin_spend(
+                    self.service.wallet_state_manager,
+                    self.service.wallet_state_manager.main_wallet,
+                    launcher_coin.coin,
+                    did_puzzle,
+                    coin_spend,
+                    f"DID {encode_puzzle_hash(launcher_id, AddressType.DID.hrp(self.service.config))}",
+                )
+            else:
+                assert did_wallet.did_info.current_inner is not None
+                if did_wallet.did_info.current_inner.get_tree_hash() != did_puzzle.get_tree_hash():
+                    # Inner DID puzzle doesn't match, we need to update the DID info
+                    full_solution: Program = Program.from_bytes(bytes(coin_spend.solution))
+                    inner_solution: Program = full_solution.rest().rest().first()
+                    recovery_list: List[bytes32] = []
+                    backup_required: int = num_verification.as_int()
+                    if recovery_list_hash != Program.to([]).get_tree_hash():
+                        try:
+                            for did in inner_solution.rest().rest().rest().rest().rest().as_python():
+                                recovery_list.append(did[0])
+                        except Exception:
+                            # We cannot recover the recovery list, but it's okay to leave it blank
+                            pass
+                    did_info: DIDInfo = DIDInfo(
+                        did_wallet.did_info.origin_coin,
+                        recovery_list,
+                        uint64(backup_required),
+                        [],
+                        did_puzzle,
+                        None,
+                        None,
+                        None,
+                        False,
+                        json.dumps(did_wallet_puzzles.program_to_metadata(metadata)),
+                    )
+                    await did_wallet.save_info(did_info)
+                    await self.service.wallet_state_manager.update_wallet_puzzle_hashes(did_wallet.wallet_info.id)
+
+            try:
+                coins = await did_wallet.select_coins(uint64(1))
+                coin = coins.pop()
+                if coin.name() == coin_state.coin.name():
+                    return {"success": True, "latest_coin_id": coin.name().hex()}
+            except ValueError:
+                # We don't have any coin for this wallet, add the coin
+                pass
+
+            wallet_id = did_wallet.id()
+            wallet_type = WalletType(did_wallet.type())
+            assert coin_state.created_height is not None
+            coin_record: WalletCoinRecord = WalletCoinRecord(
+                coin_state.coin, uint32(coin_state.created_height), uint32(0), False, False, wallet_type, wallet_id
+            )
+            await self.service.wallet_state_manager.coin_store.add_coin_record(coin_record, coin_state.coin.name())
+            await did_wallet.coin_added(coin_state.coin, uint32(coin_state.created_height), peer)
+            return {"success": True, "latest_coin_id": coin_state.coin.name().hex()}
+
     async def did_update_metadata(self, request) -> EndpointResult:
         wallet_id = uint32(request["wallet_id"])
         wallet = self.service.wallet_state_manager.wallets[wallet_id]

chia/wallet/did_wallet/did_wallet.py

@@ -26,7 +26,7 @@ from chia.wallet.derivation_record import DerivationRecord
 from chia.wallet.derive_keys import master_sk_to_wallet_sk_unhardened
 from chia.wallet.did_wallet import did_wallet_puzzles
 from chia.wallet.did_wallet.did_info import DIDInfo
-from chia.wallet.did_wallet.did_wallet_puzzles import create_fullpuz
+from chia.wallet.did_wallet.did_wallet_puzzles import create_fullpuz, uncurry_innerpuz
 from chia.wallet.lineage_proof import LineageProof
 from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import (
     DEFAULT_HIDDEN_PUZZLE_HASH,
@@ -203,8 +203,14 @@ class DIDWallet:
         recovery_list: List[bytes32] = []
         backup_required: int = num_verification.as_int()
         if recovery_list_hash != Program.to([]).get_tree_hash():
+            try:
                 for did in inner_solution.rest().rest().rest().rest().rest().as_python():
                     recovery_list.append(did[0])
+            except Exception:
+                self.log.warning(
+                    f"DID {launch_coin.name().hex()} has a recovery list hash but missing a reveal,"
+                    " you may need to reset the recovery info."
+                )
         self.did_info = DIDInfo(
             launch_coin,
             recovery_list,
@@ -383,6 +389,7 @@ class DIDWallet:
         assert full_puzzle.get_tree_hash() == coin.puzzle_hash
         if self.did_info.temp_coin is not None:
             self.wallet_state_manager.state_changed("did_coin_added", self.wallet_info.id)
+
         new_info = DIDInfo(
             self.did_info.origin_coin,
             self.did_info.backup_ids,
@@ -1132,12 +1139,26 @@ class DIDWallet:
             did_hash
         )
         assert self.did_info.origin_coin is not None
+        assert self.did_info.current_inner is not None
+        uncurried_args = uncurry_innerpuz(self.did_info.current_inner)
+        assert uncurried_args is not None
+        old_recovery_list_hash: Optional[Program] = None
+        p2_puzzle, old_recovery_list_hash, _, _, _ = uncurried_args
+        if record is None:
+            record = await self.wallet_state_manager.puzzle_store.get_derivation_record_for_puzzle_hash(
+                p2_puzzle.get_tree_hash()
+            )
+        if not (self.did_info.num_of_backup_ids_needed > 0 and len(self.did_info.backup_ids) == 0):
+            # We have the recovery list, don't reset it
+            old_recovery_list_hash = None
+
         inner_puzzle: Program = did_wallet_puzzles.create_innerpuz(
             puzzle_for_pk(record.pubkey),
             self.did_info.backup_ids,
             self.did_info.num_of_backup_ids_needed,
             self.did_info.origin_coin.name(),
             did_wallet_puzzles.metadata_to_program(json.loads(self.did_info.metadata)),
+            old_recovery_list_hash,
         )
         return inner_puzzle
 

chia/wallet/did_wallet/did_wallet_puzzles.py

@@ -29,6 +29,7 @@ def create_innerpuz(
     num_of_backup_ids_needed: uint64,
     launcher_id: bytes32,
     metadata: Program = Program.to([]),
+    recovery_list_hash: bytes32 = None,
 ) -> Program:
     """
     Create DID inner puzzle
@@ -37,9 +38,12 @@ def create_innerpuz(
     :param num_of_backup_ids_needed: Need how many DIDs for the recovery
     :param launcher_id: ID of the launch coin
     :param metadata: DID customized metadata
+    :param recovery_list_hash: Recovery list hash
     :return: DID inner puzzle
     """
     backup_ids_hash = Program(Program.to(recovery_list)).get_tree_hash()
+    if recovery_list_hash is not None:
+        backup_ids_hash = recovery_list_hash
     singleton_struct = Program.to((SINGLETON_TOP_LAYER_MOD_HASH, (launcher_id, LAUNCHER_PUZZLE_HASH)))
     return DID_INNERPUZ_MOD.curry(p2_puzzle, backup_ids_hash, num_of_backup_ids_needed, singleton_struct, metadata)
 

chia/wallet/wallet_state_manager.py

@@ -766,33 +766,15 @@ class WalletStateManager:
         self.log.info(f"parent: {parent_coin_state.coin.name()} inner_puzzle_hash for parent is {inner_puzzle_hash}")
 
         hint_list = compute_coin_hints(coin_spend)
-        old_inner_puzhash = DID_INNERPUZ_MOD.curry(
-            p2_puzzle, recovery_list_hash, num_verification, singleton_struct, metadata
-        ).get_tree_hash()
+
         derivation_record = None
-        # Hint is required, if it doesn't have any hint then it should a bugged DID
-        is_bugged = len(hint_list) == 0
-        new_p2_puzhash: Optional[bytes32] = None
         for hint in hint_list:
             derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(bytes32(hint))
             if derivation_record is not None:
-                new_p2_puzhash = hint
                 break
-            # Check if the mismatch is because of the memo bug
-            if hint == old_inner_puzhash:
-                new_p2_puzhash = hint
-                is_bugged = True
-                break
-        if is_bugged:
-            # This is a bugged DID, check if we are owner
-            derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(bytes32(p2_puzzle))
 
         launch_id: bytes32 = bytes32(bytes(singleton_struct.rest().first())[1:])
         if derivation_record is None:
-            if new_p2_puzhash != old_inner_puzhash and p2_puzzle != new_p2_puzhash:
-                # We only delete DID when the p2 puzzle doesn't change
-                self.log.info(f"Not sure if the DID belong to us, {coin_state}. Waiting for next spend ...")
-                return wallet_id, wallet_type
             self.log.info(f"Received state for the coin that doesn't belong to us {coin_state}")
             # Check if it was owned by us
             removed_wallet_ids = []

tests/wallet/did_wallet/test_did.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 import json
 from typing import Optional
 
@@ -11,6 +12,7 @@ from chia.rpc.wallet_rpc_api import WalletRpcApi
 from chia.simulator.simulator_protocol import FarmNewBlockProtocol
 from chia.simulator.time_out_assert import time_out_assert, time_out_assert_not_none
 from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.condition_opcodes import ConditionOpcode
 from chia.types.peer_info import PeerInfo
 from chia.types.spend_bundle import SpendBundle
@@ -429,6 +431,106 @@ class TestDIDWallet:
         else:
             assert False
 
+    @pytest.mark.parametrize(
+        "trusted",
+        [True, False],
+    )
+    @pytest.mark.asyncio
+    async def test_did_find_lost_did(self, self_hostname, two_wallet_nodes, trusted):
+        num_blocks = 5
+        full_nodes, wallets, _ = two_wallet_nodes
+        full_node_api = full_nodes[0]
+        server_1 = full_node_api.server
+        wallet_node, server_2 = wallets[0]
+        wallet_node_2, server_3 = wallets[1]
+        wallet = wallet_node.wallet_state_manager.main_wallet
+        wallet2 = wallet_node_2.wallet_state_manager.main_wallet
+        api_0 = WalletRpcApi(wallet_node)
+        ph = await wallet.get_new_puzzlehash()
+        if trusted:
+            wallet_node.config["trusted_peers"] = {
+                full_node_api.full_node.server.node_id.hex(): full_node_api.full_node.server.node_id.hex()
+            }
+            wallet_node_2.config["trusted_peers"] = {
+                full_node_api.full_node.server.node_id.hex(): full_node_api.full_node.server.node_id.hex()
+            }
+        else:
+            wallet_node.config["trusted_peers"] = {}
+            wallet_node_2.config["trusted_peers"] = {}
+        await server_2.start_client(PeerInfo(self_hostname, uint16(server_1._port)), None)
+        await server_3.start_client(PeerInfo(self_hostname, uint16(server_1._port)), None)
+        for i in range(1, num_blocks):
+            await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
+
+        funds = sum(
+            [
+                calculate_pool_reward(uint32(i)) + calculate_base_farmer_reward(uint32(i))
+                for i in range(1, num_blocks - 1)
+            ]
+        )
+
+        await time_out_assert(15, wallet.get_confirmed_balance, funds)
+
+        async with wallet_node.wallet_state_manager.lock:
+            did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet(
+                wallet_node.wallet_state_manager, wallet, uint64(101)
+            )
+        spend_bundle_list = await wallet_node.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(did_wallet.id())
+
+        spend_bundle = spend_bundle_list[0].spend_bundle
+        await time_out_assert_not_none(15, full_node_api.full_node.mempool_manager.get_spendbundle, spend_bundle.name())
+        ph2 = await wallet2.get_new_puzzlehash()
+        for i in range(1, num_blocks):
+            await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph2))
+
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
+        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        # Delete the coin and wallet
+        coins = await did_wallet.select_coins(uint64(1))
+        coin = coins.pop()
+        await wallet_node.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 0)
+        await wallet_node.wallet_state_manager.user_store.delete_wallet(did_wallet.wallet_info.id)
+        wallet_node.wallet_state_manager.wallets.pop(did_wallet.wallet_info.id)
+        assert len(wallet_node.wallet_state_manager.wallets) == 1
+        # Find lost DID
+        resp = await api_0.did_find_lost_did({"coin_id": did_wallet.did_info.origin_coin.name().hex()})
+        assert resp["success"]
+        did_wallets = list(
+            filter(
+                lambda w: (w.type == WalletType.DECENTRALIZED_ID),
+                await wallet_node.wallet_state_manager.get_all_wallet_info_entries(),
+            )
+        )
+        did_wallet: Optional[DIDWallet] = wallet_node.wallet_state_manager.wallets[did_wallets[0].id]
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
+        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        # Spend DID
+        recovery_list = [bytes32.fromhex(did_wallet.get_my_DID())]
+        await did_wallet.update_recovery_list(recovery_list, uint64(1))
+        assert did_wallet.did_info.backup_ids == recovery_list
+        await did_wallet.create_update_spend()
+        spend_bundle_list = await wallet_node.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(did_wallet.id())
+        spend_bundle = spend_bundle_list[0].spend_bundle
+        await time_out_assert_not_none(5, full_node_api.full_node.mempool_manager.get_spendbundle, spend_bundle.name())
+        for i in range(1, num_blocks):
+            await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph2))
+
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
+        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        # Delete the coin and change inner puzzle
+        coins = await did_wallet.select_coins(uint64(1))
+        coin = coins.pop()
+        await wallet_node.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 0)
+        new_inner_puzzle = await did_wallet.get_new_did_innerpuz()
+        did_wallet.did_info = dataclasses.replace(did_wallet.did_info, current_inner=new_inner_puzzle)
+        # Recovery the coin
+        resp = await api_0.did_find_lost_did({"coin_id": did_wallet.did_info.origin_coin.name().hex()})
+        assert resp["success"]
+        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
+        assert did_wallet.did_info.current_inner != new_inner_puzzle
+
     @pytest.mark.parametrize(
         "trusted",
         [True, False],
@@ -690,10 +792,8 @@ class TestDIDWallet:
         await time_out_assert(15, wallet.get_unconfirmed_balance, 7999999997899)
         # Check if the DID wallet is created in the wallet2
 
-        async def num_wallets() -> int:
-            return len(await wallet_node_2.wallet_state_manager.get_all_wallet_info_entries())
-
-        await time_out_assert(30, num_wallets, 2)
+        await time_out_assert(30, get_wallet_num, 2, wallet_node_2.wallet_state_manager)
+        await time_out_assert(30, get_wallet_num, 1, wallet_node.wallet_state_manager)
         # Get the new DID wallet
         did_wallets = list(
             filter(
@@ -702,7 +802,7 @@ class TestDIDWallet:
             )
         )
         did_wallet_2: Optional[DIDWallet] = wallet_node_2.wallet_state_manager.wallets[did_wallets[0].id]
-        assert len(wallet_node.wallet_state_manager.wallets) == 2
+        assert len(wallet_node.wallet_state_manager.wallets) == 1
         assert did_wallet_1.did_info.origin_coin == did_wallet_2.did_info.origin_coin
         if with_recovery:
             assert did_wallet_1.did_info.backup_ids[0] == did_wallet_2.did_info.backup_ids[0]