py-libp2p/muxer/mplex/muxed_connection.py

import asyncio
from .utils import encode_uvarint, decode_uvarint
from .muxed_connection_interface import IMuxedConn
from .muxed_stream import MuxedStream

class MuxedConn(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()

        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
        :return: a new stream
        """
        stream = MuxedStream(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 = MuxedStream(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)
        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!')