rename muxed_conn

stream_muxer/mplex/constants.py

@@ -0,0 +1,6 @@
+HEADER_TAGS = {
+    "NEW_STREAM": 0,
+    "MESSAGE": 2,
+    "CLOSE": 4,
+    "RESET": 6
+}

stream_muxer/mplex/mplex.py

@@ -0,0 +1,121 @@
+import asyncio
+from .utils import encode_uvarint, decode_uvarint
+from .mplex_stream import MplexStream
+from ..muxed_connection_interface import IMuxedConn
+
+
+class Mplex(IMuxedConn):
+    """
+    reference: https://github.com/libp2p/go-mplex/blob/master/multiplex.go
+    """
+    def __init__(self, conn, initiator):
+        """
+        create a new muxed connection
+        :param conn: an instance of raw connection
+        :param initiator: boolean to prevent multiplex with self
+        """
+        self.raw_conn = conn
+        self.initiator = initiator
+        self.buffers = {}
+        self.streams = {}
+        self.stream_queue = asyncio.Queue()
+
+        # The initiator need not read upon construction time.
+        # It should read when the user decides that it wants to read from the constructed stream.
+        if not initiator:
+            asyncio.ensure_future(self.handle_incoming())
+
+    def close(self):
+        """
+        close the stream muxer and underlying raw connection
+        """
+        self.raw_conn.close()
+
+    def is_closed(self):
+        """
+        check connection is fully closed
+        :return: true if successful
+        """
+        pass
+
+    async def read_buffer(self, stream_id):
+        # Empty buffer or nonexistent stream
+        # TODO: propagate up timeout exception and catch
+        if stream_id not in self.buffers or not self.buffers[stream_id]:
+            await self.handle_incoming()
+
+        data = self.buffers[stream_id]
+        self.buffers[stream_id] = bytearray()
+        return data
+
+    def open_stream(self, protocol_id, stream_id, peer_id, multi_addr):
+        """
+        creates a new muxed_stream
+        :param protocol_id: protocol_id of stream
+        :param stream_id: stream_id of stream
+        :param peer_id: peer_id that stream connects to
+        :param multi_addr: multi_addr that stream connects to
+        :return: a new stream
+        """
+        stream = MplexStream(stream_id, multi_addr, self)
+        self.streams[stream_id] = stream
+        return stream
+
+    async def accept_stream(self):
+        """
+        accepts a muxed stream opened by the other end
+        :return: the accepted stream
+        """
+        # TODO update to pull out protocol_id from message
+        protocol_id = "/echo/1.0.0"
+        stream_id = await self.stream_queue.get()
+        stream = MplexStream(stream_id, False, self)
+        return stream, stream_id, protocol_id
+
+    async def send_message(self, flag, data, stream_id):
+        """
+        sends a message over the connection
+        :param header: header to use
+        :param data: data to send in the message
+        :param stream_id: stream the message is in
+        :return: True if success
+        """
+        # << by 3, then or with flag
+        header = (stream_id << 3) | flag
+        header = encode_uvarint(header)
+        if data is None:
+            data_length = encode_uvarint(0)
+            _bytes = header + data_length
+        else:
+            data_length = encode_uvarint(len(data))
+            _bytes = header + data_length + data
+
+        return await self.write_to_stream(_bytes)
+
+    async def write_to_stream(self, _bytes):
+        self.raw_conn.writer.write(_bytes)
+        await self.raw_conn.writer.drain()
+        return len(_bytes)
+
+    async def handle_incoming(self):
+        data = bytearray()
+        try:
+            chunk = await asyncio.wait_for(self.raw_conn.reader.read(1024), timeout=5)
+            data += chunk
+
+            header, end_index = decode_uvarint(data, 0)
+            length, end_index = decode_uvarint(data, end_index)
+
+            message = data[end_index:end_index + length + 1]
+
+            # Deal with other types of messages
+            flag = header & 0x07
+            stream_id = header >> 3
+
+            if stream_id not in self.buffers:
+                self.buffers[stream_id] = message
+                await self.stream_queue.put(stream_id)
+            else:
+                self.buffers[stream_id] = self.buffers[stream_id] + message
+        except asyncio.TimeoutError:
+            print('timeout!')

stream_muxer/mplex/mplex_stream.py

@@ -0,0 +1,112 @@
+from .constants import HEADER_TAGS
+from ..muxed_stream_interface import IMuxedStream
+
+
+class MplexStream(IMuxedStream):
+    """
+    reference: https://github.com/libp2p/go-mplex/blob/master/stream.go
+    """
+
+    def __init__(self, stream_id, initiator, muxed_conn):
+        """
+        create new MuxedStream in muxer
+        :param stream_id: stream stream id
+        :param initiator: boolean if this is an initiator
+        :param muxed_conn: muxed connection of this muxed_stream
+        """
+        self.stream_id = stream_id
+        self.initiator = initiator
+        self.muxed_conn = muxed_conn
+
+        self.read_deadline = None
+        self.write_deadline = None
+
+        self.local_closed = False
+        self.remote_closed = False
+
+    def get_flag(self, action):
+        """
+        get header flag based on action for mplex
+        :param action: action type in str
+        :return: int flag
+        """
+        if self.initiator:
+            return HEADER_TAGS[action]
+
+        return HEADER_TAGS[action] - 1
+
+    async def read(self):
+        """
+        read messages associated with stream from buffer til end of file
+        :return: bytes of input
+        """
+        return await self.muxed_conn.read_buffer(self.stream_id)
+
+    async def write(self, data):
+        """
+        write to stream
+        :return: number of bytes written
+        """
+        return await self.muxed_conn.send_message(self.get_flag("MESSAGE"), data, self.stream_id)
+
+    async def close(self):
+        """
+        close stream
+        :return: true if successful
+        """
+
+        if self.local_closed and self.remote_closed:
+            return True
+
+        await self.muxed_conn.send_message(self.get_flag("CLOSE"), None, self.stream_id)
+        self.muxed_conn.streams.pop(self.stream_id)
+
+        self.local_closed = True
+        self.remote_closed = True
+
+        return True
+
+    def reset(self):
+        """
+        closes both ends of the stream
+        tells this remote side to hang up
+        :return: true if successful
+        """
+        # TODO behavior not fully understood
+        pass
+        # if self.local_closed and self.remote_closed:
+        #     return True
+        #
+        # self.muxed_conn.send_message(self.get_flag("RESET"), None, self.id)
+        # self.muxed_conn.streams.pop(self.id, None)
+        #
+        # self.local_closed = True
+        # self.remote_closed = True
+        #
+        # return True
+
+    # TODO deadline not in use
+    def set_deadline(self, ttl):
+        """
+        set deadline for muxed stream
+        :return: True if successful
+        """
+        self.read_deadline = ttl
+        self.write_deadline = ttl
+        return True
+
+    def set_read_deadline(self, ttl):
+        """
+        set read deadline for muxed stream
+        :return: True if successful
+        """
+        self.read_deadline = ttl
+        return True
+
+    def set_write_deadline(self, ttl):
+        """
+        set write deadline for muxed stream
+        :return: True if successful
+        """
+        self.write_deadline = ttl
+        return True

stream_muxer/mplex/utils.py

@@ -0,0 +1,25 @@
+def encode_uvarint(number):
+    """Pack `number` into varint bytes"""
+    buf = b''
+    while True:
+        towrite = number & 0x7f
+        number >>= 7
+        if number:
+            buf += bytes((towrite | 0x80, ))
+        else:
+            buf += bytes((towrite, ))
+            break
+    return buf
+
+def decode_uvarint(buff, index):
+    shift = 0
+    result = 0
+    while True:
+        i = buff[index]
+        result |= (i & 0x7f) << shift
+        shift += 7
+        if not i & 0x80:
+            break
+        index += 1
+
+    return result, index + 1