From 6d113888075b34a410e228af19ff8776c3c5f898 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Tue, 12 Jan 2021 23:40:36 +0100
Subject: [PATCH] api/admin: add 'wait' parameter to admin.vm.Shutdown

Add support for blocking shutdown call. This adds a symmetry to
admin.vm.Start call which is blocking.
Since the admin.vm.Shutdown has established semantic already, add a
'wait' parameter. It can be combined with 'force' as 'force+wait' (or
the other way around).
---
 qubes/api/admin.py       | 12 +++++++---
 qubes/tests/api_admin.py | 48 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/qubes/api/admin.py b/qubes/api/admin.py
index 2c9c30d8..33966626 100644
--- a/qubes/api/admin.py
+++ b/qubes/api/admin.py
@@ -899,9 +899,15 @@ class QubesAdminAPI(qubes.api.AbstractQubesAPI):
         scope='local', execute=True)
     @asyncio.coroutine
     def vm_shutdown(self):
-        force = (self.arg == 'force')
-        self.fire_event_for_permission(force=force)
-        yield from self.dest.shutdown(force=force)
+        if self.arg:
+            args = self.arg.split('+')
+        else:
+            args = []
+        self.enforce(all(arg in ('force', 'wait') for arg in args))
+        force = ('force' in args)
+        wait = ('wait' in args)
+        self.fire_event_for_permission(force=force, wait=wait)
+        yield from self.dest.shutdown(force=force, wait=wait)
 
     @qubes.api.method('admin.vm.Pause', no_payload=True,
         scope='local', execute=True)
diff --git a/qubes/tests/api_admin.py b/qubes/tests/api_admin.py
index be94d916..9b5b3b54 100644
--- a/qubes/tests/api_admin.py
+++ b/qubes/tests/api_admin.py
@@ -1030,7 +1030,7 @@ netvm default=True type=vm \n'''
         self.vm.shutdown = coroutine_mock
         value = self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1')
         self.assertIsNone(value)
-        func_mock.assert_called_once_with(force=False)
+        func_mock.assert_called_once_with(force=False, wait=False)
 
     def test_231_shutdown_force(self):
         func_mock = unittest.mock.Mock()
@@ -1041,7 +1041,51 @@ netvm default=True type=vm \n'''
         self.vm.shutdown = coroutine_mock
         value = self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1', b'force')
         self.assertIsNone(value)
-        func_mock.assert_called_once_with(force=True)
+        func_mock.assert_called_once_with(force=True, wait=False)
+
+    def test_232_shutdown_wait(self):
+        func_mock = unittest.mock.Mock()
+
+        @asyncio.coroutine
+        def coroutine_mock(*args, **kwargs):
+            return func_mock(*args, **kwargs)
+        self.vm.shutdown = coroutine_mock
+        value = self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1', b'wait')
+        self.assertIsNone(value)
+        func_mock.assert_called_once_with(force=False, wait=True)
+
+    def test_233_shutdown_wait_force(self):
+        func_mock = unittest.mock.Mock()
+
+        @asyncio.coroutine
+        def coroutine_mock(*args, **kwargs):
+            return func_mock(*args, **kwargs)
+        self.vm.shutdown = coroutine_mock
+        value = self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1', b'wait+force')
+        self.assertIsNone(value)
+        func_mock.assert_called_once_with(force=True, wait=True)
+
+    def test_234_shutdown_force_wait(self):
+        func_mock = unittest.mock.Mock()
+
+        @asyncio.coroutine
+        def coroutine_mock(*args, **kwargs):
+            return func_mock(*args, **kwargs)
+        self.vm.shutdown = coroutine_mock
+        value = self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1', b'force+wait')
+        self.assertIsNone(value)
+        func_mock.assert_called_once_with(force=True, wait=True)
+
+    def test_234_shutdown_force_wait_invalid(self):
+        func_mock = unittest.mock.Mock()
+
+        @asyncio.coroutine
+        def coroutine_mock(*args, **kwargs):
+            return func_mock(*args, **kwargs)
+        self.vm.shutdown = coroutine_mock
+        with self.assertRaises(qubes.api.PermissionDenied):
+            self.call_mgmt_func(b'admin.vm.Shutdown', b'test-vm1', b'forcewait')
+        func_mock.assert_not_called()
 
     def test_240_pause(self):
         func_mock = unittest.mock.Mock()