diff --git a/window-icon-updater/icon-sender b/window-icon-updater/icon-sender index 06c97722..35a19617 100644 --- a/window-icon-updater/icon-sender +++ b/window-icon-updater/icon-sender @@ -1,5 +1,5 @@ #!/usr/bin/python3 -# -*- encoding: utf8 -*- +# -*- encoding: utf-8 -*- # # The Qubes OS Project, http://www.qubes-os.org # @@ -23,18 +23,23 @@ # # Tool for sending windows icon over qrexec to dom0/GUI domain. Icons will be -# tainted by dom0 and then attached to appropriate windows. -# Usage: qrexec-client-vm dom0 qubes.WindowIconUpdater ./icon-sender - -import xcffib -from xcffib import xproto +# tinted by dom0 and then attached to appropriate windows. import sys import struct +import asyncio +import logging +import time + +import xcffib +from xcffib import xproto ICON_MAX_SIZE = 128 +log = logging.getLogger('icon-sender') + + class NoIconError(KeyError): pass @@ -52,6 +57,9 @@ class IconRetriever(object): self.atom_net_wm_icon = self.conn.core.InternAtom( False, len("_NET_WM_ICON"), "_NET_WM_ICON").reply().atom + def disconnect(self): + log.info('disconnecting from X') + self.conn.disconnect() def watch_window(self, w): self.conn.core.ChangeWindowAttributesChecked( @@ -105,48 +113,131 @@ class IconRetriever(object): raise NoIconError() return icons - def send_icon(self, w): + def describe_icon(self, w): try: icons = self.get_icons(w) chosen_size = sorted(icons.keys())[-1] - sys.stdout.buffer.write("{}\n".format(w).encode('ascii')) - sys.stdout.buffer.write("{} {}\n".format( - chosen_size[0], chosen_size[1]).encode('ascii')) - sys.stdout.buffer.write(b''.join( + data = b'' + data += "{}\n".format(w).encode('ascii') + data += "{} {}\n".format( + chosen_size[0], chosen_size[1]).encode('ascii') + data += b''.join( [struct.pack('>I', ((b << 8) & 0xffffff00) | (b >> 24)) for b in - icons[chosen_size]])) - sys.stdout.buffer.flush() + icons[chosen_size]]) + return data except NoIconError: - pass + return None def initial_sync(self): - cookie = self.conn.core.QueryTree(self.root) - root_tree = cookie.reply() - for w in root_tree.children: - self.watch_window(w) - self.send_icon(w) - - def watch_and_send_icons(self): self.conn.core.ChangeWindowAttributesChecked( self.root, xproto.CW.EventMask, [xproto.EventMask.SubstructureNotify]) self.conn.flush() - self.initial_sync() - for ev in iter(self.conn.wait_for_event, None): + cookie = self.conn.core.QueryTree(self.root) + root_tree = cookie.reply() + for w in root_tree.children: + self.watch_window(w) + yield self.describe_icon(w) + + async def watch_and_send_icons(self): + ''' + Yield data for all icons we receive. + This is an asynchronous generator, so that we can handle reconnections + during waiting for X events. + ''' + for icon in self.initial_sync(): + yield icon + + # Emulate select() + event = asyncio.Event() + asyncio.get_event_loop().add_reader( + self.conn.get_file_descriptor(), event.set) + + while True: + await event.wait() + event.clear() + for icon in self.handle_pending_events(): + yield icon + + def handle_pending_events(self): + for ev in iter(self.conn.poll_for_event, None): if isinstance(ev, xproto.CreateNotifyEvent): self.window_queue.add(ev.window) self.watch_window(ev.window) elif isinstance(ev, xproto.MapNotifyEvent): if ev.window in self.window_queue: - self.send_icon(ev.window) + yield self.describe_icon(ev.window) self.window_queue.remove(ev.window) elif isinstance(ev, xproto.PropertyNotifyEvent): if ev.atom == self.atom_net_wm_icon: - self.send_icon(ev.window) + yield self.describe_icon(ev.window) + + +class IconSender: + async def run(self): + fail_threshold_seconds = 5 + restart_wait_seconds = 5 + restart_tries = 5 + + try_num = 0 + while True: + t = await self.run_client() + if t < fail_threshold_seconds: + try_num += 1 + if try_num == restart_tries: + log.error('giving up after %d tries', try_num) + break + log.error('process exited too soon, waiting %d seconds and retrying', + restart_wait_seconds) + await asyncio.sleep(restart_wait_seconds) + else: + log.info('process exited, trying to reconnect') + try_num = 0 + + async def run_client(self): + cmd = ['qrexec-client-vm', 'dom0', 'qubes.WindowIconUpdater'] + log.info('running: %s', cmd) + + start_time = time.time() + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE + ) + send_icons_task = asyncio.create_task(self.send_icons(proc.stdin)) + try: + await proc.wait() + if proc.returncode: + log.error('process failed with status %d', proc.returncode) + return time.time() - start_time + finally: + send_icons_task.cancel() + + async def send_icons(self, writer): + retriever = IconRetriever() + try: + async for data in retriever.watch_and_send_icons(): + if data: + writer.write(data) + await writer.drain() + except IOError: + log.exception() + except asyncio.CancelledError: + pass + finally: + retriever.disconnect() + writer.close() + + +def main(): + logging.basicConfig( + stream=sys.stderr, level=logging.INFO, + format='%(asctime)s %(name)s: %(message)s') + + sender = IconSender() + asyncio.run(sender.run()) if __name__ == '__main__': - retriever = IconRetriever() - retriever.watch_and_send_icons() + main() diff --git a/window-icon-updater/qubes-icon-sender.desktop b/window-icon-updater/qubes-icon-sender.desktop index d9892edb..f62d3c6b 100644 --- a/window-icon-updater/qubes-icon-sender.desktop +++ b/window-icon-updater/qubes-icon-sender.desktop @@ -2,7 +2,7 @@ Version=1.0 Encoding=UTF-8 Name=Window icon updater -Exec=/usr/bin/qrexec-client-vm dom0 qubes.WindowIconUpdater /usr/lib/qubes/icon-sender +Exec=/usr/lib/qubes/icon-sender Terminal=false Type=Application Categories=