| # Copyright 2023 The Fuchsia Authors. All rights reserved. | |
| # Use of this source code is governed by a BSD-style license that can be | |
| # found in the LICENSE file. | |
| import asyncio | |
| import inspect | |
| import logging | |
| import struct | |
| from collections.abc import Coroutine | |
| from inspect import getframeinfo, stack | |
| from typing import Any, Dict, Set, Type, TypeVar, overload | |
| import fuchsia_controller_py as fc | |
| from ._fidl_common import ( | |
| FIDL_EPITAPH_ORDINAL, | |
| Decodable, | |
| Encodable, | |
| EpitaphError, | |
| FidlMessage, | |
| FidlMeta, | |
| StopEventHandler, | |
| TXID_Type, | |
| parse_epitaph_value, | |
| parse_ordinal, | |
| parse_txid, | |
| ) | |
| from ._ipc import GlobalHandleWaker, HandleWaker | |
| from ._registry import get_registered_method | |
| # The active TXID (mutable). | |
| TXID: TXID_Type = 0 | |
| T = TypeVar("T", bound=Decodable) | |
| # The TXID of a FIDL event. | |
| EVENT_TXID: TXID_Type = 0 | |
| # Simple client ID. Monotonically increasing for each client. | |
| _CLIENT_ID = 0 | |
| _LOGGER = logging.getLogger("fidl.client") | |
| # FIDL transactional message header size: | |
| # txid (4B) + at-rest flags (2B) + dynamic flags (2B) + magic (1B) + reserved (1B) + ordinal (8B) | |
| _FIDL_MESSAGE_HEADER_SIZE = 16 | |
| class FidlClient(metaclass=FidlMeta): | |
| def __init__( | |
| self, | |
| channel: int | fc.Channel, | |
| channel_waker: HandleWaker | None = None, | |
| ) -> None: | |
| global _CLIENT_ID | |
| self.id = _CLIENT_ID | |
| _CLIENT_ID += 1 | |
| if isinstance(channel, int): | |
| self._channel: fc.Channel | None = fc.Channel(channel) | |
| else: | |
| self._channel = channel | |
| if channel_waker is None: | |
| self._channel_waker: HandleWaker = GlobalHandleWaker() | |
| else: | |
| self._channel_waker = channel_waker | |
| self.pending_txids: Set[TXID_Type] = set({}) | |
| self.staged_messages: Dict[TXID_Type, asyncio.Queue[FidlMessage]] = {} | |
| self.epitaph_received: EpitaphError | None = None | |
| self.epitaph_event = asyncio.Event() | |
| caller = getframeinfo(stack()[1][0]) | |
| _LOGGER.debug( | |
| f"{self} instantiated from {caller.filename}:{caller.lineno}" | |
| ) | |
| def close_cleanly(self) -> None: | |
| """Closes the underlying channel. | |
| This is so-named to avoid name conflicts with existing FIDL methods. | |
| """ | |
| if self._channel is not None: | |
| self._channel.close() | |
| # self._channel = None is NOT done here, as it may be used in pending tasks. | |
| # The handle being closed will result in errors in those tasks. | |
| def __str__(self) -> str: | |
| return f"client:{type(self).__name__}:{self.id}" | |
| async def _get_staged_message(self, txid: TXID_Type) -> FidlMessage: | |
| res = await self.staged_messages[txid].get() | |
| self.staged_messages[txid].task_done() | |
| return res | |
| def _stage_message(self, txid: TXID_Type, msg: FidlMessage) -> None: | |
| # This should only ever happen if we're a channel reading another channel's response before | |
| # it has ever made a request. | |
| if txid not in self.staged_messages: | |
| self.staged_messages[txid] = asyncio.Queue(1) | |
| self.staged_messages[txid].put_nowait(msg) | |
| def _clean_staging(self, txid: TXID_Type) -> None: | |
| self.staged_messages.pop(txid) | |
| # Events are never added to this set, since they're always pending. | |
| if txid != EVENT_TXID: | |
| self.pending_txids.remove(txid) | |
| def _decode( | |
| self, | |
| txid: TXID_Type, | |
| msg: FidlMessage, | |
| response_cls: Type[T] | None, | |
| expected_ordinal: int | None = None, | |
| ) -> T | None: | |
| self._clean_staging(txid) | |
| if expected_ordinal is not None: | |
| actual_ordinal = parse_ordinal(msg) | |
| if actual_ordinal != expected_ordinal: | |
| type_name = ( | |
| response_cls.__name__ | |
| if response_cls is not None | |
| else "None" | |
| ) | |
| raise RuntimeError( | |
| f"Ordinal mismatch when decoding {type_name}: expected {expected_ordinal}, got {actual_ordinal}" | |
| ) | |
| if response_cls is None: | |
| return None | |
| handles = msg[1] | |
| verified_handles: list[int] = [0] * len(handles) | |
| for i in range(len(handles)): | |
| hdl = handles[i] | |
| if isinstance(hdl, tuple): | |
| verified_handles[i] = hdl[1] | |
| else: | |
| verified_handles[i] = hdl.take() | |
| return response_cls.decode( | |
| msg[0][_FIDL_MESSAGE_HEADER_SIZE:], verified_handles | |
| ) | |
| async def next_event(self) -> FidlMessage | None: | |
| """Attempts to read the next FIDL event from this client. | |
| Returns: | |
| The next FIDL event. If ZX_ERR_PEER_CLOSED is received on the channel, will return None. | |
| Note: this does not check to see if the protocol supports any events, so if not this | |
| function could wait forever. | |
| Raises: | |
| Any exceptions other than ZX_ERR_PEER_CLOSED (fuchsia_controller_py.FcTransportStatus) | |
| """ | |
| # TODO(awdavies): Raise an exception if there are no events supported for this client. | |
| try: | |
| return await self._read_and_decode(0) | |
| except fc.FcTransportStatus as e: | |
| err_code = e.code() | |
| if err_code != fc.FcTransportStatus.FC_ERR_FDOMAIN: | |
| _LOGGER.warning( | |
| f"{self} received error waiting for next event: {e}" | |
| ) | |
| raise e | |
| return None | |
| def _epitaph_check(self, msg: FidlMessage) -> None: | |
| # If the epitaph is already set, no need to continue with the remaining | |
| # work. | |
| if self.epitaph_received is not None: | |
| raise self.epitaph_received | |
| ordinal = parse_ordinal(msg) | |
| if ordinal == FIDL_EPITAPH_ORDINAL: | |
| if self.epitaph_received is None: | |
| self.epitaph_received = EpitaphError(parse_epitaph_value(msg)) | |
| self.epitaph_event.set() | |
| if self.epitaph_received is not None: | |
| raise self.epitaph_received | |
| async def _epitaph_event_wait(self) -> None: | |
| await self.epitaph_event.wait() | |
| if self.epitaph_received is not None: | |
| raise self.epitaph_received | |
| async def _read_and_decode( | |
| self, | |
| txid: int, | |
| response_cls: Type[T] | None = None, | |
| expected_ordinal: int | None = None, | |
| ) -> T | FidlMessage | None: | |
| if txid not in self.staged_messages: | |
| self.staged_messages[txid] = asyncio.Queue(1) | |
| if self._channel is None: | |
| raise ValueError("Channel is already closed") | |
| msg: FidlMessage | |
| with self._channel_waker.registration(self._channel, name=str(self)): | |
| while True: | |
| if self.epitaph_received is not None: | |
| raise self.epitaph_received | |
| # 1. Try to read from the channel. | |
| try: | |
| msg = self._channel.read() | |
| self._epitaph_check(msg) | |
| recvd_txid = parse_txid(msg) | |
| if recvd_txid == txid: | |
| if txid != EVENT_TXID: | |
| return self._decode( | |
| txid, msg, response_cls, expected_ordinal | |
| ) | |
| else: | |
| self._clean_staging(txid) | |
| return msg | |
| if ( | |
| recvd_txid != EVENT_TXID | |
| and recvd_txid not in self.pending_txids | |
| ): | |
| _LOGGER.warning( | |
| f"{self} received unexpected TXID: {recvd_txid}" | |
| ) | |
| # Unexpected TXID is often a sign of a bad server or a serious bug. | |
| # We don't close the channel here, but we raise. | |
| raise RuntimeError( | |
| f"{self} received unexpected TXID: {recvd_txid}" | |
| ) | |
| self._stage_message(recvd_txid, msg) | |
| except fc.FcTransportStatus as e: | |
| if e.code() != fc.FcTransportStatus.FC_ERR_SHOULD_WAIT: | |
| raise e | |
| # 2. Wait for either: | |
| # - The channel being readable again. | |
| # - A staged message for our TXID becoming available. | |
| # - An epitaph being received. | |
| # | |
| # This is handled with a task group in the event that | |
| # there are task cancellations they apply to the other | |
| # tasks in the task group, making cleanup predictable. | |
| try: | |
| async with asyncio.TaskGroup() as tg: | |
| read_ready_task = tg.create_task( | |
| self._channel_waker.wait_ready(self._channel) | |
| ) | |
| staged_msg_task = tg.create_task( | |
| self._get_staged_message(txid) | |
| ) | |
| epitaph_task = tg.create_task(self.epitaph_event.wait()) | |
| done, pending = await asyncio.wait( | |
| [read_ready_task, staged_msg_task, epitaph_task], | |
| return_when=asyncio.FIRST_COMPLETED, | |
| ) | |
| for p in pending: | |
| p.cancel() | |
| except ExceptionGroup as eg: | |
| # If there is only a single exception, raise this as the | |
| # main exception. We care particularly if there is an | |
| # epitaph exception raised. | |
| if len(eg.exceptions) == 1: | |
| raise eg.exceptions[0] from None | |
| raise | |
| if staged_msg_task in done: | |
| # If we got a staged message, we're done. | |
| msg = staged_msg_task.result() | |
| # If read_ready_task was also done, we should "re-notify" because we didn't read. | |
| if read_ready_task in done: | |
| if self._channel is None: | |
| raise ValueError("Channel is already closed") | |
| self._channel_waker.post_ready(self._channel) | |
| if txid != EVENT_TXID: | |
| return self._decode( | |
| txid, msg, response_cls, expected_ordinal | |
| ) | |
| else: | |
| self._clean_staging(txid) | |
| return msg | |
| if epitaph_task in done: | |
| # This will raise the epitaph error. | |
| if self.epitaph_received is not None: | |
| raise self.epitaph_received | |
| return None | |
| # If only read_ready_task reached here, we just loop again and try to read(). | |
| # Overload for two-way FIDL requests with a response payload. | |
| # Enables static type checking and IDE autocomplete to infer the exact response | |
| # object type (e.g. `Coroutine[Any, Any, EchoResponse]`) without requiring | |
| # `typing.cast` or leaving the return type as `Any` or `T | None`. | |
| # See: https://mypy.readthedocs.io/en/stable/more_types.html#type-checking-calls-to-overloads | |
| @overload | |
| def _send_two_way_fidl_request( | |
| self, | |
| ordinal: int, | |
| library: str, | |
| msg_obj: Encodable | None, | |
| response_cls: Type[T], | |
| ) -> Coroutine[Any, Any, T]: | |
| ... | |
| # Overload for two-way FIDL requests with an empty response payload (-> ()). | |
| # Returns `Coroutine[..., None]` statically rather than `Coroutine[..., T | None]`, | |
| # allowing IDEs and type checkers to know the call resolves to `None`. | |
| @overload | |
| def _send_two_way_fidl_request( | |
| self, | |
| ordinal: int, | |
| library: str, | |
| msg_obj: Encodable | None, | |
| response_cls: None = None, | |
| ) -> Coroutine[Any, Any, None]: | |
| ... | |
| def _send_two_way_fidl_request( | |
| self, | |
| ordinal: int, | |
| library: str, | |
| msg_obj: Encodable | None, | |
| response_cls: Type[T] | None = None, | |
| ) -> Coroutine[Any, Any, T | None]: | |
| """Sends a two-way asynchronous FIDL request. | |
| Args: | |
| ordinal: The method ordinal (for encoding). | |
| library: The FIDL library from which this method ordinal exists. | |
| msg_obj: The object being sent. | |
| response_cls: The class of the response object. | |
| Returns: | |
| The object from the two-way function. | |
| """ | |
| global TXID | |
| TXID += 1 | |
| self.pending_txids.add(TXID) | |
| self._send_one_way_fidl_request(TXID, ordinal, library, msg_obj) | |
| async def result(txid: int) -> Any: | |
| return await self._read_and_decode( | |
| txid, response_cls, expected_ordinal=ordinal | |
| ) | |
| return result(TXID) | |
| def _send_one_way_fidl_request( | |
| self, | |
| txid: int, | |
| ordinal: int, | |
| library: str, | |
| msg_obj: Encodable | None, | |
| ) -> None: | |
| """Sends a synchronous one-way FIDL request. | |
| Args: | |
| ordinal: The method ordinal (for encoding). | |
| library: The FIDL library from which this method ordinal exists. | |
| msg_obj: The object being sent. | |
| """ | |
| if msg_obj is not None: | |
| payload_bytes, handles = msg_obj.encode() | |
| else: | |
| payload_bytes, handles = b"", [] | |
| header = struct.pack("<IHBBQ", txid, 0x02, 0x00, 0x01, ordinal) | |
| encoded_msg = header + payload_bytes | |
| if self._channel is None: | |
| raise ValueError("Channel is already closed") | |
| self._channel.write((encoded_msg, handles)) | |
| class EventHandlerBase( | |
| metaclass=FidlMeta, | |
| required_class_variables=[ | |
| ("library", str), | |
| ("method_map", dict), | |
| ], | |
| ): | |
| """Base object for doing FIDL client event handling.""" | |
| client: FidlClient | |
| method_map: Dict[int, Any] | |
| def __init__(self, client: FidlClient) -> None: | |
| self.client = client | |
| def __str__(self) -> str: | |
| return f"event:{type(self.client).__name__}:{self.client.id}" | |
| async def serve(self) -> None: | |
| while True: | |
| msg = await self.client.next_event() | |
| # msg is None if the channel has been closed. | |
| if msg is None: | |
| break | |
| if not await self._handle_request(msg): | |
| break | |
| async def _handle_request(self, msg: FidlMessage) -> bool: | |
| try: | |
| await self._handle_request_helper(msg) | |
| return True | |
| except StopEventHandler: | |
| return False | |
| async def _handle_request_helper(self, msg: FidlMessage) -> None: | |
| ordinal = parse_ordinal(msg) | |
| method = get_registered_method(ordinal) | |
| if method is None: | |
| raise RuntimeError(f"Unknown ordinal {ordinal}") | |
| _, response_cls = method # Events are registered as response_cls | |
| handles = msg[1] | |
| verified_handles: list[int] = [0] * len(handles) | |
| for i in range(len(handles)): | |
| hdl = handles[i] | |
| if isinstance(hdl, tuple): | |
| verified_handles[i] = hdl[1] | |
| else: | |
| verified_handles[i] = hdl.take() | |
| if response_cls is not None: | |
| request_obj = response_cls.decode(msg[0][16:], verified_handles) | |
| else: | |
| request_obj = None | |
| method_meta = self.method_map[ordinal] | |
| method_lambda = getattr(self, method_meta.name) | |
| if request_obj is not None: | |
| res = method_lambda(request_obj) | |
| else: | |
| res = method_lambda() | |
| if inspect.isawaitable(res): | |
| await res |