From 40a86efd663bf87de310227654b164c8dc5a43f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Mon, 10 Apr 2017 01:01:40 +0200
Subject: [PATCH] mgmt: add support for events-sending methods

Standard methods return only one value, after operation is completed,
but events-related methods may return multiple values during the method
execution time. Provide a callback for such cases.
Also, according to specification, avoid sending both event and non-event
values.

QubesOS/qubes-issues#2622
---
 qubes/mgmt.py         | 5 ++++-
 qubes/tools/qubesd.py | 7 ++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/qubes/mgmt.py b/qubes/mgmt.py
index a87a1892..fa78879e 100644
--- a/qubes/mgmt.py
+++ b/qubes/mgmt.py
@@ -96,7 +96,7 @@ class AbstractQubesMgmt(object):
     There are also two helper functions for firing events associated with API
     calls.
     '''
-    def __init__(self, app, src, method, dest, arg):
+    def __init__(self, app, src, method, dest, arg, send_event=None):
         #: :py:class:`qubes.Qubes` object
         self.app = app
 
@@ -112,6 +112,9 @@ class AbstractQubesMgmt(object):
         #: name of the method
         self.method = method.decode('ascii')
 
+        #: callback for sending events if applicable
+        self.send_event = send_event
+
         #: is this operation cancellable?
         self.cancellable = False
 
diff --git a/qubes/tools/qubesd.py b/qubes/tools/qubesd.py
index cb1d3cc1..b308d550 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.event_sent = False
         self.mgmt = None
 
     def connection_made(self, transport):
@@ -81,6 +82,7 @@ class QubesDaemonProtocol(asyncio.Protocol):
                 self.send_event)
             response = yield from self.mgmt.execute(
                 untrusted_payload=untrusted_payload)
+            assert not (self.event_sent and response)
             if self.transport is None:
                 return
 
@@ -116,7 +118,8 @@ class QubesDaemonProtocol(asyncio.Protocol):
                     src, method, dest, arg, len(untrusted_payload))
 
         else:
-            self.send_response(response)
+            if not self.event_sent:
+                self.send_response(response)
             try:
                 self.transport.write_eof()
             except NotImplementedError:
@@ -133,11 +136,13 @@ class QubesDaemonProtocol(asyncio.Protocol):
         self.transport.write(self.header.pack(*args))
 
     def send_response(self, content):
+        assert not self.event_sent
         self.send_header(0x30)
         if content is not None:
             self.transport.write(content.encode('utf-8'))
 
     def send_event(self, subject, event, **kwargs):
+        self.event_sent = True
         self.send_header(0x31)
 
         if subject is not self.app: