From 7bcab46f96e58b772a2aa82093582dd7f2941adc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Sat, 28 Oct 2017 22:40:24 +0200
Subject: [PATCH] tools/qvm-shutdown: fix handling shutdown timeout for
 multiple VMs

When some VM timeout on shutdown, the tool will try to kill all of them,
but at this point some of them may be already powered off (not all
hanged during shutdown, but only some). Handle this
situation instead of crashing. And add appropriate test.
---
 qubesadmin/tests/tools/qvm_shutdown.py | 56 ++++++++++++++++++++++++++
 qubesadmin/tools/qvm_shutdown.py       | 12 +++++-
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/qubesadmin/tests/tools/qvm_shutdown.py b/qubesadmin/tests/tools/qvm_shutdown.py
index d61c9d8..b37956b 100644
--- a/qubesadmin/tests/tools/qvm_shutdown.py
+++ b/qubesadmin/tests/tools/qvm_shutdown.py
@@ -240,3 +240,59 @@ class TC_00_qvm_shutdown(qubesadmin.tests.QubesTestCase):
                 False):
             qubesadmin.tools.qvm_shutdown.main(['--wait', '--all'], app=self.app)
         self.assertAllCalled()
+
+    @unittest.skipUnless(qubesadmin.tools.qvm_shutdown.have_events,
+        'Events not present')
+    def test_015_wait_all_kill_timeout(self):
+        '''test --wait option, with multiple VMs and killing on timeout'''
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        patch = unittest.mock.patch(
+            'qubesadmin.events.EventsDispatcher._get_events_reader')
+        mock_events = patch.start()
+        self.addCleanup(patch.stop)
+        mock_events.side_effect = qubesadmin.tests.tools.MockEventsReader([
+            b'1\0\0connection-established\0\0',
+            b'1\0sys-net\0domain-shutdown\0\0',
+            ])
+
+        self.app.expected_calls[
+            ('some-vm', 'admin.vm.Shutdown', None, None)] = \
+            b'0\x00'
+        self.app.expected_calls[
+            ('some-vm', 'admin.vm.Kill', None, None)] = \
+            b'2\x00QubesVMNotStartedError\x00\x00Domain is powered off\x00'
+        self.app.expected_calls[
+            ('other-vm', 'admin.vm.Shutdown', None, None)] = \
+            b'0\x00'
+        self.app.expected_calls[
+            ('other-vm', 'admin.vm.Kill', None, None)] = \
+            b'0\x00'
+        self.app.expected_calls[
+            ('sys-net', 'admin.vm.Shutdown', None, None)] = \
+            b'0\x00'
+        self.app.expected_calls[
+            ('sys-net', 'admin.vm.Kill', None, None)] = \
+            b'2\x00QubesVMNotStartedError\x00\x00Domain is powered off\x00'
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00' \
+            b'sys-net class=AppVM state=Running\n' \
+            b'some-vm class=AppVM state=Running\n' \
+            b'other-vm class=AppVM state=Running\n'
+        self.app.expected_calls[
+            ('some-vm', 'admin.vm.List', None, None)] = [
+            b'0\x00some-vm class=AppVM state=Running\n',
+        ]
+        self.app.expected_calls[
+            ('other-vm', 'admin.vm.List', None, None)] = [
+            b'0\x00other-vm class=AppVM state=Running\n',
+            b'0\x00other-vm class=AppVM state=Running\n',
+        ]
+        self.app.expected_calls[
+            ('sys-net', 'admin.vm.List', None, None)] = \
+            b'0\x00sys-net class=AppVM state=Halted\n'
+        qubesadmin.tools.qvm_shutdown.main(
+            ['--wait', '--all', '--timeout=1'], app=self.app)
+        self.assertAllCalled()
diff --git a/qubesadmin/tools/qvm_shutdown.py b/qubesadmin/tools/qvm_shutdown.py
index b6d24da..48687a3 100644
--- a/qubesadmin/tools/qvm_shutdown.py
+++ b/qubesadmin/tools/qvm_shutdown.py
@@ -87,7 +87,11 @@ def main(args=None, app=None):  # pylint: disable=missing-docstring
                     args.timeout))
             except asyncio.TimeoutError:
                 for vm in this_round_domains:
-                    vm.kill()
+                    try:
+                        vm.kill()
+                    except qubesadmin.exc.QubesVMNotStartedError:
+                        # already shut down
+                        pass
         else:
             timeout = args.timeout
             current_vms = list(sorted(this_round_domains))
@@ -105,7 +109,11 @@ def main(args=None, app=None):  # pylint: disable=missing-docstring
                     'Killing remaining qubes: {}'
                     .format(', '.join([str(vm) for vm in current_vms])))
             for vm in current_vms:
-                vm.kill()
+                try:
+                    vm.kill()
+                except qubesadmin.exc.QubesVMNotStartedError:
+                    # already shut down
+                    pass
 
     if args.wait:
         if have_events: