server.py

import abc
import time
import socket
import logging
import asyncio
import warnings

from types import TracebackType
from typing import TYPE_CHECKING, Optional, Collection, Generic, Type, cast
from typing import List, Tuple, Dict, Any, Callable, ContextManager, Set
from contextlib import nullcontext

import h2.config
import h2.exceptions

from multidict import MultiDict

from .utils import DeadlineWrapper, Wrapper
from .const import Status, Cardinality
from .config import Configuration
from .stream import send_message, recv_message, StreamIterator
from .stream import _RecvType, _SendType
from .events import _DispatchServerEvents
from .metadata import Deadline, encode_grpc_message, _Metadata
from .metadata import encode_metadata, decode_metadata, _MetadataLike
from .metadata import _STATUS_DETAILS_KEY, encode_bin_value
from .protocol import H2Protocol, AbstractHandler
from .exceptions import GRPCError, ProtocolError, StreamTerminatedError
from .encoding.base import GRPC_CONTENT_TYPE, CodecBase, StatusDetailsCodecBase
from .encoding.proto import ProtoCodec, ProtoStatusDetailsCodec
from .encoding.proto import _googleapis_available

from ._registry import servers as _servers

if TYPE_CHECKING:
    import ssl as _ssl  # noqa
    from . import const  # noqa
    from . import protocol  # noqa
    from ._typing import IServable  # noqa


log = logging.getLogger(__name__)

_Headers = List[Tuple[str, str]]


class Stream(StreamIterator[_RecvType], Generic[_RecvType, _SendType]):
    """
    Represents gRPC method call – HTTP/2 request/stream, and everything you
    need to communicate with client in order to handle this request.

    As you can see, every method handler accepts single positional argument -
    stream:

    .. code-block:: python3

        async def MakeLatte(self, stream: grpclib.server.Stream):
            task: cafe_pb2.LatteOrder = await stream.recv_message()
            ...
            await stream.send_message(empty_pb2.Empty())

    This is true for every gRPC method type.
    """
    # state
    _send_initial_metadata_done = False
    _send_message_done = False
    _send_trailing_metadata_done = False
    _cancel_done = False

    # stats
    _messages_sent = 0
    _messages_received = 0

    def __init__(
        self,
        stream: 'protocol.Stream',
        method_name: str,
        cardinality: Cardinality,
        recv_type: Type[_RecvType],
        send_type: Type[_SendType],
        *,
        codec: CodecBase,
        status_details_codec: Optional[StatusDetailsCodecBase],
        dispatch: _DispatchServerEvents,
        deadline: Optional[Deadline] = None,
        user_agent: Optional[str] = None,
    ):
        self._stream = stream
        self._method_name = method_name
        self._cardinality = cardinality
        self._recv_type = recv_type
        self._send_type = send_type
        self._codec = codec
        self._status_details_codec = status_details_codec
        self._dispatch = dispatch
        #: :py:class:`~grpclib.metadata.Deadline` of the current request
        self.deadline = deadline
        #: Invocation metadata, received with headers from the client.
        #: Represented as a multi-dict object.
        self.metadata: Optional[_Metadata] = None
        #: Client's user-agent
        self.user_agent = user_agent
        #: Connection's peer info of type :py:class:`~grpclib.protocol.Peer`
        self.peer = self._stream.connection.get_peer()

    @property
    def _content_type(self) -> str:
        return GRPC_CONTENT_TYPE + '+' + self._codec.__content_subtype__

    async def recv_message(self) -> Optional[_RecvType]:
        """Coroutine to receive incoming message from the client.

        If client sends UNARY request, then you can call this coroutine
        only once. If client sends STREAM request, then you should call this
        coroutine several times, until it returns None when the client has
        ended the stream. To simplify your code in this case,
        :py:class:`Stream` class implements async iteration protocol, so
        you can use it like this:

        .. code-block:: python3

            async for message in stream:
                do_smth_with(message)

        or even like this:

        .. code-block:: python3

            messages = [msg async for msg in stream]

        HTTP/2 has flow control mechanism, so server will acknowledge received
        DATA frames as a message only after user consumes this coroutine.

        :returns: message
        """
        message = await recv_message(self._stream, self._codec, self._recv_type)
        if message is not None:
            message, = await self._dispatch.recv_message(message)
            self._messages_received += 1
            self._stream.connection.messages_received += 1
            self._stream.connection.last_message_received = time.monotonic()
            return message
        else:
            return None

    async def send_initial_metadata(
        self,
        *,
        metadata: Optional[_MetadataLike] = None,
    ) -> None:
        """Coroutine to send headers with initial metadata to the client.

        In gRPC you can send initial metadata as soon as possible, because
        gRPC doesn't use `:status` pseudo header to indicate success or failure
        of the current request. gRPC uses trailers for this purpose, and
        trailers are sent during :py:meth:`send_trailing_metadata` call, which
        should be called in the end.

        .. note:: This coroutine will be called implicitly during first
            :py:meth:`send_message` coroutine call, if not called before
            explicitly.

        :param metadata: custom initial metadata, dict or list of pairs
        """
        if self._send_initial_metadata_done:
            raise ProtocolError('Initial metadata was already sent')

        headers = [
            (':status', '200'),
            ('content-type', self._content_type),
        ]
        metadata = MultiDict(metadata or ())
        metadata, = await self._dispatch.send_initial_metadata(metadata)
        headers.extend(encode_metadata(cast(_Metadata, metadata)))

        await self._stream.send_headers(headers)
        self._send_initial_metadata_done = True

    async def send_message(self, message: _SendType) -> None:
        """Coroutine to send message to the client.

        If server sends UNARY response, then you should call this coroutine only
        once. If server sends STREAM response, then you can call this coroutine
        as many times as you need.

        :param message: message object
        """
        if not self._send_initial_metadata_done:
            await self.send_initial_metadata()

        if not self._cardinality.server_streaming:
            if self._send_message_done:
                raise ProtocolError('Message was already sent')

        message, = await self._dispatch.send_message(message)
        await send_message(self._stream, self._codec, message, self._send_type)
        self._send_message_done = True
        self._messages_sent += 1
        self._stream.connection.messages_sent += 1
        self._stream.connection.last_message_sent = time.monotonic()

    async def send_trailing_metadata(
        self,
        *,
        status: Status = Status.OK,
        status_message: Optional[str] = None,
        status_details: Any = None,
        metadata: Optional[_MetadataLike] = None,
    ) -> None:
        """Coroutine to send trailers with trailing metadata to the client.

        This coroutine allows sending trailers-only responses, in case of some
        failure conditions during handling current request, i.e. when
        ``status is not OK``.

        .. note:: This coroutine will be called implicitly at exit from
            request handler, with appropriate status code, if not called
            explicitly during handler execution.

        :param status: resulting status of this coroutine call
        :param status_message: description for a status
        :param metadata: custom trailing metadata, dict or list of pairs
        """
        if self._send_trailing_metadata_done:
            raise ProtocolError('Trailing metadata was already sent')

        if (
            not self._cardinality.server_streaming
            and not self._send_message_done
            and status is Status.OK
        ):
            raise ProtocolError('Unary response with OK status requires '
                                'a single message to be sent')

        if self._send_initial_metadata_done:
            headers: _Headers = []
        else:
            # trailers-only response
            headers = [
                (':status', '200'),
                ('content-type', self._content_type),
            ]

        headers.append(('grpc-status', str(status.value)))
        if status_message is not None:
            headers.append(('grpc-message',
                            encode_grpc_message(status_message)))
        if (
            status_details is not None
            and self._status_details_codec is not None
        ):
            status_details_bin = (
                encode_bin_value(self._status_details_codec.encode(
                    status, status_message, status_details,
                )).decode('ascii')
            )
            headers.append((_STATUS_DETAILS_KEY, status_details_bin))

        metadata = MultiDict(metadata or ())
        metadata, = await self._dispatch.send_trailing_metadata(
            metadata,
            status=status,
            status_message=status_message,
            status_details=status_details,
        )
        headers.extend(encode_metadata(cast(_Metadata, metadata)))

        await self._stream.send_headers(headers, end_stream=True)
        self._send_trailing_metadata_done = True

        if status != Status.OK and self._stream.closable:
            self._stream.reset_nowait()

    async def cancel(self) -> None:
        """Coroutine to cancel this request/stream.

        Server will send RST_STREAM frame to the client, so it will be
        explicitly informed that there is nothing to expect from the server
        regarding this request/stream.
        """
        if self._cancel_done:
            raise ProtocolError('Stream was already cancelled')

        await self._stream.reset()  # TODO: specify error code
        self._cancel_done = True

    async def __aenter__(self) -> 'Stream[_RecvType, _SendType]':
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> Optional[bool]:
        if (
            self._send_trailing_metadata_done
            or self._cancel_done
            or self._stream._transport.is_closing()
        ):
            # to suppress exception propagation
            return True

        protocol_error = None
        if exc_val is not None:
            # This error should be logged by ``request_handler``, here we
            # have to convert it into trailers and send to the client using
            # ``send_trailing_metadata`` method.
            if isinstance(exc_val, GRPCError):
                status = exc_val.status
                status_message = exc_val.message
                status_details = exc_val.details
            elif isinstance(exc_val, Exception):
                status = Status.UNKNOWN
                status_message = 'Internal Server Error'
                status_details = None
            else:
                # propagate exception
                return None
        elif (
            # There is a possibility of a ``ProtocolError`` in the
            # ``send_trailing_metadata`` method, so we are checking for such
            # errors here
            not self._cardinality.server_streaming
            and not self._send_message_done
        ):
            status = Status.UNKNOWN
            status_message = 'Internal Server Error'
            status_details = None
            protocol_error = ('Unary response with OK status requires '
                              'a single message to be sent: {!r}'
                              .format(self._method_name))
        else:
            status = Status.OK
            status_message = None
            status_details = None

        try:
            await self.send_trailing_metadata(status=status,
                                              status_message=status_message,
                                              status_details=status_details)
        except h2.exceptions.StreamClosedError:
            pass

        if protocol_error is not None:
            raise ProtocolError(protocol_error)

        # to suppress exception propagation
        return True


async def _abort(
    h2_stream: 'protocol.Stream',
    h2_status: int,
    grpc_status: Optional[Status] = None,
    grpc_message: Optional[str] = None,
) -> None:
    headers = [(':status', str(h2_status))]
    if grpc_status is not None:
        headers.append(('grpc-status', str(grpc_status.value)))
    if grpc_message is not None:
        headers.append(('grpc-message', grpc_message))
    await h2_stream.send_headers(headers, end_stream=True)
    if h2_stream.closable:
        h2_stream.reset_nowait()


async def request_handler(
    mapping: Dict[str, 'const.Handler'],
    _stream: 'protocol.Stream',
    headers: _Headers,
    codec: CodecBase,
    status_details_codec: Optional[StatusDetailsCodecBase],
    dispatch: _DispatchServerEvents,
    release_stream: Callable[[], Any],
) -> None:
    try:
        headers_map = dict(headers)

        if headers_map[':method'] != 'POST':
            await _abort(_stream, 405)
            return

        content_type = headers_map.get('content-type')
        if content_type is None:
            await _abort(_stream, 415, Status.UNKNOWN,
                         'Missing content-type header')
            return

        base_content_type, _, sub_type = content_type.partition('+')
        sub_type = sub_type or ProtoCodec.__content_subtype__
        if (
            base_content_type != GRPC_CONTENT_TYPE
            or sub_type != codec.__content_subtype__
        ):
            await _abort(_stream, 415, Status.UNKNOWN,
                         'Unacceptable content-type header')
            return

        if headers_map.get('te') != 'trailers':
            await _abort(_stream, 400, Status.UNKNOWN,
                         'Required "te: trailers" header is missing')
            return

        method_name = headers_map[':path']
        method = mapping.get(method_name)
        if method is None:
            await _abort(_stream, 200, Status.UNIMPLEMENTED,
                         'Method not found')
            return

        try:
            deadline = Deadline.from_headers(headers)
        except ValueError:
            await _abort(_stream, 200, Status.UNKNOWN,
                         'Invalid grpc-timeout header')
            return

        metadata = decode_metadata(headers)
        user_agent = headers_map.get('user-agent')

        async with Stream(
            _stream, method_name, method.cardinality,
            method.request_type, method.reply_type,
            codec=codec, status_details_codec=status_details_codec,
            dispatch=dispatch, deadline=deadline, user_agent=user_agent,
        ) as stream:
            deadline_wrapper: 'ContextManager[Any]'
            if deadline is None:
                wrapper = _stream.wrapper = Wrapper()
                deadline_wrapper = nullcontext()
            else:
                wrapper = _stream.wrapper = DeadlineWrapper()
                deadline_wrapper = wrapper.start(deadline)
            try:
                with deadline_wrapper, wrapper:
                    stream.metadata, method_func = await dispatch.recv_request(
                        metadata,
                        method.func,
                        method_name=method_name,
                        deadline=deadline,
                        content_type=content_type,
                        user_agent=user_agent,
                        peer=stream.peer,
                    )
                    await method_func(stream)
            except GRPCError:
                raise
            except asyncio.TimeoutError:
                if wrapper.cancel_failed:
                    log.exception('Failed to handle cancellation')
                    raise GRPCError(Status.DEADLINE_EXCEEDED)
                elif wrapper.cancelled:
                    log.info('Deadline exceeded')
                    raise GRPCError(Status.DEADLINE_EXCEEDED)
                else:
                    log.exception('Timeout occurred')
                    raise
            except StreamTerminatedError as err:
                if wrapper.cancel_failed:
                    log.exception('Failed to handle cancellation')
                    raise
                else:
                    assert wrapper.cancelled
                    log.info('Request was cancelled: %s', err)
                    raise
            except Exception:
                log.exception('Application error')
                raise
    except ProtocolError:
        log.exception('Application error')
    except Exception:
        log.exception('Server error')
    finally:
        release_stream()


class _GC(abc.ABC):
    _gc_counter = 0

    @property
    @abc.abstractmethod
    def __gc_interval__(self) -> int:
        raise NotImplementedError

    @abc.abstractmethod
    def __gc_collect__(self) -> None:
        pass

    def __gc_step__(self) -> None:
        self._gc_counter += 1
        if not (self._gc_counter % self.__gc_interval__):
            self.__gc_collect__()


class Handler(_GC, AbstractHandler):
    __gc_interval__ = 10

    closing = False

    def __init__(
        self,
        mapping: Dict[str, 'const.Handler'],
        codec: CodecBase,
        status_details_codec: Optional[StatusDetailsCodecBase],
        dispatch: _DispatchServerEvents,
    ) -> None:
        self.mapping = mapping
        self.codec = codec
        self.status_details_codec = status_details_codec
        self.dispatch = dispatch
        self.loop = asyncio.get_event_loop()
        self._tasks: Dict['protocol.Stream', 'asyncio.Task[None]'] = {}
        self._cancelled: Set['asyncio.Task[None]'] = set()

    def __gc_collect__(self) -> None:
        self._tasks = {s: t for s, t in self._tasks.items()
                       if not t.done()}
        self._cancelled = {t for t in self._cancelled
                           if not t.done()}

    def accept(
        self,
        stream: 'protocol.Stream',
        headers: _Headers,
        release_stream: Callable[[], Any],
    ) -> None:
        self.__gc_step__()
        self._tasks[stream] = self.loop.create_task(request_handler(
            self.mapping, stream, headers, self.codec,
            self.status_details_codec, self.dispatch, release_stream,
        ))

    def cancel(self, stream: 'protocol.Stream') -> None:
        task = self._tasks.pop(stream)
        task.cancel()
        self._cancelled.add(task)

    def close(self) -> None:
        for task in self._tasks.values():
            task.cancel()
        self._cancelled.update(self._tasks.values())
        self.closing = True

    async def wait_closed(self) -> None:
        if self._cancelled:
            await asyncio.wait(self._cancelled)

    def check_closed(self) -> bool:
        self.__gc_collect__()
        return not self._tasks and not self._cancelled


class Server(_GC):
    """
    HTTP/2 server, which uses gRPC service handlers to handle requests.

    Handler is a subclass of the abstract base class, which was generated
    from .proto file:

    .. code-block:: python3

        class CoffeeMachine(cafe_grpc.CoffeeMachineBase):

            async def MakeLatte(self, stream):
                task: cafe_pb2.LatteOrder = await stream.recv_message()
                ...
                await stream.send_message(empty_pb2.Empty())

        server = Server([CoffeeMachine()])
    """
    __gc_interval__ = 10

    def __init__(
        self,
        handlers: Collection['IServable'],
        *,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        codec: Optional[CodecBase] = None,
        status_details_codec: Optional[StatusDetailsCodecBase] = None,
        config: Optional[Configuration] = None,
    ) -> None:
        """
        :param handlers: list of handlers

        :param loop: (deprecated) asyncio-compatible event loop

        :param codec: instance of a codec to encode and decode messages,
            if omitted ``ProtoCodec`` is used by default

        :param status_details_codec: instance of a status details codec to
            encode error details in a trailing metadata, if omitted
            ``ProtoStatusDetailsCodec`` is used by default
        """
        if loop:
            warnings.warn("The loop argument is deprecated and scheduled "
                          "for removal in grpclib 0.5",
                          DeprecationWarning, stacklevel=2)

        mapping: Dict[str, 'const.Handler'] = {}
        for handler in handlers:
            mapping.update(handler.__mapping__())

        self._mapping = mapping
        self._loop = loop or asyncio.get_event_loop()

        if codec is None:
            codec = ProtoCodec()
            if status_details_codec is None and _googleapis_available():
                status_details_codec = ProtoStatusDetailsCodec()

        self._codec = codec
        self._status_details_codec = status_details_codec

        self._h2_config = h2.config.H2Configuration(
            client_side=False,
            header_encoding='ascii',
            validate_inbound_headers=False,
            validate_outbound_headers=False,
            normalize_inbound_headers=False,
            normalize_outbound_headers=False,
        )

        config = Configuration() if config is None else config
        self._config = config.__for_server__()

        self._server: Optional[asyncio.AbstractServer] = None
        self._server_closed_fut: Optional[asyncio.Future[None]] = None
        self._handlers: Set[Handler] = set()

        self.__dispatch__ = _DispatchServerEvents()
        _servers.add(self)

    def __gc_collect__(self) -> None:
        self._handlers = {h for h in self._handlers
                          if not (h.closing and h.check_closed())}

    def _protocol_factory(self) -> H2Protocol:
        self.__gc_step__()
        handler = Handler(
            self._mapping, self._codec, self._status_details_codec,
            self.__dispatch__,
        )
        self._handlers.add(handler)
        return H2Protocol(handler, self._config, self._h2_config)

    async def start(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        *,
        path: Optional[str] = None,
        family: 'socket.AddressFamily' = socket.AF_UNSPEC,
        flags: 'socket.AddressInfo' = socket.AI_PASSIVE,
        sock: Optional[socket.socket] = None,
        backlog: int = 100,
        ssl: Optional['_ssl.SSLContext'] = None,
        reuse_address: Optional[bool] = None,
        reuse_port: Optional[bool] = None,
    ) -> None:
        """Coroutine to start the server.

        :param host: can be a string, containing IPv4/v6 address or domain name.
            If host is None, server will be bound to all available interfaces.

        :param port: port number.

        :param path: UNIX domain socket path. If specified, host and port should
            be omitted (must be None).

        :param family: can be set to either :py:data:`python:socket.AF_INET` or
            :py:data:`python:socket.AF_INET6` to force the socket to use IPv4 or
            IPv6. If not set it will be determined from host.

        :param flags: is a bitmask for
            :py:meth:`~python:asyncio.AbstractEventLoop.getaddrinfo`.

        :param sock: sock can optionally be specified in order to use a
            preexisting socket object. If specified, host and port should be
            omitted (must be None).

        :param backlog: is the maximum number of queued connections passed to
            listen().

        :param ssl: can be set to an :py:class:`~python:ssl.SSLContext`
            to enable SSL over the accepted connections.

        :param reuse_address: tells the kernel to reuse a local socket in
            TIME_WAIT state, without waiting for its natural timeout to expire.

        :param reuse_port: tells the kernel to allow this endpoint to be bound
            to the same port as other existing endpoints are bound to,
            so long as they all set this flag when being created.
        """
        if path is not None and (host is not None or port is not None):
            raise ValueError("The 'path' parameter can not be used with the "
                             "'host' or 'port' parameters.")

        if self._server is not None:
            raise RuntimeError('Server is already started')

        if path is not None:
            self._server = await self._loop.create_unix_server(
                self._protocol_factory, path, sock=sock, backlog=backlog,
                ssl=ssl
            )
        else:
            # FIXME: Not all union combinations were tried because there are
            #  too many unions
            self._server = await self._loop.create_server(  # type: ignore
                self._protocol_factory, host,
                port,  # type: ignore
                family=family, flags=flags,
                sock=sock,  # type: ignore
                backlog=backlog, ssl=ssl,
                reuse_address=reuse_address, reuse_port=reuse_port
            )
        self._server_closed_fut = self._loop.create_future()

    def close(self) -> None:
        """Stops accepting new connections, cancels all currently running
        requests. Request handlers are able to handle `CancelledError` and
        exit properly.
        """
        if self._server is None or self._server_closed_fut is None:
            raise RuntimeError('Server is not started')
        self._server.close()
        if not self._server_closed_fut.done():
            self._server_closed_fut.set_result(None)
        for handler in self._handlers:
            handler.close()

    async def wait_closed(self) -> None:
        """Coroutine to wait until all existing request handlers will exit
        properly.
        """
        if self._server is None or self._server_closed_fut is None:
            raise RuntimeError('Server is not started')
        await self._server_closed_fut
        await self._server.wait_closed()
        if self._handlers:
            await asyncio.wait({
                self._loop.create_task(h.wait_closed()) for h in self._handlers
            })

    async def __aenter__(self) -> 'Server':
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        self.close()
        await self.wait_closed()