from __future__ import annotations import asyncio import gc import logging import signal import sqlite3 from collections.abc import AsyncGenerator, AsyncIterator, Iterator from contextlib import asynccontextmanager, contextmanager from pathlib import Path from types import FrameType from typing import Any from chia_rs import ConsensusConstants from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import uint16 from chia.cmds.init_funcs import init from chia.consensus.constants import replace_str_to_bytes from chia.daemon.server import WebSocketServer, daemon_launch_lock_path from chia.farmer.farmer_service import FarmerService from chia.farmer.start_farmer import create_farmer_service from chia.full_node.full_node_service import FullNodeService from chia.full_node.start_full_node import create_full_node_service from chia.harvester.harvester_service import HarvesterService from chia.harvester.start_harvester import create_harvester_service from chia.introducer.introducer_service import IntroducerService from chia.introducer.start_introducer import create_introducer_service from chia.protocols.outbound_message import NodeType from chia.protocols.shared_protocol import Capability, default_capabilities from chia.seeder.crawler_service import CrawlerService from chia.seeder.dns_server import DNSServer, create_dns_server_service from chia.seeder.start_crawler import create_full_node_crawler_service from chia.server.resolve_peer_info import set_peer_info from chia.server.signal_handlers import SignalHandlers from chia.simulator.block_tools import BlockTools, test_constants from chia.simulator.keyring import TempKeyring from chia.simulator.ssl_certs import get_next_nodes_certs_and_keys, get_next_private_ca_cert_and_key from chia.simulator.start_simulator import SimulatorFullNodeService, create_full_node_simulator_service from chia.solver.solver_service import SolverService from chia.solver.start_solver import create_solver_service from chia.ssl.create_ssl import create_all_ssl from chia.timelord.start_timelord import create_timelord_service from chia.timelord.timelord_launcher import VDFClientProcessMgr, find_vdf_client, spawn_process from chia.timelord.timelord_service import TimelordService from chia.types.peer_info import UnresolvedPeerInfo from chia.util.bech32m import encode_puzzle_hash from chia.util.config import config_path_for_filename, load_config, lock_and_load_config, save_config from chia.util.db_wrapper import generate_in_memory_db_uri from chia.util.keychain import bytes_to_mnemonic from chia.util.lock import Lockfile from chia.util.task_referencer import create_referenced_task from chia.wallet.start_wallet import create_wallet_service from chia.wallet.wallet_service import WalletService log = logging.getLogger(__name__) @contextmanager def create_lock_and_load_config(certs_path: Path, root_path: Path) -> Iterator[dict[str, Any]]: init(None, root_path) init(certs_path, root_path) path = config_path_for_filename(root_path=root_path, filename="config.yaml") # Using localhost leads to flakiness on CI path.write_text(path.read_text().replace("localhost", "127.0.0.1")) with lock_and_load_config(root_path, "config.yaml") as config: yield config def get_capability_overrides(node_type: NodeType, disabled_capabilities: list[Capability]) -> list[tuple[uint16, str]]: return [ ( capability if Capability(int(capability[0])) not in disabled_capabilities or Capability(int(capability[0])) == Capability.BASE else (capability[0], "0") ) for capability in default_capabilities[node_type] ] @asynccontextmanager async def setup_daemon(btools: BlockTools) -> AsyncGenerator[WebSocketServer, None]: root_path = btools.root_path config = btools.config assert "daemon_port" in config crt_path = root_path / config["daemon_ssl"]["private_crt"] key_path = root_path / config["daemon_ssl"]["private_key"] ca_crt_path = root_path / config["private_ssl_ca"]["crt"] ca_key_path = root_path / config["private_ssl_ca"]["key"] with Lockfile.create(daemon_launch_lock_path(root_path)): ws_server = WebSocketServer(root_path, ca_crt_path, ca_key_path, crt_path, key_path) async with ws_server.run(): yield ws_server @asynccontextmanager async def setup_full_node( consensus_constants: ConsensusConstants, db_name: str, self_hostname: str, local_bt: BlockTools, introducer_port: int | None = None, simulator: bool = False, send_uncompact_interval: int = 0, sanitize_weight_proof_only: bool = False, connect_to_daemon: bool = False, db_version: int = 1, disable_capabilities: list[Capability] | None = None, *, reuse_db: bool = False, ) -> AsyncGenerator[FullNodeService | SimulatorFullNodeService, None]: if reuse_db: db_path: str | Path = local_bt.root_path / f"{db_name}" uri = False else: db_path = generate_in_memory_db_uri() uri = True if not reuse_db and db_version > 1: with sqlite3.connect(db_path, uri=uri) as connection: connection.execute("CREATE TABLE database_version(version int)") connection.execute("INSERT INTO database_version VALUES (?)", (db_version,)) connection.commit() if connect_to_daemon: assert local_bt.config["daemon_port"] is not None config = local_bt.config service_config = config["full_node"] service_config["database_path"] = db_name service_config["testing"] = True service_config["send_uncompact_interval"] = send_uncompact_interval service_config["target_uncompact_proofs"] = 30 service_config["peer_connect_interval"] = 50 service_config["sanitize_weight_proof_only"] = sanitize_weight_proof_only if introducer_port is not None: service_config["introducer_peer"]["host"] = self_hostname service_config["introducer_peer"]["port"] = introducer_port else: service_config["introducer_peer"] = None service_config["dns_servers"] = [] service_config["port"] = 0 service_config["rpc_port"] = 0 config["simulator"]["auto_farm"] = False # Disable Auto Farm for tests config["simulator"]["use_current_time"] = False # Disable Real timestamps when running tests overrides = service_config["network_overrides"]["constants"][service_config["selected_network"]] updated_constants = replace_str_to_bytes(consensus_constants, **overrides) local_bt.change_config(config) override_capabilities = ( None if disable_capabilities is None else get_capability_overrides(NodeType.FULL_NODE, disable_capabilities) ) service: FullNodeService | SimulatorFullNodeService if simulator: service = await create_full_node_simulator_service( local_bt.root_path, config, local_bt, connect_to_daemon=connect_to_daemon, override_capabilities=override_capabilities, ) else: service = await create_full_node_service( local_bt.root_path, config, updated_constants, connect_to_daemon=connect_to_daemon, override_capabilities=override_capabilities, ) async with service.manage(): yield service @asynccontextmanager async def setup_crawler( root_path_populated_with_config: Path, database_uri: str, start_crawler_loop: bool = True ) -> AsyncGenerator[CrawlerService, None]: create_all_ssl( root_path=root_path_populated_with_config, private_ca_crt_and_key=get_next_private_ca_cert_and_key().collateral.cert_and_key, node_certs_and_keys=get_next_nodes_certs_and_keys().collateral.certs_and_keys, ) config = load_config(root_path_populated_with_config, "config.yaml") service_config = config["seeder"] service_config["selected_network"] = "testnet0" service_config["port"] = 0 service_config["crawler"]["start_rpc_server"] = False service_config["other_peers_port"] = 58444 service_config["crawler_db_path"] = database_uri overrides = service_config["network_overrides"]["constants"][service_config["selected_network"]] updated_constants = replace_str_to_bytes(test_constants, **overrides) service = create_full_node_crawler_service( root_path_populated_with_config, config, updated_constants, connect_to_daemon=False, start_crawler_loop=start_crawler_loop, ) async with service.manage(): if not service_config["crawler"]["start_rpc_server"]: # otherwise the loops don't work. service._node.state_changed_callback = lambda x, y: None yield service @asynccontextmanager async def setup_seeder(root_path_populated_with_config: Path, database_uri: str) -> AsyncGenerator[DNSServer, None]: config = load_config(root_path_populated_with_config, "config.yaml") service_config = config["seeder"] service_config["selected_network"] = "testnet0" if service_config["domain_name"].endswith("."): # remove the trailing . so that we can test that logic. service_config["domain_name"] = service_config["domain_name"][:-1] service_config["dns_port"] = 0 service_config["crawler_db_path"] = database_uri service = create_dns_server_service( config, root_path_populated_with_config, ) async with service.run(): yield service # Note: convert these setup functions to fixtures, or push it one layer up, # keeping these usable independently? @asynccontextmanager async def setup_wallet_node( self_hostname: str, consensus_constants: ConsensusConstants, local_bt: BlockTools, spam_filter_after_n_txs: int | None = 200, xch_spam_amount: int = 1000000, full_node_port: uint16 | None = None, introducer_port: uint16 | None = None, key_seed: bytes | None = None, initial_num_public_keys: int = 5, ) -> AsyncGenerator[WalletService, None]: with TempKeyring(populate=True) as keychain: config = local_bt.config service_config = config["wallet"] service_config["testing"] = True service_config["port"] = 0 service_config["rpc_port"] = 0 service_config["initial_num_public_keys"] = initial_num_public_keys service_config["spam_filter_after_n_txs"] = spam_filter_after_n_txs service_config["xch_spam_amount"] = xch_spam_amount entropy = bytes32.secret() if key_seed is None: key_seed = entropy keychain.add_key(bytes_to_mnemonic(key_seed)) first_pk = keychain.get_first_public_key() assert first_pk is not None db_path_key_suffix = str(first_pk.get_fingerprint()) db_name = f"test-wallet-db-{full_node_port}-KEY.sqlite" db_path_replaced: str = db_name.replace("KEY", db_path_key_suffix) db_path = local_bt.root_path / db_path_replaced if db_path.exists(): # TODO: remove (maybe) when fixed https://github.com/python/cpython/issues/97641 gc.collect() db_path.unlink() service_config["database_path"] = str(db_name) service_config["testing"] = True service_config["introducer_peer"]["host"] = self_hostname if introducer_port is not None: service_config["introducer_peer"]["port"] = introducer_port service_config["peer_connect_interval"] = 10 else: service_config["introducer_peer"] = None if full_node_port is not None: service_config.pop("full_node_peer", None) service_config["full_node_peers"] = [ { "host": self_hostname, "port": full_node_port, }, ] else: service_config.pop("full_node_peer", None) service_config.pop("full_node_peers", None) service = create_wallet_service( local_bt.root_path, config, consensus_constants, keychain, connect_to_daemon=False, ) try: async with service.manage(): yield service finally: if db_path.exists(): # TODO: remove (maybe) when fixed https://github.com/python/cpython/issues/97641 # 3.11 switched to using functools.lru_cache for the statement cache. # See #87028. This introduces a reference cycle involving the connection # object, so the connection object no longer gets immediately # deallocated, not until, for example, gc.collect() is called to break # the cycle. gc.collect() for _ in range(10): try: db_path.unlink() break except PermissionError as e: print(f"db_path.unlink(): {e}") await asyncio.sleep(0.1) # filesystem operations are async on windows # [WinError 32] The process cannot access the file because it is # being used by another process keychain.delete_all_keys() @asynccontextmanager async def setup_harvester( b_tools: BlockTools, root_path: Path, farmer_peer: UnresolvedPeerInfo | None, consensus_constants: ConsensusConstants, start_service: bool = True, ) -> AsyncGenerator[HarvesterService, None]: with create_lock_and_load_config(b_tools.root_path / "config" / "ssl" / "ca", root_path) as config: config["logging"]["log_stdout"] = True config["selected_network"] = "testnet0" config["harvester"]["selected_network"] = "testnet0" config["harvester"]["port"] = 0 config["harvester"]["rpc_port"] = 0 config["harvester"]["plot_directories"] = [str(b_tools.plot_dir.resolve())] # CI doesn't like GPU compressed farming config["harvester"]["parallel_decompressor_count"] = 0 save_config(root_path, "config.yaml", config) service = create_harvester_service( root_path, config, consensus_constants, farmer_peers={farmer_peer} if farmer_peer is not None else set(), connect_to_daemon=False, ) async with service.manage(start=start_service): yield service @asynccontextmanager async def setup_farmer( b_tools: BlockTools, root_path: Path, self_hostname: str, consensus_constants: ConsensusConstants, full_node_port: uint16 | None = None, start_service: bool = True, port: uint16 = uint16(0), solver_peer: UnresolvedPeerInfo | None = None, ) -> AsyncGenerator[FarmerService, None]: with create_lock_and_load_config(b_tools.root_path / "config" / "ssl" / "ca", root_path) as root_config: root_config["logging"]["log_stdout"] = True root_config["selected_network"] = "testnet0" root_config["farmer"]["selected_network"] = "testnet0" save_config(root_path, "config.yaml", root_config) service_config = root_config["farmer"] config_pool = root_config["pool"] service_config["xch_target_address"] = encode_puzzle_hash(b_tools.farmer_ph, "xch") service_config["pool_public_keys"] = [bytes(pk).hex() for pk in b_tools.pool_pubkeys] service_config["port"] = port service_config["rpc_port"] = uint16(0) config_pool["xch_target_address"] = encode_puzzle_hash(b_tools.pool_ph, "xch") if full_node_port: service_config.pop("full_node_peer", None) service_config["full_node_peers"] = [ { "host": self_hostname, "port": full_node_port, }, ] else: service_config.pop("full_node_peer", None) service_config.pop("full_node_peers", None) if solver_peer: service_config["solver_peers"] = [ { "host": solver_peer.host, "port": solver_peer.port, }, ] else: service_config.pop("solver_peers", None) service = create_farmer_service( root_path, root_config, config_pool, consensus_constants, b_tools.local_keychain, connect_to_daemon=False, solver_peer=solver_peer, ) async with service.manage(start=start_service): yield service @asynccontextmanager async def setup_introducer(bt: BlockTools, port: int) -> AsyncGenerator[IntroducerService, None]: service = create_introducer_service( bt.root_path, bt.config, advertised_port=port, connect_to_daemon=False, ) async with service.manage(): yield service @asynccontextmanager async def setup_vdf_client(bt: BlockTools, self_hostname: str, port: int) -> AsyncIterator[None]: find_vdf_client() # raises FileNotFoundError if not found process_mgr = VDFClientProcessMgr() vdf_task_1 = create_referenced_task( spawn_process(self_hostname, port, 1, process_mgr, prefer_ipv6=bt.config.get("prefer_ipv6", False)), name="vdf_client_1", ) async def stop( signal_: signal.Signals, stack_frame: FrameType | None, loop: asyncio.AbstractEventLoop, ) -> None: await process_mgr.kill_processes() async with SignalHandlers.manage() as signal_handlers: signal_handlers.setup_async_signal_handler(handler=stop) try: yield finally: await process_mgr.kill_processes() vdf_task_1.cancel() try: await vdf_task_1 except (Exception, asyncio.CancelledError): pass @asynccontextmanager async def setup_vdf_clients(bt: BlockTools, self_hostname: str, port: int) -> AsyncIterator[None]: find_vdf_client() # raises FileNotFoundError if not found process_mgr = VDFClientProcessMgr() tasks = [] prefer_ipv6 = bt.config.get("prefer_ipv6", False) for i in range(1, 4): tasks.append( create_referenced_task( spawn_process( host=self_hostname, port=port, counter=i, process_mgr=process_mgr, prefer_ipv6=prefer_ipv6 ), name=f"vdf_client_{i}", ) ) async def stop( signal_: signal.Signals, stack_frame: FrameType | None, loop: asyncio.AbstractEventLoop, ) -> None: await process_mgr.kill_processes() signal_handlers = SignalHandlers() async with signal_handlers.manage(): signal_handlers.setup_async_signal_handler(handler=stop) try: yield finally: await process_mgr.kill_processes() for task in tasks: task.cancel() try: await task except (Exception, asyncio.CancelledError): pass @asynccontextmanager async def setup_timelord( full_node_port: int, sanitizer: bool, consensus_constants: ConsensusConstants, config: dict[str, Any], root_path: Path, vdf_port: uint16 = uint16(0), ) -> AsyncGenerator[TimelordService, None]: service_config = config["timelord"] set_peer_info(service_config, peer_type=NodeType.FULL_NODE, peer_port=full_node_port) service_config["bluebox_mode"] = sanitizer service_config["fast_algorithm"] = False service_config["vdf_server"]["port"] = vdf_port service_config["start_rpc_server"] = True service_config["rpc_port"] = uint16(0) service = create_timelord_service( root_path, config, consensus_constants, connect_to_daemon=False, ) async with service.manage(): yield service @asynccontextmanager async def setup_solver( root_path: Path, b_tools: BlockTools, consensus_constants: ConsensusConstants, start_service: bool = True, ) -> AsyncGenerator[SolverService, None]: with create_lock_and_load_config(b_tools.root_path / "config" / "ssl" / "ca", root_path) as config: config["logging"]["log_stdout"] = True config["solver"]["enable_upnp"] = True config["solver"]["selected_network"] = "testnet0" config["solver"]["port"] = 0 config["solver"]["rpc_port"] = 0 config["solver"]["num_threads"] = 1 save_config(root_path, "config.yaml", config) service = create_solver_service( root_path, config, consensus_constants, ) async with service.manage(start=start_service): yield service