chia-blockchain 2.5.0rc2__py3-none-any.whl → 2.5.1rc2__py3-none-any.whl

chia/server/node_discovery.py

@@ -9,7 +9,7 @@ from logging import Logger
 from pathlib import Path
 from random import Random
 from secrets import randbits
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Optional
 
 import dns.asyncresolver
 
@@ -24,7 +24,10 @@ from chia.server.ws_connection import WSChiaConnection
 from chia.types.peer_info import PeerInfo, TimestampedPeerInfo, UnresolvedPeerInfo
 from chia.util.hash import std_hash
 from chia.util.ints import uint16, uint64
-from chia.util.network import IPAddress, resolve
+from chia.util.ip_address import IPAddress
+from chia.util.network import resolve
+from chia.util.safe_cancel_task import cancel_task_safe
+from chia.util.task_referencer import create_referenced_task
 
 MAX_PEERS_RECEIVED_PER_REQUEST = 1000
 MAX_TOTAL_PEERS_RECEIVED = 3000
@@ -46,8 +49,8 @@ class FullNodeDiscovery:
         server: ChiaServer,
         target_outbound_count: int,
         peers_file_path: Path,
-        introducer_info: Optional[Dict[str, Any]],
-        dns_servers: List[str],
+        introducer_info: Optional[dict[str, Any]],
+        dns_servers: list[str],
         peer_connect_interval: int,
         selected_network: str,
         default_port: Optional[int],
@@ -68,10 +71,10 @@ class FullNodeDiscovery:
             self.enable_private_networks = False
         self.peer_connect_interval = peer_connect_interval
         self.log = log
-        self.relay_queue: Optional[asyncio.Queue[Tuple[TimestampedPeerInfo, int]]] = None
+        self.relay_queue: Optional[asyncio.Queue[tuple[TimestampedPeerInfo, int]]] = None
         self.address_manager: Optional[AddressManager] = None
-        self.connection_time_pretest: Dict[str, Any] = {}
-        self.received_count_from_peers: Dict[str, Any] = {}
+        self.connection_time_pretest: dict[str, Any] = {}
+        self.received_count_from_peers: dict[str, Any] = {}
         self.lock = asyncio.Lock()
         self.connect_peers_task: Optional[asyncio.Task[None]] = None
         self.serialize_task: Optional[asyncio.Task[None]] = None
@@ -82,8 +85,8 @@ class FullNodeDiscovery:
         except Exception:
             self.resolver = None
             self.log.exception("Error initializing asyncresolver")
-        self.pending_outbound_connections: Set[str] = set()
-        self.pending_tasks: Set[asyncio.Task[None]] = set()
+        self.pending_outbound_connections: set[str] = set()
+        self.pending_tasks: set[asyncio.Task[None]] = set()
         self.default_port: Optional[int] = default_port
         if default_port is None and selected_network in NETWORK_ID_DEFAULT_PORTS:
             self.default_port = NETWORK_ID_DEFAULT_PORTS[selected_network]
@@ -96,27 +99,20 @@ class FullNodeDiscovery:
 
     async def start_tasks(self) -> None:
         random = Random()
-        self.connect_peers_task = asyncio.create_task(self._connect_to_peers(random))
-        self.serialize_task = asyncio.create_task(self._periodically_serialize(random))
-        self.cleanup_task = asyncio.create_task(self._periodically_cleanup())
+        self.connect_peers_task = create_referenced_task(self._connect_to_peers(random))
+        self.serialize_task = create_referenced_task(self._periodically_serialize(random))
+        self.cleanup_task = create_referenced_task(self._periodically_cleanup())
 
     async def _close_common(self) -> None:
         self.is_closed = True
-        self.cancel_task_safe(self.connect_peers_task)
-        self.cancel_task_safe(self.serialize_task)
-        self.cancel_task_safe(self.cleanup_task)
+        cancel_task_safe(self.connect_peers_task, self.log)
+        cancel_task_safe(self.serialize_task, self.log)
+        cancel_task_safe(self.cleanup_task, self.log)
         for t in self.pending_tasks:
-            self.cancel_task_safe(t)
+            cancel_task_safe(t, self.log)
         if len(self.pending_tasks) > 0:
             await asyncio.wait(self.pending_tasks)
 
-    def cancel_task_safe(self, task: Optional[asyncio.Task[None]]) -> None:
-        if task is not None:
-            try:
-                task.cancel()
-            except Exception as e:
-                self.log.error(f"Error while canceling task.{e} {task}")
-
     async def on_connect(self, peer: WSChiaConnection) -> None:
         if (
             peer.is_outbound is False
@@ -201,7 +197,7 @@ class FullNodeDiscovery:
                 self.log.warning("Skipping DNS query: asyncresolver not initialized.")
                 return
             for rdtype in ["A", "AAAA"]:
-                peers: List[TimestampedPeerInfo] = []
+                peers: list[TimestampedPeerInfo] = []
                 result = await self.resolver.resolve(qname=dns_address, rdtype=rdtype, lifetime=30)
                 for ip in result:
                     peers.append(
@@ -402,7 +398,7 @@ class FullNodeDiscovery:
                         self.log.debug("Max concurrent outbound connections reached. waiting")
                         await asyncio.wait(self.pending_tasks, return_when=asyncio.FIRST_COMPLETED)
                     self.pending_tasks.add(
-                        asyncio.create_task(self.start_client_async(addr, disconnect_after_handshake))
+                        create_referenced_task(self.start_client_async(addr, disconnect_after_handshake))
                     )
 
                 await asyncio.sleep(connect_peer_interval)
@@ -448,7 +444,7 @@ class FullNodeDiscovery:
         return port in NETWORK_ID_DEFAULT_PORTS.values() and port != self.default_port
 
     async def _add_peers_common(
-        self, peer_list: List[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
+        self, peer_list: list[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
     ) -> None:
         # Check if we got the peers from a full node or from the introducer.
         peers_adjusted_timestamp = []
@@ -503,8 +499,8 @@ class FullNodePeers(FullNodeDiscovery):
         server: ChiaServer,
         target_outbound_count: int,
         peers_file_path: Path,
-        introducer_info: Dict[str, Any],
-        dns_servers: List[str],
+        introducer_info: dict[str, Any],
+        dns_servers: list[str],
         peer_connect_interval: int,
         selected_network: str,
         default_port: Optional[int],
@@ -522,19 +518,19 @@ class FullNodePeers(FullNodeDiscovery):
             log,
         )
         self.relay_queue = asyncio.Queue()
-        self.neighbour_known_peers: Dict[PeerInfo, Set[str]] = {}
+        self.neighbour_known_peers: dict[PeerInfo, set[str]] = {}
         self.key = randbits(256)
 
     async def start(self) -> None:
         await self.initialize_address_manager()
-        self.self_advertise_task = asyncio.create_task(self._periodically_self_advertise_and_clean_data())
-        self.address_relay_task = asyncio.create_task(self._address_relay())
+        self.self_advertise_task = create_referenced_task(self._periodically_self_advertise_and_clean_data())
+        self.address_relay_task = create_referenced_task(self._address_relay())
         await self.start_tasks()
 
     async def close(self) -> None:
         await self._close_common()
-        self.cancel_task_safe(self.self_advertise_task)
-        self.cancel_task_safe(self.address_relay_task)
+        cancel_task_safe(self.self_advertise_task, self.log)
+        cancel_task_safe(self.address_relay_task, self.log)
 
     async def _periodically_self_advertise_and_clean_data(self) -> None:
         while not self.is_closed:
@@ -571,7 +567,7 @@ class FullNodePeers(FullNodeDiscovery):
                 self.log.error(f"Exception in self advertise: {e}")
                 self.log.error(f"Traceback: {traceback.format_exc()}")
 
-    async def add_peers_neighbour(self, peers: List[TimestampedPeerInfo], neighbour_info: PeerInfo) -> None:
+    async def add_peers_neighbour(self, peers: list[TimestampedPeerInfo], neighbour_info: PeerInfo) -> None:
         async with self.lock:
             for peer in peers:
                 if neighbour_info not in self.neighbour_known_peers:
@@ -604,7 +600,7 @@ class FullNodePeers(FullNodeDiscovery):
             return None
 
     async def add_peers(
-        self, peer_list: List[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
+        self, peer_list: list[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
     ) -> None:
         try:
             await self._add_peers_common(peer_list, peer_src, is_full_node)
@@ -683,8 +679,8 @@ class WalletPeers(FullNodeDiscovery):
         server: ChiaServer,
         target_outbound_count: int,
         peers_file_path: Path,
-        introducer_info: Dict[str, Any],
-        dns_servers: List[str],
+        introducer_info: dict[str, Any],
+        dns_servers: list[str],
         peer_connect_interval: int,
         selected_network: str,
         default_port: Optional[int],
@@ -713,6 +709,6 @@ class WalletPeers(FullNodeDiscovery):
         await self._close_common()
 
     async def add_peers(
-        self, peer_list: List[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
+        self, peer_list: list[TimestampedPeerInfo], peer_src: Optional[PeerInfo], is_full_node: bool
     ) -> None:
         await self._add_peers_common(peer_list, peer_src, is_full_node)

chia/server/rate_limit_numbers.py

@@ -4,14 +4,16 @@ from __future__ import annotations
 import copy
 import dataclasses
 import functools
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from chia.protocols.protocol_message_types import ProtocolMessageTypes
 from chia.protocols.shared_protocol import Capability
 
-compose_rate_limits_cache: Dict[int, Dict[str, Any]] = {}
+compose_rate_limits_cache: dict[int, dict[str, Any]] = {}
 
 
+# this class is used to configure the *rate* limit for a message type. The
+# limits are counts and size per 60 seconds.
 @dataclasses.dataclass(frozen=True)
 class RLSettings:
     frequency: int  # Max request per time period (ie 1 min)
@@ -19,7 +21,15 @@ class RLSettings:
     max_total_size: Optional[int] = None  # Max cumulative size of all requests in that period
 
 
-def get_rate_limits_to_use(our_capabilities: List[Capability], peer_capabilities: List[Capability]) -> Dict[str, Any]:
+# this class is used to indicate that a message type is not subject to a rate
+# limit, but just a per-message size limit. This may be appropriate for response
+# messages that are implicitly limited by their corresponding request message
+@dataclasses.dataclass(frozen=True)
+class Unlimited:
+    max_size: int  # Max size of each request
+
+
+def get_rate_limits_to_use(our_capabilities: list[Capability], peer_capabilities: list[Capability]) -> dict[str, Any]:
     # This will use the newest possible rate limits that both peers support. At this time there are only two
     # options, v1 and v2.
 
@@ -35,14 +45,14 @@ def get_rate_limits_to_use(our_capabilities: List[Capability], peer_capabilities
         return rate_limits[1]
 
 
-def compose_rate_limits(old_rate_limits: Dict[str, Any], new_rate_limits: Dict[str, Any]) -> Dict[str, Any]:
+def compose_rate_limits(old_rate_limits: dict[str, Any], new_rate_limits: dict[str, Any]) -> dict[str, Any]:
     # Composes two rate limits dicts, so that the newer values override the older values
-    final_rate_limits: Dict[str, Any] = copy.deepcopy(new_rate_limits)
-    categories: List[str] = ["rate_limits_tx", "rate_limits_other"]
-    all_new_msgs_lists: List[List[ProtocolMessageTypes]] = [
+    final_rate_limits: dict[str, Any] = copy.deepcopy(new_rate_limits)
+    categories: list[str] = ["rate_limits_tx", "rate_limits_other"]
+    all_new_msgs_lists: list[list[ProtocolMessageTypes]] = [
         list(new_rate_limits[category].keys()) for category in categories
     ]
-    all_new_msgs: List[ProtocolMessageTypes] = functools.reduce(lambda a, b: a + b, all_new_msgs_lists)
+    all_new_msgs: list[ProtocolMessageTypes] = functools.reduce(lambda a, b: a + b, all_new_msgs_lists)
     for old_cat, mapping in old_rate_limits.items():
         if old_cat in categories:
             for old_protocol_msg, old_rate_limit_value in mapping.items():
@@ -94,11 +104,11 @@ rate_limits = {
             ProtocolMessageTypes.request_proof_of_weight: RLSettings(5, 100),
             ProtocolMessageTypes.respond_proof_of_weight: RLSettings(5, 50 * 1024 * 1024, 100 * 1024 * 1024),
             ProtocolMessageTypes.request_block: RLSettings(200, 100),
-            ProtocolMessageTypes.reject_block: RLSettings(200, 100),
+            ProtocolMessageTypes.reject_block: Unlimited(100),
             ProtocolMessageTypes.request_blocks: RLSettings(500, 100),
-            ProtocolMessageTypes.respond_blocks: RLSettings(100, 50 * 1024 * 1024, 5 * 50 * 1024 * 1024),
-            ProtocolMessageTypes.reject_blocks: RLSettings(100, 100),
-            ProtocolMessageTypes.respond_block: RLSettings(200, 2 * 1024 * 1024, 10 * 2 * 1024 * 1024),
+            ProtocolMessageTypes.respond_blocks: Unlimited(50 * 1024 * 1024),
+            ProtocolMessageTypes.reject_blocks: Unlimited(100),
+            ProtocolMessageTypes.respond_block: Unlimited(2 * 1024 * 1024),
             ProtocolMessageTypes.new_unfinished_block: RLSettings(200, 100),
             ProtocolMessageTypes.request_unfinished_block: RLSettings(200, 100),
             ProtocolMessageTypes.new_unfinished_block2: RLSettings(200, 100),
@@ -144,10 +154,10 @@ rate_limits = {
             ProtocolMessageTypes.plot_sync_done: RLSettings(1000, 100 * 1024 * 1024),
             ProtocolMessageTypes.plot_sync_response: RLSettings(3000, 100 * 1024 * 1024),
             ProtocolMessageTypes.coin_state_update: RLSettings(1000, 100 * 1024 * 1024),
-            ProtocolMessageTypes.register_interest_in_puzzle_hash: RLSettings(1000, 100 * 1024 * 1024),
-            ProtocolMessageTypes.respond_to_ph_update: RLSettings(1000, 100 * 1024 * 1024),
-            ProtocolMessageTypes.register_interest_in_coin: RLSettings(1000, 100 * 1024 * 1024),
-            ProtocolMessageTypes.respond_to_coin_update: RLSettings(1000, 100 * 1024 * 1024),
+            ProtocolMessageTypes.register_for_ph_updates: RLSettings(1000, 100 * 1024 * 1024),
+            ProtocolMessageTypes.respond_to_ph_updates: RLSettings(1000, 100 * 1024 * 1024),
+            ProtocolMessageTypes.register_for_coin_updates: RLSettings(1000, 100 * 1024 * 1024),
+            ProtocolMessageTypes.respond_to_coin_updates: RLSettings(1000, 100 * 1024 * 1024),
             ProtocolMessageTypes.request_remove_puzzle_subscriptions: RLSettings(1000, 100 * 1024 * 1024),
             ProtocolMessageTypes.respond_remove_puzzle_subscriptions: RLSettings(1000, 100 * 1024 * 1024),
             ProtocolMessageTypes.request_remove_coin_subscriptions: RLSettings(1000, 100 * 1024 * 1024),

chia/server/rate_limits.py

@@ -4,12 +4,12 @@ import dataclasses
 import logging
 import time
 from collections import Counter
-from typing import List
+from typing import Optional
 
 from chia.protocols.protocol_message_types import ProtocolMessageTypes
 from chia.protocols.shared_protocol import Capability
 from chia.server.outbound_message import Message
-from chia.server.rate_limit_numbers import RLSettings, get_rate_limits_to_use
+from chia.server.rate_limit_numbers import RLSettings, Unlimited, get_rate_limits_to_use
 
 log = logging.getLogger(__name__)
 
@@ -43,10 +43,12 @@ class RateLimiter:
         self.non_tx_cumulative_size = 0
 
     def process_msg_and_check(
-        self, message: Message, our_capabilities: List[Capability], peer_capabilities: List[Capability]
-    ) -> bool:
+        self, message: Message, our_capabilities: list[Capability], peer_capabilities: list[Capability]
+    ) -> Optional[str]:
         """
-        Returns True if message can be processed successfully, false if a rate limit is passed.
+        Returns a string indicating which limit was hit if a rate limit is
+        exceeded, and the message should be blocked. Returns None if the limit was not
+        hit and the message is good to be sent or received.
         """
 
         current_minute = int(time.time() // self.reset_seconds)
@@ -60,7 +62,7 @@ class RateLimiter:
             message_type = ProtocolMessageTypes(message.type)
         except Exception as e:
             log.warning(f"Invalid message: {message.type}, {e}")
-            return True
+            return None
 
         new_message_counts: int = self.message_counts[message_type] + 1
         new_cumulative_size: int = self.message_cumulative_sizes[message_type] + len(message.data)
@@ -77,30 +79,68 @@ class RateLimiter:
                 limits = rate_limits["rate_limits_tx"][message_type]
             elif message_type in rate_limits["rate_limits_other"]:
                 limits = rate_limits["rate_limits_other"][message_type]
-                non_tx_freq = rate_limits["non_tx_freq"]
-                non_tx_max_total_size = rate_limits["non_tx_max_total_size"]
-                new_non_tx_count = self.non_tx_message_counts + 1
-                new_non_tx_size = self.non_tx_cumulative_size + len(message.data)
-                if new_non_tx_count > non_tx_freq * proportion_of_limit:
-                    return False
-                if new_non_tx_size > non_tx_max_total_size * proportion_of_limit:
-                    return False
-            else:
-                log.warning(f"Message type {message_type} not found in rate limits")
+                if isinstance(limits, RLSettings):
+                    non_tx_freq = rate_limits["non_tx_freq"]
+                    non_tx_max_total_size = rate_limits["non_tx_max_total_size"]
+                    new_non_tx_count = self.non_tx_message_counts + 1
+                    new_non_tx_size = self.non_tx_cumulative_size + len(message.data)
+                    if new_non_tx_count > non_tx_freq * proportion_of_limit:
+                        return " ".join(
+                            [
+                                f"non-tx count: {new_non_tx_count}",
+                                f"> {non_tx_freq * proportion_of_limit}",
+                                f"(scale factor: {proportion_of_limit})",
+                            ]
+                        )
+                    if new_non_tx_size > non_tx_max_total_size * proportion_of_limit:
+                        return " ".join(
+                            [
+                                f"non-tx size: {new_non_tx_size}",
+                                f"> {non_tx_max_total_size * proportion_of_limit}",
+                                f"(scale factor: {proportion_of_limit})",
+                            ]
+                        )
+            else:  # pragma: no cover
+                log.warning(
+                    f"Message type {message_type} not found in rate limits (scale factor: {proportion_of_limit})",
+                )
 
-            if limits.max_total_size is None:
-                limits = dataclasses.replace(limits, max_total_size=limits.frequency * limits.max_size)
-            assert limits.max_total_size is not None
+            if isinstance(limits, Unlimited):
+                # this message type is not rate limited. This is used for
+                # response messages and must be combined with banning peers
+                # sending unsolicited responses of this type
+                if len(message.data) > limits.max_size:
+                    return f"message size: {len(message.data)} > {limits.max_size}"
+                ret = True
+                return None
+            elif isinstance(limits, RLSettings):
+                if limits.max_total_size is None:
+                    limits = dataclasses.replace(limits, max_total_size=limits.frequency * limits.max_size)
+                assert limits.max_total_size is not None
 
-            if new_message_counts > limits.frequency * proportion_of_limit:
-                return False
-            if len(message.data) > limits.max_size:
-                return False
-            if new_cumulative_size > limits.max_total_size * proportion_of_limit:
-                return False
+                if new_message_counts > limits.frequency * proportion_of_limit:
+                    return " ".join(
+                        [
+                            f"message count: {new_message_counts}"
+                            f"> {limits.frequency * proportion_of_limit}"
+                            f"(scale factor: {proportion_of_limit})"
+                        ]
+                    )
+                if len(message.data) > limits.max_size:
+                    return f"message size: {len(message.data)} > {limits.max_size}"
+                if new_cumulative_size > limits.max_total_size * proportion_of_limit:
+                    return " ".join(
+                        [
+                            f"cumulative size: {new_cumulative_size}",
+                            f"> {limits.max_total_size * proportion_of_limit}",
+                            f"(scale factor: {proportion_of_limit})",
+                        ]
+                    )
 
-            ret = True
-            return True
+                ret = True
+                return None
+            else:  # pragma: no cover
+                return f"Internal Error, unknown rate limit for message: {message_type}, limit: {limits}"
         finally:
             if self.incoming or ret:
                 # now that we determined that it's OK to send the message, commit the

chia/server/server.py

@@ -5,10 +5,11 @@ import logging
 import ssl
 import time
 import traceback
+from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from ipaddress import IPv4Network, IPv6Network, ip_network
 from pathlib import Path
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Callable, Optional, Union, cast
 
 from aiohttp import (
     ClientResponseError,
@@ -39,6 +40,7 @@ from chia.util.ints import uint16
 from chia.util.network import WebServer, is_in_network, is_localhost, is_trusted_peer
 from chia.util.ssl_check import verify_ssl_certs_and_keys
 from chia.util.streamable import Streamable
+from chia.util.task_referencer import create_referenced_task
 
 max_message_size = 50 * 1024 * 1024  # 50MB
 
@@ -116,7 +118,7 @@ def calculate_node_id(cert_path: Path) -> bytes32:
 class ChiaServer:
     _port: Optional[int]
     _local_type: NodeType
-    _local_capabilities_for_handshake: List[Tuple[uint16, str]]
+    _local_capabilities_for_handshake: list[tuple[uint16, str]]
     _ping_interval: int
     _network_id: str
     _inbound_rate_limit_percent: int
@@ -124,13 +126,14 @@ class ChiaServer:
     api: ApiProtocol
     node: Any
     root_path: Path
-    config: Dict[str, Any]
+    config: dict[str, Any]
     log: logging.Logger
     ssl_context: ssl.SSLContext
     ssl_client_context: ssl.SSLContext
     node_id: bytes32
-    exempt_peer_networks: List[Union[IPv4Network, IPv6Network]]
-    all_connections: Dict[bytes32, WSChiaConnection] = field(default_factory=dict)
+    exempt_peer_networks: list[Union[IPv4Network, IPv6Network]]
+    class_for_type: dict[NodeType, type[ApiProtocol]]
+    all_connections: dict[bytes32, WSChiaConnection] = field(default_factory=dict)
     on_connect: Optional[ConnectionCallback] = None
     shut_down_event: asyncio.Event = field(default_factory=asyncio.Event)
     introducer_peers: Optional[IntroducerPeers] = None
@@ -138,7 +141,7 @@ class ChiaServer:
     webserver: Optional[WebServer] = None
     connection_close_task: Optional[asyncio.Task[None]] = None
     received_message_callback: Optional[ConnectionCallback] = None
-    banned_peers: Dict[str, float] = field(default_factory=dict)
+    banned_peers: dict[str, float] = field(default_factory=dict)
     invalid_protocol_ban_seconds = INVALID_PROTOCOL_BAN_SECONDS
 
     @classmethod
@@ -152,11 +155,12 @@ class ChiaServer:
         network_id: str,
         inbound_rate_limit_percent: int,
         outbound_rate_limit_percent: int,
-        capabilities: List[Tuple[uint16, str]],
+        capabilities: list[tuple[uint16, str]],
         root_path: Path,
-        config: Dict[str, Any],
-        private_ca_crt_key: Tuple[Path, Path],
-        chia_ca_crt_key: Tuple[Path, Path],
+        config: dict[str, Any],
+        private_ca_crt_key: tuple[Path, Path],
+        chia_ca_crt_key: tuple[Path, Path],
+        class_for_type: dict[NodeType, type[ApiProtocol]],
         name: str = __name__,
     ) -> ChiaServer:
         log = logging.getLogger(name)
@@ -232,6 +236,7 @@ class ChiaServer:
             node_id=calculate_node_id(node_id_cert_path),
             exempt_peer_networks=[ip_network(net, strict=False) for net in config.get("exempt_peer_networks", [])],
             introducer_peers=IntroducerPeers() if local_type is NodeType.INTRODUCER else None,
+            class_for_type=class_for_type,
         )
 
     def set_received_message_callback(self, callback: ConnectionCallback) -> None:
@@ -245,12 +250,12 @@ class ChiaServer:
         is_crawler = getattr(self.node, "crawl", None)
         while True:
             await asyncio.sleep(600 if is_crawler is None else 2)
-            to_remove: List[WSChiaConnection] = []
+            to_remove: list[WSChiaConnection] = []
             for connection in self.all_connections.values():
                 if connection.closed:
                     to_remove.append(connection)
                 elif (
-                    self._local_type == NodeType.FULL_NODE or self._local_type == NodeType.WALLET
+                    self._local_type in {NodeType.FULL_NODE, NodeType.WALLET}
                 ) and connection.connection_type == NodeType.FULL_NODE:
                     if is_crawler is not None:
                         if time.time() - connection.creation_time > 5:
@@ -281,7 +286,7 @@ class ChiaServer:
         if self.webserver is not None:
             raise RuntimeError("ChiaServer already started")
         if self.gc_task is None:
-            self.gc_task = asyncio.create_task(self.garbage_collect_connections_task())
+            self.gc_task = create_referenced_task(self.garbage_collect_connections_task())
 
         if self._port is not None:
             self.on_connect = on_connect
@@ -332,6 +337,7 @@ class ChiaServer:
                 inbound_rate_limit_percent=self._inbound_rate_limit_percent,
                 outbound_rate_limit_percent=self._outbound_rate_limit_percent,
                 local_capabilities_for_handshake=self._local_capabilities_for_handshake,
+                class_for_type=self.class_for_type,
             )
             await connection.perform_handshake(self._network_id, self.get_port(), self._local_type)
             assert connection.connection_type is not None, "handshake failed to set connection type, still None"
@@ -482,6 +488,7 @@ class ChiaServer:
                 inbound_rate_limit_percent=self._inbound_rate_limit_percent,
                 outbound_rate_limit_percent=self._outbound_rate_limit_percent,
                 local_capabilities_for_handshake=self._local_capabilities_for_handshake,
+                class_for_type=self.class_for_type,
                 session=session,
             )
             await connection.perform_handshake(self._network_id, server_port, self._local_type)
@@ -492,29 +499,42 @@ class ChiaServer:
             connection_type_str = ""
             if connection.connection_type is not None:
                 connection_type_str = connection.connection_type.name.lower()
-            self.log.info(f"Connected with {connection_type_str} {target_node}")
-            if is_feeler:
-                asyncio.create_task(connection.close())
+            if not is_feeler:
+                self.log.info(f"Connected with {connection_type_str} {target_node}")
+            else:
+                self.log.debug(f"Successful feeler connection with {connection_type_str} {target_node}")
+                create_referenced_task(connection.close(), known_unreferenced=True)
             return True
         except client_exceptions.ClientConnectorError as e:
-            self.log.info(f"{e}")
+            if is_feeler:
+                self.log.debug(f"Feeler connection error. {e}")
+            else:
+                self.log.info(f"{e}")
         except ProtocolError as e:
             if connection is not None:
                 await connection.close(self.invalid_protocol_ban_seconds, WSCloseCode.PROTOCOL_ERROR, e.code)
             if e.code == Err.INVALID_HANDSHAKE:
-                self.log.warning(f"Invalid handshake with peer {target_node}. Maybe the peer is running old software.")
+                self.log.warning(
+                    f"Invalid handshake with peer {target_node}{' during feeler connection' if is_feeler else ''}"
+                    f". Maybe the peer is running old software."
+                )
             elif e.code == Err.INCOMPATIBLE_NETWORK_ID:
-                self.log.warning("Incompatible network ID. Maybe the peer is on another network")
+                self.log.warning(
+                    f"Incompatible network ID{' during feeler connection' if is_feeler else ''}"
+                    f". Maybe the peer is on another network"
+                )
             elif e.code == Err.SELF_CONNECTION:
                 pass
             else:
                 error_stack = traceback.format_exc()
-                self.log.error(f"Exception {e}, exception Stack: {error_stack}")
+                self.log.error(
+                    f"{'Feeler connection ' if is_feeler else ''}Exception {e}, exception Stack: {error_stack}"
+                )
         except Exception as e:
             if connection is not None:
                 await connection.close(self.invalid_protocol_ban_seconds, WSCloseCode.PROTOCOL_ERROR, Err.UNKNOWN)
             error_stack = traceback.format_exc()
-            self.log.error(f"Exception {e}, exception Stack: {error_stack}")
+            self.log.error(f"{'Feeler connection ' if is_feeler else ''}Exception {e}, exception Stack: {error_stack}")
         finally:
             if session is not None:
                 await session.close()
@@ -535,8 +555,9 @@ class ChiaServer:
             ban_until: float = time.time() + ban_time
             self.log.warning(f"Banning {connection.peer_info.host} for {ban_time} seconds")
             if connection.peer_info.host in self.banned_peers:
-                if ban_until > self.banned_peers[connection.peer_info.host]:
-                    self.banned_peers[connection.peer_info.host] = ban_until
+                self.banned_peers[connection.peer_info.host] = max(
+                    ban_until, self.banned_peers[connection.peer_info.host]
+                )
             else:
                 self.banned_peers[connection.peer_info.host] = ban_until
 
@@ -558,7 +579,7 @@ class ChiaServer:
             if on_disconnect is not None:
                 await on_disconnect(connection)
 
-    async def validate_broadcast_message_type(self, messages: List[Message], node_type: NodeType) -> None:
+    async def validate_broadcast_message_type(self, messages: list[Message], node_type: NodeType) -> None:
         for message in messages:
             if message_requires_reply(ProtocolMessageTypes(message.type)):
                 # Internal protocol logic error - we will raise, blocking messages to all peers
@@ -574,7 +595,7 @@ class ChiaServer:
 
     async def send_to_all(
         self,
-        messages: List[Message],
+        messages: list[Message],
         node_type: NodeType,
         exclude: Optional[bytes32] = None,
     ) -> None:
@@ -586,7 +607,7 @@ class ChiaServer:
 
     async def send_to_all_if(
         self,
-        messages: List[Message],
+        messages: list[Message],
         node_type: NodeType,
         predicate: Callable[[WSChiaConnection], bool],
         exclude: Optional[bytes32] = None,
@@ -597,7 +618,7 @@ class ChiaServer:
                 for message in messages:
                     await connection.send_message(message)
 
-    async def send_to_specific(self, messages: List[Message], node_id: bytes32) -> None:
+    async def send_to_specific(self, messages: list[Message], node_id: bytes32) -> None:
         if node_id in self.all_connections:
             connection = self.all_connections[node_id]
             for message in messages:
@@ -614,7 +635,7 @@ class ChiaServer:
 
     def get_connections(
         self, node_type: Optional[NodeType] = None, *, outbound: Optional[bool] = None
-    ) -> List[WSChiaConnection]:
+    ) -> list[WSChiaConnection]:
         result = []
         for _, connection in self.all_connections.items():
             node_type_match = node_type is None or connection.connection_type == node_type
@@ -631,7 +652,7 @@ class ChiaServer:
                 self.log.error(f"Exception while closing connection {e}")
 
     def close_all(self) -> None:
-        self.connection_close_task = asyncio.create_task(self.close_all_connections())
+        self.connection_close_task = create_referenced_task(self.close_all_connections())
         if self.webserver is not None:
             self.webserver.close()
 
@@ -707,7 +728,7 @@ class ChiaServer:
             return inbound_count < cast(int, self.config.get("max_inbound_timelord", 5))
         return True
 
-    def is_trusted_peer(self, peer: WSChiaConnection, trusted_peers: Dict[str, Any]) -> bool:
+    def is_trusted_peer(self, peer: WSChiaConnection, trusted_peers: dict[str, Any]) -> bool:
         return is_trusted_peer(
             host=peer.peer_info.host,
             node_id=peer.peer_node_id,
@@ -716,5 +737,5 @@ class ChiaServer:
             testing=self.config.get("testing", False),
         )
 
-    def set_capabilities(self, capabilities: List[Tuple[uint16, str]]) -> None:
+    def set_capabilities(self, capabilities: list[tuple[uint16, str]]) -> None:
         self._local_capabilities_for_handshake = capabilities