Add delim_encode and delim_read

- Add `StreamCommunicator` and `RawConnectionCommunicator`, read/write
messages with delim codec, with `IMuxedStream` and `IRawConnection`
respectively.
- Use it in `Multiselect` and `MultiselectClient`.

libp2p/protocol_muxer/multiselect.py

@@ -1,8 +1,7 @@
 from typing import Dict, Tuple
 
-from libp2p.typing import NegotiableTransport, StreamHandlerFn, TProtocol
+from libp2p.typing import StreamHandlerFn, TProtocol
 
-from .multiselect_communicator import MultiselectCommunicator
 from .multiselect_communicator_interface import IMultiselectCommunicator
 from .multiselect_muxer_interface import IMultiselectMuxer
 
@@ -31,7 +30,7 @@ class Multiselect(IMultiselectMuxer):
         self.handlers[protocol] = handler
 
     async def negotiate(
-        self, stream: NegotiableTransport
+        self, communicator: IMultiselectCommunicator
     ) -> Tuple[TProtocol, StreamHandlerFn]:
         """
         Negotiate performs protocol selection
@@ -39,8 +38,6 @@ class Multiselect(IMultiselectMuxer):
         :return: selected protocol name, handler function
         :raise Exception: negotiation failed exception
         """
-        # Create a communicator to handle all communication across the stream
-        communicator = MultiselectCommunicator(stream)
 
         # Perform handshake to ensure multiselect protocol IDs match
         await self.handshake(communicator)
@@ -48,7 +45,7 @@ class Multiselect(IMultiselectMuxer):
         # Read and respond to commands until a valid protocol ID is sent
         while True:
             # Read message
-            command = await communicator.read_stream_until_eof()
+            command = await communicator.read()
 
             # Command is ls or a protocol
             if command == "ls":
@@ -78,7 +75,7 @@ class Multiselect(IMultiselectMuxer):
         await communicator.write(MULTISELECT_PROTOCOL_ID)
 
         # Read in the protocol ID from other party
-        handshake_contents = await communicator.read_stream_until_eof()
+        handshake_contents = await communicator.read()
 
         # Confirm that the protocols are the same
         if not validate_handshake(handshake_contents):

libp2p/protocol_muxer/multiselect_client.py

@@ -1,10 +1,8 @@
 from typing import Sequence
 
-from libp2p.stream_muxer.abc import IMuxedStream
-from libp2p.typing import NegotiableTransport, TProtocol
+from libp2p.typing import TProtocol
 
 from .multiselect_client_interface import IMultiselectClient
-from .multiselect_communicator import MultiselectCommunicator
 from .multiselect_communicator_interface import IMultiselectCommunicator
 
 MULTISELECT_PROTOCOL_ID = "/multistream/1.0.0"
@@ -31,7 +29,7 @@ class MultiselectClient(IMultiselectClient):
         await communicator.write(MULTISELECT_PROTOCOL_ID)
 
         # Read in the protocol ID from other party
-        handshake_contents = await communicator.read_stream_until_eof()
+        handshake_contents = await communicator.read()
 
         # Confirm that the protocols are the same
         if not validate_handshake(handshake_contents):
@@ -40,7 +38,7 @@ class MultiselectClient(IMultiselectClient):
         # Handshake succeeded if this point is reached
 
     async def select_protocol_or_fail(
-        self, protocol: TProtocol, stream: IMuxedStream
+        self, protocol: TProtocol, communicator: IMultiselectCommunicator
     ) -> TProtocol:
         """
         Send message to multiselect selecting protocol
@@ -49,9 +47,6 @@ class MultiselectClient(IMultiselectClient):
         :param stream: stream to communicate with multiselect over
         :return: selected protocol
         """
-        # Create a communicator to handle all communication across the stream
-        communicator = MultiselectCommunicator(stream)
-
         # Perform handshake to ensure multiselect protocol IDs match
         await self.handshake(communicator)
 
@@ -61,7 +56,7 @@ class MultiselectClient(IMultiselectClient):
         return selected_protocol
 
     async def select_one_of(
-        self, protocols: Sequence[TProtocol], stream: NegotiableTransport
+        self, protocols: Sequence[TProtocol], communicator: IMultiselectCommunicator
     ) -> TProtocol:
         """
         For each protocol, send message to multiselect selecting protocol
@@ -71,10 +66,6 @@ class MultiselectClient(IMultiselectClient):
         :param stream: stream to communicate with multiselect over
         :return: selected protocol
         """
-
-        # Create a communicator to handle all communication across the stream
-        communicator = MultiselectCommunicator(stream)
-
         # Perform handshake to ensure multiselect protocol IDs match
         await self.handshake(communicator)
 
@@ -105,7 +96,7 @@ class MultiselectClient(IMultiselectClient):
         await communicator.write(protocol)
 
         # Get what counterparty says in response
-        response = await communicator.read_stream_until_eof()
+        response = await communicator.read()
 
         # Return protocol if response is equal to protocol or raise error
         if response == protocol:

libp2p/protocol_muxer/multiselect_client_interface.py

@@ -1,7 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import Sequence
 
-from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.protocol_muxer.multiselect_communicator_interface import (
+    IMultiselectCommunicator,
+)
 from libp2p.typing import TProtocol
 
 
@@ -13,7 +15,7 @@ class IMultiselectClient(ABC):
 
     @abstractmethod
     async def select_protocol_or_fail(
-        self, protocol: TProtocol, stream: IMuxedStream
+        self, protocol: TProtocol, communicator: IMultiselectCommunicator
     ) -> TProtocol:
         """
         Send message to multiselect selecting protocol
@@ -25,7 +27,7 @@ class IMultiselectClient(ABC):
 
     @abstractmethod
     async def select_one_of(
-        self, protocols: Sequence[TProtocol], stream: IMuxedStream
+        self, protocols: Sequence[TProtocol], communicator: IMultiselectCommunicator
     ) -> TProtocol:
         """
         For each protocol, send message to multiselect selecting protocol

libp2p/protocol_muxer/multiselect_communicator.py

@@ -1,35 +1,47 @@
-from libp2p.typing import NegotiableTransport
+from libp2p.network.connection.raw_connection_interface import IRawConnection
+from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.stream_muxer.mplex.utils import decode_uvarint_from_stream, encode_uvarint
+from libp2p.typing import StreamReader
 
 from .multiselect_communicator_interface import IMultiselectCommunicator
 
 
-class MultiselectCommunicator(IMultiselectCommunicator):
-    """
-    Communicator helper class that ensures both the client
-    and multistream module will follow the same multistream protocol,
-    which is necessary for them to work
-    """
+def delim_encode(msg_str: str) -> bytes:
+    msg_bytes = msg_str.encode()
+    varint_len_msg = encode_uvarint(len(msg_bytes) + 1)
+    return varint_len_msg + msg_bytes + b"\n"
 
-    reader_writer: NegotiableTransport
 
-    def __init__(self, reader_writer: NegotiableTransport) -> None:
-        """
-        MultistreamCommunicator expects a reader_writer object that has
-        an async read and an async write function (this could be a stream,
-        raw connection, or other object implementing those functions)
-        """
-        self.reader_writer = reader_writer
+async def delim_read(reader: StreamReader, timeout: int = 10) -> str:
+    len_msg = await decode_uvarint_from_stream(reader, timeout)
+    msg_bytes = await reader.read(len_msg)
+    return msg_bytes.decode().rstrip()
+
+
+class RawConnectionCommunicator(IMultiselectCommunicator):
+    conn: IRawConnection
+
+    def __init__(self, conn: IRawConnection) -> None:
+        self.conn = conn
 
     async def write(self, msg_str: str) -> None:
-        """
-        Write message to reader_writer
-        :param msg_str: message to write
-        """
-        await self.reader_writer.write(msg_str.encode())
+        msg_bytes = delim_encode(msg_str)
+        self.conn.writer.write(msg_bytes)
+        await self.conn.writer.drain()
 
-    async def read_stream_until_eof(self) -> str:
-        """
-        Reads message from reader_writer until EOF
-        """
-        read_str = (await self.reader_writer.read()).decode()
-        return read_str
+    async def read(self) -> str:
+        return await delim_read(self.conn.reader)
+
+
+class StreamCommunicator(IMultiselectCommunicator):
+    stream: IMuxedStream
+
+    def __init__(self, stream: IMuxedStream) -> None:
+        self.stream = stream
+
+    async def write(self, msg_str: str) -> None:
+        msg_bytes = delim_encode(msg_str)
+        await self.stream.write(msg_bytes)
+
+    async def read(self) -> str:
+        return await delim_read(self.stream)

libp2p/protocol_muxer/multiselect_communicator_interface.py

@@ -16,7 +16,7 @@ class IMultiselectCommunicator(ABC):
         """
 
     @abstractmethod
-    async def read_stream_until_eof(self) -> str:
+    async def read(self) -> str:
         """
         Reads message from stream until EOF
         """

libp2p/protocol_muxer/multiselect_muxer_interface.py

@@ -1,7 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Tuple
 
-from libp2p.typing import NegotiableTransport, StreamHandlerFn, TProtocol
+from libp2p.typing import StreamHandlerFn, TProtocol
+
+from .multiselect_communicator_interface import IMultiselectCommunicator
 
 
 class IMultiselectMuxer(ABC):
@@ -23,7 +25,7 @@ class IMultiselectMuxer(ABC):
 
     @abstractmethod
     async def negotiate(
-        self, stream: NegotiableTransport
+        self, communicator: IMultiselectCommunicator
     ) -> Tuple[TProtocol, StreamHandlerFn]:
         """
         Negotiate performs protocol selection