diff --git a/chia/_tests/wallet/test_wallet.py b/chia/_tests/wallet/test_wallet.py index e95f2ac4dc..f5f60499a5 100644 --- a/chia/_tests/wallet/test_wallet.py +++ b/chia/_tests/wallet/test_wallet.py @@ -541,6 +541,15 @@ class TestWalletSimulator: wallet_environments.tx_config, ) assert len(resp.transaction_ids) == 1 + # just test that an extra submission doesn't do anything + await env.rpc_client.spend_clawback_coins( + SpendClawbackCoins( + coin_ids=[merkle_coin.name()], + fee=uint64(test_fee), + push=True, + ), + wallet_environments.tx_config, + ) await wallet_environments.process_pending_states( [ diff --git a/chia/wallet/clawback_manager.py b/chia/wallet/clawback_manager.py new file mode 100644 index 0000000000..cd2b4f38dc --- /dev/null +++ b/chia/wallet/clawback_manager.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import logging +import time +from collections.abc import Awaitable, Callable +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass, replace + +from chia_rs import Coin, CoinSpend, G2Element +from chia_rs.sized_bytes import bytes32 +from chia_rs.sized_ints import uint32, uint64 + +from chia.types.mempool_inclusion_status import MempoolInclusionStatus +from chia.util.hash import std_hash +from chia.util.streamable import UInt32Range, UInt64Range +from chia.wallet.conditions import ( + AssertCoinAnnouncement, + Condition, + CreateCoin, + CreateCoinAnnouncement, + parse_timelock_info, +) +from chia.wallet.puzzles.clawback.drivers import generate_clawback_spend_bundle +from chia.wallet.puzzles.clawback.metadata import ClawbackMetadata +from chia.wallet.transaction_record import TransactionRecord +from chia.wallet.util.compute_memos import compute_memos +from chia.wallet.util.transaction_type import TransactionType +from chia.wallet.util.tx_config import TXConfig +from chia.wallet.util.wallet_types import CoinType, WalletType +from chia.wallet.wallet import Wallet +from chia.wallet.wallet_action_scope import WalletActionScope +from chia.wallet.wallet_blockchain import WalletBlockchain +from chia.wallet.wallet_coin_store import WalletCoinStore +from chia.wallet.wallet_puzzle_store import WalletPuzzleStore +from chia.wallet.wallet_spend_bundle import WalletSpendBundle +from chia.wallet.wallet_transaction_store import WalletTransactionStore + + +@dataclass(frozen=True, kw_only=True) +class ClawbackManager: + log: logging.Logger + blockchain: WalletBlockchain + coin_store: WalletCoinStore + puzzle_store: WalletPuzzleStore + transaction_store: WalletTransactionStore + xch_wallet: Wallet + auto_claim_tx_fee: uint64 + auto_claim_batch_size: int + timestamp_for_height: Callable[[uint32], Awaitable[uint64]] + puzzle_hash_encoder: Callable[[bytes32], str] + action_scope_sandbox: Callable[[TXConfig, bool], AbstractAsyncContextManager[WalletActionScope]] + + async def auto_claim_coins(self, action_scope: WalletActionScope) -> None: + # Get unspent clawback coin + current_timestamp = self.blockchain.get_latest_timestamp() + clawback_coins: dict[Coin, ClawbackMetadata] = {} + unspent_coins = await self.coin_store.get_coin_records( + coin_type=CoinType.CLAWBACK, + wallet_type=WalletType.STANDARD_WALLET, + spent_range=UInt32Range(stop=uint32(0)), + amount_range=UInt64Range( + start=action_scope.config.tx_config.coin_selection_config.min_coin_amount, + stop=action_scope.config.tx_config.coin_selection_config.max_coin_amount, + ), + ) + + for coin in unspent_coins.records: + try: + metadata = coin.parsed_metadata() + assert isinstance(metadata, ClawbackMetadata) + if await metadata.is_recipient(self.puzzle_store): + coin_timestamp = await self.timestamp_for_height(coin.confirmed_block_height) + if current_timestamp - coin_timestamp >= metadata.time_lock: + clawback_coins[coin.coin] = metadata + if len(clawback_coins) >= self.auto_claim_batch_size: + await self.spend_clawback_coins(clawback_coins, self.auto_claim_tx_fee, action_scope) + clawback_coins = {} + except Exception as e: + self.log.error(f"Failed to claim clawback coin {coin.coin.name().hex()}: %s", e) + if len(clawback_coins) > 0: + await self.spend_clawback_coins(clawback_coins, self.auto_claim_tx_fee, action_scope) + + async def spend_clawback_coins( + self, + clawback_coins: dict[Coin, ClawbackMetadata], + fee: uint64, + action_scope: WalletActionScope, + force: bool = False, + extra_conditions: tuple[Condition, ...] = tuple(), + ) -> None: + assert len(clawback_coins) > 0 + coin_spends: list[CoinSpend] = [] + message = std_hash(b"".join([c.name() for c in clawback_coins.keys()])) + derivation_record = None + amount = uint64(0) + for coin, metadata in clawback_coins.items(): + try: + self.log.info(f"Claiming clawback coin {coin.name().hex()}") + # Get incoming tx + incoming_tx = await self.transaction_store.get_transaction_record(coin.name()) + assert incoming_tx is not None, f"Cannot find incoming tx for clawback coin {coin.name().hex()}" + if incoming_tx.sent > 0 and not force: + self.log.error( + f"Clawback coin {coin.name().hex()} is already in a pending spend bundle. {incoming_tx}" + ) + continue + + recipient_puzhash = metadata.recipient_puzzle_hash + sender_puzhash = metadata.sender_puzzle_hash + is_recipient: bool = await metadata.is_recipient(self.puzzle_store) + if is_recipient: + derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(recipient_puzhash) + else: + derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(sender_puzhash) + assert derivation_record is not None + amount = uint64(amount + coin.amount) + # Remove the clawback hint since it is unnecessary for the XCH coin + memos: list[bytes] = [] if len(incoming_tx.memos) == 0 else next(iter(incoming_tx.memos.items()))[1][1:] + inner_puzzle = self.xch_wallet.puzzle_for_pk(derivation_record.pubkey) + inner_solution = self.xch_wallet.make_solution( + primaries=[ + CreateCoin( + derivation_record.puzzle_hash, + uint64(coin.amount), + memos, # Forward memo of the first coin + ) + ], + conditions=( + extra_conditions + if len(coin_spends) > 0 or fee == 0 + else (*extra_conditions, CreateCoinAnnouncement(message)) + ), + ) + coin_spend: CoinSpend = generate_clawback_spend_bundle(coin, metadata, inner_puzzle, inner_solution) + coin_spends.append(coin_spend) + # Update incoming tx to prevent double spend and mark it is pending + await self.transaction_store.increment_sent(incoming_tx.name, "", MempoolInclusionStatus.PENDING, None) + except Exception as e: + self.log.error(f"Failed to create clawback spend bundle for {coin.name().hex()}: {e}") + if len(coin_spends) == 0: + return + spend_bundle = WalletSpendBundle(coin_spends, G2Element()) + if fee > 0: + async with self.action_scope_sandbox(action_scope.config.tx_config, False) as inner_action_scope: + async with action_scope.use() as interface: + async with inner_action_scope.use() as inner_interface: + inner_interface.side_effects.selected_coins = interface.side_effects.selected_coins + await self.xch_wallet.create_tandem_xch_tx( + fee, + inner_action_scope, + extra_conditions=( + AssertCoinAnnouncement(asserted_id=coin_spends[0].coin.name(), asserted_msg=message), + ), + ) + async with inner_action_scope.use() as inner_interface: + # This should not be looked to for best practice. + # Ideally, the two spend bundles can exist separately on each tx record until they are pushed. + # This is not very supported behavior at the moment + # so to avoid any potential backwards compatibility issues, + # we're moving the spend bundle from this TX to the main + interface.side_effects.transactions.extend( + [replace(tx, spend_bundle=None) for tx in inner_interface.side_effects.transactions] + ) + interface.side_effects.selected_coins.extend(inner_interface.side_effects.selected_coins) + spend_bundle = WalletSpendBundle.aggregate( + [ + spend_bundle, + *( + tx.spend_bundle + for tx in inner_action_scope.side_effects.transactions + if tx.spend_bundle is not None + ), + ] + ) + assert derivation_record is not None + tx_record = TransactionRecord( + confirmed_at_height=uint32(0), + created_at_time=uint64(time.time()), + to_puzzle_hash=derivation_record.puzzle_hash, + to_address=self.puzzle_hash_encoder(derivation_record.puzzle_hash), + amount=amount, + fee_amount=uint64(fee), + confirmed=False, + sent=uint32(0), + spend_bundle=spend_bundle, + additions=spend_bundle.additions(), + removals=spend_bundle.removals(), + wallet_id=uint32(1), + sent_to=[], + trade_id=None, + type=uint32(TransactionType.OUTGOING_CLAWBACK), + name=spend_bundle.name(), + memos=compute_memos(spend_bundle), + valid_times=parse_timelock_info(extra_conditions), + ) + async with action_scope.use() as interface: + interface.side_effects.transactions.append(tx_record) diff --git a/chia/wallet/wallet_node.py b/chia/wallet/wallet_node.py index e2cd00603e..13e983b523 100644 --- a/chia/wallet/wallet_node.py +++ b/chia/wallet/wallet_node.py @@ -323,6 +323,12 @@ class WalletNode: if "auto_claim" not in self.config or self.config["auto_claim"] != auto_claim_config_json: # Update in memory config self.config["auto_claim"] = auto_claim_config_json + if self._wallet_state_manager is not None: + self.wallet_state_manager.clawback_manager = dataclasses.replace( + self.wallet_state_manager.clawback_manager, + auto_claim_tx_fee=auto_claim_config.tx_fee, + auto_claim_batch_size=auto_claim_config.batch_size, + ) # Update config file with lock_and_load_config(self.root_path, "config.yaml") as config: config["wallet"]["auto_claim"] = self.config["auto_claim"] @@ -1299,7 +1305,10 @@ class WalletNode: # Check if any coin needs auto spending if self.config.get("auto_claim", {}).get("enabled", False): - await self.wallet_state_manager.auto_claim_coins() + async with self.wallet_state_manager.new_action_scope( + self.wallet_state_manager.tx_config, push=True + ) as action_scope: + await self.wallet_state_manager.clawback_manager.auto_claim_coins(action_scope) if new_peak_hb.foliage_transaction_block is not None: await self._retry_fee_failed_transactions() diff --git a/chia/wallet/wallet_rpc_api.py b/chia/wallet/wallet_rpc_api.py index 327a0c871b..c60272e08e 100644 --- a/chia/wallet/wallet_rpc_api.py +++ b/chia/wallet/wallet_rpc_api.py @@ -1500,7 +1500,7 @@ class WalletRpcApi: batch_size = ( request.batch_size if request.batch_size is not None - else self.service.wallet_state_manager.config.get("auto_claim", {}).get("batch_size", 50) + else self.service.wallet_state_manager.clawback_manager.auto_claim_batch_size ) records_list = list(coin_records.coin_id_to_record.values()) for i in range(0, len(records_list), batch_size): @@ -1511,7 +1511,7 @@ class WalletRpcApi: except WalletCoinRecordMetadataParsingError as e: log.error("Failed to spend clawback coin: %s", e) continue - await self.service.wallet_state_manager.spend_clawback_coins( + await self.service.wallet_state_manager.clawback_manager.spend_clawback_coins( # Semantically, we're guaranteed the right type here, but the typing isn't there coin_batch, # type: ignore[arg-type] request.fee, diff --git a/chia/wallet/wallet_state_manager.py b/chia/wallet/wallet_state_manager.py index 13e0f44d83..048598e1a6 100644 --- a/chia/wallet/wallet_state_manager.py +++ b/chia/wallet/wallet_state_manager.py @@ -48,12 +48,11 @@ from chia.wallet.cat_wallet.cat_info import CATCoinData, CATInfo, CRCATInfo from chia.wallet.cat_wallet.cat_utils import CAT_MOD, CAT_MOD_HASH, construct_cat_puzzle, match_cat_puzzle from chia.wallet.cat_wallet.cat_wallet import CATWallet from chia.wallet.cat_wallet.r_cat_wallet import RCATWallet +from chia.wallet.clawback_manager import ClawbackManager from chia.wallet.conditions import ( - AssertCoinAnnouncement, Condition, ConditionValidTimes, CreateCoin, - CreateCoinAnnouncement, parse_timelock_info, ) from chia.wallet.db_wallet.db_wallet_puzzles import MIRROR_PUZZLE_HASH @@ -88,7 +87,7 @@ from chia.wallet.outer_puzzles import AssetType from chia.wallet.plotnft_wallet.plotnft_store import PlotNFTStore from chia.wallet.plotnft_wallet.plotnft_wallet import PlotNFT2Wallet from chia.wallet.puzzle_drivers import PuzzleInfo -from chia.wallet.puzzles.clawback.drivers import generate_clawback_spend_bundle, match_clawback_puzzle +from chia.wallet.puzzles.clawback.drivers import match_clawback_puzzle 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 @@ -122,7 +121,7 @@ from chia.wallet.vc_wallet.vc_wallet import VCWallet from chia.wallet.wallet import Wallet from chia.wallet.wallet_action_scope import PlotNFTTargetStateInfo, WalletActionScope, new_wallet_action_scope from chia.wallet.wallet_blockchain import WalletBlockchain -from chia.wallet.wallet_coin_record import MetadataTypes, WalletCoinRecord +from chia.wallet.wallet_coin_record import WalletCoinRecord from chia.wallet.wallet_coin_store import CoinRecordOrder, WalletCoinStore from chia.wallet.wallet_info import WalletInfo from chia.wallet.wallet_interested_store import WalletInterestedStore @@ -206,6 +205,7 @@ class WalletStateManager: initial_num_public_keys: int decorator_manager: PuzzleDecoratorManager signer: WalletSigner + clawback_manager: ClawbackManager @staticmethod async def create( @@ -297,6 +297,20 @@ class WalletStateManager: self.main_wallet = await Wallet.create(self, main_wallet_info) + self.clawback_manager = ClawbackManager( + log=self.log, + blockchain=self.blockchain, + coin_store=self.coin_store, + puzzle_store=self.puzzle_store, + transaction_store=self.tx_store, + xch_wallet=self.main_wallet, + auto_claim_tx_fee=uint64(self.config.get("auto_claim", {}).get("tx_fee", 0)), + auto_claim_batch_size=self.config.get("auto_claim", {}).get("batch_size", 50), + timestamp_for_height=self.wallet_node.get_timestamp_for_height, + puzzle_hash_encoder=self.encode_puzzle_hash, + action_scope_sandbox=self.new_action_scope, + ) + self.wallets = {main_wallet_info.id: self.main_wallet} self.asset_to_wallet_map = { @@ -1036,157 +1050,6 @@ class WalletStateManager: logged_in_fingerprint=self.root_pubkey.get_fingerprint(), ) - async def auto_claim_coins(self) -> None: - # Get unspent clawback coin - current_timestamp = self.blockchain.get_latest_timestamp() - clawback_coins: dict[Coin, ClawbackMetadata] = {} - tx_fee = uint64(self.config.get("auto_claim", {}).get("tx_fee", 0)) - unspent_coins = await self.coin_store.get_coin_records( - coin_type=CoinType.CLAWBACK, - wallet_type=WalletType.STANDARD_WALLET, - spent_range=UInt32Range(stop=uint32(0)), - amount_range=UInt64Range( - start=self.tx_config.coin_selection_config.min_coin_amount, - stop=self.tx_config.coin_selection_config.max_coin_amount, - ), - ) - async with self.new_action_scope(self.tx_config, push=True) as action_scope: - for coin in unspent_coins.records: - try: - metadata: MetadataTypes = coin.parsed_metadata() - assert isinstance(metadata, ClawbackMetadata) - if await metadata.is_recipient(self.puzzle_store): - coin_timestamp = await self.wallet_node.get_timestamp_for_height(coin.confirmed_block_height) - if current_timestamp - coin_timestamp >= metadata.time_lock: - clawback_coins[coin.coin] = metadata - if len(clawback_coins) >= self.config.get("auto_claim", {}).get("batch_size", 50): - await self.spend_clawback_coins(clawback_coins, tx_fee, action_scope) - clawback_coins = {} - except Exception as e: - self.log.error(f"Failed to claim clawback coin {coin.coin.name().hex()}: %s", e) - if len(clawback_coins) > 0: - await self.spend_clawback_coins(clawback_coins, tx_fee, action_scope) - - async def spend_clawback_coins( - self, - clawback_coins: dict[Coin, ClawbackMetadata], - fee: uint64, - action_scope: WalletActionScope, - force: bool = False, - extra_conditions: tuple[Condition, ...] = tuple(), - ) -> None: - assert len(clawback_coins) > 0 - coin_spends: list[CoinSpend] = [] - message: bytes32 = std_hash(b"".join([c.name() for c in clawback_coins.keys()])) - now: uint64 = uint64(time.time()) - derivation_record: DerivationRecord | None = None - amount: uint64 = uint64(0) - for coin, metadata in clawback_coins.items(): - try: - self.log.info(f"Claiming clawback coin {coin.name().hex()}") - # Get incoming tx - incoming_tx = await self.tx_store.get_transaction_record(coin.name()) - assert incoming_tx is not None, f"Cannot find incoming tx for clawback coin {coin.name().hex()}" - if incoming_tx.sent > 0 and not force: - self.log.error( - f"Clawback coin {coin.name().hex()} is already in a pending spend bundle. {incoming_tx}" - ) - continue - - recipient_puzhash: bytes32 = metadata.recipient_puzzle_hash - sender_puzhash: bytes32 = metadata.sender_puzzle_hash - is_recipient: bool = await metadata.is_recipient(self.puzzle_store) - if is_recipient: - derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(recipient_puzhash) - else: - derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(sender_puzhash) - assert derivation_record is not None - amount = uint64(amount + coin.amount) - # Remove the clawback hint since it is unnecessary for the XCH coin - memos: list[bytes] = [] if len(incoming_tx.memos) == 0 else next(iter(incoming_tx.memos.items()))[1][1:] - inner_puzzle: Program = self.main_wallet.puzzle_for_pk(derivation_record.pubkey) - inner_solution: Program = self.main_wallet.make_solution( - primaries=[ - CreateCoin( - derivation_record.puzzle_hash, - uint64(coin.amount), - memos, # Forward memo of the first coin - ) - ], - conditions=( - extra_conditions - if len(coin_spends) > 0 or fee == 0 - else (*extra_conditions, CreateCoinAnnouncement(message)) - ), - ) - coin_spend: CoinSpend = generate_clawback_spend_bundle(coin, metadata, inner_puzzle, inner_solution) - coin_spends.append(coin_spend) - # Update incoming tx to prevent double spend and mark it is pending - await self.tx_store.increment_sent(incoming_tx.name, "", MempoolInclusionStatus.PENDING, None) - except Exception as e: - self.log.error(f"Failed to create clawback spend bundle for {coin.name().hex()}: {e}") - if len(coin_spends) == 0: - return - spend_bundle = WalletSpendBundle(coin_spends, G2Element()) - if fee > 0: - async with self.new_action_scope(action_scope.config.tx_config, push=False) as inner_action_scope: - async with action_scope.use() as interface: - async with inner_action_scope.use() as inner_interface: - inner_interface.side_effects.selected_coins = interface.side_effects.selected_coins - await self.main_wallet.create_tandem_xch_tx( - fee, - inner_action_scope, - extra_conditions=( - AssertCoinAnnouncement(asserted_id=coin_spends[0].coin.name(), asserted_msg=message), - ), - ) - async with inner_action_scope.use() as inner_interface: - # This should not be looked to for best practice. - # Ideally, the two spend bundles can exist separately on each tx record until they are pushed. - # This is not very supported behavior at the moment - # so to avoid any potential backwards compatibility issues, - # we're moving the spend bundle from this TX to the main - interface.side_effects.transactions.extend( - [ - dataclasses.replace(tx, spend_bundle=None) - for tx in inner_interface.side_effects.transactions - ] - ) - interface.side_effects.selected_coins.extend(inner_interface.side_effects.selected_coins) - spend_bundle = WalletSpendBundle.aggregate( - [ - spend_bundle, - *( - tx.spend_bundle - for tx in inner_action_scope.side_effects.transactions - if tx.spend_bundle is not None - ), - ] - ) - assert derivation_record is not None - tx_record = TransactionRecord( - confirmed_at_height=uint32(0), - created_at_time=now, - to_puzzle_hash=derivation_record.puzzle_hash, - to_address=self.encode_puzzle_hash(derivation_record.puzzle_hash), - amount=amount, - fee_amount=uint64(fee), - confirmed=False, - sent=uint32(0), - spend_bundle=spend_bundle, - additions=spend_bundle.additions(), - removals=spend_bundle.removals(), - wallet_id=uint32(1), - sent_to=[], - trade_id=None, - type=uint32(TransactionType.OUTGOING_CLAWBACK), - name=spend_bundle.name(), - memos=compute_memos(spend_bundle), - valid_times=parse_timelock_info(extra_conditions), - ) - async with action_scope.use() as interface: - interface.side_effects.transactions.append(tx_record) - async def filter_spam(self, new_coin_state: list[CoinState]) -> list[CoinState]: xch_spam_amount = self.config.get("xch_spam_amount", 1000000)