from __future__ import ( annotations, ) import base64 from collections.abc import ( Callable, KeysView, ) import functools import hashlib import logging import random import time from typing import ( NamedTuple, Protocol, cast, ) import multibase import trio from libp2p.abc import ( IHost, INetStream, IPubsub, IPubsubRouter, ISubscriptionAPI, ) from libp2p.crypto.keys import ( PrivateKey, ) from libp2p.custom_types import ( AsyncValidatorFn, SyncValidatorFn, TProtocol, ValidatorFn, ) from libp2p.encoding_config import get_default_encoding from libp2p.exceptions import ( ParseError, ValidationError, ) from libp2p.io.exceptions import ( IncompleteReadError, ) from libp2p.network.exceptions import ( SwarmException, ) from libp2p.network.stream.exceptions import ( StreamClosed, StreamEOF, StreamError, StreamReset, ) from libp2p.peer.id import ( ID, ) from libp2p.peer.peerdata import ( PeerDataError, ) from libp2p.peer.peerstore import env_to_send_in_RPC from libp2p.pubsub.extensions import ( ExtensionsState, ) from libp2p.pubsub.utils import maybe_consume_signed_record from libp2p.tools.anyio_service import ( Service, ) from libp2p.tools.timed_cache.last_seen_cache import ( LastSeenCache, ) from libp2p.utils import ( encode_varint_prefixed, read_varint_prefixed_bytes, ) from libp2p.utils.varint import encode_uvarint from .pb import ( rpc_pb2, ) from .pubsub_notifee import ( PubsubNotifee, ) from .rpc_queue import ( RpcQueue, drop_rpc, ) from .subscription import ( TrioSubscriptionAPI, ) from .validators import ( PUBSUB_SIGNING_PREFIX, signature_validator, ) # GossipSub v1.3+ protocol IDs. Extensions Control Message is only sent when # negotiating one of these protocols (per spec: extensions in first message). _MESHSUB_V13_PLUS = frozenset( ( TProtocol("/meshsub/1.3.0"), TProtocol("/meshsub/1.4.0"), TProtocol("/meshsub/2.0.0"), ) ) class _RouterWithExtensions(Protocol): """Protocol for a router that supports GossipSub v1.3 extensions.""" extensions_state: ExtensionsState # Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/40e1c94708658b155f30cf99e4574f384756d83c/topic.go#L97 # noqa: E501 SUBSCRIPTION_CHANNEL_SIZE = 32 _ANNOUNCE_RETRY_MIN_DELAY_MS = 1 _ANNOUNCE_RETRY_JITTER_MS = 1000 _ANNOUNCE_RETRY_MAX_ATTEMPTS = 10 # Peer-registration retry policy. The network ``connected`` notifee can fire # before the muxer handshake completes, so a one-shot ``new_stream`` may fail; # pubsub keeps retrying with capped exponential backoff while the peer remains # connected so a peer is never silently left unregistered. _PEER_STREAM_BACKOFF_INITIAL = 0.5 _PEER_STREAM_BACKOFF_MAX = 10.0 logger = logging.getLogger(__name__) def get_peer_and_seqno_msg_id(msg: rpc_pb2.Message) -> bytes: # Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/ab876fc71c34e89a7f0c8f4e361720ca9fa8588a/pubsub.go#L1327-L1330 # noqa: E501 return msg.from_id + msg.seqno def get_content_addressed_msg_id( msg: rpc_pb2.Message, encoding: str | None = None ) -> bytes: """ Generate content-addressed message ID using multibase encoding. :param msg: Pubsub message :param encoding: Encoding to use. When *None* the process-wide default from :mod:`libp2p.encoding_config` is used. :return: Multibase-encoded message ID """ if encoding is None: encoding = get_default_encoding() digest = hashlib.sha256(msg.data).digest() return multibase.encode(encoding, digest) def get_topic_aware_msg_id(msg: rpc_pb2.Message) -> bytes: """ Generate message ID that includes topic information for better deduplication across topics. Useful for v1.4 multi-topic scenarios. """ # Include topics in the hash for better separation topic_str = "|".join(sorted(msg.topicIDs)) combined = msg.seqno + msg.from_id + topic_str.encode() return hashlib.sha256(combined).digest() def get_timestamp_msg_id(msg: rpc_pb2.Message) -> bytes: """ Generate message ID that includes timestamp for time-based deduplication. Useful for v1.4 time-sensitive applications. """ import time timestamp = int(time.time() * 1000).to_bytes(8, byteorder="big") return msg.seqno + msg.from_id + timestamp def get_secure_msg_id(msg: rpc_pb2.Message) -> bytes: """ Generate cryptographically secure message ID using HMAC. Useful for v1.4 security-sensitive applications. """ # Use a combination of message content for HMAC key = msg.from_id + msg.seqno content = msg.data + "|".join(msg.topicIDs).encode() return hashlib.sha256(key + content).digest() class MessageIDGenerator: """ Abstract base class for message ID generators in GossipSub v1.4. Allows for more sophisticated message ID generation strategies that can maintain state or use external configuration. """ def generate_id(self, msg: rpc_pb2.Message) -> bytes: """Generate a message ID for the given message.""" raise NotImplementedError def __call__(self, msg: rpc_pb2.Message) -> bytes: """Make the generator callable like a function.""" return self.generate_id(msg) class CustomMessageIDGenerator(MessageIDGenerator): """ Customizable message ID generator that allows users to provide their own ID generation function. """ def __init__(self, id_fn: Callable[[rpc_pb2.Message], bytes]): self.id_fn = id_fn def generate_id(self, msg: rpc_pb2.Message) -> bytes: return self.id_fn(msg) class PeerAndSeqnoMessageIDGenerator(MessageIDGenerator): """Standard peer+seqno message ID generator.""" def generate_id(self, msg: rpc_pb2.Message) -> bytes: return msg.from_id + msg.seqno class ContentAddressedMessageIDGenerator(MessageIDGenerator): """Content-addressed message ID generator using SHA256.""" def generate_id(self, msg: rpc_pb2.Message) -> bytes: return base64.b64encode(hashlib.sha256(msg.data).digest()) class TopicValidator(NamedTuple): validator: ValidatorFn is_async: bool class ValidationResult(NamedTuple): """Result of message validation with caching metadata.""" is_valid: bool timestamp: float error_message: str | None = None class ValidationCache: """Cache for validation results to avoid redundant validation.""" def __init__(self, ttl: int = 300, max_size: int = 1000): """ Initialize validation cache. :param ttl: Time-to-live for cache entries in seconds :param max_size: Maximum number of entries to cache """ self.ttl = ttl self.max_size = max_size self.cache: dict[bytes, ValidationResult] = {} self.access_order: list[bytes] = [] # For LRU eviction def get(self, msg_id: bytes) -> ValidationResult | None: """Get cached validation result if still valid.""" if msg_id not in self.cache: return None result = self.cache[msg_id] current_time = time.time() # Check if result is still valid if current_time - result.timestamp > self.ttl: self._remove(msg_id) return None # Update access order for LRU if msg_id in self.access_order: self.access_order.remove(msg_id) self.access_order.append(msg_id) return result def put(self, msg_id: bytes, result: ValidationResult) -> None: """Cache a validation result.""" # Evict old entries if cache is full while len(self.cache) >= self.max_size and self.access_order: oldest = self.access_order.pop(0) self.cache.pop(oldest, None) self.cache[msg_id] = result if msg_id in self.access_order: self.access_order.remove(msg_id) self.access_order.append(msg_id) def _remove(self, msg_id: bytes) -> None: """Remove entry from cache.""" self.cache.pop(msg_id, None) if msg_id in self.access_order: self.access_order.remove(msg_id) def clear_expired(self) -> None: """Clear expired entries from cache.""" current_time = time.time() expired_keys = [ msg_id for msg_id, result in self.cache.items() if current_time - result.timestamp > self.ttl ] for msg_id in expired_keys: self._remove(msg_id) MAX_CONCURRENT_VALIDATORS = 10 class GossipsubEvent: peer_id: str | None = None # Inbound publish: bool = False subopts: bool = False control: bool = False message_size: int | None = None # Outbound / lifecycle publish_out: bool = False subscription_change: bool = False topic: str | None = None action: str | None = None peers_sent: int | None = None class Pubsub(Service, IPubsub): host: IHost router: IPubsubRouter peer_receive_channel: trio.MemoryReceiveChannel[ID] dead_peer_receive_channel: trio.MemoryReceiveChannel[ID] _validator_semaphore: trio.Semaphore seen_messages: LastSeenCache subscribed_topics_send: dict[str, trio.MemorySendChannel[rpc_pb2.Message]] subscribed_topics_receive: dict[str, TrioSubscriptionAPI] peer_topics: dict[str, set[ID]] peers: dict[ID, INetStream] peer_queues: dict[ID, RpcQueue] topic_validators: dict[str, TopicValidator] validation_cache: ValidationCache validation_timeout: float # Timeout for async validators in seconds counter: int # uint64 # Indicate if we should enforce signature verification strict_signing: bool sign_key: PrivateKey | None # Set of blacklisted peer IDs blacklisted_peers: set[ID] event_handle_peer_queue_started: trio.Event event_handle_dead_peer_queue_started: trio.Event _msg_id_constructor: Callable[[rpc_pb2.Message], bytes] _pending_announce_retries: set[tuple[ID, str, bool]] # topics whose message-cache window we already replayed, per peer _replayed_recent_topics: dict[ID, set[str]] def __init__( self, host: IHost, router: IPubsubRouter, cache_size: int | None = None, seen_ttl: int = 120, sweep_interval: int = 60, strict_signing: bool = True, msg_id_constructor: Callable[[rpc_pb2.Message], bytes] | MessageIDGenerator = get_peer_and_seqno_msg_id, max_concurrent_validator_count: int = MAX_CONCURRENT_VALIDATORS, validation_cache_ttl: int = 300, validation_cache_size: int = 1000, validation_timeout: float = 5.0, ) -> None: """ Construct a new Pubsub object, which is responsible for handling all Pubsub-related messages and relaying messages as appropriate to the Pubsub router (which is responsible for choosing who to send messages to). Since the logic for choosing peers to send pubsub messages to is in the router, the same Pubsub impl can back floodsub, gossipsub, etc. """ self.host = host self.router = router # Support both callable functions and MessageIDGenerator objects if isinstance(msg_id_constructor, MessageIDGenerator): self._msg_id_constructor = msg_id_constructor.generate_id else: self._msg_id_constructor = msg_id_constructor # Attach this new Pubsub object to the router self.router.attach(self) peer_send, peer_receive = trio.open_memory_channel[ID](0) dead_peer_send, dead_peer_receive = trio.open_memory_channel[ID](0) # Only keep the receive channels in `Pubsub`. # Therefore, we can only close from the receive side. self.peer_receive_channel = peer_receive self.dead_peer_receive_channel = dead_peer_receive self._validator_semaphore = trio.Semaphore(max_concurrent_validator_count) # Register a notifee self.host.get_network().register_notifee( PubsubNotifee(peer_send, dead_peer_send) ) # Register stream handlers for each pubsub router protocol to handle # the pubsub streams opened on those protocols for protocol in router.get_protocols(): self.host.set_stream_handler(protocol, self.stream_handler) # keeps track of seen messages as LRU cache if cache_size is None: self.cache_size = 128 else: self.cache_size = cache_size self.strict_signing = strict_signing if strict_signing: self.sign_key = self.host.get_private_key() else: self.sign_key = None self.seen_messages = LastSeenCache(seen_ttl, sweep_interval) # Map of topics we are subscribed to blocking queues # for when the given topic receives a message self.subscribed_topics_send = {} self.subscribed_topics_receive = {} # Map of topic to peers to keep track of what peers are subscribed to self.peer_topics = {} # Create peers map, which maps peer_id (as string) to stream (to a given peer) self.peers = {} # Per-peer outbound RPC queues self.peer_queues = {} # Map of topic to topic validator self.topic_validators = {} # Enhanced validation features (v2.0) self.validation_cache = ValidationCache( validation_cache_ttl, validation_cache_size ) self.validation_timeout = validation_timeout self.counter = int(time.time()) # Set of blacklisted peer IDs self.blacklisted_peers = set() # Event-based waiting: maps for trio.Event instances # Used by wait_for_peer / wait_for_subscription to avoid busy-waiting self._peer_added_events: dict[ID, trio.Event] = {} self._subscription_events: dict[tuple[ID, str], trio.Event] = {} self._pending_announce_retries = set() self._replayed_recent_topics = {} # Peers with a background stream-registration retry task already in # flight, so concurrent stream failures cannot pile up duplicate tasks. self._peer_stream_retries_pending: set[ID] = set() self.event_handle_peer_queue_started = trio.Event() self.event_handle_dead_peer_queue_started = trio.Event() async def run(self) -> None: self.manager.run_daemon_task(self.handle_peer_queue) self.manager.run_daemon_task(self.handle_dead_peer_queue) self.manager.run_daemon_task(self._validation_cache_cleanup) await self.manager.wait_finished() @property def my_id(self) -> ID: return self.host.get_id() @property def protocols(self) -> tuple[TProtocol, ...]: return tuple(self.router.get_protocols()) @property def topic_ids(self) -> KeysView[str]: return self.subscribed_topics_receive.keys() def get_hello_packet(self) -> rpc_pb2.RPC: """ Generate subscription message with all topics we are subscribed to only send hello packet if we have subscribed topics. """ packet = rpc_pb2.RPC() for topic_id in self.topic_ids: packet.subscriptions.extend( [rpc_pb2.RPC.SubOpts(subscribe=True, topicid=topic_id)] ) # Add the sender's signedRecord in the RPC message envelope_bytes, _ = env_to_send_in_RPC(self.host) packet.senderRecord = envelope_bytes return packet async def continuously_read_stream(self, stream: INetStream) -> None: """ Read from input stream in an infinite loop. Process messages from other nodes. :param stream: stream to continously read from """ peer_id = stream.muxed_conn.peer_id try: while self.manager.is_running: incoming: bytes = await read_varint_prefixed_bytes(stream) rpc_incoming: rpc_pb2.RPC = rpc_pb2.RPC() rpc_incoming.ParseFromString(incoming) # Process the sender's signed-record if sent if not maybe_consume_signed_record(rpc_incoming, self.host, peer_id): logger.error( "Received an invalid-signed-record, ignoring the incoming msg" ) continue event = GossipsubEvent() event.peer_id = peer_id.pretty() event.message_size = len(incoming) if rpc_incoming.publish: # deal with RPC.publish event.publish = True for msg in rpc_incoming.publish: if not self._is_subscribed_to_msg(msg): continue logger.debug( "received `publish` message %s from peer %s", msg, peer_id ) # Only schedule task if service is still running if self.manager.is_running: self.manager.run_task(self.push_msg, peer_id, msg) if rpc_incoming.subscriptions: # deal with RPC.subscriptions # We don't need to relay the subscription to our # peers because a given node only needs its peers # to know that it is subscribed to the topic (doesn't # need everyone to know) event.subopts = True for message in rpc_incoming.subscriptions: logger.debug( "received `subscriptions` message %s from peer %s", message, peer_id, ) self.handle_subscription(peer_id, message) # NOTE: Check if `rpc_incoming.control` is set through `HasField`. # This is necessary because `control` is an optional field in pb2. # Ref: https://developers.google.com/protocol-buffers/docs/reference/python-generated#singular-fields-proto2 # noqa: E501 if rpc_incoming.HasField("control"): event.control = True # Pass rpc to router so router could perform custom logic logger.debug( "received `control` message %s from peer %s", rpc_incoming.control, peer_id, ) await self.router.handle_rpc(rpc_incoming, peer_id) self.host.get_event_bus().emit(event) except StreamEOF: logger.debug( f"Stream closed for peer {peer_id}, exiting read loop cleanly." ) def set_topic_validator( self, topic: str, validator: ValidatorFn, is_async_validator: bool ) -> None: """ Register a validator under the given topic. One topic can only have one validtor. :param topic: the topic to register validator under :param validator: the validator used to validate messages published to the topic :param is_async_validator: indicate if the validator is an asynchronous validator """ # noqa: E501 self.topic_validators[topic] = TopicValidator(validator, is_async_validator) def remove_topic_validator(self, topic: str) -> None: """ Remove the validator from the given topic. :param topic: the topic to remove validator from """ self.topic_validators.pop(topic, None) def get_msg_validators(self, msg: rpc_pb2.Message) -> tuple[TopicValidator, ...]: """ Get all validators corresponding to the topics in the message. :param msg: the message published to the topic """ return tuple( self.topic_validators[topic] for topic in msg.topicIDs if topic in self.topic_validators ) def add_to_blacklist(self, peer_id: ID) -> None: """ Add a peer to the blacklist. When a peer is blacklisted: - Any existing connection to that peer is immediately closed and removed - The peer is removed from all topic subscription mappings - Future connection attempts from this peer will be rejected - Messages forwarded by or originating from this peer will be dropped - The peer will not be able to participate in pubsub communication :param peer_id: the peer ID to blacklist """ self.blacklisted_peers.add(peer_id) logger.debug("Added peer %s to blacklist", peer_id) self.manager.run_task(self._teardown_if_connected, peer_id) async def _teardown_if_connected(self, peer_id: ID) -> None: """Close their stream and remove them if connected""" stream = self.peers.get(peer_id) if stream is not None: try: await stream.reset() except Exception: pass del self.peers[peer_id] # Also remove from any subscription maps: self._forget_all_subscriptions(peer_id) def remove_from_blacklist(self, peer_id: ID) -> None: """ Remove a peer from the blacklist. Once removed from the blacklist: - The peer can establish new connections to this node - Messages from this peer will be processed normally - The peer can participate in topic subscriptions and message forwarding :param peer_id: the peer ID to remove from blacklist """ self.blacklisted_peers.discard(peer_id) logger.debug("Removed peer %s from blacklist", peer_id) def is_peer_blacklisted(self, peer_id: ID) -> bool: """ Check if a peer is blacklisted. :param peer_id: the peer ID to check :return: True if peer is blacklisted, False otherwise """ return peer_id in self.blacklisted_peers def clear_blacklist(self) -> None: """ Clear all peers from the blacklist. This removes all blacklist restrictions, allowing previously blacklisted peers to: - Establish new connections - Send and forward messages - Participate in topic subscriptions """ self.blacklisted_peers.clear() logger.debug("Cleared all peers from blacklist") def get_blacklisted_peers(self) -> set[ID]: """ Get a copy of the current blacklisted peers. Returns a snapshot of all currently blacklisted peer IDs. These peers are completely isolated from pubsub communication - their connections are rejected and their messages are dropped. :return: a set containing all blacklisted peer IDs """ return self.blacklisted_peers.copy() async def stream_handler(self, stream: INetStream) -> None: """ Stream handler for pubsub. Gets invoked whenever a new stream is created on one of the supported pubsub protocols. :param stream: newly created stream """ peer_id = stream.muxed_conn.peer_id try: await self.continuously_read_stream(stream) except (StreamError, ParseError, IncompleteReadError) as error: logger.debug( "fail to read from peer %s, error=%s," "closing the stream and remove the peer from record", peer_id, error, ) await stream.reset() self._handle_dead_peer(peer_id) async def wait_until_ready(self) -> None: await self.event_handle_peer_queue_started.wait() await self.event_handle_dead_peer_queue_started.wait() async def wait_for_peer(self, peer_id: ID, timeout: float = 5.0) -> None: """ Wait until a pubsub stream with the given peer has been established. This method blocks until the given peer has been added to the pubsub peers map, indicating that a pubsub protocol stream exists. Use this instead of arbitrary trio.sleep() calls to avoid race conditions. Uses an event-based approach: the task blocks until ``_handle_new_peer`` fires the corresponding :class:`trio.Event`, consuming zero CPU while waiting. :param peer_id: the peer ID to wait for :param timeout: maximum time to wait in seconds (default: 5.0) :raises trio.TooSlowError: if the peer stream is not established within the timeout Example:: await connect(host1, host2) await pubsub1.wait_for_peer(host2.get_id()) # Now safe to publish or check peer_topics """ if peer_id in self.peers: return event = self._peer_added_events.setdefault(peer_id, trio.Event()) with trio.fail_after(timeout): await event.wait() async def wait_for_subscription( self, peer_id: ID, topic_id: str, timeout: float = 5.0 ) -> None: """ Wait until a specific peer has subscribed to a topic. This method blocks until the given peer appears in the peer_topics map for the specified topic, indicating that they have sent a subscription message. Use this instead of arbitrary trio.sleep() calls to avoid race conditions. Uses an event-based approach: the task blocks until ``handle_subscription`` fires the corresponding :class:`trio.Event`, consuming zero CPU while waiting. :param peer_id: the peer ID to wait for :param topic_id: the topic to check subscription for :param timeout: maximum time to wait in seconds (default: 5.0) :raises trio.TooSlowError: if the peer does not subscribe within the timeout Example:: await connect(host1, host2) await pubsub1.wait_for_subscription(host2.get_id(), "my-topic") # Now safe to assert subscription state """ if topic_id in self.peer_topics and peer_id in self.peer_topics[topic_id]: return key = (peer_id, topic_id) event = self._subscription_events.setdefault(key, trio.Event()) with trio.fail_after(timeout): await event.wait() async def ensure_peer_stream(self, peer_id: ID, timeout: float = 15.0) -> bool: """ Ensure a pubsub stream with *peer_id* is open (idempotent). Useful when the application established the connection out-of-band — e.g. ``connect_peer`` reusing an existing mDNS connection does not fire a fresh ``connected`` notifee, so without this the peer would never be registered with pubsub. Opens the stream and registers the peer with the router, retrying while the peer stays connected until it succeeds or *timeout* seconds elapse. :param peer_id: the peer to register with pubsub :param timeout: maximum seconds to keep retrying (default 15.0) :return: True if the peer is registered with pubsub, False otherwise """ if peer_id in self.peers: return True if self.is_peer_blacklisted(peer_id): return False await self._handle_new_peer_with_retry(peer_id, timeout) return peer_id in self.peers async def _handle_new_peer(self, peer_id: ID) -> None: # Check if we already have a pubsub stream with this peer to avoid duplicates if peer_id in self.peers: logger.debug("Peer %s already has pubsub stream, skipping", peer_id) return if self.is_peer_blacklisted(peer_id): logger.debug("Rejecting blacklisted peer %s", peer_id) return try: stream: INetStream = await self.host.new_stream(peer_id, self.protocols) except SwarmException as error: # The `connected` notifee can fire before the muxer handshake # completes, so a single stream open may legitimately fail. Raise so # callers (e.g. `_handle_new_peer_with_retry`) can retry while the # peer stays connected instead of silently never registering it. logger.debug( "fail to open pubsub stream to peer %s, error %s", peer_id, error ) raise # Build hello packet. hello = self.get_hello_packet() # GossipSub v1.3 – Extensions Control Message injection. # Per spec: "If a peer supports any extension, the Extensions control # message MUST be included in the first message on the stream." # Only inject when we negotiated v1.3+; peers on v1.1/v1.2 must not # receive extension fields. negotiated_protocol = stream.get_protocol() router = self.router if ( negotiated_protocol in _MESHSUB_V13_PLUS and hasattr(router, "extensions_state") and hasattr(router, "supports_v13_features") ): # We pass the peer_id because extensions_state needs to track # "sent_extensions" per peer for the at-most-once rule. # cast() tells static type-checkers the narrowed type without # creating a runtime dependency on gossipsub.py from pubsub.py. v13_router = cast(_RouterWithExtensions, router) hello = v13_router.extensions_state.build_hello_extensions(peer_id, hello) try: await stream.write(encode_varint_prefixed(hello.SerializeToString())) except StreamClosed: logger.debug("Fail to add new peer %s: stream closed", peer_id) raise try: self.router.add_peer(peer_id, negotiated_protocol) except Exception as error: logger.debug("fail to add new peer %s, error %s", peer_id, error) return self.peers[peer_id] = stream # Create per-peer outbound queue and spawn sending task queue = RpcQueue() self.peer_queues[peer_id] = queue self.manager.run_task(self.handle_sending_messages, peer_id, stream, queue) # Notify anyone waiting in wait_for_peer() if peer_id in self._peer_added_events: self._peer_added_events.pop(peer_id).set() # Flush any messages that were queued while this peer's protocol # identification was still in progress (identify-aware publishing). try: await self.router.flush_pending_messages(peer_id) except Exception as error: logger.debug( "failed to flush pending messages for peer %s: %s", peer_id, error, ) await self._send_recent_messages_to_new_peer(peer_id) logger.debug("added new peer %s", peer_id) async def _send_recent_messages_to_new_peer(self, peer_id: ID) -> None: """ Replay recent messages to a peer whose subscriptions we already hold. The peer's subscriptions can land before ``_handle_new_peer`` registers the outbound stream, and ``handle_subscription``'s catch-up is a no-op in that ordering, so the replay has to run again once we can write. :param peer_id: the peer that just became writable """ subscribed_topics = [ topic for topic, peers in self.peer_topics.items() if peer_id in peers ] for topic in subscribed_topics: await self._replay_recent_messages(peer_id, topic) async def _replay_recent_messages(self, peer_id: ID, topic: str) -> None: """ Replay the router's recent messages for a topic, once per subscription. The gate keeps a peer that reconnects or re-announces from being handed the whole message-cache window again. It is dropped when the peer unsubscribes, disconnects, is blacklisted, or when the replay fails. :param peer_id: the peer to replay messages to :param topic: the topic to replay messages for """ if peer_id not in self.peers: # Not writable yet, so the replay would be a no-op. Once the # outbound stream is registered, `_handle_new_peer` replays. return replayed = self._replayed_recent_topics.setdefault(peer_id, set()) if topic in replayed: return replayed.add(topic) try: await self.router.send_recent_messages(peer_id, topic) except Exception as error: # Un-spend the gate so a later announcement can retry. replayed.discard(topic) logger.debug( "failed to send recent messages for topic %s to peer %s: %s", topic, peer_id, error, ) async def _handle_new_peer_safe(self, peer_id: ID) -> None: """ Safely handle new peer with exception handling. This wrapper ensures that any exceptions during peer negotiation don't crash the entire pubsub service. Kept for backward compatibility and tests; the live peer queue now uses ``_handle_new_peer_with_retry`` which retries transient failures. """ try: await self._handle_new_peer(peer_id) except Exception as error: logger.info(f"Protocol negotiation failed for peer {peer_id}: {error}") async def _handle_new_peer_with_retry( self, peer_id: ID, timeout: float | None = None ) -> None: """ Open a pubsub stream to *peer_id*, retrying while it stays connected. The network ``connected`` notifee fires as soon as the transport connection is established, which can be before the muxer handshake completes. A single ``new_stream`` attempt at that point frequently fails; if pubsub never retried, the peer would silently never be registered and messages to it would be dropped forever. Retries with capped exponential backoff while the peer remains connected (unbounded by design when *timeout* is ``None``, matching go-libp2p's persistent mesh maintenance). Registration is idempotent (``_handle_new_peer`` is a no-op once the peer is in ``self.peers``), so concurrent retry tasks for the same peer are safe. If the peer disconnects we give up — a fresh ``connected`` notifee will restart the process on the next connection. :param peer_id: the peer to register with pubsub :param timeout: optional maximum seconds to keep retrying (``None`` for unbounded retries while the peer stays connected) """ delay = _PEER_STREAM_BACKOFF_INITIAL deadline = None if timeout is None else trio.current_time() + timeout while True: try: await self._handle_new_peer(peer_id) return except (SwarmException, StreamClosed) as error: if not self._peer_is_connected(peer_id): logger.debug( "peer %s is no longer connected; giving up on pubsub " "stream registration: %s", peer_id, error, ) return if deadline is not None and trio.current_time() >= deadline: logger.debug( "timed out registering pubsub stream with peer %s", peer_id ) return logger.debug( "failed to open pubsub stream to peer %s (retrying in %.1fs): %s", peer_id, delay, error, ) await trio.sleep(delay) delay = min(delay * 2, _PEER_STREAM_BACKOFF_MAX) except Exception as error: # Non-retryable registration failure; do not spin forever. logger.debug("failed to register pubsub peer %s: %s", peer_id, error) return def _peer_is_connected(self, peer_id: ID) -> bool: """ Return True if the peer currently has an open (non-closed) connection. A connection whose muxer handshake has not completed yet (``muxed_conn`` not set) still counts as connected: the retry loop must keep trying during that window instead of giving up. """ try: connections = self.host.get_network().get_connections(peer_id) except Exception: return False for conn in connections: muxed_conn = getattr(conn, "muxed_conn", None) if muxed_conn is None or not getattr(muxed_conn, "is_closed", False): return True return False def _schedule_peer_stream_retry(self, peer_id: ID) -> None: """ Schedule a background stream-registration retry for *peer_id*. Ensures at most one retry task per peer is in flight, so concurrent stream failures (multiple connections churning) do not pile up duplicate tasks. """ if not self.manager.is_running: return if peer_id in self._peer_stream_retries_pending: return self._peer_stream_retries_pending.add(peer_id) self.manager.run_task(self._peer_stream_retry_task, peer_id) async def _peer_stream_retry_task(self, peer_id: ID) -> None: try: await self._handle_new_peer_with_retry(peer_id) finally: self._peer_stream_retries_pending.discard(peer_id) def _handle_dead_peer(self, peer_id: ID) -> None: # Runs before the `peers` check: subscriptions arrive on the peer's # inbound stream, so a half-registered peer has state to clean up here. self._clear_pending_announce_retries_for_peer(peer_id) self._forget_all_subscriptions(peer_id) if peer_id not in self.peers: # A stream failed on a peer that never finished registering. If the # peer is still connected, try again; otherwise there is nothing to # clean up. if self._peer_is_connected(peer_id): self._schedule_peer_stream_retry(peer_id) return del self.peers[peer_id] # Close the outbound queue so the sending task exits if peer_id in self.peer_queues: self.peer_queues.pop(peer_id).close() self.router.remove_peer(peer_id) # The pubsub stream died but the peer is still connected — e.g. the # stream was opened on a connection whose muxer handshake failed while # a healthy connection exists (common when mDNS auto-connect races an # explicit dial, producing multiple simultaneous connections). # Re-establish the stream so messaging with this peer does not silently # die. if self._peer_is_connected(peer_id): logger.debug( "peer %s still connected after stream close; re-establishing " "pubsub stream", peer_id, ) self._schedule_peer_stream_retry(peer_id) return logger.debug("removed dead peer %s", peer_id) def _forget_all_subscriptions(self, peer_id: ID) -> None: for peers in self.peer_topics.values(): peers.discard(peer_id) self._replayed_recent_topics.pop(peer_id, None) def _forget_subscription(self, peer_id: ID, topic: str) -> None: if topic in self.peer_topics: self.peer_topics[topic].discard(peer_id) replayed = self._replayed_recent_topics.get(peer_id) if replayed is not None: replayed.discard(topic) def _clear_pending_announce_retries_for_peer(self, peer_id: ID) -> None: # This is O(n) over pending retry keys. Keep this representation because # retries are bounded and deduplicated per (peer, topic, subscribe). self._pending_announce_retries = { key for key in self._pending_announce_retries if key[0] != peer_id } async def handle_peer_queue(self) -> None: """ Continuously read from peer queue and each time a new peer is found, open a stream to the peer using a supported pubsub protocol pubsub protocols we support. """ async with self.peer_receive_channel: self.event_handle_peer_queue_started.set() async for peer_id in self.peer_receive_channel: # Add Peer - retry while connected so a registration that races # the muxer handshake is not silently dropped. self.manager.run_task(self._handle_new_peer_with_retry, peer_id) async def handle_dead_peer_queue(self) -> None: """ Continuously read from dead peer channel and close the stream between that peer and remove peer info from pubsub and pubsub router. Only removes the peer if there are no remaining active connections. """ async with self.dead_peer_receive_channel: self.event_handle_dead_peer_queue_started.set() async for peer_id in self.dead_peer_receive_channel: # Check if peer still has active connections before removing # This prevents premature removal when multiple connections exist network = self.host.get_network() remaining_connections = network.get_connections(peer_id) if remaining_connections: # Filter out closed connections. A connection whose muxer # handshake has not completed (``muxed_conn`` not set yet) # counts as active. active_connections = [ c for c in remaining_connections if getattr(c, "muxed_conn", None) is None or not getattr(c.muxed_conn, "is_closed", False) ] if active_connections: logger.debug( "Peer %s still has %d active connections, not removing", peer_id, len(active_connections), ) continue # Remove Peer - no more active connections self._handle_dead_peer(peer_id) async def handle_sending_messages( self, peer_id: ID, stream: INetStream, queue: RpcQueue ) -> None: """ Per-peer sending loop: pops RPCs from *queue*, splits them if needed, and writes each chunk to *stream*. Runs as a task spawned by :meth:`_handle_new_peer`. Exits when the queue is closed (peer disconnected) or the stream errors. """ try: while True: rpc = await queue.pop() if rpc is None: # Queue was closed return ok = await self.write_msg(stream, rpc) if not ok: return except Exception: logger.debug("sending loop for %s terminated with error", peer_id) self._handle_dead_peer(peer_id) def handle_subscription( self, origin_id: ID, sub_message: rpc_pb2.RPC.SubOpts ) -> None: """ Handle an incoming subscription message from a peer. Update internal mapping to mark the peer as subscribed or unsubscribed to topics as defined in the subscription message. :param origin_id: id of the peer who subscribe to the message :param sub_message: RPC.SubOpts """ if sub_message.subscribe: was_newly_added = False if sub_message.topicid not in self.peer_topics: self.peer_topics[sub_message.topicid] = {origin_id} was_newly_added = True elif origin_id not in self.peer_topics[sub_message.topicid]: # Add peer to topic self.peer_topics[sub_message.topicid].add(origin_id) was_newly_added = True if was_newly_added: # Notify anyone waiting in wait_for_subscription() key = (origin_id, sub_message.topicid) if key in self._subscription_events: self._subscription_events.pop(key).set() # Both hooks are async while `handle_subscription` is sync, so # they have to be spawned. if self.manager.is_running: # Flush any messages that were queued while waiting for this # peer's subscription (identify-aware publishing). self.manager.run_task( self.router.flush_pending_messages, origin_id, ) # Also send recent messages from mcache for this topic. # This handles the case where messages were published before # this peer was even in pubsub.peers (race during connection # setup). self.manager.run_task( self._replay_recent_messages, origin_id, sub_message.topicid, ) else: self._forget_subscription(origin_id, sub_message.topicid) def notify_subscriptions(self, publish_message: rpc_pb2.Message) -> None: """ Put incoming message from a peer onto my blocking queue. :param publish_message: RPC.Message format """ # Check if this message has any topics that we are subscribed to for topic in publish_message.topicIDs: if topic in self.topic_ids: # we are subscribed to a topic this message was sent for, # so add message to the subscription output queue # for each topic try: self.subscribed_topics_send[topic].send_nowait(publish_message) except trio.WouldBlock: # Channel is full, ignore this message. logger.warning( "fail to deliver message to subscription for topic %s", topic ) def _build_announce_rpc( self, topic_id: str, subscribe: bool ) -> tuple[rpc_pb2.RPC, rpc_pb2.RPC.SubOpts]: packet = rpc_pb2.RPC() subopt = rpc_pb2.RPC.SubOpts(subscribe=subscribe, topicid=topic_id) packet.subscriptions.extend([subopt]) envelope_bytes, _ = env_to_send_in_RPC(self.host) packet.senderRecord = envelope_bytes return packet, subopt def _announce_state_matches(self, topic_id: str, subscribe: bool) -> bool: is_currently_subscribed = topic_id in self.subscribed_topics_receive return subscribe == is_currently_subscribed async def subscribe(self, topic_id: str) -> ISubscriptionAPI: """ Subscribe ourself to a topic. :param topic_id: topic_id to subscribe to """ logger.debug("subscribing to topic %s", topic_id) # Already subscribed if topic_id in self.topic_ids: return self.subscribed_topics_receive[topic_id] send_channel, receive_channel = trio.open_memory_channel[rpc_pb2.Message]( SUBSCRIPTION_CHANNEL_SIZE ) subscription = TrioSubscriptionAPI( receive_channel, unsubscribe_fn=functools.partial(self.unsubscribe, topic_id), ) self.subscribed_topics_send[topic_id] = send_channel self.subscribed_topics_receive[topic_id] = subscription # Create subscribe announcement packet, subopt = self._build_announce_rpc(topic_id, subscribe=True) # Send out subscribe message to all peers await self.message_all_peers(packet.SerializeToString(), announce=subopt) # Tell router we are joining this topic await self.router.join(topic_id) # Emit subscription-change metric event event = GossipsubEvent() event.peer_id = self.my_id.pretty() event.subscription_change = True event.topic = topic_id event.action = "subscribe" self.host.get_event_bus().emit(event) # Return the subscription for messages on this topic return subscription async def unsubscribe(self, topic_id: str) -> None: """ Unsubscribe ourself from a topic. :param topic_id: topic_id to unsubscribe from """ logger.debug("unsubscribing from topic %s", topic_id) # Return if we already unsubscribed from the topic if topic_id not in self.topic_ids: return # Remove topic_id from the maps before yielding send_channel = self.subscribed_topics_send[topic_id] del self.subscribed_topics_send[topic_id] del self.subscribed_topics_receive[topic_id] # Only close the send side await send_channel.aclose() # Create unsubscribe announcement packet, subopt = self._build_announce_rpc(topic_id, subscribe=False) # Send out unsubscribe message to all peers await self.message_all_peers(packet.SerializeToString(), announce=subopt) # Tell router we are leaving this topic await self.router.leave(topic_id) # Emit subscription-change metric event event = GossipsubEvent() event.peer_id = self.my_id.pretty() event.subscription_change = True event.topic = topic_id event.action = "unsubscribe" self.host.get_event_bus().emit(event) async def message_all_peers( self, raw_msg: bytes, announce: rpc_pb2.RPC.SubOpts | None = None ) -> None: """ Broadcast a message to peers. :param raw_msg: raw contents of the message to broadcast """ rpc_msg: rpc_pb2.RPC | None = None # Broadcast message via per-peer outbound queues to preserve # queue back-pressure/drop semantics. for peer_id in tuple(self.peers): queue = self.peer_queues.get(peer_id) if queue is None: logger.debug("No outbound queue for peer %s", peer_id) continue # Fast path for small RPCs: avoid split/clone overhead and # enqueue the parsed RPC directly. if len(raw_msg) <= queue.max_message_size: if rpc_msg is None: rpc_msg = rpc_pb2.RPC() rpc_msg.ParseFromString(raw_msg) self._enqueue_or_retry_announce(peer_id, queue, rpc_msg, announce) continue if rpc_msg is None: rpc_msg = rpc_pb2.RPC() rpc_msg.ParseFromString(raw_msg) for part in queue.split_rpc(rpc_msg): if part.ByteSize() > queue.max_message_size: # Intentional asymmetry: only queue-full drops schedule # announce retries. Oversized chunks are terminal here, # matching _run_announce_retry, which also bails out when # the announce RPC itself exceeds max_message_size. drop_rpc(peer_id, part) continue ok = queue.push(part) if not ok: drop_rpc(peer_id, part) if announce is not None: self._schedule_announce_retry(peer_id, announce) break def _enqueue_or_retry_announce( self, peer_id: ID, queue: RpcQueue, rpc_msg: rpc_pb2.RPC, announce: rpc_pb2.RPC.SubOpts | None, ) -> None: ok = queue.push(rpc_msg) if ok: return drop_rpc(peer_id, rpc_msg) if announce is not None: self._schedule_announce_retry(peer_id, announce) def _schedule_announce_retry( self, peer_id: ID, announce: rpc_pb2.RPC.SubOpts ) -> None: if not self.manager.is_running: return key = (peer_id, announce.topicid, announce.subscribe) if key in self._pending_announce_retries: return self._pending_announce_retries.add(key) self.manager.run_task( self._run_announce_retry, peer_id, announce.topicid, announce.subscribe ) async def _run_announce_retry( self, peer_id: ID, topic_id: str, subscribe: bool ) -> None: key = (peer_id, topic_id, subscribe) try: for _ in range(_ANNOUNCE_RETRY_MAX_ATTEMPTS): if not self.manager.is_running: return delay_ms = _ANNOUNCE_RETRY_MIN_DELAY_MS + random.randint( 0, _ANNOUNCE_RETRY_JITTER_MS - 1 ) await trio.sleep(delay_ms / 1000) if not self._announce_state_matches(topic_id, subscribe): return queue = self.peer_queues.get(peer_id) if queue is None: return # Rebuild the announce RPC on every attempt so senderRecord and # any host-address-derived record data stays fresh if listen # addresses changed since the previous attempt. retry_rpc, _ = self._build_announce_rpc(topic_id, subscribe) if retry_rpc.ByteSize() > queue.max_message_size: drop_rpc(peer_id, retry_rpc) return ok = queue.push(retry_rpc) if ok: return drop_rpc(peer_id, retry_rpc) finally: self._pending_announce_retries.discard(key) async def publish(self, topic_id: str | list[str], data: bytes) -> None: """ Publish data to a topic or multiple topics. :param topic_id: topic (str) or topics (list[str]) to publish the data to :param data: data which we are publishing """ # Handle both single topic (str) and multiple topics (list[str]) if isinstance(topic_id, str): topic_ids = [topic_id] else: topic_ids = topic_id msg = rpc_pb2.Message( data=data, topicIDs=topic_ids, # Origin is ourself. from_id=self.my_id.to_bytes(), seqno=self._next_seqno(), ) if self.strict_signing: priv_key = self.sign_key if priv_key is None: raise PeerDataError("private key not found") signature = priv_key.sign( PUBSUB_SIGNING_PREFIX.encode() + msg.SerializeToString() ) msg.key = self.host.get_public_key().serialize() msg.signature = signature await self.push_msg(self.my_id, msg) logger.debug("successfully published message %s", msg) # Emit publish (outbound) metric event event = GossipsubEvent() event.peer_id = self.my_id.pretty() event.publish_out = True event.topic = ",".join(topic_ids) event.message_size = len(msg.data) event.peers_sent = len(self.peers) self.host.get_event_bus().emit(event) async def validate_msg( self, msg_forwarder: ID, msg: rpc_pb2.Message, ) -> None: """ Validate the received message with caching and timeout support. :param msg_forwarder: the peer who forward us the message. :param msg: the message. """ # Check validation cache first msg_id = self._msg_id_constructor(msg) cached_result = self.validation_cache.get(msg_id) if cached_result is not None: if not cached_result.is_valid: error_msg = cached_result.error_message or "unknown error" raise ValidationError( f"Cached validation failed for msg={msg}: {error_msg}" ) return sync_topic_validators: list[SyncValidatorFn] = [] async_topic_validators: list[AsyncValidatorFn] = [] for topic_validator in self.get_msg_validators(msg): if topic_validator.is_async: async_topic_validators.append( cast(AsyncValidatorFn, topic_validator.validator) ) else: sync_topic_validators.append( cast(SyncValidatorFn, topic_validator.validator) ) validation_error = None try: # Run synchronous validators first for validator in sync_topic_validators: if not validator(msg_forwarder, msg): validation_error = "Synchronous validation failed" raise ValidationError(f"Validation failed for msg={msg}") # Run asynchronous validators with timeout if len(async_topic_validators) > 0: try: with trio.move_on_after(self.validation_timeout) as cancel_scope: # Appends to lists are thread safe in CPython results: list[bool] = [] async with trio.open_nursery() as nursery: for async_validator in async_topic_validators: nursery.start_soon( self._run_async_validator, async_validator, msg_forwarder, msg, results, ) if not all(results): validation_error = "Asynchronous validation failed" raise ValidationError(f"Validation failed for msg={msg}") if cancel_scope.cancelled_caught: validation_error = "Validation timeout" raise ValidationError(f"Validation timeout for msg={msg}") except ValidationError: raise except Exception as e: validation_error = f"Validation error: {e}" raise ValidationError(f"Validation error for msg={msg}: {e}") # Cache successful validation self.validation_cache.put( msg_id, ValidationResult(is_valid=True, timestamp=time.time()) ) except ValidationError: # Cache failed validation self.validation_cache.put( msg_id, ValidationResult( is_valid=False, timestamp=time.time(), error_message=validation_error, ), ) raise async def _run_async_validator( self, func: AsyncValidatorFn, msg_forwarder: ID, msg: rpc_pb2.Message, results: list[bool], ) -> None: async with self._validator_semaphore: result = await func(msg_forwarder, msg) results.append(result) async def push_msg(self, msg_forwarder: ID, msg: rpc_pb2.Message) -> None: """ Push a pubsub message to others. :param msg_forwarder: the peer who forward us the message. :param msg: the message we are going to push out. """ logger.debug("attempting to publish message %s", msg) # Check if the message forwarder (source) is in the blacklist. If yes, reject. if self.is_peer_blacklisted(msg_forwarder): logger.debug( "Rejecting message from blacklisted source peer %s", msg_forwarder ) return # Check if the message originator (from) is in the blacklist. If yes, reject. msg_from_peer = ID(msg.from_id) if self.is_peer_blacklisted(msg_from_peer): logger.debug( "Rejecting message from blacklisted originator peer %s", msg_from_peer ) return # If the message is processed before, return(i.e., don't further process the message) # noqa: E501 if self._is_msg_seen(msg): return try: scorer = getattr(self.router, "scorer", None) if scorer is not None: if not scorer.allow_publish(msg_forwarder, list(msg.topicIDs)): logger.debug( "Rejecting message from %s by publish score gate", msg_forwarder ) return except Exception: # Router may not support scoring; ignore gracefully pass # Check if signing is required and if so validate the signature if self.strict_signing: # Validate the signature of the message if not signature_validator(msg): logger.debug("Signature validation failed for msg: %s", msg) return # Validate the message with registered topic validators. # If the validation failed, return(i.e., don't further process the message). try: await self.validate_msg(msg_forwarder, msg) except ValidationError: # Scoring: count invalid messages try: scorer = getattr(self.router, "scorer", None) if scorer is not None: for topic in msg.topicIDs: scorer.on_invalid_message(msg_forwarder, topic) except Exception: pass logger.debug( "Topic validation failed: sender %s sent data %s under topic IDs: %s %s:%s", # noqa: E501 msg_forwarder, msg.data.hex(), msg.topicIDs, ID(msg.from_id).to_base58(), msg.seqno.hex(), ) return self._mark_msg_seen(msg) # Scoring: first delivery for this sender per topic try: scorer = getattr(self.router, "scorer", None) if scorer is not None: for topic in msg.topicIDs: scorer.on_first_delivery(msg_forwarder, topic) except Exception: pass # reject messages claiming to be from ourselves but not locally published self_id = self.host.get_id() if ID(msg.from_id) == self_id and msg_forwarder != self_id: logger.debug( "dropping message claiming to be from self but forwarded from %s", msg_forwarder, ) return self.notify_subscriptions(msg) await self.router.publish(msg_forwarder, msg) def _next_seqno(self) -> bytes: """Make the next message sequence id.""" self.counter += 1 return self.counter.to_bytes(8, "big") def _is_msg_seen(self, msg: rpc_pb2.Message) -> bool: msg_id = self._msg_id_constructor(msg) return self.seen_messages.has(msg_id) def _mark_msg_seen(self, msg: rpc_pb2.Message) -> None: msg_id = self._msg_id_constructor(msg) self.seen_messages.add(msg_id) def _is_subscribed_to_msg(self, msg: rpc_pb2.Message) -> bool: return any(topic in self.topic_ids for topic in msg.topicIDs) def get_message_id(self, msg: rpc_pb2.Message) -> bytes: """ Get the message ID for a given message using the configured message ID constructor. This method provides a public interface for external components (like routers) to access message ID construction functionality. :param msg: the message to get the ID for :return: the message ID as bytes """ return self._msg_id_constructor(msg) async def write_msg(self, stream: INetStream, rpc_msg: rpc_pb2.RPC) -> bool: """ Write an RPC message to a stream with proper error handling. Implements WriteMsg similar to go-msgio which is used in go-libp2p Ref: https://github.com/libp2p/go-msgio/blob/master/protoio/uvarint_writer.go#L56 :param stream: stream to write the message to :param rpc_msg: RPC message to write :return: True if successful, False if stream was closed (StreamClosed) or reset (StreamReset) """ try: # Calculate message size first msg_bytes = rpc_msg.SerializeToString() msg_size = len(msg_bytes) # Calculate varint size and allocate exact buffer size needed varint_bytes = encode_uvarint(msg_size) varint_size = len(varint_bytes) # Allocate buffer with exact size (like Go's pool.Get()) buf = bytearray(varint_size + msg_size) # Write varint length prefix to buffer (like Go's binary.PutUvarint()) buf[:varint_size] = varint_bytes # Write serialized message after varint (like Go's rpc.MarshalTo()) buf[varint_size:] = msg_bytes # Single write operation (like Go's s.Write(buf)) await stream.write(bytes(buf)) return True except (StreamClosed, StreamReset): peer_id = stream.muxed_conn.peer_id logger.debug("Fail to write message to %s: stream closed or reset", peer_id) self._handle_dead_peer(peer_id) return False async def _validation_cache_cleanup(self) -> None: """ Periodically clean up expired validation cache entries. """ while self.manager.is_running: await trio.sleep(60) # Clean up every minute try: self.validation_cache.clear_expired() except Exception as e: logger.debug("Error during validation cache cleanup: %s", e)