icon-sender: reconnect on failure

This makes icon-sender handle GUI restart, same as qubes-gui and
audio. However, in this case we are not running raw vchan, but
Qubes RPC, so the procedure is more complicated: we start a
qrexec-client-vm subprocess, and restart it if it breaks.

Needs icon-receiver to be a service with wait-for-session
(QubesOS/qubes-gui-daemon#37), and fix for socket services in dom0
(QubesOS/qubes-core-qrexec#42).

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()

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=