From bd9f1d2e7c7ce499e4ef523ff2276ca61ef34f07 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Mon, 10 Apr 2017 00:54:29 +0200
Subject: [PATCH] mgmt: add support for cancellable operations

Allow method handler to decide if operation could be cancelled. If yes,
when connection to the qubesd is terminated (and
protocol.connection_lost get called) the operation is cancelled using
standard asyncio method - in which case asyncio.CancelledError is thrown
inside method handler. This needs to be explicitly enabled, because
cancellable methods are much harder to write, to maintain consistent
system state.
Caveat: protocol.connection_lost is called only when trying to send some
data to it (and it's already terminated). Which makes this whole
mechanism useful only for events. Otherwise, when sending some data (and
possibly detecting that connection is broken), the operation is already
completed.

QubesOS/qubes-issues#2622
---
 qubes/mgmt.py         | 21 ++++++++++++++++++++-
 qubes/tools/qubesd.py | 23 ++++++++++++++++++-----
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/qubes/mgmt.py b/qubes/mgmt.py
index b94a0218..a87a1892 100644
--- a/qubes/mgmt.py
+++ b/qubes/mgmt.py
@@ -112,6 +112,9 @@ class AbstractQubesMgmt(object):
         #: name of the method
         self.method = method.decode('ascii')
 
+        #: is this operation cancellable?
+        self.cancellable = False
+
         untrusted_candidates = []
         for attr in dir(self):
             untrusted_func = getattr(self, attr)
@@ -135,9 +138,25 @@ class AbstractQubesMgmt(object):
             'multiple candidates for method {!r}'.format(self.method)
 
         #: the method to execute
-        self.execute = untrusted_candidates[0]
+        self._handler = untrusted_candidates[0]
+        self._running_handler = None
         del untrusted_candidates
 
+    def execute(self, *, untrusted_payload):
+        '''Execute management operation.
+
+        This method is a coroutine.
+        '''
+        self._running_handler = asyncio.ensure_future(self._handler(
+            untrusted_payload=untrusted_payload))
+        return self._running_handler
+
+    def cancel(self):
+        '''If operation is cancellable, interrupt it'''
+        if self.cancellable and self._running_handler is not None:
+            self._running_handler.cancel()
+
+
     def fire_event_for_permission(self, **kwargs):
         '''Fire an event on the source qube to check for permission'''
         return self.src.fire_event_pre('mgmt-permission:{}'.format(self.method),
diff --git a/qubes/tools/qubesd.py b/qubes/tools/qubesd.py
index 91f18732..cb1d3cc1 100644
--- a/qubes/tools/qubesd.py
+++ b/qubes/tools/qubesd.py
@@ -32,6 +32,7 @@ class QubesDaemonProtocol(asyncio.Protocol):
         self.len_untrusted_buffer = 0
         self.transport = None
         self.debug = debug
+        self.mgmt = None
 
     def connection_made(self, transport):
         print('connection_made()')
@@ -40,6 +41,10 @@ class QubesDaemonProtocol(asyncio.Protocol):
     def connection_lost(self, exc):
         print('connection_lost(exc={!r})'.format(exc))
         self.untrusted_buffer.close()
+        # for cancellable operation, interrupt it, otherwise it will do nothing
+        if self.mgmt is not None:
+            self.mgmt.cancel()
+        self.transport = None
 
     def data_received(self, untrusted_data):  # pylint: disable=arguments-differ
         print('data_received(untrusted_data={!r})'.format(untrusted_data))
@@ -72,9 +77,12 @@ class QubesDaemonProtocol(asyncio.Protocol):
     @asyncio.coroutine
     def respond(self, src, method, dest, arg, *, untrusted_payload):
         try:
-            mgmt = self.handler(self.app, src, method, dest, arg)
-            response = yield from mgmt.execute(
+            self.mgmt = self.handler(self.app, src, method, dest, arg,
+                self.send_event)
+            response = yield from self.mgmt.execute(
                 untrusted_payload=untrusted_payload)
+            if self.transport is None:
+                return
 
         # except clauses will fall through to transport.abort() below
 
@@ -91,9 +99,14 @@ class QubesDaemonProtocol(asyncio.Protocol):
                     method, arg, src, dest, len(untrusted_payload))
 
         except qubes.exc.QubesException as err:
-            self.send_exception(err)
-            self.transport.write_eof()
-            self.transport.close()
+            self.app.log.exception(
+                'error while calling '
+                'src=%r method=%r dest=%r arg=%r len(untrusted_payload)=%d',
+                src, method, dest, arg, len(untrusted_payload))
+            if self.transport is not None:
+                self.send_exception(err)
+                self.transport.write_eof()
+                self.transport.close()
             return
 
         except Exception:  # pylint: disable=broad-except