[CHIA-307] (long-lived) Vault

chia/_tests/clvm/test_p2_singletons.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import pytest
+from chia_rs import G2Element
+
+from chia.clvm.spend_sim import CostLogger, sim_and_client
+from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.coin_spend import make_spend
+from chia.types.mempool_inclusion_status import MempoolInclusionStatus
+from chia.types.spend_bundle import SpendBundle
+from chia.util.errors import Err
+from chia.util.ints import uint64
+from chia.wallet.conditions import CreateCoin
+from chia.wallet.puzzles import p2_singleton_via_delegated_puzzle_safe as dp_safe
+from chia.wallet.util.curry_and_treehash import shatree_atom
+
+ACS = Program.to(1)
+ACS_HASH = ACS.get_tree_hash()
+
+MOCK_SINGLETON_MOD = Program.to([2, 5, 7])  # (a 5 11) - (mod (_ PUZZLE . solution) (a PUZZLE solution))
+MOCK_SINGLETON_MOD_HASH = MOCK_SINGLETON_MOD.get_tree_hash()
+MOCK_SINGLETON_LAUNCHER_ID = bytes32([0] * 32)
+MOCK_SINGLETON_LAUNCHER_HASH = bytes32([1] * 32)
+MOCK_SINGLETON = MOCK_SINGLETON_MOD.curry(
+    Program.to((MOCK_SINGLETON_MOD_HASH, (MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_LAUNCHER_HASH))),
+    ACS,
+)
+MOCK_SINGLETON_HASH = MOCK_SINGLETON.get_tree_hash()
+dp_safe.PRE_HASHED_HASHES[MOCK_SINGLETON_MOD_HASH] = shatree_atom(MOCK_SINGLETON_MOD_HASH)
+dp_safe.PRE_HASHED_HASHES[MOCK_SINGLETON_LAUNCHER_HASH] = shatree_atom(MOCK_SINGLETON_LAUNCHER_HASH)
+
+
+@pytest.mark.anyio
+async def test_dp_safe_lifecycle(cost_logger: CostLogger) -> None:
+    P2_SINGLETON = dp_safe.construct(MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_MOD_HASH, MOCK_SINGLETON_LAUNCHER_HASH)
+    P2_SINGLETON_HASH = dp_safe.construct_hash(
+        MOCK_SINGLETON_LAUNCHER_ID, MOCK_SINGLETON_MOD_HASH, MOCK_SINGLETON_LAUNCHER_HASH
+    )
+    assert dp_safe.match(P2_SINGLETON) is not None
+    assert dp_safe.match(ACS) is None
+    assert dp_safe.match(MOCK_SINGLETON) is None
+
+    async with sim_and_client() as (sim, sim_client):
+        await sim.farm_block(P2_SINGLETON_HASH)
+        await sim.farm_block(MOCK_SINGLETON_HASH)
+        p2_singleton = (await sim_client.get_coin_records_by_puzzle_hash(P2_SINGLETON_HASH, include_spent_coins=False))[
+            0
+        ].coin
+        singleton = (await sim_client.get_coin_records_by_puzzle_hash(MOCK_SINGLETON_HASH, include_spent_coins=False))[
+            0
+        ].coin
+
+        dp = Program.to(1)
+        dp_hash = dp.get_tree_hash()
+        bundle = cost_logger.add_cost(
+            "p2_singleton_w_mock_singleton",
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, p2_singleton.name()).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            ),
+        )
+        result = await sim_client.push_tx(bundle)
+        assert result == (MempoolInclusionStatus.SUCCESS, None)
+        checkpoint = sim.block_height
+        await sim.farm_block()
+
+        assert len(await sim_client.get_coin_records_by_puzzle_hash(bytes32([0] * 32), include_spent_coins=False)) == 1
+
+        await sim.rewind(checkpoint)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            bytes32([0] * 32),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, p2_singleton.name()).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_MY_COIN_ID_FAILED)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_ANNOUNCE_CONSUMED_FAILED)
+
+        result = await sim_client.push_tx(
+            SpendBundle(
+                [
+                    make_spend(
+                        p2_singleton,
+                        P2_SINGLETON,
+                        dp_safe.solve(
+                            ACS_HASH,
+                            dp,
+                            Program.to([CreateCoin(bytes32([0] * 32), uint64(0)).to_program()]),
+                            p2_singleton.name(),
+                        ),
+                    ),
+                    make_spend(
+                        singleton,
+                        MOCK_SINGLETON,
+                        Program.to([dp_safe.required_announcement(dp_hash, bytes32([0] * 32)).to_program()]),
+                    ),
+                ],
+                G2Element(),
+            )
+        )
+        assert result == (MempoolInclusionStatus.FAILED, Err.ASSERT_ANNOUNCE_CONSUMED_FAILED)

chia/_tests/core/daemon/test_daemon.py

@@ -32,7 +32,7 @@
 from chia.simulator.setup_services import setup_full_node
 from chia.util.config import load_config
 from chia.util.json_util import dict_to_json_str
-from chia.util.keychain import Keychain, KeyData, supports_os_passphrase_storage
+from chia.util.keychain import Keychain, KeyData, KeyTypes, supports_os_passphrase_storage
 from chia.util.keyring_wrapper import DEFAULT_PASSPHRASE_IF_NO_MASTER_PASSPHRASE, KeyringWrapper
 from chia.util.ws_message import create_payload, create_payload_dict
 from chia.wallet.derive_keys import master_sk_to_farmer_sk, master_sk_to_pool_sk
@@ -237,6 +237,7 @@ async def get_keys_for_plotting(self, request: Dict[str, Any]) -> Dict[str, Any]
 def add_private_key_response_data(fingerprint: int) -> Dict[str, object]:
     return {
         "success": True,
+        "key_type": KeyTypes.G1_ELEMENT.value,
         "fingerprint": fingerprint,
     }
 

chia/_tests/core/daemon/test_keychain_proxy.py

@@ -5,6 +5,7 @@
 from typing import Any, AsyncGenerator
 
 import pytest
+from chia_rs import G1Element
 
 from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate
 from chia.simulator.block_tools import BlockTools
@@ -47,21 +48,22 @@ async def test_add_private_key(keychain_proxy: KeychainProxy) -> None:
 @pytest.mark.anyio
 async def test_add_public_key(keychain_proxy: KeychainProxy) -> None:
     keychain = keychain_proxy
-    await keychain.add_key(bytes(TEST_KEY_3.public_key).hex(), TEST_KEY_3.label, private=False)
+    assert isinstance(TEST_KEY_3.observation_root, G1Element)
+    await keychain.add_key(TEST_KEY_3.public_key.hex(), TEST_KEY_3.label, private=False)
     with pytest.raises(Exception, match="already exists"):
-        await keychain.add_key(bytes(TEST_KEY_3.public_key).hex(), "", private=False)
+        await keychain.add_key(TEST_KEY_3.public_key.hex(), "", private=False)
     key = await keychain.get_key(TEST_KEY_3.fingerprint, include_secrets=False)
     assert key is not None
-    assert key.public_key == TEST_KEY_3.public_key
+    assert key.observation_root == TEST_KEY_3.observation_root
     assert key.secrets is None
 
     pk = await keychain.get_key_for_fingerprint(TEST_KEY_3.fingerprint, private=False)
     assert pk is not None
-    assert pk == TEST_KEY_3.public_key
+    assert pk == TEST_KEY_3.observation_root
 
     pk = await keychain.get_key_for_fingerprint(None, private=False)
     assert pk is not None
-    assert pk == TEST_KEY_3.public_key
+    assert pk == TEST_KEY_3.observation_root
 
     with pytest.raises(KeychainKeyNotFound):
         pk = await keychain.get_key_for_fingerprint(1234567890, private=False)
@@ -82,8 +84,8 @@ async def test_get_key_for_fingerprint(keychain_proxy: KeychainProxy) -> None:
     with pytest.raises(KeychainIsEmpty):
         await keychain.get_key_for_fingerprint(None, private=False)
     await keychain_proxy.add_key(TEST_KEY_1.mnemonic_str(), TEST_KEY_1.label)
-    assert await keychain.get_key_for_fingerprint(TEST_KEY_1.fingerprint, private=False) == TEST_KEY_1.public_key
-    assert await keychain.get_key_for_fingerprint(None, private=False) == TEST_KEY_1.public_key
+    assert await keychain.get_key_for_fingerprint(TEST_KEY_1.fingerprint, private=False) == TEST_KEY_1.observation_root
+    assert await keychain.get_key_for_fingerprint(None, private=False) == TEST_KEY_1.observation_root
     with pytest.raises(KeychainKeyNotFound):
         await keychain.get_key_for_fingerprint(1234567890, private=False)
 

chia/_tests/core/data_layer/test_data_rpc.py

@@ -68,8 +68,8 @@
 from chia.util.timing import adjusted_timeout, backoff_times
 from chia.wallet.trading.offer import Offer as TradingOffer
 from chia.wallet.transaction_record import TransactionRecord
-from chia.wallet.wallet import Wallet
 from chia.wallet.wallet_node import WalletNode
+from chia.wallet.wallet_protocol import MainWalletProtocol
 
 pytestmark = pytest.mark.data_layer
 nodes = Tuple[WalletNode, FullNodeSimulator]
@@ -813,7 +813,7 @@ async def offer_setup_fixture(
     [full_node_service], wallet_services, bt = two_wallet_nodes_services
     enable_batch_autoinsertion_settings = getattr(request, "param", (True, True))
     full_node_api = full_node_service._api
-    wallets: List[Wallet] = []
+    wallets: List[MainWalletProtocol] = []
     for wallet_service in wallet_services:
         wallet_node = wallet_service._node
         assert wallet_node.server is not None
@@ -2308,7 +2308,7 @@ async def test_wallet_log_in_changes_active_fingerprint(
 
     mnemonic = create_mnemonic()
     assert wallet_rpc_api.service.local_keychain is not None
-    private_key = wallet_rpc_api.service.local_keychain.add_key(mnemonic_or_pk=mnemonic)
+    private_key, _ = wallet_rpc_api.service.local_keychain.add_key(mnemonic_or_pk=mnemonic)
     secondary_fingerprint: int = private_key.get_g1().get_fingerprint()
 
     await wallet_rpc_api.log_in(request={"fingerprint": primary_fingerprint})

chia/_tests/core/mempool/test_mempool_manager.py

@@ -53,9 +53,9 @@
 from chia.wallet.conditions import AssertCoinAnnouncement
 from chia.wallet.payment import Payment
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG
-from chia.wallet.wallet import Wallet
 from chia.wallet.wallet_coin_record import WalletCoinRecord
 from chia.wallet.wallet_node import WalletNode
+from chia.wallet.wallet_protocol import MainWalletProtocol
 
 IDENTITY_PUZZLE = SerializedProgram.to(1)
 IDENTITY_PUZZLE_HASH = IDENTITY_PUZZLE.get_tree_hash()
@@ -1620,7 +1620,7 @@ async def farm_a_block(full_node_api: FullNodeSimulator, wallet_node: WalletNode
 
     async def make_setup_and_coins(
         full_node_api: FullNodeSimulator, wallet_node: WalletNode
-    ) -> Tuple[Wallet, list[WalletCoinRecord], bytes32]:
+    ) -> Tuple[MainWalletProtocol, list[WalletCoinRecord], bytes32]:
         wallet = wallet_node.wallet_state_manager.main_wallet
         ph = await wallet.get_new_puzzlehash()
         phs = [await wallet.get_new_puzzlehash() for _ in range(3)]

chia/_tests/core/util/test_keychain.py

@@ -3,7 +3,7 @@
 import json
 import random
 from dataclasses import replace
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 import importlib_resources
 import pytest
@@ -25,12 +25,14 @@
     Keychain,
     KeyData,
     KeyDataSecrets,
+    KeyTypes,
     bytes_from_mnemonic,
     bytes_to_mnemonic,
     generate_mnemonic,
     mnemonic_from_short_words,
     mnemonic_to_seed,
 )
+from chia.util.observation_root import ObservationRoot
 
 mnemonic = (
     "rapid this oven common drive ribbon bulb urban uncover napkin kitten usage enforce uncle unveil scene "
@@ -81,6 +83,9 @@ def test_basic_add_delete(self, empty_temp_file_keyring: TempKeyring, seeded_ran
         assert kc._get_free_private_key_index() == 2
         assert len(kc.get_all_private_keys()) == 2
         assert len(kc.get_all_public_keys()) == 2
+        all_pks: List[G1Element] = kc.get_all_public_keys_of_type(G1Element)
+        assert len(all_pks) == 2
+
         assert kc.get_all_private_keys()[0] == kc.get_first_private_key()
         assert kc.get_all_public_keys()[0] == kc.get_first_public_key()
 
@@ -97,7 +102,7 @@ def test_basic_add_delete(self, empty_temp_file_keyring: TempKeyring, seeded_ran
         assert len(kc.get_all_private_keys()) == 0
 
         kc.add_key(bech32_pubkey, label=None, private=False)
-        all_pks = kc.get_all_public_keys()
+        all_pks = kc.get_all_public_keys_of_type(G1Element)
         assert len(all_pks) == 1
         assert all_pks[0] == public_key
         kc.delete_all_keys()
@@ -155,7 +160,7 @@ def test_bip39_eip2333_test_vector(self, empty_temp_file_keyring: TempKeyring):
 
         mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
         print("entropy to seed:", mnemonic_to_seed(mnemonic).hex())
-        master_sk = kc.add_key(mnemonic)
+        master_sk, _ = kc.add_key(mnemonic)
         tv_master_int = 8075452428075949470768183878078858156044736575259233735633523546099624838313
         tv_child_int = 18507161868329770878190303689452715596635858303241878571348190917018711023613
         assert master_sk == PrivateKey.from_bytes(tv_master_int.to_bytes(32, "big"))
@@ -234,7 +239,7 @@ def test_key_data_generate(label: Optional[str]) -> None:
     key_data = KeyData.generate(label)
     assert key_data.private_key == AugSchemeMPL.key_gen(mnemonic_to_seed(key_data.mnemonic_str()))
     assert key_data.entropy == bytes_from_mnemonic(key_data.mnemonic_str())
-    assert key_data.public_key == key_data.private_key.get_g1()
+    assert key_data.observation_root == key_data.private_key.get_g1()
     assert key_data.fingerprint == key_data.private_key.get_g1().get_fingerprint()
     assert key_data.label == label
 
@@ -246,7 +251,7 @@ def test_key_data_generate(label: Optional[str]) -> None:
 def test_key_data_creation(input_data: object, from_method: Callable[..., KeyData], label: Optional[str]) -> None:
     key_data = from_method(input_data, label)
     assert key_data.fingerprint == fingerprint
-    assert key_data.public_key == public_key
+    assert key_data.observation_root == public_key
     assert key_data.mnemonic == mnemonic.split()
     assert key_data.mnemonic_str() == mnemonic
     assert key_data.entropy == entropy
@@ -255,7 +260,7 @@ def test_key_data_creation(input_data: object, from_method: Callable[..., KeyDat
 
 
 def test_key_data_without_secrets() -> None:
-    key_data = KeyData(fingerprint, public_key, None, None)
+    key_data = KeyData(fingerprint, bytes(public_key), None, None, KeyTypes.G1_ELEMENT.value)
     assert key_data.secrets is None
 
     with pytest.raises(KeychainSecretsMissing):
@@ -287,12 +292,21 @@ def test_key_data_secrets_post_init(input_data: Tuple[List[str], bytes, PrivateK
 @pytest.mark.parametrize(
     "input_data, data_type",
     [
-        ((fingerprint, G1Element(), None, KeyDataSecrets(mnemonic.split(), entropy, private_key)), "public_key"),
-        ((fingerprint, G1Element(), None, None), "fingerprint"),
+        (
+            (
+                fingerprint,
+                bytes(G1Element()),
+                None,
+                KeyDataSecrets(mnemonic.split(), entropy, private_key),
+                KeyTypes.G1_ELEMENT.value,
+            ),
+            "public_key",
+        ),
+        ((fingerprint, bytes(G1Element()), None, None, KeyTypes.G1_ELEMENT.value), "fingerprint"),
     ],
 )
 def test_key_data_post_init(
-    input_data: Tuple[uint32, G1Element, Optional[str], Optional[KeyDataSecrets]], data_type: str
+    input_data: Tuple[uint32, bytes, Optional[str], Optional[KeyDataSecrets], str], data_type: str
 ) -> None:
     with pytest.raises(KeychainKeyDataMismatch, match=data_type):
         KeyData(*input_data)
@@ -465,3 +479,17 @@ async def test_delete_drops_labels(get_temp_keyring: Keychain, delete_all: bool)
         for key_data in keys:
             keychain.delete_key_by_fingerprint(key_data.fingerprint)
             assert keychain.keyring_wrapper.get_label(key_data.fingerprint) is None
+
+
+@pytest.mark.parametrize("key_type", [e.value for e in KeyTypes])
+def test_key_type_support(key_type: str) -> None:
+    """
+    The purpose of this test is to make sure that whenever KeyTypes is updated, all relevant functionality is
+    also updated with it.
+    """
+    generate_test_key_for_key_type: Dict[str, Tuple[int, bytes, ObservationRoot]] = {
+        KeyTypes.G1_ELEMENT.value: (G1Element().get_fingerprint(), bytes(G1Element()), G1Element())
+    }
+    obr_fingerprint, obr_bytes, obr = generate_test_key_for_key_type[key_type]
+    assert KeyData(uint32(obr_fingerprint), obr_bytes, None, None, key_type).observation_root == obr
+    assert KeyTypes.parse_observation_root(obr_bytes, KeyTypes(key_type)) == obr