Feat/fix mplex stream muxer issues

.gitignore

@@ -30,6 +30,7 @@ downloads/
 wheels/
 MANIFEST
 pip-wheel-metadata
+.ruff_cache
 
 # Installer logs
 pip-log.txt

libp2p/abc.py

@@ -1490,6 +1490,103 @@ async def listen_close(self, network: "INetwork", multiaddr: Multiaddr) -> None:
         """
 
 
+class IMultiselectCommunicator(ABC):
+    """
+    Communicator helper for multiselect.
+
+    Ensures that both the client and multistream module follow the same
+    multistream protocol.
+    """
+
+    @abstractmethod
+    async def write(self, msg_str: str) -> None:
+        """
+        Write a message to the stream.
+
+        Parameters
+        ----------
+        msg_str : str
+            The message string to write.
+
+        """
+
+    @abstractmethod
+    async def read(self) -> str:
+        """
+        Read a message from the stream until EOF.
+
+        Returns
+        -------
+        str
+            The message read from the stream.
+
+        """
+
+
+# -------------------------- multiselect_muxer interface.py --------------------------
+
+
+class IMultiselectMuxer(ABC):
+    """
+    Multiselect module for protocol negotiation.
+
+    Responsible for responding to a multiselect client by selecting a protocol
+    and its corresponding handler for communication.
+    """
+
+    handlers: dict[TProtocol | None, StreamHandlerFn | None]
+
+    @abstractmethod
+    def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
+        """
+        Store a handler for the specified protocol.
+
+        Parameters
+        ----------
+        protocol : TProtocol
+            The protocol name.
+        handler : StreamHandlerFn
+            The handler function associated with the protocol.
+
+        """
+
+    def get_protocols(self) -> tuple[TProtocol | None, ...]:
+        """
+        Retrieve the protocols for which handlers have been registered.
+
+        Returns
+        -------
+        tuple[TProtocol, ...]
+            A tuple of registered protocol names.
+
+        """
+        return tuple(self.handlers.keys())
+
+    @abstractmethod
+    async def negotiate(
+        self, communicator: IMultiselectCommunicator
+    ) -> tuple[TProtocol | None, StreamHandlerFn | None]:
+        """
+        Negotiate a protocol selection with a multiselect client.
+
+        Parameters
+        ----------
+        communicator : IMultiselectCommunicator
+            The communicator used to negotiate the protocol.
+
+        Returns
+        -------
+        tuple[TProtocol, StreamHandlerFn]
+            A tuple containing the selected protocol and its handler.
+
+        Raises
+        ------
+        Exception
+            If negotiation fails.
+
+        """
+
+
 # -------------------------- host interface.py --------------------------
 
 
@@ -1551,13 +1648,14 @@ def get_network(self) -> INetworkService:
         """
 
     @abstractmethod
-    def get_mux(self) -> "Multiselect":
+    def get_mux(self) -> IMultiselectMuxer:
+
         """
         Retrieve the muxer instance for the host.
 
         Returns
         -------
-        Any
+        IMultiselectMuxer
             The muxer instance of the host.
 
         """
@@ -2020,39 +2118,6 @@ def is_expired(self) -> bool:
 # ------------------ multiselect_communicator interface.py ------------------
 
 
-class IMultiselectCommunicator(ABC):
-    """
-    Communicator helper for multiselect.
-
-    Ensures that both the client and multistream module follow the same
-    multistream protocol.
-    """
-
-    @abstractmethod
-    async def write(self, msg_str: str) -> None:
-        """
-        Write a message to the stream.
-
-        Parameters
-        ----------
-        msg_str : str
-            The message string to write.
-
-        """
-
-    @abstractmethod
-    async def read(self) -> str:
-        """
-        Read a message from the stream until EOF.
-
-        Returns
-        -------
-        str
-            The message read from the stream.
-
-        """
-
-
 # -------------------------- multiselect_client interface.py --------------------------
 
 
@@ -2197,8 +2262,6 @@ async def negotiate(
             If negotiation fails.
 
         """
-
-
 # -------------------------- routing interface.py --------------------------
 
 

libp2p/host/basic_host.py

@@ -16,6 +16,7 @@
 
 from libp2p.abc import (
     IHost,
+    IMultiselectMuxer,
     INetConn,
     INetStream,
     INetworkService,
@@ -130,7 +131,7 @@ def get_peerstore(self) -> IPeerStore:
         """
         return self.peerstore
 
-    def get_mux(self) -> Multiselect:
+    def get_mux(self) -> IMultiselectMuxer:
         """
         :return: mux instance of host
         """
@@ -274,8 +275,13 @@ async def close(self) -> None:
 
     # Reference: `BasicHost.newStreamHandler` in Go.
     async def _swarm_stream_handler(self, net_stream: INetStream) -> None:
+        """
+        Handles incoming network streams by performing protocol negotiation
+        and dispatching to the appropriate handler.
+        """
         # Perform protocol muxing to determine protocol to use
         try:
+            # The protocol returned here can now be None
             protocol, handler = await self.multiselect.negotiate(
                 MultiselectCommunicator(net_stream), self.negotiate_timeout
             )
@@ -286,6 +292,24 @@ async def _swarm_stream_handler(self, net_stream: INetStream) -> None:
             )
             await net_stream.reset()
             return
+
+        # Handle case where protocol is None
+        if protocol is None:
+            peer_id = net_stream.muxed_conn.peer_id
+            logger.debug(
+                "No protocol selected by peer %s during negotiation. Resetting stream.",
+                peer_id,
+            )
+            await net_stream.reset()
+            # The BasicHost analysis suggested raising StreamFailure here.
+            # However, the current structure of _swarm_stream_handler
+            # just returns on failure, so let's maintain that pattern
+            # for now, unless further analysis suggests a raise is better.
+            # For strict adherence to the analysis, it might be:
+            # raise StreamFailure(f"No protocol selected from peer {peer_id}")
+            # But the 'return' is consistent with the `except` block's handling.
+            return
+
         net_stream.set_protocol(protocol)
         if handler is None:
             logger.debug(
@@ -322,4 +346,4 @@ def get_peer_connection_info(self, peer_id: ID) -> INetConn | None:
         :param peer_id: ID of the peer to get info for
         :return: Connection object if peer is connected, None otherwise
         """
-        return self._network.connections.get(peer_id)
+        return self._network.connections.get(peer_id)

libp2p/protocol_muxer/multiselect.py

@@ -48,16 +48,15 @@ def add_handler(
         """
         self.handlers[protocol] = handler
 
-    # FIXME: Make TProtocol Optional[TProtocol] to keep types consistent
     async def negotiate(
         self,
         communicator: IMultiselectCommunicator,
         negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
-    ) -> tuple[TProtocol, StreamHandlerFn | None]:
+    ) -> tuple[TProtocol | None, StreamHandlerFn | None]:
         """
-        Negotiate performs protocol selection.
+        Negotiate performs protocol selection with a multiselect client.
 
-        :param stream: stream to negotiate on
+        :param communicator: The communicator used to negotiate the protocol.
         :param negotiate_timeout: timeout for negotiation
         :return: selected protocol name, handler function
         :raise MultiselectError: raised when negotiation failed
@@ -144,4 +143,4 @@ def is_valid_handshake(handshake_contents: str) -> bool:
     :param handshake_contents: contents of handshake message
     :return: true if handshake is complete, false otherwise
     """
-    return handshake_contents == MULTISELECT_PROTOCOL_ID
+    return handshake_contents == MULTISELECT_PROTOCOL_ID

libp2p/security/security_multistream.py

@@ -17,6 +17,7 @@
 from libp2p.peer.id import (
     ID,
 )
+from libp2p.protocol_muxer.exceptions import MultiselectError
 from libp2p.protocol_muxer.multiselect import (
     Multiselect,
 )
@@ -26,6 +27,7 @@
 from libp2p.protocol_muxer.multiselect_communicator import (
     MultiselectCommunicator,
 )
+from libp2p.transport.exceptions import SecurityUpgradeFailure
 
 """
 Represents a secured connection object, which includes a connection and details about
@@ -96,23 +98,36 @@ async def secure_outbound(self, conn: IRawConnection, peer_id: ID) -> ISecureCon
     async def select_transport(
         self, conn: IRawConnection, is_initiator: bool
     ) -> ISecureTransport:
-        """
-        Select a transport that both us and the node on the other end of conn
-        support and agree on.
-
-        :param conn: conn to choose a transport over
-        :param is_initiator: true if we are the initiator, false otherwise
-        :return: selected secure transport
-        """
-        protocol: TProtocol
+        # Note: protocol is TProtocol | None here due to negotiate's new type hint
+        protocol: TProtocol | None  # <--- UPDATE TYPE HINT FOR 'protocol' VARIABLE
         communicator = MultiselectCommunicator(conn)
-        if is_initiator:
-            # Select protocol if initiator
-            protocol = await self.multiselect_client.select_one_of(
-                list(self.transports.keys()), communicator
+
+        # Use a try-except block to catch MultiselectError from negotiate
+        try:
+            if is_initiator:
+                # Select protocol if initiator (multiselect_client.select_one_of should
+                # raise if no protocol)
+                protocol = await self.multiselect_client.select_one_of(
+                    list(self.transports.keys()), communicator
+                )
+            else:
+                # Select protocol if non-initiator
+                # protocol can now be None if negotiate doesn't find a suitable one
+                protocol, _ = await self.multiselect.negotiate(communicator)
+        except MultiselectError as error:
+            # Catch errors from both select_one_of and negotiate, and re-raise as
+            # SecurityUpgradeFailure
+            raise SecurityUpgradeFailure(
+                "failed to negotiate security protocol"
+            ) from error
+
+        # --- NEW CODE: Handle case where protocol is None after negotiation ---
+        if protocol is None:
+            raise SecurityUpgradeFailure(
+                "No security protocol selected during negotiation"
             )
-        else:
-            # Select protocol if non-initiator
-            protocol, _ = await self.multiselect.negotiate(communicator)
-        # Return transport from protocol
+        # --- END NEW CODE ---
+
+        # protocol is guaranteed to be TProtocol here, so no TProtocol(protocol) cast
+        # needed
         return self.transports[protocol]