py-libp2p/libp2p/stream_muxer/mplex/mplex.py

import math
import asyncio
import logging
from typing import Any  # noqa: F401
from typing import Awaitable, Dict, List, Optional, Tuple

import trio
from async_service import Service

from libp2p.exceptions import ParseError
from libp2p.io.exceptions import IncompleteReadError
from libp2p.network.connection.exceptions import RawConnError
from libp2p.peer.id import ID
from libp2p.security.secure_conn_interface import ISecureConn
from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
from libp2p.typing import TProtocol
from libp2p.utils import (
    decode_uvarint_from_stream,
    encode_uvarint,
    encode_varint_prefixed,
    read_varint_prefixed_bytes,
)

from .constants import HeaderTags
from .datastructures import StreamID
from .exceptions import MplexUnavailable
from .mplex_stream import MplexStream

MPLEX_PROTOCOL_ID = TProtocol("/mplex/6.7.0")

logger = logging.getLogger("libp2p.stream_muxer.mplex.mplex")


class Mplex(IMuxedConn, Service):
    """
    reference: https://github.com/libp2p/go-mplex/blob/master/multiplex.go
    """

    secured_conn: ISecureConn
    peer_id: ID
    next_channel_id: int
    streams: Dict[StreamID, MplexStream]
    streams_lock: trio.Lock
    streams_msg_channels: Dict[StreamID, "trio.MemorySendChannel[bytes]"]
    new_stream_send_channel: "trio.MemorySendChannel[IMuxedStream]"
    new_stream_receive_channel: "trio.MemoryReceiveChannel[IMuxedStream]"

    event_shutting_down: trio.Event
    event_closed: trio.Event

    def __init__(self, secured_conn: ISecureConn, peer_id: ID) -> None:
        """
        create a new muxed connection.

        :param secured_conn: an instance of ``ISecureConn``
        :param generic_protocol_handler: generic protocol handler
        for new muxed streams
        :param peer_id: peer_id of peer the connection is to
        """
        self.secured_conn = secured_conn

        self.next_channel_id = 0

        # Set peer_id
        self.peer_id = peer_id

        # Mapping from stream ID -> buffer of messages for that stream
        self.streams = {}
        self.streams_lock = trio.Lock()
        self.streams_msg_channels = {}
        send_channel, receive_channel = trio.open_memory_channel(math.inf)
        self.new_stream_send_channel = send_channel
        self.new_stream_receive_channel = receive_channel
        self.event_shutting_down = trio.Event()
        self.event_closed = trio.Event()

    async def run(self):
        self.manager.run_task(self.handle_incoming)
        await self.manager.wait_finished()

    @property
    def is_initiator(self) -> bool:
        return self.secured_conn.is_initiator

    async def close(self) -> None:
        """close the stream muxer and underlying secured connection."""
        if self.event_shutting_down.is_set():
            return
        # Set the `event_shutting_down`, to allow graceful shutdown.
        self.event_shutting_down.set()
        await self.secured_conn.close()
        # Blocked until `close` is finally set.
        await self.event_closed.wait()

    def is_closed(self) -> bool:
        """
        check connection is fully closed.

        :return: true if successful
        """
        return self.event_closed.is_set()

    def _get_next_channel_id(self) -> int:
        """
        Get next available stream id.

        :return: next available stream id for the connection
        """
        next_id = self.next_channel_id
        self.next_channel_id += 1
        return next_id

    async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream:
        # Use an unbounded buffer, to avoid `handle_incoming` being blocked when doing
        # `send_channel.send`.
        send_channel, receive_channel = trio.open_memory_channel(math.inf)
        stream = MplexStream(name, stream_id, self, receive_channel)
        async with self.streams_lock:
            self.streams[stream_id] = stream
            self.streams_msg_channels[stream_id] = send_channel
        return stream

    async def open_stream(self) -> IMuxedStream:
        """
        creates a new muxed_stream.

        :return: a new ``MplexStream``
        """
        channel_id = self._get_next_channel_id()
        stream_id = StreamID(channel_id=channel_id, is_initiator=True)
        # Default stream name is the `channel_id`
        name = str(channel_id)
        stream = await self._initialize_stream(stream_id, name)
        await self.send_message(HeaderTags.NewStream, name.encode(), stream_id)
        return stream

    async def accept_stream(self) -> IMuxedStream:
        """accepts a muxed stream opened by the other end."""
        try:
            return await self.new_stream_receive_channel.receive()
        except (trio.ClosedResourceError, trio.EndOfChannel):
            raise MplexUnavailable

    async def send_message(
        self, flag: HeaderTags, data: Optional[bytes], stream_id: StreamID
    ) -> int:
        """
        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
        """
        print(
            f"!@# send_message: {self._id}: flag={flag}, data={data}, stream_id={stream_id}"
        )
        # << by 3, then or with flag
        header = encode_uvarint((stream_id.channel_id << 3) | flag.value)

        if data is None:
            data = b""

        _bytes = header + encode_varint_prefixed(data)

        return await self.write_to_stream(_bytes)

    async def write_to_stream(self, _bytes: bytes) -> int:
        """
        writes a byte array to a secured connection.

        :param _bytes: byte array to write
        :return: length written
        """
        await self.secured_conn.write(_bytes)
        return len(_bytes)

    async def handle_incoming(self) -> None:
        """Read a message off of the secured connection and add it to the
        corresponding message buffer."""

        while self.manager.is_running:
            try:
                print(
                    f"!@# handle_incoming: {self._id}: before _handle_incoming_message"
                )
                await self._handle_incoming_message()
                print(
                    f"!@# handle_incoming: {self._id}: after _handle_incoming_message"
                )
            except MplexUnavailable as e:
                logger.debug("mplex unavailable while waiting for incoming: %s", e)
                print(f"!@# handle_incoming: {self._id}: MplexUnavailable: {e}")
                break

        print(f"!@# handle_incoming: {self._id}: leaving")
        # If we enter here, it means this connection is shutting down.
        # We should clean things up.
        await self._cleanup()

    async def read_message(self) -> Tuple[int, int, bytes]:
        """
        Read a single message off of the secured connection.

        :return: stream_id, flag, message contents
        """

        try:
            header = await decode_uvarint_from_stream(self.secured_conn)
        except (ParseError, RawConnError, IncompleteReadError) as error:
            raise MplexUnavailable(
                f"failed to read the header correctly from the underlying connection: {error}"
            )
        try:
            message = await read_varint_prefixed_bytes(self.secured_conn)
        except (ParseError, RawConnError, IncompleteReadError) as error:
            raise MplexUnavailable(
                "failed to read the message body correctly from the underlying connection: "
                f"{error}"
            )

        flag = header & 0x07
        channel_id = header >> 3

        return channel_id, flag, message

    @property
    def _id(self) -> int:
        return 0 if self.is_initiator else 1

    async def _handle_incoming_message(self) -> None:
        """
        Read and handle a new incoming message.

        :raise MplexUnavailable: `Mplex` encounters fatal error or is shutting down.
        """
        print(f"!@# _handle_incoming_message: {self._id}: before reading")
        channel_id, flag, message = await self.read_message()
        print(
            f"!@# _handle_incoming_message: {self._id}: channel_id={channel_id}, flag={flag}, message={message}"
        )
        stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
        print(f"!@# _handle_incoming_message: {self._id}: 2")

        if flag == HeaderTags.NewStream.value:
            print(f"!@# _handle_incoming_message: {self._id}: 3")
            await self._handle_new_stream(stream_id, message)
            print(f"!@# _handle_incoming_message: {self._id}: 4")
        elif flag in (
            HeaderTags.MessageInitiator.value,
            HeaderTags.MessageReceiver.value,
        ):
            print(f"!@# _handle_incoming_message: {self._id}: 5")
            await self._handle_message(stream_id, message)
            print(f"!@# _handle_incoming_message: {self._id}: 6")
        elif flag in (HeaderTags.CloseInitiator.value, HeaderTags.CloseReceiver.value):
            print(f"!@# _handle_incoming_message: {self._id}: 7")
            await self._handle_close(stream_id)
            print(f"!@# _handle_incoming_message: {self._id}: 8")
        elif flag in (HeaderTags.ResetInitiator.value, HeaderTags.ResetReceiver.value):
            print(f"!@# _handle_incoming_message: {self._id}: 9")
            await self._handle_reset(stream_id)
            print(f"!@# _handle_incoming_message: {self._id}: 10")
        else:
            print(f"!@# _handle_incoming_message: {self._id}: 11")
            # Receives messages with an unknown flag
            # TODO: logging
            async with self.streams_lock:
                print(f"!@# _handle_incoming_message: {self._id}: 12")
                if stream_id in self.streams:
                    print(f"!@# _handle_incoming_message: {self._id}: 13")
                    stream = self.streams[stream_id]
                    await stream.reset()
            print(f"!@# _handle_incoming_message: {self._id}: 14")

    async def _handle_new_stream(self, stream_id: StreamID, message: bytes) -> None:
        async with self.streams_lock:
            if stream_id in self.streams:
                # `NewStream` for the same id is received twice...
                raise MplexUnavailable(
                    f"received NewStream message for existing stream: {stream_id}"
                )
        mplex_stream = await self._initialize_stream(stream_id, message.decode())
        try:
            await self.new_stream_send_channel.send(mplex_stream)
        except (trio.BrokenResourceError, trio.EndOfChannel):
            raise MplexUnavailable

    async def _handle_message(self, stream_id: StreamID, message: bytes) -> None:
        print(
            f"!@# _handle_message: {self._id}: stream_id={stream_id}, message={message}"
        )
        async with self.streams_lock:
            print(f"!@# _handle_message: {self._id}: 1")
            if stream_id not in self.streams:
                # We receive a message of the stream `stream_id` which is not accepted
                #   before. It is abnormal. Possibly disconnect?
                # TODO: Warn and emit logs about this.
                print(f"!@# _handle_message: {self._id}: 2")
                return
            print(f"!@# _handle_message: {self._id}: 3")
            stream = self.streams[stream_id]
            send_channel = self.streams_msg_channels[stream_id]
        async with stream.close_lock:
            print(f"!@# _handle_message: {self._id}: 4")
            if stream.event_remote_closed.is_set():
                print(f"!@# _handle_message: {self._id}: 5")
                # TODO: Warn "Received data from remote after stream was closed by them. (len = %d)"  # noqa: E501
                return
        print(f"!@# _handle_message: {self._id}: 6")
        await send_channel.send(message)
        print(f"!@# _handle_message: {self._id}: 7")

    async def _handle_close(self, stream_id: StreamID) -> None:
        print(f"!@# _handle_close: {self._id}: step=0")
        async with self.streams_lock:
            if stream_id not in self.streams:
                # Ignore unmatched messages for now.
                return
            stream = self.streams[stream_id]
            send_channel = self.streams_msg_channels[stream_id]
        print(f"!@# _handle_close: {self._id}: step=1")
        await send_channel.aclose()
        print(f"!@# _handle_close: {self._id}: step=2")
        # NOTE: If remote is already closed, then return: Technically a bug
        #   on the other side. We should consider killing the connection.
        async with stream.close_lock:
            if stream.event_remote_closed.is_set():
                return
        print(f"!@# _handle_close: {self._id}: step=3")
        is_local_closed: bool
        async with stream.close_lock:
            stream.event_remote_closed.set()
            is_local_closed = stream.event_local_closed.is_set()
        print(f"!@# _handle_close: {self._id}: step=4")
        # If local is also closed, both sides are closed. Then, we should clean up
        #   the entry of this stream, to avoid others from accessing it.
        if is_local_closed:
            async with self.streams_lock:
                if stream_id in self.streams:
                    del self.streams[stream_id]
        print(f"!@# _handle_close: {self._id}: step=5")

    async def _handle_reset(self, stream_id: StreamID) -> None:
        async with self.streams_lock:
            if stream_id not in self.streams:
                # This is *ok*. We forget the stream on reset.
                return
            stream = self.streams[stream_id]
            send_channel = self.streams_msg_channels[stream_id]
        await send_channel.aclose()
        async with stream.close_lock:
            if not stream.event_remote_closed.is_set():
                stream.event_reset.set()
                stream.event_remote_closed.set()
            # If local is not closed, we should close it.
            if not stream.event_local_closed.is_set():
                stream.event_local_closed.set()
        async with self.streams_lock:
            if stream_id in self.streams:
                del self.streams[stream_id]
                del self.streams_msg_channels[stream_id]

    async def _cleanup(self) -> None:
        if not self.event_shutting_down.is_set():
            self.event_shutting_down.set()
        async with self.streams_lock:
            for stream_id, stream in self.streams.items():
                async with stream.close_lock:
                    if not stream.event_remote_closed.is_set():
                        stream.event_remote_closed.set()
                        stream.event_reset.set()
                        stream.event_local_closed.set()
                send_channel = self.streams_msg_channels[stream_id]
                await send_channel.aclose()
            self.streams = None
        self.event_closed.set()
        await self.new_stream_send_channel.aclose()
        await self.new_stream_receive_channel.aclose()