From 402236dfdb9b74f743663d4c5b2419cc005118cc Mon Sep 17 00:00:00 2001 From: Matt Hauff Date: Wed, 5 Aug 2026 06:07:48 +1200 Subject: [PATCH] [LABS-488] Extract signing functionality to its own module (#21159) * [LABS-480] Remove interested_ph_cache and initialize coin cache in init * [LABS-481] Remove unnecessary `server` field from WSM * [LABS-486] Remove useless `add_transaction` from WSM * [LABS-488] Extract signing functionality to its own module * Test coverage --- .../wallet/cat_wallet/test_cat_wallet.py | 4 +- chia/_tests/wallet/cat_wallet/test_trades.py | 26 +- .../_tests/wallet/db_wallet/test_dl_offers.py | 4 +- .../wallet/nft_wallet/test_nft_1_offers.py | 14 +- .../wallet/nft_wallet/test_nft_offers.py | 20 +- chia/_tests/wallet/test_sign_coin_spends.py | 52 +-- chia/_tests/wallet/test_signer_protocol.py | 34 +- chia/_tests/wallet/test_wallet.py | 2 +- chia/wallet/wallet.py | 202 +--------- chia/wallet/wallet_rpc_api.py | 12 +- chia/wallet/wallet_signer.py | 357 ++++++++++++++++++ chia/wallet/wallet_state_manager.py | 180 +-------- 12 files changed, 469 insertions(+), 438 deletions(-) create mode 100644 chia/wallet/wallet_signer.py diff --git a/chia/_tests/wallet/cat_wallet/test_cat_wallet.py b/chia/_tests/wallet/cat_wallet/test_cat_wallet.py index 1174b5928a..cd7715bf89 100644 --- a/chia/_tests/wallet/cat_wallet/test_cat_wallet.py +++ b/chia/_tests/wallet/cat_wallet/test_cat_wallet.py @@ -1454,7 +1454,7 @@ async def test_cat_change_detection(wallet_environments: WalletTestFramework, wa construct_cat_puzzle(CAT_MOD, Program.NIL.get_tree_hash(), our_puzzle).get_tree_hash(), cat_amount_0, ) - eve_spend, _ = await wsm.sign_bundle( + eve_spend, _ = await wsm.signer.sign_bundle( [ make_spend( cat_coin, @@ -1670,7 +1670,7 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non ) ], ) - signed_spend, _ = await env.wallet_state_manager.sign_bundle(new_spend.coin_spends) + signed_spend, _ = await env.wallet_state_manager.signer.sign_bundle(new_spend.coin_spends) await env.rpc_client.push_tx(PushTX(spend_bundle=signed_spend)) await time_out_assert(10, simulator.tx_id_in_mempool, True, signed_spend.name()) diff --git a/chia/_tests/wallet/cat_wallet/test_trades.py b/chia/_tests/wallet/cat_wallet/test_trades.py index d23aebb2ca..d690d33732 100644 --- a/chia/_tests/wallet/cat_wallet/test_trades.py +++ b/chia/_tests/wallet/cat_wallet/test_trades.py @@ -514,7 +514,7 @@ async def test_cat_trades( assert trade_make is not None peer = wallet_node_taker.get_full_node_peer() - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -731,7 +731,7 @@ async def test_cat_trades( assert success is True assert trade_make is not None - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -860,7 +860,7 @@ async def test_cat_trades( assert error is None assert success is True assert trade_make is not None - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1061,7 +1061,7 @@ async def test_cat_trades( assert success is True assert trade_make is not None - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1318,7 +1318,7 @@ async def test_cat_trades( assert error is None assert success is True assert trade_make is not None - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1451,7 +1451,7 @@ async def test_cat_trades( assert success is True assert trade_make is not None - [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1733,7 +1733,7 @@ async def test_trade_cancellation(wallet_environments: WalletTestFramework, wall # Due to current mempool rules, trying to force a take out of the mempool with a cancel will not work. # Uncomment this when/if it does - # [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.sign_offers( + # [maker_offer], signing_response = await wallet_node_maker.wallet_state_manager.signer.sign_offers( # [Offer.from_bytes(trade_make.offer)] # ) # trade_take = await trade_manager_taker.respond_to_offer( @@ -2037,7 +2037,7 @@ async def test_trade_conflict(wallet_environments: WalletTestFramework, wallet_t assert trade_make is not None peer = env_taker.node.get_full_node_peer() offer = Offer.from_bytes(trade_make.offer) - [offer], signing_response = await env_maker.wallet_state_manager.sign_offers([offer]) + [offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers([offer]) async with trade_manager_taker.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response ) as action_scope: @@ -2055,7 +2055,7 @@ async def test_trade_conflict(wallet_environments: WalletTestFramework, wallet_t await trade_manager_taker.respond_to_offer(offer, peer, action_scope, fee=fee) await time_out_assert(15, get_trade_and_status, TradeStatus.PENDING_CONFIRM, trade_manager_taker, tr1) # pushing into mempool while already in it should fail - [offer], signing_response = await env_maker.wallet_state_manager.sign_offers([offer]) + [offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers([offer]) async with trade_manager_trader.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response ) as action_scope: @@ -2352,7 +2352,9 @@ async def test_trade_high_fee(wallet_environments: WalletTestFramework, wallet_t assert success is True assert trade_make is not None peer = env_taker.node.get_full_node_peer() - [offer], signing_response = await env_maker.wallet_state_manager.sign_offers([Offer.from_bytes(trade_make.offer)]) + [offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( + [Offer.from_bytes(trade_make.offer)] + ) fee = uint64(1_000_000_000_000) async with trade_manager_taker.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response @@ -2515,10 +2517,10 @@ async def test_aggregated_trade_state(wallet_environments: WalletTestFramework, assert success is True assert trade_make_2 is not None - [offer_1], signing_response_1 = await env_maker.node.wallet_state_manager.sign_offers( + [offer_1], signing_response_1 = await env_maker.node.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make_1.offer)] ) - [offer_2], signing_response_2 = await env_maker.node.wallet_state_manager.sign_offers( + [offer_2], signing_response_2 = await env_maker.node.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make_2.offer)] ) agg_offer = Offer.aggregate([offer_1, offer_2]) diff --git a/chia/_tests/wallet/db_wallet/test_dl_offers.py b/chia/_tests/wallet/db_wallet/test_dl_offers.py index 5f1435b56c..95bd9b1fe4 100644 --- a/chia/_tests/wallet/db_wallet/test_dl_offers.py +++ b/chia/_tests/wallet/db_wallet/test_dl_offers.py @@ -195,7 +195,7 @@ async def test_dl_offers(wallet_environments: WalletTestFramework) -> None: ] ) - [_maker_offer], signing_response = await wsm_maker.sign_offers([Offer.from_bytes(offer_maker.offer)]) + [_maker_offer], signing_response = await wsm_maker.signer.sign_offers([Offer.from_bytes(offer_maker.offer)]) async with trade_manager_taker.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response ) as action_scope: @@ -750,7 +750,7 @@ async def test_multiple_dl_offers(wallet_environments: WalletTestFramework) -> N assert success is True assert offer_maker is not None - [_maker_offer], signing_response = await wsm_maker.sign_offers([Offer.from_bytes(offer_maker.offer)]) + [_maker_offer], signing_response = await wsm_maker.signer.sign_offers([Offer.from_bytes(offer_maker.offer)]) async with trade_manager_taker.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response ) as action_scope: diff --git a/chia/_tests/wallet/nft_wallet/test_nft_1_offers.py b/chia/_tests/wallet/nft_wallet/test_nft_1_offers.py index 2d4535a9e6..5e991c71ca 100644 --- a/chia/_tests/wallet/nft_wallet/test_nft_1_offers.py +++ b/chia/_tests/wallet/nft_wallet/test_nft_1_offers.py @@ -178,7 +178,7 @@ async def test_nft_offer_sell_nft(wallet_environments: WalletTestFramework, zero taker_fee = 1 peer = env_taker.node.get_full_node_peer() - [_maker_offer], signing_response = await env_maker.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -410,7 +410,7 @@ async def test_nft_offer_request_nft(wallet_environments: WalletTestFramework, z taker_fee = 1 peer = env_taker.node.get_full_node_peer() - [_maker_offer], signing_response = await env_maker.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -700,7 +700,7 @@ async def test_nft_offer_sell_did_to_did(wallet_environments: WalletTestFramewor taker_fee = 1 peer = env_taker.node.get_full_node_peer() - [_maker_offer], signing_response = await env_maker.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1006,7 +1006,7 @@ async def test_nft_offer_sell_nft_for_cat( taker_fee = 1 peer = env_taker.node.get_full_node_peer() - [_maker_offer], signing_response = await env_maker.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1362,7 +1362,7 @@ async def test_nft_offer_request_nft_for_cat( taker_fee = 1 peer = env_taker.node.get_full_node_peer() - [_maker_offer], signing_response = await env_maker.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_maker.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -2030,7 +2030,7 @@ async def test_complex_nft_offer( assert success assert trade_make is not None - [maker_offer], signing_response = await wsm_maker.sign_offers([Offer.from_bytes(trade_make.offer)]) + [maker_offer], signing_response = await wsm_maker.signer.sign_offers([Offer.from_bytes(trade_make.offer)]) if royalty_basis_pts_maker == 10000: with pytest.raises(ValueError): async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -2239,7 +2239,7 @@ async def test_complex_nft_offer( assert success assert trade_make is not None - [maker_offer], signing_response = await wsm_maker.sign_offers([Offer.from_bytes(trade_make.offer)]) + [maker_offer], signing_response = await wsm_maker.signer.sign_offers([Offer.from_bytes(trade_make.offer)]) async with trade_manager_taker.wallet_state_manager.new_action_scope( wallet_environments.tx_config, push=True, additional_signing_responses=signing_response ) as action_scope: diff --git a/chia/_tests/wallet/nft_wallet/test_nft_offers.py b/chia/_tests/wallet/nft_wallet/test_nft_offers.py index baa2798088..3257fad0b4 100644 --- a/chia/_tests/wallet/nft_wallet/test_nft_offers.py +++ b/chia/_tests/wallet/nft_wallet/test_nft_offers.py @@ -136,7 +136,7 @@ async def test_nft_offer_with_fee(wallet_environments: WalletTestFramework) -> N taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) peer = env_1.node.get_full_node_peer() @@ -235,7 +235,7 @@ async def test_nft_offer_with_fee(wallet_environments: WalletTestFramework) -> N taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -622,7 +622,7 @@ async def test_nft_offer_with_metadata_update(wallet_environments: WalletTestFra taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) peer = env_1.node.get_full_node_peer() @@ -850,7 +850,7 @@ async def test_nft_offer_nft_for_cat(wallet_environments: WalletTestFramework, w taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) peer = env_1.node.get_full_node_peer() @@ -980,7 +980,7 @@ async def test_nft_offer_nft_for_cat(wallet_environments: WalletTestFramework, w taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) async with trade_manager_taker.wallet_state_manager.new_action_scope( @@ -1259,7 +1259,7 @@ async def test_nft_offer_nft_for_nft(wallet_environments: WalletTestFramework) - taker_fee = uint64(1) - [_maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers( + [_maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( [Offer.from_bytes(trade_make.offer)] ) peer = env_1.node.get_full_node_peer() @@ -1490,7 +1490,9 @@ async def test_nft_offer_nft0_and_xch_for_cat( assert error is None assert trade_make is not None - [maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers([Offer.from_bytes(trade_make.offer)]) + [maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( + [Offer.from_bytes(trade_make.offer)] + ) taker_fee = uint64(1) @@ -1619,7 +1621,9 @@ async def test_nft_offer_nft0_and_xch_for_cat( assert error is None assert trade_make is not None - [maker_offer], signing_response = await env_0.wallet_state_manager.sign_offers([Offer.from_bytes(trade_make.offer)]) + [maker_offer], signing_response = await env_0.wallet_state_manager.signer.sign_offers( + [Offer.from_bytes(trade_make.offer)] + ) taker_fee = uint64(1) diff --git a/chia/_tests/wallet/test_sign_coin_spends.py b/chia/_tests/wallet/test_sign_coin_spends.py index 66ea39a2da..8a269917d9 100644 --- a/chia/_tests/wallet/test_sign_coin_spends.py +++ b/chia/_tests/wallet/test_sign_coin_spends.py @@ -22,10 +22,9 @@ from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import ( puzzle_hash_for_synthetic_public_key, ) from chia.wallet.util.wallet_types import WalletType -from chia.wallet.wallet import Wallet from chia.wallet.wallet_puzzle_store import WalletPuzzleStore -from chia.wallet.wallet_state_manager import WalletStateManager -from chia.wallet.wallet_user_store import WalletUserStore +from chia.wallet.wallet_signer import WalletSigner +from chia.wallet.wallet_spend_bundle import WalletSpendBundle top_sk: PrivateKey = PrivateKey.from_bytes(bytes([1] * 32)) sk1_h: PrivateKey = master_sk_to_wallet_sk(top_sk, uint32(1)) @@ -65,28 +64,39 @@ spend_u: CoinSpend = make_spend( ) +async def _noop_push(_bundle: WalletSpendBundle) -> None: ... + + @pytest.mark.anyio -async def test_wsm_sign_transaction() -> None: +async def test_wallet_signer() -> None: async with manage_connection("file:temp.db?mode=memory&cache=shared", uri=True, name="writer") as writer_conn: async with manage_connection("file:temp.db?mode=memory&cache=shared", uri=True, name="reader") as reader_conn: - wsm = WalletStateManager() db = DBWrapper2(writer_conn) await db.add_connection(reader_conn) - wsm.puzzle_store = await WalletPuzzleStore.create(db) - wsm.constants = DEFAULT_CONSTANTS - wsm.private_key = top_sk - wsm.root_pubkey = top_sk.get_g1() - wsm.user_store = await WalletUserStore.create(db) - wallet_info = await wsm.user_store.get_wallet_by_id(1) - assert wallet_info is not None - wsm.main_wallet = await Wallet.create(wsm, wallet_info) + with pytest.raises(ValueError, match=re.escape("root_private_key does not match self.root_pubkey")): + WalletSigner( + root_pubkey=G1Element(), + root_private_key=top_sk, + puzzle_store=await WalletPuzzleStore.create(db), + max_block_cost_clvm=DEFAULT_CONSTANTS.MAX_BLOCK_COST_CLVM, + agg_sig_me_additional_data=DEFAULT_CONSTANTS.AGG_SIG_ME_ADDITIONAL_DATA, + spend_bundle_push=_noop_push, + ) + signer = WalletSigner( + root_pubkey=top_sk.get_g1(), + root_private_key=top_sk, + puzzle_store=await WalletPuzzleStore.create(db), + max_block_cost_clvm=DEFAULT_CONSTANTS.MAX_BLOCK_COST_CLVM, + agg_sig_me_additional_data=DEFAULT_CONSTANTS.AGG_SIG_ME_ADDITIONAL_DATA, + spend_bundle_push=_noop_push, + ) with pytest.raises( ValueError, match=re.escape(f"Pubkey {pk1_h.get_fingerprint()} not found (or path/sum hinted to)") ): - await wsm.sign_bundle([spend_h]) + await signer.sign_bundle([spend_h]) - await wsm.puzzle_store.add_derivation_paths( + await signer.puzzle_store.add_derivation_paths( [ DerivationRecord( uint32(1), @@ -99,7 +109,7 @@ async def test_wsm_sign_transaction() -> None: ] ) - await wsm.puzzle_store.add_derivation_paths( + await signer.puzzle_store.add_derivation_paths( [ DerivationRecord( uint32(2), @@ -112,7 +122,7 @@ async def test_wsm_sign_transaction() -> None: ] ) - signature: G2Element = ((await wsm.sign_bundle([spend_h]))[0]).aggregated_signature + signature: G2Element = ((await signer.sign_bundle([spend_h]))[0]).aggregated_signature assert signature == AugSchemeMPL.aggregate( [ AugSchemeMPL.sign(sk1_h, msg1), @@ -123,9 +133,9 @@ async def test_wsm_sign_transaction() -> None: with pytest.raises( ValueError, match=re.escape(f"Pubkey {pk1_u.get_fingerprint()} not found (or path/sum hinted to)") ): - await wsm.sign_bundle([spend_u]) + await signer.sign_bundle([spend_u]) - await wsm.puzzle_store.add_derivation_paths( + await signer.puzzle_store.add_derivation_paths( [ DerivationRecord( uint32(1), @@ -138,7 +148,7 @@ async def test_wsm_sign_transaction() -> None: ] ) - await wsm.puzzle_store.add_derivation_paths( + await signer.puzzle_store.add_derivation_paths( [ DerivationRecord( uint32(2), @@ -150,7 +160,7 @@ async def test_wsm_sign_transaction() -> None: ) ] ) - signature2: G2Element = ((await wsm.sign_bundle([spend_u]))[0]).aggregated_signature + signature2: G2Element = ((await signer.sign_bundle([spend_u]))[0]).aggregated_signature assert signature2 == AugSchemeMPL.aggregate( [ AugSchemeMPL.sign(sk1_u, msg1), diff --git a/chia/_tests/wallet/test_signer_protocol.py b/chia/_tests/wallet/test_signer_protocol.py index e4f422a3f4..5664d59efb 100644 --- a/chia/_tests/wallet/test_signer_protocol.py +++ b/chia/_tests/wallet/test_signer_protocol.py @@ -221,9 +221,9 @@ async def test_p2dohp_wallet_signer_protocol(wallet_environments: WalletTestFram assert not_our_utx.signing_instructions.targets[1].message == bytes(not_our_message) not_our_signing_instructions: SigningInstructions = not_our_utx.signing_instructions with pytest.raises(ValueError, match=r"not found \(or path/sum hinted to\)"): - await wallet_state_manager.execute_signing_instructions(not_our_signing_instructions) + await wallet_state_manager.signer.execute_signing_instructions(not_our_signing_instructions) with pytest.raises(ValueError, match=r"No pubkey found \(or path hinted to\) for fingerprint"): - await wallet_state_manager.execute_signing_instructions( + await wallet_state_manager.signer.execute_signing_instructions( dataclasses.replace( not_our_signing_instructions, key_hints=dataclasses.replace( @@ -236,7 +236,7 @@ async def test_p2dohp_wallet_signer_protocol(wallet_environments: WalletTestFram ) ) with pytest.raises(ValueError, match="No root pubkey for fingerprint"): - await wallet_state_manager.execute_signing_instructions( + await wallet_state_manager.signer.execute_signing_instructions( dataclasses.replace( not_our_signing_instructions, key_hints=dataclasses.replace( @@ -248,7 +248,7 @@ async def test_p2dohp_wallet_signer_protocol(wallet_environments: WalletTestFram ), ) ) - signing_responses_2 = await wallet_state_manager.execute_signing_instructions( + signing_responses_2 = await wallet_state_manager.signer.execute_signing_instructions( not_our_signing_instructions, partial_allowed=True ) assert len(signing_responses_2) == 2 @@ -324,7 +324,7 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle # Test just a path hint test_name: bytes32 = std_hash(b"path hint only") child_sk: PrivateKey = _derive_path_unhardened(root_sk, [uint64(1), uint64(2), uint64(3), uint64(4)]) - signing_responses: list[SigningResponse] = await wallet.execute_signing_instructions( + signing_responses: list[SigningResponse] = await wallet.wallet_state_manager.signer.execute_signing_instructions( SigningInstructions( KeyHints( [], @@ -347,7 +347,9 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle [SigningTarget(sum_pk.get_fingerprint().to_bytes(4, "big"), test_name, test_name)], ) for partial_allowed in (True, False): - signing_responses = await wallet.execute_signing_instructions(signing_instructions, partial_allowed) + signing_responses = await wallet.wallet_state_manager.signer.execute_signing_instructions( + signing_instructions, partial_allowed + ) assert signing_responses == [ SigningResponse( bytes( @@ -369,7 +371,9 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle SigningTarget(b"random fingerprint", test_name, test_name), ], ) - signing_responses = await wallet.execute_signing_instructions(signing_instructions, partial_allowed=True) + signing_responses = await wallet.wallet_state_manager.signer.execute_signing_instructions( + signing_instructions, partial_allowed=True + ) assert signing_responses == [ SigningResponse( bytes( @@ -398,7 +402,9 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle [SigningTarget(sum_pk.get_fingerprint().to_bytes(4, "big"), test_name, test_name)], ) for partial_allowed in (True, False): - signing_responses = await wallet.execute_signing_instructions(signing_instructions, partial_allowed) + signing_responses = await wallet.wallet_state_manager.signer.execute_signing_instructions( + signing_instructions, partial_allowed + ) assert signing_responses == [ SigningResponse( bytes( @@ -423,7 +429,7 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle other_sk_2 = PrivateKey.from_bytes(test_name_2) sum_pk = child_sk.get_g1() + other_sk.get_g1() sum_pk_2 = child_sk_2.get_g1() + other_sk_2.get_g1() - signing_responses = await wallet.execute_signing_instructions( + signing_responses = await wallet.wallet_state_manager.signer.execute_signing_instructions( SigningInstructions( KeyHints( [ @@ -482,18 +488,18 @@ async def test_p2blsdohp_execute_signing_instructions(wallet_environments: Walle [SigningTarget(b"unknown fingerprint", b"", std_hash(b"some hook"))], ) with pytest.raises(ValueError, match="No root pubkey for fingerprint"): - await wallet.execute_signing_instructions(unknown_path_hint) + await wallet.wallet_state_manager.signer.execute_signing_instructions(unknown_path_hint) with pytest.raises(ValueError, match="No pubkey found"): - await wallet.execute_signing_instructions(unknown_sum_hint) + await wallet.wallet_state_manager.signer.execute_signing_instructions(unknown_sum_hint) with pytest.raises(ValueError, match="not found"): - await wallet.execute_signing_instructions(unknown_target) + await wallet.wallet_state_manager.signer.execute_signing_instructions(unknown_target) # Test no private key partial sign sum hint - wallet.wallet_state_manager.private_key = None + wallet.wallet_state_manager.signer = dataclasses.replace(wallet.wallet_state_manager.signer, root_private_key=None) test_name = std_hash(b"sum hint partial no private key") other_sk = PrivateKey.from_bytes(test_name) sum_pk = other_sk.get_g1() + root_pk - signing_responses = await wallet.execute_signing_instructions( + signing_responses = await wallet.wallet_state_manager.signer.execute_signing_instructions( SigningInstructions( KeyHints( [SumHint([root_fingerprint], test_name, bytes(sum_pk))], diff --git a/chia/_tests/wallet/test_wallet.py b/chia/_tests/wallet/test_wallet.py index 002ec14fc1..e95f2ac4dc 100644 --- a/chia/_tests/wallet/test_wallet.py +++ b/chia/_tests/wallet/test_wallet.py @@ -1871,7 +1871,7 @@ class TestWalletSimulator: assert stolen_cs is not None # get a legit signature - stolen_sb, _ = await wallet.wallet_state_manager.sign_bundle([stolen_cs]) + stolen_sb, _ = await wallet.wallet_state_manager.signer.sign_bundle([stolen_cs]) stolen_tx = wallet.wallet_state_manager.new_outgoing_transaction( wallet_id=wallet.id(), puzzle_hash=bytes32.zeros, diff --git a/chia/wallet/wallet.py b/chia/wallet/wallet.py index 767e2e7e80..6435caa937 100644 --- a/chia/wallet/wallet.py +++ b/chia/wallet/wallet.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, ClassVar, cast -from chia_rs import AugSchemeMPL, CoinSpend, G1Element, G2Element, PrivateKey +from chia_rs import CoinSpend, G1Element, G2Element, PrivateKey from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import uint32, uint64, uint128 from typing_extensions import Unpack @@ -20,34 +20,16 @@ from chia.wallet.conditions import ( CreateCoin, CreateCoinAnnouncement, ) -from chia.wallet.derivation_record import DerivationRecord -from chia.wallet.derive_keys import ( - MAX_POOL_WALLETS, - _derive_path, - _derive_path_unhardened, - master_sk_to_singleton_owner_sk, -) +from chia.wallet.derive_keys import _derive_path from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import ( DEFAULT_HIDDEN_PUZZLE_HASH, - calculate_synthetic_offset, calculate_synthetic_public_key, calculate_synthetic_secret_key, puzzle_for_pk, puzzle_hash_for_pk, - puzzle_hash_for_synthetic_public_key, solution_for_conditions, ) from chia.wallet.puzzles.puzzle_utils import make_reserve_fee_condition -from chia.wallet.signer_protocol import ( - PathHint, - Signature, - SignedTransaction, - SigningInstructions, - SigningResponse, - Spend, - SumHint, - TransactionInfo, -) from chia.wallet.transaction_record import TransactionRecord from chia.wallet.util.puzzle_decorator import PuzzleDecoratorManager from chia.wallet.util.transaction_type import CLAWBACK_INCOMING_TRANSACTION_TYPES @@ -498,183 +480,3 @@ class Wallet: def hardened_pubkey_for_path(self, path: list[int]) -> G1Element: return _derive_path(self.wallet_state_manager.get_master_private_key(), path).get_g1() - - async def sum_hint_for_pubkey(self, pk: bytes) -> SumHint | None: - pk_parsed: G1Element = G1Element.from_bytes(pk) - dr: DerivationRecord | None = await self.wallet_state_manager.puzzle_store.record_for_puzzle_hash( - puzzle_hash_for_synthetic_public_key(pk_parsed) - ) - if dr is None: - return None - return SumHint( - [dr.pubkey.get_fingerprint().to_bytes(4, "big")], - calculate_synthetic_offset(dr.pubkey, DEFAULT_HIDDEN_PUZZLE_HASH).to_bytes(32, "big"), - pk, - ) - - async def path_hint_for_pubkey(self, pk: bytes) -> PathHint | None: - pk_parsed: G1Element = G1Element.from_bytes(pk) - index: uint32 | None = await self.wallet_state_manager.puzzle_store.index_for_pubkey(pk_parsed) - if index is None: - index = await self.wallet_state_manager.puzzle_store.index_for_puzzle_hash( - puzzle_hash_for_synthetic_public_key(pk_parsed) - ) - root_pubkey: bytes = self.wallet_state_manager.root_pubkey.get_fingerprint().to_bytes(4, "big") - if index is None: - # Pool wallet may have a secret key here - if self.wallet_state_manager.private_key is not None: - for pool_wallet_index in range(MAX_POOL_WALLETS): - try_owner_sk = master_sk_to_singleton_owner_sk( - self.wallet_state_manager.private_key, uint32(pool_wallet_index) - ) - if try_owner_sk.get_g1() == pk_parsed: - return PathHint( - root_pubkey, - [uint64(12381), uint64(8444), uint64(5), uint64(pool_wallet_index)], - ) - return None - return PathHint( - root_pubkey, - [uint64(12381), uint64(8444), uint64(2), uint64(index)], - ) - - async def execute_signing_instructions( - self, signing_instructions: SigningInstructions, partial_allowed: bool = False - ) -> list[SigningResponse]: - root_pubkey: G1Element = self.wallet_state_manager.root_pubkey - pk_lookup: dict[int, G1Element] = ( - {root_pubkey.get_fingerprint(): root_pubkey} if self.wallet_state_manager.private_key is not None else {} - ) - sk_lookup: dict[int, PrivateKey] = ( - {root_pubkey.get_fingerprint(): self.wallet_state_manager.get_master_private_key()} - if self.wallet_state_manager.private_key is not None - else {} - ) - aggregate_responses_at_end: bool = True - responses: list[SigningResponse] = [] - - # TODO: expand path hints and sum hints recursively (a sum hint can give a new key to path hint) - # Next, expand our pubkey set with path hints - if self.wallet_state_manager.private_key is not None: - for path_hint in signing_instructions.key_hints.path_hints: - if int.from_bytes(path_hint.root_fingerprint, "big") != root_pubkey.get_fingerprint(): - if not partial_allowed: - raise ValueError(f"No root pubkey for fingerprint {root_pubkey.get_fingerprint()}") - else: - continue - else: - path = [int(step) for step in path_hint.path] - derive_child_sk = _derive_path(self.wallet_state_manager.get_master_private_key(), path) - derive_child_sk_unhardened = _derive_path_unhardened( - self.wallet_state_manager.get_master_private_key(), path - ) - derive_child_pk = derive_child_sk.get_g1() - derive_child_pk_unhardened = derive_child_sk_unhardened.get_g1() - pk_lookup[derive_child_pk.get_fingerprint()] = derive_child_pk - pk_lookup[derive_child_pk_unhardened.get_fingerprint()] = derive_child_pk_unhardened - sk_lookup[derive_child_pk.get_fingerprint()] = derive_child_sk - sk_lookup[derive_child_pk_unhardened.get_fingerprint()] = derive_child_sk_unhardened - - # Next, expand our pubkey set with sum hints - sum_hint_lookup: dict[int, list[int]] = {} - for sum_hint in signing_instructions.key_hints.sum_hints: - fingerprints_we_have: list[int] = [] - for fingerprint in sum_hint.fingerprints: - fingerprint_as_int = int.from_bytes(fingerprint, "big") - if fingerprint_as_int not in pk_lookup: - if not partial_allowed: - raise ValueError( - f"No pubkey found (or path hinted to) for fingerprint {int.from_bytes(fingerprint, 'big')}" - ) - else: - aggregate_responses_at_end = False - else: - fingerprints_we_have.append(fingerprint_as_int) - - # Add any synthetic offsets as keys we "have" - offset_sk = PrivateKey.from_bytes(sum_hint.synthetic_offset) - offset_pk = offset_sk.get_g1() - pk_lookup[offset_pk.get_fingerprint()] = offset_pk - sk_lookup[offset_pk.get_fingerprint()] = offset_sk - final_pubkey: G1Element = G1Element.from_bytes(sum_hint.final_pubkey) - final_fingerprint: int = final_pubkey.get_fingerprint() - pk_lookup[final_fingerprint] = final_pubkey - sum_hint_lookup[final_fingerprint] = [*fingerprints_we_have, offset_pk.get_fingerprint()] - - for target in signing_instructions.targets: - pk_fingerprint: int = int.from_bytes(target.fingerprint, "big") - if pk_fingerprint not in sk_lookup and pk_fingerprint not in sum_hint_lookup: - if not partial_allowed: - raise ValueError(f"Pubkey {pk_fingerprint} not found (or path/sum hinted to)") - else: - aggregate_responses_at_end = False - continue - elif pk_fingerprint in sk_lookup: - responses.append( - SigningResponse( - bytes(AugSchemeMPL.sign(sk_lookup[pk_fingerprint], target.message)), - target.hook, - ) - ) - else: # Implicit if pk_fingerprint in sum_hint_lookup - signatures: list[G2Element] = [] - for partial_fingerprint in sum_hint_lookup[pk_fingerprint]: - signatures.append( - AugSchemeMPL.sign(sk_lookup[partial_fingerprint], target.message, pk_lookup[pk_fingerprint]) - ) - if partial_allowed: - # In multisig scenarios, we return everything as a component signature - for sig in signatures: - responses.append( - SigningResponse( - bytes(sig), - target.hook, - ) - ) - else: - # In the scenario where we are the only signer, we can collapse many responses into one - responses.append( - SigningResponse( - bytes(AugSchemeMPL.aggregate(signatures)), - target.hook, - ) - ) - - # If we have the full set of signing responses for the instructions, aggregate them as much as possible - if aggregate_responses_at_end: - new_responses: list[SigningResponse] = [] - grouped_responses: dict[bytes32, list[SigningResponse]] = {} - for response in responses: - grouped_responses.setdefault(response.hook, []) - grouped_responses[response.hook].append(response) - for hook, group in grouped_responses.items(): - new_responses.append( - SigningResponse( - bytes(AugSchemeMPL.aggregate([G2Element.from_bytes(res.signature) for res in group])), - hook, - ) - ) - responses = new_responses - - return responses - - async def apply_signatures( - self, spends: list[Spend], signing_responses: list[SigningResponse] - ) -> SignedTransaction: - signing_responses_set = set(signing_responses) - return SignedTransaction( - TransactionInfo(spends), - [ - Signature( - "bls_12381_aug_scheme", - bytes( - AugSchemeMPL.aggregate( - [ - G2Element.from_bytes(signing_response.signature) - for signing_response in signing_responses_set - ] - ) - ), - ) - ], - ) diff --git a/chia/wallet/wallet_rpc_api.py b/chia/wallet/wallet_rpc_api.py index 78b5c308fe..327a0c871b 100644 --- a/chia/wallet/wallet_rpc_api.py +++ b/chia/wallet/wallet_rpc_api.py @@ -430,7 +430,7 @@ def tx_endpoint( # deferring to parent action scope return response - unsigned_txs = await self.service.wallet_state_manager.gather_signing_info_for_txs( + unsigned_txs = await self.service.wallet_state_manager.signer.gather_signing_info_for_txs( action_scope.side_effects.transactions ) @@ -3535,7 +3535,7 @@ class WalletRpcApi: request: GatherSigningInfo, ) -> GatherSigningInfoResponse: return GatherSigningInfoResponse( - signing_instructions=await self.service.wallet_state_manager.gather_signing_info(request.spends) + signing_instructions=await self.service.wallet_state_manager.signer.gather_signing_info(request.spends) ) async def apply_signatures( @@ -3544,7 +3544,9 @@ class WalletRpcApi: ) -> ApplySignaturesResponse: return ApplySignaturesResponse( signed_transactions=[ - await self.service.wallet_state_manager.apply_signatures(request.spends, request.signing_responses) + await self.service.wallet_state_manager.signer.apply_signatures( + request.spends, request.signing_responses + ) ] ) @@ -3553,7 +3555,7 @@ class WalletRpcApi: request: SubmitTransactions, ) -> SubmitTransactionsResponse: return SubmitTransactionsResponse( - mempool_ids=await self.service.wallet_state_manager.submit_transactions(request.signed_transactions) + mempool_ids=await self.service.wallet_state_manager.signer.submit_transactions(request.signed_transactions) ) async def execute_signing_instructions( @@ -3561,7 +3563,7 @@ class WalletRpcApi: request: ExecuteSigningInstructions, ) -> ExecuteSigningInstructionsResponse: return ExecuteSigningInstructionsResponse( - signing_responses=await self.service.wallet_state_manager.execute_signing_instructions( + signing_responses=await self.service.wallet_state_manager.signer.execute_signing_instructions( request.signing_instructions, request.partial_allowed ) ) diff --git a/chia/wallet/wallet_signer.py b/chia/wallet/wallet_signer.py new file mode 100644 index 0000000000..2a147d3b85 --- /dev/null +++ b/chia/wallet/wallet_signer.py @@ -0,0 +1,357 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, replace + +from chia_rs import AugSchemeMPL, CoinSpend, G1Element, G2Element, PrivateKey +from chia_rs.sized_bytes import bytes32 +from chia_rs.sized_ints import uint32, uint64 + +from chia.consensus.condition_tools import conditions_dict_for_solution, pkm_pairs_for_conditions_dict +from chia.util.hash import std_hash +from chia.wallet.derive_keys import ( + MAX_POOL_WALLETS, + _derive_path, + _derive_path_unhardened, + master_sk_to_singleton_owner_sk, +) +from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import ( + DEFAULT_HIDDEN_PUZZLE_HASH, + calculate_synthetic_offset, + puzzle_hash_for_synthetic_public_key, +) +from chia.wallet.signer_protocol import ( + KeyHints, + PathHint, + Signature, + SignedTransaction, + SigningInstructions, + SigningResponse, + SigningTarget, + Spend, + SumHint, + TransactionInfo, + UnsignedTransaction, +) +from chia.wallet.trading.offer import Offer +from chia.wallet.transaction_record import TransactionRecord +from chia.wallet.wallet_puzzle_store import WalletPuzzleStore +from chia.wallet.wallet_spend_bundle import WalletSpendBundle + + +@dataclass(frozen=True, kw_only=True) +class WalletSigner: + root_pubkey: G1Element + root_private_key: PrivateKey | None + puzzle_store: WalletPuzzleStore + max_block_cost_clvm: int + agg_sig_me_additional_data: bytes32 + spend_bundle_push: Callable[[WalletSpendBundle], Awaitable[None]] + + def __post_init__(self) -> None: + if self.root_private_key is not None and self.root_private_key.get_g1() != self.root_pubkey: + raise ValueError("root_private_key does not match self.root_pubkey") + + async def sum_hint_for_pubkey(self, pk: bytes) -> SumHint | None: + pk_parsed: G1Element = G1Element.from_bytes(pk) + dr = await self.puzzle_store.record_for_puzzle_hash(puzzle_hash_for_synthetic_public_key(pk_parsed)) + if dr is None: + return None + return SumHint( + [dr.pubkey.get_fingerprint().to_bytes(4, "big")], + calculate_synthetic_offset(dr.pubkey, DEFAULT_HIDDEN_PUZZLE_HASH).to_bytes(32, "big"), + pk, + ) + + async def path_hint_for_pubkey(self, pk: bytes) -> PathHint | None: + pk_parsed = G1Element.from_bytes(pk) + index = await self.puzzle_store.index_for_pubkey(pk_parsed) + if index is None: + index = await self.puzzle_store.index_for_puzzle_hash(puzzle_hash_for_synthetic_public_key(pk_parsed)) + root_pubkey: bytes = self.root_pubkey.get_fingerprint().to_bytes(4, "big") + if index is None: + # Pool wallet may have a secret key here + if self.root_private_key is not None: + for pool_wallet_index in range(MAX_POOL_WALLETS): + try_owner_sk = master_sk_to_singleton_owner_sk(self.root_private_key, uint32(pool_wallet_index)) + if try_owner_sk.get_g1() == pk_parsed: + return PathHint( + root_pubkey, + [uint64(12381), uint64(8444), uint64(5), uint64(pool_wallet_index)], + ) + return None + return PathHint( + root_pubkey, + [uint64(12381), uint64(8444), uint64(2), uint64(index)], + ) + + async def key_hints_for_pubkeys(self, pks: list[bytes]) -> KeyHints: + return KeyHints( + [sum_hint for pk in pks if (sum_hint := await self.sum_hint_for_pubkey(pk)) is not None], + [path_hint for pk in pks if (path_hint := await self.path_hint_for_pubkey(pk)) is not None], + ) + + async def gather_signing_info(self, spends: list[Spend]) -> SigningInstructions: + pks: list[bytes] = [] + signing_targets: list[SigningTarget] = [] + for spend in spends: + coin_spend = spend.as_coin_spend() + # Get AGG_SIG conditions + conditions_dict = conditions_dict_for_solution( + coin_spend.puzzle_reveal, + coin_spend.solution, + self.max_block_cost_clvm, + ) + # Create signature + for pk, msg in pkm_pairs_for_conditions_dict( + conditions_dict, coin_spend.coin, self.agg_sig_me_additional_data + ): + pk_bytes = bytes(pk) + pks.append(pk_bytes) + fingerprint: bytes = pk.get_fingerprint().to_bytes(4, "big") + signing_targets.append(SigningTarget(fingerprint, msg, std_hash(pk_bytes + msg))) + + return SigningInstructions( + await self.key_hints_for_pubkeys(pks), + signing_targets, + ) + + async def gather_signing_info_for_bundles(self, bundles: list[WalletSpendBundle]) -> list[UnsignedTransaction]: + utxs: list[UnsignedTransaction] = [] + for bundle in bundles: + signer_protocol_spends: list[Spend] = [Spend.from_coin_spend(spend) for spend in bundle.coin_spends] + utxs.append( + UnsignedTransaction( + TransactionInfo(signer_protocol_spends), + await self.gather_signing_info(signer_protocol_spends), + ) + ) + + return utxs + + async def gather_signing_info_for_txs(self, txs: list[TransactionRecord]) -> list[UnsignedTransaction]: + return await self.gather_signing_info_for_bundles( + [tx.spend_bundle for tx in txs if tx.spend_bundle is not None] + ) + + async def gather_signing_info_for_trades(self, offers: list[Offer]) -> list[UnsignedTransaction]: + return await self.gather_signing_info_for_bundles([offer._bundle for offer in offers]) + + async def execute_signing_instructions( + self, signing_instructions: SigningInstructions, partial_allowed: bool = False + ) -> list[SigningResponse]: + pk_lookup: dict[int, G1Element] = ( + {self.root_pubkey.get_fingerprint(): self.root_pubkey} if self.root_private_key is not None else {} + ) + sk_lookup: dict[int, PrivateKey] = ( + {self.root_pubkey.get_fingerprint(): self.root_private_key} if self.root_private_key is not None else {} + ) + aggregate_responses_at_end: bool = True + responses: list[SigningResponse] = [] + + # TODO: expand path hints and sum hints recursively (a sum hint can give a new key to path hint) + # Next, expand our pubkey set with path hints + if self.root_private_key is not None: + for path_hint in signing_instructions.key_hints.path_hints: + if int.from_bytes(path_hint.root_fingerprint, "big") != self.root_pubkey.get_fingerprint(): + if not partial_allowed: + raise ValueError(f"No root pubkey for fingerprint {self.root_pubkey.get_fingerprint()}") + else: + continue + else: + path = [int(step) for step in path_hint.path] + derive_child_sk = _derive_path(self.root_private_key, path) + derive_child_sk_unhardened = _derive_path_unhardened(self.root_private_key, path) + derive_child_pk = derive_child_sk.get_g1() + derive_child_pk_unhardened = derive_child_sk_unhardened.get_g1() + pk_lookup[derive_child_pk.get_fingerprint()] = derive_child_pk + pk_lookup[derive_child_pk_unhardened.get_fingerprint()] = derive_child_pk_unhardened + sk_lookup[derive_child_pk.get_fingerprint()] = derive_child_sk + sk_lookup[derive_child_pk_unhardened.get_fingerprint()] = derive_child_sk_unhardened + + # Next, expand our pubkey set with sum hints + sum_hint_lookup: dict[int, list[int]] = {} + for sum_hint in signing_instructions.key_hints.sum_hints: + fingerprints_we_have: list[int] = [] + for fingerprint in sum_hint.fingerprints: + fingerprint_as_int = int.from_bytes(fingerprint, "big") + if fingerprint_as_int not in pk_lookup: + if not partial_allowed: + raise ValueError( + f"No pubkey found (or path hinted to) for fingerprint {int.from_bytes(fingerprint, 'big')}" + ) + else: + aggregate_responses_at_end = False + else: + fingerprints_we_have.append(fingerprint_as_int) + + # Add any synthetic offsets as keys we "have" + offset_sk = PrivateKey.from_bytes(sum_hint.synthetic_offset) + offset_pk = offset_sk.get_g1() + pk_lookup[offset_pk.get_fingerprint()] = offset_pk + sk_lookup[offset_pk.get_fingerprint()] = offset_sk + final_pubkey: G1Element = G1Element.from_bytes(sum_hint.final_pubkey) + final_fingerprint: int = final_pubkey.get_fingerprint() + pk_lookup[final_fingerprint] = final_pubkey + sum_hint_lookup[final_fingerprint] = [*fingerprints_we_have, offset_pk.get_fingerprint()] + + for target in signing_instructions.targets: + pk_fingerprint: int = int.from_bytes(target.fingerprint, "big") + if pk_fingerprint not in sk_lookup and pk_fingerprint not in sum_hint_lookup: + if not partial_allowed: + raise ValueError(f"Pubkey {pk_fingerprint} not found (or path/sum hinted to)") + else: + aggregate_responses_at_end = False + continue + elif pk_fingerprint in sk_lookup: + responses.append( + SigningResponse( + bytes(AugSchemeMPL.sign(sk_lookup[pk_fingerprint], target.message)), + target.hook, + ) + ) + else: # Implicit if pk_fingerprint in sum_hint_lookup + signatures: list[G2Element] = [] + for partial_fingerprint in sum_hint_lookup[pk_fingerprint]: + signatures.append( + AugSchemeMPL.sign(sk_lookup[partial_fingerprint], target.message, pk_lookup[pk_fingerprint]) + ) + if partial_allowed: + # In multisig scenarios, we return everything as a component signature + for sig in signatures: + responses.append( + SigningResponse( + bytes(sig), + target.hook, + ) + ) + else: + # In the scenario where we are the only signer, we can collapse many responses into one + responses.append( + SigningResponse( + bytes(AugSchemeMPL.aggregate(signatures)), + target.hook, + ) + ) + + # If we have the full set of signing responses for the instructions, aggregate them as much as possible + if aggregate_responses_at_end: + new_responses: list[SigningResponse] = [] + grouped_responses: dict[bytes32, list[SigningResponse]] = {} + for response in responses: + grouped_responses.setdefault(response.hook, []) + grouped_responses[response.hook].append(response) + for hook, group in grouped_responses.items(): + new_responses.append( + SigningResponse( + bytes(AugSchemeMPL.aggregate([G2Element.from_bytes(res.signature) for res in group])), + hook, + ) + ) + responses = new_responses + + return responses + + async def apply_signatures( + self, spends: list[Spend], signing_responses: list[SigningResponse] + ) -> SignedTransaction: + signing_responses_set = set(signing_responses) + return SignedTransaction( + TransactionInfo(spends), + [ + Signature( + "bls_12381_aug_scheme", + bytes( + AugSchemeMPL.aggregate( + [ + G2Element.from_bytes(signing_response.signature) + for signing_response in signing_responses_set + ] + ) + ), + ) + ], + ) + + def signed_tx_to_spendbundle(self, signed_tx: SignedTransaction) -> WalletSpendBundle: + if len([_ for _ in signed_tx.signatures if _.type != "bls_12381_aug_scheme"]) > 0: + raise ValueError("Unable to handle signatures that are not bls_12381_aug_scheme") # pragma: no cover + return WalletSpendBundle( + [spend.as_coin_spend() for spend in signed_tx.transaction_info.spends], + AugSchemeMPL.aggregate([G2Element.from_bytes(sig.signature) for sig in signed_tx.signatures]), + ) + + async def sign_transactions( + self, + tx_records: list[TransactionRecord], + additional_signing_responses: list[SigningResponse] = [], + partial_allowed: bool = False, + ) -> tuple[list[TransactionRecord], list[SigningResponse]]: + unsigned_txs: list[UnsignedTransaction] = await self.gather_signing_info_for_txs(tx_records) + new_txs: list[TransactionRecord] = [] + all_signing_responses = additional_signing_responses.copy() + for unsigned_tx, tx in zip( + unsigned_txs, [tx_record for tx_record in tx_records if tx_record.spend_bundle is not None] + ): + signing_responses: list[SigningResponse] = await self.execute_signing_instructions( + unsigned_tx.signing_instructions, partial_allowed=partial_allowed + ) + all_signing_responses.extend(signing_responses) + new_bundle = self.signed_tx_to_spendbundle( + await self.apply_signatures( + unsigned_tx.transaction_info.spends, + [*additional_signing_responses, *signing_responses], + ) + ) + new_txs.append(replace(tx, spend_bundle=new_bundle, name=new_bundle.name())) + new_txs.extend([tx_record for tx_record in tx_records if tx_record.spend_bundle is None]) + return new_txs, all_signing_responses + + async def sign_offers( + self, + offers: list[Offer], + additional_signing_responses: list[SigningResponse] = [], + partial_allowed: bool = False, + ) -> tuple[list[Offer], list[SigningResponse]]: + unsigned_txs: list[UnsignedTransaction] = await self.gather_signing_info_for_trades(offers) + new_offers: list[Offer] = [] + all_signing_responses = additional_signing_responses.copy() + for unsigned_tx, offer in zip(unsigned_txs, [offer for offer in offers]): + signing_responses: list[SigningResponse] = await self.execute_signing_instructions( + unsigned_tx.signing_instructions, partial_allowed=partial_allowed + ) + all_signing_responses.extend(signing_responses) + new_bundle = self.signed_tx_to_spendbundle( + await self.apply_signatures( + unsigned_tx.transaction_info.spends, + [*additional_signing_responses, *signing_responses], + ) + ) + new_offers.append(Offer(offer.requested_payments, new_bundle, offer.driver_dict)) + return new_offers, all_signing_responses + + async def sign_bundle( + self, + coin_spends: list[CoinSpend], + additional_signing_responses: list[SigningResponse] = [], + partial_allowed: bool = False, + ) -> tuple[WalletSpendBundle, list[SigningResponse]]: + [unsigned_tx] = await self.gather_signing_info_for_bundles([WalletSpendBundle(coin_spends, G2Element())]) + signing_responses: list[SigningResponse] = await self.execute_signing_instructions( + unsigned_tx.signing_instructions, partial_allowed=partial_allowed + ) + return ( + self.signed_tx_to_spendbundle( + await self.apply_signatures( + unsigned_tx.transaction_info.spends, + [*additional_signing_responses, *signing_responses], + ) + ), + signing_responses, + ) + + async def submit_transactions(self, signed_txs: list[SignedTransaction]) -> list[bytes32]: + bundles: list[WalletSpendBundle] = [self.signed_tx_to_spendbundle(tx) for tx in signed_txs] + for bundle in bundles: + await self.spend_bundle_push(bundle) + return [bundle.name() for bundle in bundles] diff --git a/chia/wallet/wallet_state_manager.py b/chia/wallet/wallet_state_manager.py index 0d362fdbb2..2db02a0bb9 100644 --- a/chia/wallet/wallet_state_manager.py +++ b/chia/wallet/wallet_state_manager.py @@ -16,13 +16,12 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast import aiosqlite -from chia_rs import AugSchemeMPL, CoinRecord, CoinSpend, CoinState, ConsensusConstants, G1Element, G2Element, PrivateKey +from chia_rs import CoinRecord, CoinSpend, CoinState, ConsensusConstants, G1Element, G2Element, PrivateKey from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import uint8, uint16, uint32, uint64, uint128 from chia.consensus.block_rewards import calculate_base_farmer_reward, calculate_pool_reward from chia.consensus.coinbase import farmer_parent_id, pool_parent_id -from chia.consensus.condition_tools import conditions_dict_for_solution, pkm_pairs_for_conditions_dict from chia.data_layer.data_layer_wallet import DataLayerWallet from chia.data_layer.dl_wallet_store import DataLayerStore from chia.data_layer.singleton_record import SingletonRecord @@ -93,22 +92,10 @@ from chia.wallet.puzzles.clawback.drivers import generate_clawback_spend_bundle, from chia.wallet.puzzles.clawback.metadata import ClawbackMetadata, ClawbackVersion from chia.wallet.remote_wallet.remote_coin_store import RemoteCoinStore from chia.wallet.remote_wallet.remote_wallet import RemoteWallet -from chia.wallet.signer_protocol import ( - KeyHints, - PathHint, - SignedTransaction, - SigningInstructions, - SigningResponse, - SigningTarget, - Spend, - SumHint, - TransactionInfo, - UnsignedTransaction, -) +from chia.wallet.signer_protocol import SigningResponse from chia.wallet.singleton import SINGLETON_LAUNCHER_PUZZLE_HASH as SINGLETON_LAUNCHER_HASH from chia.wallet.singleton import create_singleton_puzzle, get_inner_puzzle_from_singleton from chia.wallet.trade_manager import TradeManager -from chia.wallet.trading.offer import Offer from chia.wallet.trading.trade_status import TradeStatus from chia.wallet.transaction_record import LightTransactionRecord, TransactionRecord from chia.wallet.uncurried_puzzle import uncurry_puzzle @@ -149,6 +136,7 @@ from chia.wallet.wallet_protocol import ( ) from chia.wallet.wallet_puzzle_store import WalletPuzzleStore from chia.wallet.wallet_retry_store import WalletRetryStore +from chia.wallet.wallet_signer import WalletSigner from chia.wallet.wallet_spend_bundle import WalletSpendBundle from chia.wallet.wallet_transaction_store import WalletTransactionStore from chia.wallet.wallet_user_store import WalletUserStore @@ -217,6 +205,7 @@ class WalletStateManager: asset_to_wallet_map: dict[AssetType, Any] initial_num_public_keys: int decorator_manager: PuzzleDecoratorManager + signer: WalletSigner @staticmethod async def create( @@ -293,6 +282,15 @@ class WalletStateManager: assert root_pubkey == calculated_root_public_key self.root_pubkey = calculated_root_public_key + self.signer = WalletSigner( + root_pubkey=self.root_pubkey, + root_private_key=self.private_key, + puzzle_store=self.puzzle_store, + max_block_cost_clvm=self.constants.MAX_BLOCK_COST_CLVM, + agg_sig_me_additional_data=self.constants.AGG_SIG_ME_ADDITIONAL_DATA, + spend_bundle_push=self.wallet_node.push_tx, + ) + fingerprint = self.root_pubkey.get_fingerprint() puzzle_decorators = self.config.get("puzzle_decorators", {}).get(fingerprint, []) self.decorator_manager = PuzzleDecoratorManager.create(puzzle_decorators) @@ -2461,7 +2459,7 @@ class WalletStateManager: for i, tx in enumerate(tx_records) ] if sign: - tx_records, _ = await self.sign_transactions( + tx_records, _ = await self.signer.sign_transactions( tx_records, [] if additional_signing_responses is None else additional_signing_responses, additional_signing_responses != [] and additional_signing_responses is not None, @@ -2801,156 +2799,6 @@ class WalletStateManager: return vc_wallet - async def sum_hint_for_pubkey(self, pk: bytes) -> SumHint | None: - return await self.main_wallet.sum_hint_for_pubkey(pk) - - async def path_hint_for_pubkey(self, pk: bytes) -> PathHint | None: - return await self.main_wallet.path_hint_for_pubkey(pk) - - async def key_hints_for_pubkeys(self, pks: list[bytes]) -> KeyHints: - return KeyHints( - [sum_hint for pk in pks for sum_hint in (await self.sum_hint_for_pubkey(pk),) if sum_hint is not None], - [path_hint for pk in pks for path_hint in (await self.path_hint_for_pubkey(pk),) if path_hint is not None], - ) - - async def gather_signing_info(self, spends: list[Spend]) -> SigningInstructions: - pks: list[bytes] = [] - signing_targets: list[SigningTarget] = [] - for spend in spends: - coin_spend = spend.as_coin_spend() - # Get AGG_SIG conditions - conditions_dict = conditions_dict_for_solution( - Program.from_serialized(coin_spend.puzzle_reveal), - Program.from_serialized(coin_spend.solution), - self.constants.MAX_BLOCK_COST_CLVM, - ) - # Create signature - for pk, msg in pkm_pairs_for_conditions_dict( - conditions_dict, coin_spend.coin, self.constants.AGG_SIG_ME_ADDITIONAL_DATA - ): - pk_bytes = bytes(pk) - pks.append(pk_bytes) - fingerprint: bytes = pk.get_fingerprint().to_bytes(4, "big") - signing_targets.append(SigningTarget(fingerprint, msg, std_hash(pk_bytes + msg))) - - return SigningInstructions( - await self.key_hints_for_pubkeys(pks), - signing_targets, - ) - - async def gather_signing_info_for_bundles(self, bundles: list[WalletSpendBundle]) -> list[UnsignedTransaction]: - utxs: list[UnsignedTransaction] = [] - for bundle in bundles: - signer_protocol_spends: list[Spend] = [Spend.from_coin_spend(spend) for spend in bundle.coin_spends] - utxs.append( - UnsignedTransaction( - TransactionInfo(signer_protocol_spends), await self.gather_signing_info(signer_protocol_spends) - ) - ) - - return utxs - - async def gather_signing_info_for_txs(self, txs: list[TransactionRecord]) -> list[UnsignedTransaction]: - return await self.gather_signing_info_for_bundles( - [tx.spend_bundle for tx in txs if tx.spend_bundle is not None] - ) - - async def gather_signing_info_for_trades(self, offers: list[Offer]) -> list[UnsignedTransaction]: - return await self.gather_signing_info_for_bundles([offer._bundle for offer in offers]) - - async def execute_signing_instructions( - self, signing_instructions: SigningInstructions, partial_allowed: bool = False - ) -> list[SigningResponse]: - return await self.main_wallet.execute_signing_instructions(signing_instructions, partial_allowed) - - async def apply_signatures( - self, spends: list[Spend], signing_responses: list[SigningResponse] - ) -> SignedTransaction: - return await self.main_wallet.apply_signatures(spends, signing_responses) - - def signed_tx_to_spendbundle(self, signed_tx: SignedTransaction) -> WalletSpendBundle: - if len([_ for _ in signed_tx.signatures if _.type != "bls_12381_aug_scheme"]) > 0: - raise ValueError("Unable to handle signatures that are not bls_12381_aug_scheme") # pragma: no cover - return WalletSpendBundle( - [spend.as_coin_spend() for spend in signed_tx.transaction_info.spends], - AugSchemeMPL.aggregate([G2Element.from_bytes(sig.signature) for sig in signed_tx.signatures]), - ) - - async def sign_transactions( - self, - tx_records: list[TransactionRecord], - additional_signing_responses: list[SigningResponse] = [], - partial_allowed: bool = False, - ) -> tuple[list[TransactionRecord], list[SigningResponse]]: - unsigned_txs: list[UnsignedTransaction] = await self.gather_signing_info_for_txs(tx_records) - new_txs: list[TransactionRecord] = [] - all_signing_responses = additional_signing_responses.copy() - for unsigned_tx, tx in zip( - unsigned_txs, [tx_record for tx_record in tx_records if tx_record.spend_bundle is not None] - ): - signing_responses: list[SigningResponse] = await self.execute_signing_instructions( - unsigned_tx.signing_instructions, partial_allowed=partial_allowed - ) - all_signing_responses.extend(signing_responses) - new_bundle = self.signed_tx_to_spendbundle( - await self.apply_signatures( - unsigned_tx.transaction_info.spends, - [*additional_signing_responses, *signing_responses], - ) - ) - new_txs.append(dataclasses.replace(tx, spend_bundle=new_bundle, name=new_bundle.name())) - new_txs.extend([tx_record for tx_record in tx_records if tx_record.spend_bundle is None]) - return new_txs, all_signing_responses - - async def sign_offers( - self, - offers: list[Offer], - additional_signing_responses: list[SigningResponse] = [], - partial_allowed: bool = False, - ) -> tuple[list[Offer], list[SigningResponse]]: - unsigned_txs: list[UnsignedTransaction] = await self.gather_signing_info_for_trades(offers) - new_offers: list[Offer] = [] - all_signing_responses = additional_signing_responses.copy() - for unsigned_tx, offer in zip(unsigned_txs, [offer for offer in offers]): - signing_responses: list[SigningResponse] = await self.execute_signing_instructions( - unsigned_tx.signing_instructions, partial_allowed=partial_allowed - ) - all_signing_responses.extend(signing_responses) - new_bundle = self.signed_tx_to_spendbundle( - await self.apply_signatures( - unsigned_tx.transaction_info.spends, - [*additional_signing_responses, *signing_responses], - ) - ) - new_offers.append(Offer(offer.requested_payments, new_bundle, offer.driver_dict)) - return new_offers, all_signing_responses - - async def sign_bundle( - self, - coin_spends: list[CoinSpend], - additional_signing_responses: list[SigningResponse] = [], - partial_allowed: bool = False, - ) -> tuple[WalletSpendBundle, list[SigningResponse]]: - [unsigned_tx] = await self.gather_signing_info_for_bundles([WalletSpendBundle(coin_spends, G2Element())]) - signing_responses: list[SigningResponse] = await self.execute_signing_instructions( - unsigned_tx.signing_instructions, partial_allowed=partial_allowed - ) - return ( - self.signed_tx_to_spendbundle( - await self.apply_signatures( - unsigned_tx.transaction_info.spends, - [*additional_signing_responses, *signing_responses], - ) - ), - signing_responses, - ) - - async def submit_transactions(self, signed_txs: list[SignedTransaction]) -> list[bytes32]: - bundles: list[WalletSpendBundle] = [self.signed_tx_to_spendbundle(tx) for tx in signed_txs] - for bundle in bundles: - await self.wallet_node.push_tx(bundle) - return [bundle.name() for bundle in bundles] - @contextlib.asynccontextmanager async def new_action_scope( self,