import array from contextlib import contextmanager import errno from itertools import count import logging from typing import Optional try: from contextlib import asynccontextmanager # Python 3.7 except ImportError: from async_generator import asynccontextmanager # Backport for Python 3.6 from outcome import Value, Error import trio from trio.abc import Channel from jeepney.auth import Authenticator, BEGIN from jeepney.bus import get_bus from jeepney.fds import FileDescriptor, fds_buf_size from jeepney.low_level import Parser, MessageType, Message from jeepney.wrappers import ProxyBase, unwrap_msg from jeepney.bus_messages import message_bus from .common import ( MessageFilters, FilterHandle, ReplyMatcher, RouterClosed, check_replyable, ) log = logging.getLogger(__name__) __all__ = [ 'open_dbus_connection', 'open_dbus_router', 'Proxy', ] # The function below is copied from trio, which is under the MIT license: # Permission is hereby granted, free of charge, to any person obtaining # a copy of this software and associated documentation files (the # "Software"), to deal in the Software without restriction, including # without limitation the rights to use, copy, modify, merge, publish, # distribute, sublicense, and/or sell copies of the Software, and to # permit persons to whom the Software is furnished to do so, subject to # the following conditions: # # The above copyright notice and this permission notice shall be # included in all copies or substantial portions of the Software. @contextmanager def _translate_socket_errors_to_stream_errors(): try: yield except OSError as exc: if exc.errno in {errno.EBADF, errno.ENOTSOCK}: # EBADF on Unix, ENOTSOCK on Windows raise trio.ClosedResourceError("this socket was already closed") from None else: raise trio.BrokenResourceError( "socket connection broken: {}".format(exc) ) from exc class DBusConnection(Channel): """A plain D-Bus connection with no matching of replies. This doesn't run any separate tasks: sending and receiving are done in the task that calls those methods. It's suitable for implementing servers: several worker tasks can receive requests and send replies. For a typical client pattern, see :class:`DBusRouter`. Implements trio's channel interface for Message objects. """ def __init__(self, socket, enable_fds=False): self.socket = socket self.enable_fds = enable_fds self.parser = Parser() self.outgoing_serial = count(start=1) self.unique_name = None self.send_lock = trio.Lock() self.recv_lock = trio.Lock() self._leftover_to_send = None # type: Optional[memoryview] async def send(self, message: Message, *, serial=None): """Serialise and send a :class:`~.Message` object""" async with self.send_lock: if serial is None: serial = next(self.outgoing_serial) fds = array.array('i') if self.enable_fds else None data = message.serialise(serial, fds=fds) await self._send_data(data, fds) # _send_data is copied & modified from trio's SocketStream.send_all() . # See above for the MIT license. async def _send_data(self, data: bytes, fds): if self.socket.did_shutdown_SHUT_WR: raise trio.ClosedResourceError("can't send data after sending EOF") with _translate_socket_errors_to_stream_errors(): if self._leftover_to_send: # A previous message was partly sent - finish sending it now. await self._send_remainder(self._leftover_to_send) with memoryview(data) as data: if fds: sent = await self.socket.sendmsg([data], [( trio.socket.SOL_SOCKET, trio.socket.SCM_RIGHTS, fds )]) else: sent = await self.socket.send(data) await self._send_remainder(data, sent) async def _send_remainder(self, data: memoryview, already_sent=0): try: while already_sent < len(data): with data[already_sent:] as remaining: sent = await self.socket.send(remaining) already_sent += sent self._leftover_to_send = None except trio.Cancelled: # Sending cancelled mid-message. Keep track of the remaining data # so it can be sent before the next message, otherwise the next # message won't be recognised. self._leftover_to_send = data[already_sent:] raise async def receive(self) -> Message: """Return the next available message from the connection""" async with self.recv_lock: while True: msg = self.parser.get_next_message() if msg is not None: return msg # Once data is read, it must be given to the parser with no # checkpoints (where the task could be cancelled). b, fds = await self._read_data() if not b: raise trio.EndOfChannel("Socket closed at the other end") self.parser.add_data(b, fds) async def _read_data(self): if self.enable_fds: nbytes = self.parser.bytes_desired() with _translate_socket_errors_to_stream_errors(): data, ancdata, flags, _ = await self.socket.recvmsg( nbytes, fds_buf_size() ) if flags & getattr(trio.socket, 'MSG_CTRUNC', 0): self._close() raise RuntimeError("Unable to receive all file descriptors") return data, FileDescriptor.from_ancdata(ancdata) else: # not self.enable_fds with _translate_socket_errors_to_stream_errors(): data = await self.socket.recv(4096) return data, [] def _close(self): self.socket.close() self._leftover_to_send = None # Our closing is currently sync, but AsyncResource objects must have aclose async def aclose(self): """Close the D-Bus connection""" self._close() @asynccontextmanager async def router(self): """Temporarily wrap this connection as a :class:`DBusRouter` To be used like:: async with conn.router() as req: reply = await req.send_and_get_reply(msg) While the router is running, you shouldn't use :meth:`receive`. Once the router is closed, you can use the plain connection again. """ async with trio.open_nursery() as nursery: router = DBusRouter(self) await router.start(nursery) try: yield router finally: await router.aclose() async def open_dbus_connection(bus='SESSION', *, enable_fds=False) -> DBusConnection: """Open a plain D-Bus connection :return: :class:`DBusConnection` """ bus_addr = get_bus(bus) sock : trio.SocketStream = await trio.open_unix_socket(bus_addr) # Authentication authr = Authenticator(enable_fds=enable_fds) for req_data in authr: await sock.send_all(req_data) authr.feed(await sock.receive_some()) await sock.send_all(BEGIN) conn = DBusConnection(sock.socket, enable_fds=enable_fds) # Say *Hello* to the message bus - this must be the first message, and the # reply gives us our unique name. async with conn.router() as router: reply = await router.send_and_get_reply(message_bus.Hello()) conn.unique_name = reply.body[0] return conn class TrioFilterHandle(FilterHandle): def __init__(self, filters: MessageFilters, rule, send_chn, recv_chn): super().__init__(filters, rule, recv_chn) self.send_channel = send_chn @property def receive_channel(self): return self.queue async def aclose(self): self.close() await self.send_channel.aclose() async def __aenter__(self): return self.queue async def __aexit__(self, exc_type, exc_val, exc_tb): await self.aclose() class Future: """A very simple Future for trio based on `trio.Event`.""" def __init__(self): self._outcome = None self._event = trio.Event() def set_result(self, result): self._outcome = Value(result) self._event.set() def set_exception(self, exc): self._outcome = Error(exc) self._event.set() async def get(self): await self._event.wait() return self._outcome.unwrap() class DBusRouter: """A client D-Bus connection which can wait for replies. This runs a separate receiver task and dispatches received messages. """ _nursery_mgr = None _rcv_cancel_scope = None def __init__(self, conn: DBusConnection): self._conn = conn self._replies = ReplyMatcher() self._filters = MessageFilters() @property def unique_name(self): return self._conn.unique_name async def send(self, message, *, serial=None): """Send a message, don't wait for a reply """ await self._conn.send(message, serial=serial) async def send_and_get_reply(self, message) -> Message: """Send a method call message and wait for the reply Returns the reply message (method return or error message type). """ check_replyable(message) if self._rcv_cancel_scope is None: raise RouterClosed("This DBusRouter has stopped") serial = next(self._conn.outgoing_serial) with self._replies.catch(serial, Future()) as reply_fut: await self.send(message, serial=serial) return (await reply_fut.get()) def filter(self, rule, *, channel: Optional[trio.MemorySendChannel]=None, bufsize=1): """Create a filter for incoming messages Usage:: async with router.filter(rule) as receive_channel: matching_msg = await receive_channel.receive() # OR: send_chan, recv_chan = trio.open_memory_channel(1) async with router.filter(rule, channel=send_chan): matching_msg = await recv_chan.receive() If the channel fills up, The sending end of the channel is closed when leaving the ``async with`` block, whether or not it was passed in. :param jeepney.MatchRule rule: Catch messages matching this rule :param trio.MemorySendChannel channel: Send matching messages here :param int bufsize: If no channel is passed in, create one with this size """ if channel is None: channel, recv_channel = trio.open_memory_channel(bufsize) else: recv_channel = None return TrioFilterHandle(self._filters, rule, channel, recv_channel) # Task management ------------------------------------------- async def start(self, nursery: trio.Nursery): if self._rcv_cancel_scope is not None: raise RuntimeError("DBusRouter receiver task is already running") self._rcv_cancel_scope = await nursery.start(self._receiver) async def aclose(self): """Stop the sender & receiver tasks""" # It doesn't matter if we receive a partial message - the connection # should ensure that whatever is received is fed to the parser. if self._rcv_cancel_scope is not None: self._rcv_cancel_scope.cancel() self._rcv_cancel_scope = None # Ensure trio checkpoint await trio.sleep(0) # Code to run in receiver task ------------------------------------ def _dispatch(self, msg: Message): """Handle one received message""" if self._replies.dispatch(msg): return for filter in self._filters.matches(msg): try: filter.send_channel.send_nowait(msg) except trio.WouldBlock: pass async def _receiver(self, task_status=trio.TASK_STATUS_IGNORED): """Receiver loop - runs in a separate task""" with trio.CancelScope() as cscope: self.is_running = True task_status.started(cscope) try: while True: msg = await self._conn.receive() self._dispatch(msg) finally: self.is_running = False # Send errors to any tasks still waiting for a message. self._replies.drop_all() # Closing a memory channel can't block, but it only has an # async close method, so we need to shield it from cancellation. with trio.move_on_after(3) as cleanup_scope: for filter in self._filters.filters.values(): cleanup_scope.shield = True await filter.send_channel.aclose() class Proxy(ProxyBase): """A trio proxy for calling D-Bus methods You can call methods on the proxy object, such as ``await bus_proxy.Hello()`` to make a method call over D-Bus and wait for a reply. It will either return a tuple of returned data, or raise :exc:`.DBusErrorResponse`. The methods available are defined by the message generator you wrap. :param msggen: A message generator object. :param ~trio.DBusRouter router: Router to send and receive messages. """ def __init__(self, msggen, router): super().__init__(msggen) if not isinstance(router, DBusRouter): raise TypeError("Proxy can only be used with DBusRequester") self._router = router def _method_call(self, make_msg): async def inner(*args, **kwargs): msg = make_msg(*args, **kwargs) assert msg.header.message_type is MessageType.method_call reply = await self._router.send_and_get_reply(msg) return unwrap_msg(reply) return inner @asynccontextmanager async def open_dbus_router(bus='SESSION', *, enable_fds=False): """Open a D-Bus 'router' to send and receive messages. Use as an async context manager:: async with open_dbus_router() as req: ... :param str bus: 'SESSION' or 'SYSTEM' or a supported address. :return: :class:`DBusRouter` This is a shortcut for:: conn = await open_dbus_connection() async with conn: async with conn.router() as req: ... """ conn = await open_dbus_connection(bus, enable_fds=enable_fds) async with conn: async with conn.router() as rtr: yield rtr