feat: simplify upgrader with GenericMultistreamSelector, fix multistream_client bug (#313)

aniruddha1295 · aniruddha1295 · commit 0017ebc1b763 · 2026-02-18T23:38:09.000Z
diff --git a/libp2p/protocol_muxer/__init__.py b/libp2p/protocol_muxer/__init__.py
@@ -0,0 +1,5 @@
+from .generic_selector import (
+    GenericMultistreamSelector,
+)
+
+__all__ = ["GenericMultistreamSelector"]
diff --git a/libp2p/protocol_muxer/generic_selector.py b/libp2p/protocol_muxer/generic_selector.py
@@ -0,0 +1,108 @@
+import logging
+from collections import (
+    OrderedDict,
+)
+from typing import (
+    Generic,
+    TypeVar,
+)
+
+from libp2p.custom_types import (
+    TProtocol,
+)
+from libp2p.io.abc import (
+    ReadWriteCloser,
+)
+
+from .exceptions import (
+    MultiselectError,
+)
+from .multiselect import (
+    DEFAULT_NEGOTIATE_TIMEOUT,
+    Multiselect,
+)
+from .multiselect_client import (
+    MultiselectClient,
+)
+from .multiselect_communicator import (
+    MultiselectCommunicator,
+)
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class GenericMultistreamSelector(Generic[T]):
+    """
+    A generic multistream-select protocol negotiator.
+
+    Consolidates the duplicated negotiation logic from SecurityMultistream
+    and MuxerMultistream into a single reusable class. Both security and
+    muxer layers use this selector for negotiation only; the caller is
+    responsible for applying the chosen handler (e.g. securing the connection
+    or instantiating the muxer).
+    """
+
+    handlers: "OrderedDict[TProtocol, T]"
+    multiselect: Multiselect
+    multiselect_client: MultiselectClient
+
+    def __init__(self) -> None:
+        self.handlers = OrderedDict()
+        self.multiselect = Multiselect()
+        self.multiselect_client = MultiselectClient()
+
+    def add_handler(self, protocol: TProtocol, handler: T) -> None:
+        """
+        Add a protocol and its handler.
+
+        Re-adding an existing protocol moves it to the end, updating its
+        precedence in negotiation.
+
+        :param protocol: the protocol name negotiated via multistream-select.
+        :param handler: the handler or factory associated with the protocol.
+        """
+        self.handlers.pop(protocol, None)
+        self.handlers[protocol] = handler
+        self.multiselect.add_handler(protocol, None)
+
+    def get_protocols(self) -> list[TProtocol]:
+        """
+        Return the list of registered protocols in precedence order.
+
+        :return: list of protocol names.
+        """
+        return list(self.handlers.keys())
+
+    async def select(
+        self,
+        conn: ReadWriteCloser,
+        is_initiator: bool,
+        negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
+    ) -> "tuple[TProtocol, T]":
+        """
+        Negotiate and select a protocol with the remote peer.
+
+        :param conn: connection supporting read/write (ReadWriteCloser).
+        :param is_initiator: True if we initiated the connection.
+        :param negotiate_timeout: timeout for negotiation in seconds.
+        :return: tuple of (selected_protocol, handler).
+        :raises MultiselectError: if negotiation fails or no protocol is selected.
+        """
+        communicator = MultiselectCommunicator(conn)
+        protocol: TProtocol | None
+
+        if is_initiator:
+            protocol = await self.multiselect_client.select_one_of(
+                tuple(self.handlers.keys()), communicator, negotiate_timeout
+            )
+        else:
+            protocol, _ = await self.multiselect.negotiate(
+                communicator, negotiate_timeout
+            )
+
+        if protocol is None:
+            raise MultiselectError("Failed to negotiate: no protocol selected")
+
+        return protocol, self.handlers[protocol]
diff --git a/libp2p/security/security_multistream.py b/libp2p/security/security_multistream.py
@@ -21,15 +21,15 @@
 from libp2p.protocol_muxer.exceptions import (
     MultiselectError,
 )
+from libp2p.protocol_muxer.generic_selector import (
+    GenericMultistreamSelector,
+)
 from libp2p.protocol_muxer.multiselect import (
     Multiselect,
 )
 from libp2p.protocol_muxer.multiselect_client import (
     MultiselectClient,
 )
-from libp2p.protocol_muxer.multiselect_communicator import (
-    MultiselectCommunicator,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -48,18 +48,26 @@ class SecurityMultistream(ABC):
     Go implementation: github.com/libp2p/go-conn-security-multistream/ssms.go
     """
 
-    transports: "OrderedDict[TProtocol, ISecureTransport]"
-    multiselect: Multiselect
-    multiselect_client: MultiselectClient
+    _selector: "GenericMultistreamSelector[ISecureTransport]"
 
     def __init__(self, secure_transports_by_protocol: TSecurityOptions) -> None:
-        self.transports = OrderedDict()
-        self.multiselect = Multiselect()
-        self.multiselect_client = MultiselectClient()
+        self._selector = GenericMultistreamSelector()
 
         for protocol, transport in secure_transports_by_protocol.items():
             self.add_transport(protocol, transport)
 
+    @property
+    def transports(self) -> "OrderedDict[TProtocol, ISecureTransport]":
+        return self._selector.handlers
+
+    @property
+    def multiselect(self) -> Multiselect:
+        return self._selector.multiselect
+
+    @property
+    def multiselect_client(self) -> MultiselectClient:
+        return self._selector.multiselect_client
+
     def add_transport(self, protocol: TProtocol, transport: ISecureTransport) -> None:
         """
         Add a protocol and its corresponding transport to multistream-
@@ -69,12 +77,7 @@ def add_transport(self, protocol: TProtocol, transport: ISecureTransport) -> Non
         :param protocol: the protocol name, which is negotiated in multiselect.
         :param transport: the corresponding transportation to the ``protocol``.
         """
-        # If protocol is already added before, remove it and add it again.
-        self.transports.pop(protocol, None)
-        self.transports[protocol] = transport
-        # Note: None is added as the handler for the given protocol since
-        # we only care about selecting the protocol, not any handler function
-        self.multiselect.add_handler(protocol, None)
+        self._selector.add_handler(protocol, transport)
 
     async def secure_inbound(self, conn: IRawConnection) -> ISecureConn:
         """
@@ -113,19 +116,10 @@ async def select_transport(
         :param is_initiator: true if we are the initiator, false otherwise
         :return: selected secure transport
         """
-        protocol: TProtocol | None
-        communicator = MultiselectCommunicator(conn)
-        if is_initiator:
-            # Select protocol if initiator
-            protocol = await self.multiselect_client.select_one_of(
-                list(self.transports.keys()), communicator
-            )
-        else:
-            # Select protocol if non-initiator
-            protocol, _ = await self.multiselect.negotiate(communicator)
-        if protocol is None:
+        try:
+            _, transport = await self._selector.select(conn, is_initiator)
+        except MultiselectError:
             raise MultiselectError(
                 "Failed to negotiate a security protocol: no protocol selected"
             )
-        # Return transport from protocol
-        return self.transports[protocol]
+        return transport
diff --git a/libp2p/stream_muxer/muxer_multistream.py b/libp2p/stream_muxer/muxer_multistream.py
@@ -21,16 +21,16 @@
 from libp2p.protocol_muxer.exceptions import (
     MultiselectError,
 )
+from libp2p.protocol_muxer.generic_selector import (
+    GenericMultistreamSelector,
+)
 from libp2p.protocol_muxer.multiselect import (
     DEFAULT_NEGOTIATE_TIMEOUT,
     Multiselect,
 )
 from libp2p.protocol_muxer.multiselect_client import (
     MultiselectClient,
 )
-from libp2p.protocol_muxer.multiselect_communicator import (
-    MultiselectCommunicator,
-)
 from libp2p.stream_muxer.yamux.yamux import (
     PROTOCOL_ID,
     Yamux,
@@ -46,24 +46,31 @@ class MuxerMultistream:
     go implementation: github.com/libp2p/go-stream-muxer-multistream/multistream.go
     """
 
-    # NOTE: Can be changed to `typing.OrderedDict` since Python 3.7.2.
-    transports: "OrderedDict[TProtocol, TMuxerClass]"
-    multiselect: Multiselect
-    multiselect_client: MultiselectClient
+    _selector: "GenericMultistreamSelector[TMuxerClass]"
     negotiate_timeout: int
 
     def __init__(
         self,
         muxer_transports_by_protocol: TMuxerOptions,
         negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
     ) -> None:
-        self.transports = OrderedDict()
-        self.multiselect = Multiselect()
-        self.multistream_client = MultiselectClient()
+        self._selector = GenericMultistreamSelector()
         self.negotiate_timeout = negotiate_timeout
         for protocol, transport in muxer_transports_by_protocol.items():
             self.add_transport(protocol, transport)
 
+    @property
+    def transports(self) -> "OrderedDict[TProtocol, TMuxerClass]":
+        return self._selector.handlers
+
+    @property
+    def multiselect(self) -> Multiselect:
+        return self._selector.multiselect
+
+    @property
+    def multiselect_client(self) -> MultiselectClient:
+        return self._selector.multiselect_client
+
     def add_transport(self, protocol: TProtocol, transport: TMuxerClass) -> None:
         """
         Add a protocol and its corresponding transport to multistream-
@@ -73,10 +80,7 @@ def add_transport(self, protocol: TProtocol, transport: TMuxerClass) -> None:
         :param protocol: the protocol name, which is negotiated in multiselect.
         :param transport: the corresponding transportation to the ``protocol``.
         """
-        # If protocol is already added before, remove it and add it again.
-        self.transports.pop(protocol, None)
-        self.transports[protocol] = transport
-        self.multiselect.add_handler(protocol, None)
+        self._selector.add_handler(protocol, transport)
 
     async def select_transport(self, conn: IRawConnection) -> TMuxerClass:
         """
@@ -86,48 +90,26 @@ async def select_transport(self, conn: IRawConnection) -> TMuxerClass:
         :param conn: conn to choose a transport over
         :return: selected muxer transport
         """
-        protocol: TProtocol | None
-        communicator = MultiselectCommunicator(conn)
-        if conn.is_initiator:
-            protocol = await self.multiselect_client.select_one_of(
-                tuple(self.transports.keys()), communicator, self.negotiate_timeout
-            )
-        else:
-            protocol, _ = await self.multiselect.negotiate(
-                communicator, self.negotiate_timeout
+        try:
+            _, transport = await self._selector.select(
+                conn, conn.is_initiator, self.negotiate_timeout
             )
-        if protocol is None:
+        except MultiselectError:
             raise MultiselectError(
                 "Fail to negotiate a stream muxer protocol: no protocol selected"
             )
-        return self.transports[protocol]
+        return transport
 
     async def new_conn(self, conn: ISecureConn, peer_id: ID) -> IMuxedConn:
-        communicator = MultiselectCommunicator(conn)
         logger.debug(
             "MuxerMultistream: muxer negotiation peer=%s initiator=%s",
             peer_id,
             conn.is_initiator,
         )
-
-        # Use appropriate multiselect based on role:
-        # - Initiator (client/dialer) uses multiselect_client to select a protocol
-        # - Non-initiator (server/listener) uses multiselect to negotiate with client
-        if conn.is_initiator:
-            protocol = await self.multistream_client.select_one_of(
-                tuple(self.transports.keys()), communicator, self.negotiate_timeout
-            )
-        else:
-            negotiated_protocol, _ = await self.multiselect.negotiate(
-                communicator, self.negotiate_timeout
-            )
-            if negotiated_protocol is None:
-                raise MultiselectError(
-                    "Fail to negotiate a stream muxer protocol: no protocol selected"
-                )
-            protocol = negotiated_protocol
+        protocol, transport_class = await self._selector.select(
+            conn, conn.is_initiator, self.negotiate_timeout
+        )
         logger.debug("MuxerMultistream new_conn: negotiated protocol %s", protocol)
-        transport_class = self.transports[protocol]
         if protocol == PROTOCOL_ID:
             async with trio.open_nursery():
 
diff --git a/newsfragments/313.internal.rst b/newsfragments/313.internal.rst
@@ -0,0 +1 @@
+Simplified upgrader by consolidating duplicated multistream-select negotiation logic from SecurityMultistream and MuxerMultistream into a reusable GenericMultistreamSelector class. Fixed attribute naming bug (multistream_client → multiselect_client) in MuxerMultistream.
diff --git a/tests/core/protocol_muxer/test_generic_selector.py b/tests/core/protocol_muxer/test_generic_selector.py
diff --git a/tests/core/stream_muxer/test_muxer_multistream.py b/tests/core/stream_muxer/test_muxer_multistream.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Simplified upgrader by consolidating duplicated multistream-select negotiation logic from SecurityMultistream and MuxerMultistream into a reusable GenericMultistreamSelector class. Fixed attribute naming bug (multistream_client → multiselect_client) in MuxerMultistream.`