hivemind/hivemind/p2p/p2p_daemon_bindings/control.py

"""
Originally taken from: https://github.com/mhchia/py-libp2p-daemon-bindings
Licence: MIT
Author: Kevin Mai-Husan Chia
"""

import asyncio
from contextlib import asynccontextmanager, closing
from typing import AsyncIterator, Awaitable, Callable, Dict, Iterable, Optional, Sequence, Tuple
from uuid import UUID, uuid4

from multiaddr import Multiaddr, protocols

from hivemind.p2p.p2p_daemon_bindings.datastructures import PeerID, PeerInfo, StreamInfo
from hivemind.p2p.p2p_daemon_bindings.utils import DispatchFailure, raise_if_failed, read_pbmsg_safe, write_pbmsg
from hivemind.proto import p2pd_pb2 as p2pd_pb
from hivemind.utils.logging import get_logger

StreamHandler = Callable[[StreamInfo, asyncio.StreamReader, asyncio.StreamWriter], Awaitable[None]]

SUPPORT_CONN_PROTOCOLS = (
    protocols.P_IP4,
    # protocols.P_IP6,
    protocols.P_UNIX,
)
SUPPORTED_PROTOS = (protocols.protocol_with_code(proto) for proto in SUPPORT_CONN_PROTOCOLS)
logger = get_logger(__name__)

DEFAULT_MAX_MSG_SIZE = 4 * 1024 ** 2


def parse_conn_protocol(maddr: Multiaddr) -> int:
    proto_codes = set(proto.code for proto in maddr.protocols())
    proto_cand = proto_codes.intersection(SUPPORT_CONN_PROTOCOLS)
    if len(proto_cand) != 1:
        raise ValueError(
            f"connection protocol should be only one protocol out of {SUPPORTED_PROTOS}" f", maddr={maddr}"
        )
    return tuple(proto_cand)[0]


class DaemonConnector:
    DEFAULT_CONTROL_MADDR = "/unix/tmp/p2pd.sock"

    def __init__(self, control_maddr: Multiaddr = Multiaddr(DEFAULT_CONTROL_MADDR)) -> None:
        self.control_maddr = control_maddr
        self.proto_code = parse_conn_protocol(self.control_maddr)

    async def open_connection(self) -> (asyncio.StreamReader, asyncio.StreamWriter):
        if self.proto_code == protocols.P_UNIX:
            control_path = self.control_maddr.value_for_protocol(protocols.P_UNIX)
            return await asyncio.open_unix_connection(control_path)
        elif self.proto_code == protocols.P_IP4:
            host = self.control_maddr.value_for_protocol(protocols.P_IP4)
            port = int(self.control_maddr.value_for_protocol(protocols.P_TCP))
            return await asyncio.open_connection(host, port)
        else:
            raise ValueError(f"Protocol not supported: {protocols.protocol_with_code(self.proto_code)}")

    async def open_persistent_connection(self) -> (asyncio.StreamReader, asyncio.StreamWriter):
        """
        Open connection to daemon and upgrade it to a persistent one
        """
        reader, writer = await self.open_connection()
        req = p2pd_pb.Request(type=p2pd_pb.Request.PERSISTENT_CONN_UPGRADE)
        await write_pbmsg(writer, req)

        response = p2pd_pb.Response()
        await read_pbmsg_safe(reader, response)

        if response.type == "ERROR":
            raise P2PDaemonError(response.error.msg)

        return reader, writer


TUnaryHandler = Callable[[bytes, PeerID], Awaitable[bytes]]
CallID = UUID


class ControlClient:
    DEFAULT_LISTEN_MADDR = "/unix/tmp/p2pclient.sock"

    def __init__(
        self,
        daemon_connector: DaemonConnector,
        listen_maddr: Multiaddr = Multiaddr(DEFAULT_LISTEN_MADDR),
        *,
        _initialized_with_create: bool = False,
        persistent_conn_max_msg_size: int = DEFAULT_MAX_MSG_SIZE,
    ) -> None:
        assert _initialized_with_create, "Please use ControlClient.create coroutine to spawn new control instances"

        self.persistent_conn_max_msg_size = persistent_conn_max_msg_size

        self.listen_maddr = listen_maddr
        self.daemon_connector = daemon_connector
        self.handlers: Dict[str, StreamHandler] = {}

        self.unary_handlers: Dict[str, TUnaryHandler] = {}

        self._pending_messages: asyncio.Queue[p2pd_pb.PersistentConnectionRequest] = asyncio.Queue()
        self._pending_calls: Dict[CallID, asyncio.Future[bytes]] = {}
        self._handler_tasks: Dict[CallID, asyncio.Task] = {}

        self._read_task: Optional[asyncio.Task] = None
        self._write_task: Optional[asyncio.Task] = None

    @classmethod
    async def create(
        cls,
        daemon_connector: DaemonConnector,
        listen_maddr: Multiaddr = Multiaddr(DEFAULT_LISTEN_MADDR),
        use_persistent_conn: bool = True,
        persistent_conn_max_msg_size=2 << 22,
    ) -> "ControlClient":
        control = cls(
            daemon_connector,
            listen_maddr,
            _initialized_with_create=True,
            persistent_conn_max_msg_size=persistent_conn_max_msg_size,
        )

        if use_persistent_conn:
            await control._ensure_persistent_conn()

        return control

    def close(self) -> None:
        if self._read_task is not None:
            self._read_task.cancel()
        if self._write_task is not None:
            self._write_task.cancel()

    def __del__(self):
        self.close()

    async def _handler(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        pb_stream_info = p2pd_pb.StreamInfo()  # type: ignore
        await read_pbmsg_safe(reader, pb_stream_info)
        stream_info = StreamInfo.from_protobuf(pb_stream_info)
        try:
            handler = self.handlers[stream_info.proto]
        except KeyError as e:
            # should never enter here... daemon should reject the stream for us.
            writer.close()
            raise DispatchFailure(e)
        await handler(stream_info, reader, writer)

    @asynccontextmanager
    async def listen(self) -> AsyncIterator["ControlClient"]:
        proto_code = parse_conn_protocol(self.listen_maddr)
        if proto_code == protocols.P_UNIX:
            listen_path = self.listen_maddr.value_for_protocol(protocols.P_UNIX)
            server = await asyncio.start_unix_server(self._handler, path=listen_path)
        elif proto_code == protocols.P_IP4:
            host = self.listen_maddr.value_for_protocol(protocols.P_IP4)
            port = int(self.listen_maddr.value_for_protocol(protocols.P_TCP))
            server = await asyncio.start_server(self._handler, port=port, host=host)
        else:
            raise ValueError(f"Protocol not supported: {protocols.protocol_with_code(proto_code)}")

        async with server:
            yield self

    async def _read_from_persistent_conn(self, reader: asyncio.StreamReader):
        while True:
            resp = p2pd_pb.PersistentConnectionResponse()
            try:
                await read_pbmsg_safe(reader, resp)
            except asyncio.IncompleteReadError:
                break

            call_id = UUID(bytes=resp.callId)

            if resp.HasField("callUnaryResponse"):
                if call_id in self._pending_calls and resp.callUnaryResponse.HasField("response"):
                    self._pending_calls[call_id].set_result(resp.callUnaryResponse.response)
                elif call_id in self._pending_calls and resp.callUnaryResponse.HasField("error"):
                    remote_exc = P2PHandlerError(resp.callUnaryResponse.error.decode(errors="ignore"))
                    self._pending_calls[call_id].set_exception(remote_exc)
                else:
                    logger.debug(f"Received unexpected unary call: {resp}")

            elif resp.HasField("requestHandling"):
                handler_task = asyncio.create_task(self._handle_persistent_request(call_id, resp.requestHandling))
                self._handler_tasks[call_id] = handler_task

            elif call_id in self._handler_tasks and resp.HasField("cancel"):
                self._handler_tasks[call_id].cancel()

            elif call_id in self._pending_calls and resp.HasField("daemonError"):
                daemon_exc = P2PDaemonError(resp.daemonError.message)
                self._pending_calls[call_id].set_exception(daemon_exc)

            elif call_id in self._pending_calls:
                self._pending_calls[call_id].set_result(None)

            else:
                logger.debug(f"Received unexpected response from daemon: {resp}")

    async def _write_to_persistent_conn(self, writer: asyncio.StreamWriter):
        with closing(writer):
            while True:
                msg = await self._pending_messages.get()
                await write_pbmsg(writer, msg)

    async def _handle_persistent_request(self, call_id: UUID, request: p2pd_pb.CallUnaryRequest):
        if request.proto not in self.unary_handlers:
            logger.warning(f"Protocol {request.proto} not supported")
            return

        try:
            remote_id = PeerID(request.peer)
            response_payload: bytes = await self.unary_handlers[request.proto](request.data, remote_id)
            response = p2pd_pb.CallUnaryResponse(response=response_payload)

        except Exception as e:
            response = p2pd_pb.CallUnaryResponse(error=repr(e).encode())

        payload = p2pd_pb.PersistentConnectionRequest(callId=call_id.bytes, unaryResponse=response)
        if payload.ByteSize() <= self.persistent_conn_max_msg_size:
            await self._pending_messages.put(payload)
        else:
            error_msg = p2pd_pb.PersistentConnectionRequest(
                callId=call_id.bytes,
                callUnaryResponse=p2pd_pb.CallUnaryResponse(
                    error=b"response size exceeds message size limit",
                ),
            )
            await self._pending_messages.put(error_msg)

        self._handler_tasks.pop(call_id)

    async def _cancel_unary_call(self, call_id: UUID):
        await self._pending_messages.put(
            p2pd_pb.PersistentConnectionRequest(
                callId=call_id.bytes,
                cancel=p2pd_pb.Cancel(),
            ),
        )

    async def _ensure_persistent_conn(self):
        reader, writer = await self.daemon_connector.open_persistent_connection()

        self._read_task = asyncio.create_task(self._read_from_persistent_conn(reader))
        self._write_task = asyncio.create_task(self._write_to_persistent_conn(writer))

    async def add_unary_handler(self, proto: str, handler: TUnaryHandler):
        call_id = uuid4()

        add_unary_handler_req = p2pd_pb.AddUnaryHandlerRequest(proto=proto)
        req = p2pd_pb.PersistentConnectionRequest(callId=call_id.bytes, addUnaryHandler=add_unary_handler_req)

        if self.unary_handlers.get(proto):
            raise P2PDaemonError(f"Handler for protocol {proto} already registered")
        self.unary_handlers[proto] = handler

        self._pending_calls[call_id] = asyncio.Future()
        await self._pending_messages.put(req)
        await self._pending_calls[call_id]

    async def call_unary_handler(self, peer_id: PeerID, proto: str, data: bytes) -> bytes:
        call_id = uuid4()
        call_unary_req = p2pd_pb.CallUnaryRequest(
            peer=peer_id.to_bytes(),
            proto=proto,
            data=data,
        )
        req = p2pd_pb.PersistentConnectionRequest(
            callId=call_id.bytes,
            callUnary=call_unary_req,
        )

        if req.ByteSize() > self.persistent_conn_max_msg_size:
            raise P2PDaemonError(f"Message size exceeds set limit {self.persistent_conn_max_msg_size}")

        try:
            self._pending_calls[call_id] = asyncio.Future()
            await self._pending_messages.put(req)
            return await self._pending_calls[call_id]

        except asyncio.CancelledError:
            await self._cancel_unary_call(call_id)
            raise

        finally:
            self._pending_calls.pop(call_id, None)

    async def identify(self) -> Tuple[PeerID, Tuple[Multiaddr, ...]]:
        reader, writer = await self.daemon_connector.open_connection()
        req = p2pd_pb.Request(type=p2pd_pb.Request.IDENTIFY)
        await write_pbmsg(writer, req)

        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        writer.close()

        raise_if_failed(resp)
        peer_id_bytes = resp.identify.id
        maddrs_bytes = resp.identify.addrs

        maddrs = tuple(Multiaddr(maddr_bytes) for maddr_bytes in maddrs_bytes)
        peer_id = PeerID(peer_id_bytes)

        return peer_id, maddrs

    async def connect(self, peer_id: PeerID, maddrs: Iterable[Multiaddr]) -> None:
        reader, writer = await self.daemon_connector.open_connection()

        maddrs_bytes = [i.to_bytes() for i in maddrs]
        connect_req = p2pd_pb.ConnectRequest(peer=peer_id.to_bytes(), addrs=maddrs_bytes)
        req = p2pd_pb.Request(type=p2pd_pb.Request.CONNECT, connect=connect_req)
        await write_pbmsg(writer, req)

        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        writer.close()
        raise_if_failed(resp)

    async def list_peers(self) -> Tuple[PeerInfo, ...]:
        req = p2pd_pb.Request(type=p2pd_pb.Request.LIST_PEERS)
        reader, writer = await self.daemon_connector.open_connection()
        await write_pbmsg(writer, req)
        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        writer.close()
        raise_if_failed(resp)

        peers = tuple(PeerInfo.from_protobuf(pinfo) for pinfo in resp.peers)
        return peers

    async def disconnect(self, peer_id: PeerID) -> None:
        disconnect_req = p2pd_pb.DisconnectRequest(peer=peer_id.to_bytes())
        req = p2pd_pb.Request(type=p2pd_pb.Request.DISCONNECT, disconnect=disconnect_req)
        reader, writer = await self.daemon_connector.open_connection()
        await write_pbmsg(writer, req)
        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        writer.close()
        raise_if_failed(resp)

    async def stream_open(
        self, peer_id: PeerID, protocols: Sequence[str]
    ) -> Tuple[StreamInfo, asyncio.StreamReader, asyncio.StreamWriter]:
        reader, writer = await self.daemon_connector.open_connection()

        stream_open_req = p2pd_pb.StreamOpenRequest(peer=peer_id.to_bytes(), proto=list(protocols))
        req = p2pd_pb.Request(type=p2pd_pb.Request.STREAM_OPEN, streamOpen=stream_open_req)
        await write_pbmsg(writer, req)

        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        raise_if_failed(resp)

        pb_stream_info = resp.streamInfo
        stream_info = StreamInfo.from_protobuf(pb_stream_info)

        return stream_info, reader, writer

    async def stream_handler(self, proto: str, handler_cb: StreamHandler) -> None:
        reader, writer = await self.daemon_connector.open_connection()

        listen_path_maddr_bytes = self.listen_maddr.to_bytes()
        stream_handler_req = p2pd_pb.StreamHandlerRequest(addr=listen_path_maddr_bytes, proto=[proto])
        req = p2pd_pb.Request(type=p2pd_pb.Request.STREAM_HANDLER, streamHandler=stream_handler_req)
        await write_pbmsg(writer, req)

        resp = p2pd_pb.Response()  # type: ignore
        await read_pbmsg_safe(reader, resp)
        writer.close()
        raise_if_failed(resp)

        # if success, add the handler to the dict
        self.handlers[proto] = handler_cb


class P2PHandlerError(Exception):
    """
    Raised if remote handled a request with an exception
    """


class P2PDaemonError(Exception):
    """
    Raised if daemon failed to handle request
    """