Easy dependency switch if only preload is in chain

Fixes: QubesOS/qubes-issues#10227
For: QubesOS/qubes-issues#1512

Author: ben-grande

qubes/api/internal.py

@@ -267,7 +267,7 @@ async def suspend_pre(self):
 
         preload_templates = qubes.vm.dispvm.get_preload_templates(self.app)
         for qube in preload_templates:
-            qube.remove_preload_excess(0)
+            qube.remove_preload_excess(0, reason="system wants to suspend")
 
         # first keep track of VMs which were paused before suspending
         previously_paused = [

qubes/app.py

@@ -1632,6 +1632,7 @@ def _domain_event_callback(self, _conn, domain, event, _detail, _opaque):
     @qubes.events.handler("domain-pre-delete")
     def on_domain_pre_deleted(self, event, vm):
         # pylint: disable=unused-argument
+        preloads = []
         for obj in itertools.chain(self.domains, (self,)):
             if obj is vm:
                 # allow removed VM to reference itself
@@ -1642,6 +1643,11 @@ def on_domain_pre_deleted(self, event, vm):
                         isinstance(prop, qubes.vm.VMProperty)
                         and getattr(obj, prop.__name__) == vm
                     ):
+                        if getattr(obj, "is_preload") and (
+                            prop.__name__ in ["default_dispvm", "template"]
+                        ):
+                            preloads.append(obj)
+                            continue
                         self.log.error(
                             "Cannot remove %s, used by %s.%s",
                             vm,
@@ -1664,6 +1670,9 @@ def on_domain_pre_deleted(self, event, vm):
                 vm, "VM has devices assigned to other VMs: " + desc
             )
 
+        if preloads:
+            vm.remove_preload_excess(0, reason="domain will be deleted")
+
     @qubes.events.handler("domain-delete")
     def on_domain_deleted(self, event, vm):
         # pylint: disable=unused-argument

qubes/ext/audio.py

@@ -69,10 +69,12 @@ def set_tag_and_qubesdb_entry(self, subject, event, newvalue=None):
     @qubes.ext.handler("domain-pre-shutdown")
     def on_domain_pre_shutdown(self, vm, event, **kwargs):
         attached_vms = [
-            domain for domain in self.attached_vms(vm) if domain.is_running()
+            domain
+            for domain in self.attached_vms(vm)
+            if domain.is_running() and not getattr(domain, "is_preload")
         ]
         if attached_vms and not kwargs.get("force", False):
-            raise qubes.exc.QubesVMError(
+            raise qubes.exc.QubesVMInUseError(
                 self,
                 "There are running VMs using this VM as AudioVM: "
                 "{}".format(", ".join(vm.name for vm in attached_vms)),

qubes/tests/app.py

@@ -1096,6 +1096,44 @@ def test_205_remove_appvm_dispvm(self):
             with self.assertRaises(qubes.exc.QubesVMInUseError):
                 del self.app.domains[appvm]
 
+    def test_205_remove_appvm_dispvm_preload(self):
+        appvm = self.app.add_new_vm(
+            "AppVM",
+            name="test-appvm",
+            template=self.template,
+            template_for_dispvms=True,
+            label="red",
+        )
+        dispvm = self.app.add_new_vm(
+            "DispVM", name="test-dispvm", template=appvm, label="red"
+        )
+        dispvm_alt = self.app.add_new_vm(
+            "DispVM", name="test-dispvm-alt", template=appvm, label="red"
+        )
+        with mock.patch.object(self.app, "vmm"):
+            with self.assertRaises(qubes.exc.QubesVMInUseError):
+                del self.app.domains[appvm]
+
+        with mock.patch.object(self.app, "vmm"):
+            with mock.patch.object(appvm, "fire_event_async"):
+                appvm.features["preload-dispvm-max"] = "1"
+                appvm.features["preload-dispvm"] = str(dispvm.name)
+                with self.assertRaises(qubes.exc.QubesVMInUseError):
+                    del self.app.domains[appvm]
+
+        with mock.patch.object(self.app, "vmm"):
+            with (
+                mock.patch.object(appvm, "fire_event_async"),
+                mock.patch.object(appvm, "remove_preload_excess") as mock_remove
+            ):
+                appvm.features["preload-dispvm-max"] = "2"
+                appvm.features["preload-dispvm"] = dispvm.name
+                appvm.features["preload-dispvm"] = (
+                    dispvm.name + " " + dispvm_alt.name
+                )
+                del self.app.domains[appvm]
+                mock_remove.assert_called_with(0, reason=mock.ANY)
+
     def test_206_remove_attached(self):
         # See also qubes.tests.api_admin.
         vm = self.app.add_new_vm(

qubes/tests/ext.py

@@ -21,6 +21,7 @@
 import os
 import unittest.mock
 
+import qubes.ext.audio
 import qubes.ext.core_features
 import qubes.ext.custom_persist
 import qubes.ext.services
@@ -2579,3 +2580,58 @@ def test_015_feature_set_path_with_colon_without_options(self):
         self.vm.untrusted_qdb.write.assert_called_with(
             "/persist/test", "/var/test:dir:with:colon"
         )
+
+
+class TC_60_Audio(qubes.tests.QubesTestCase):
+    def setUp(self):
+        super().setUp()
+        self.ext = qubes.ext.audio.AUDIO()
+        self.audiovm = mock.MagicMock()
+        self.audiovm.name = "sys-audio"
+        self.client = mock.MagicMock()
+        self.client.name = "client"
+        self.client_alt = mock.MagicMock()
+        self.client_alt.name = "client"
+
+    def test_000_shutdown(self):
+        self.ext.on_domain_pre_shutdown(
+            self.audiovm,
+            "domain-pre-shutdown",
+        )
+
+    def test_000_shutdown_used(self):
+        with unittest.mock.patch.object(
+            self.ext, "attached_vms", return_value=[self.client]
+        ), unittest.mock.patch.object(
+            self.client, "is_running", return_value=True
+        ):
+            self.client.is_preload = False
+            with self.assertRaises(qubes.exc.QubesVMInUseError):
+                self.ext.on_domain_pre_shutdown(
+                    self.audiovm,
+                    "domain-pre-shutdown",
+                )
+
+            self.client.is_preload = True
+            self.ext.on_domain_pre_shutdown(
+                self.audiovm,
+                "domain-pre-shutdown",
+            )
+
+    def test_000_shutdown_used_by_some(self):
+        with unittest.mock.patch.object(
+            self.ext,
+            "attached_vms",
+            return_value=[self.client, self.client_alt],
+        ), unittest.mock.patch.object(
+            self.client, "is_running", return_value=True
+        ), unittest.mock.patch.object(
+            self.client_alt, "is_running", return_value=True
+        ):
+            self.client.is_preload = False
+            self.client_alt.is_preload = True
+            with self.assertRaises(qubes.exc.QubesVMInUseError):
+                self.ext.on_domain_pre_shutdown(
+                    self.audiovm,
+                    "domain-pre-shutdown",
+                )

qubes/tests/vm/mix/dvmtemplate.py

@@ -69,6 +69,11 @@ def setUp(self):
         self.template = self.app.add_new_vm(
             qubes.vm.templatevm.TemplateVM, name="test-template", label="red"
         )
+        self.template_alt = self.app.add_new_vm(
+            qubes.vm.templatevm.TemplateVM,
+            name="test-template-alt",
+            label="red",
+        )
         self.appvm = self.app.add_new_vm(
             qubes.vm.appvm.AppVM,
             name="test-vm",
@@ -98,9 +103,14 @@ def cleanup_dispvm(self):
         if hasattr(self, "dispvm"):
             self.dispvm.close()
             del self.dispvm
+        if hasattr(self, "dispvm_alt"):
+            self.dispvm_alt.close()
+            del self.dispvm_alt
         self.template.close()
+        self.template_alt.close()
         self.appvm.close()
         del self.template
+        del self.template_alt
         del self.appvm
         self.app.domains.clear()
         self.app.pools.clear()
@@ -241,7 +251,7 @@ def test_011_dvm_preload_del_max(self, mock_remove_preload_excess):
         del self.adminvm.features["preload-dispvm-max"]
         self.appvm.features["preload-dispvm-max"] = ""
         del self.appvm.features["preload-dispvm-max"]
-        mock_remove_preload_excess.assert_called_once_with(0)
+        mock_remove_preload_excess.assert_called_once_with(0, reason=mock.ANY)
 
     @mock.patch("qubes.events.Emitter.fire_event_async")
     def test_012_dvm_preload_set_max(self, mock_events):
@@ -251,6 +261,11 @@ def test_012_dvm_preload_set_max(self, mock_events):
             "domain-preload-dispvm-start", reason=mock.ANY
         )
 
+        mock_events.reset_mock()
+        self.appvm.template_for_dispvms = False
+        self.appvm.features["preload-dispvm-max"] = "2"
+        mock_events.assert_not_called()
+
     def test_013_dvm_preload_get_treshold(self):
         cases = [None, False, "0", "2", "1000"]
         self.assertEqual(self.appvm.get_feat_preload_threshold(), 0)
@@ -260,6 +275,101 @@ def test_013_dvm_preload_get_treshold(self):
                 threshold = self.appvm.get_feat_preload_threshold()
                 self.assertEqual(threshold, int(value or 0) * 1024**2)
 
+    @mock.patch("qubes.events.Emitter.fire_event_async")
+    @mock.patch(
+        "qubes.vm.mix.dvmtemplate.DVMTemplateMixin.remove_preload_excess"
+    )
+    def test_030_dvm_preload_set_template(self, mock_remove, mock_events):
+        # Don't try to preload if max is not set.
+        mock_events.side_effect = self.mock_coro
+        self.appvm.template = self.template_alt
+        mock_events.assert_not_called()
+        mock_remove.assert_called_once_with(0, reason=mock.ANY)
+
+        # Try to remove and preload if max is set and template has changed.
+        mock_remove.reset_mock()
+        self.appvm.features["preload-dispvm-max"] = "1"
+        mock_events.reset_mock()
+        self.appvm.template = self.template
+        mock_remove.assert_called_once_with(0, reason=mock.ANY)
+        mock_events.assert_called_once_with(
+            "domain-preload-dispvm-start", reason=mock.ANY
+        )
+
+        # Don't change anything if template hasn't changed.
+        mock_remove.reset_mock()
+        mock_events.reset_mock()
+        self.appvm.template = self.template
+        mock_remove.assert_not_called()
+        mock_events.assert_not_called()
+
+    @mock.patch("qubes.events.Emitter.fire_event_async")
+    @mock.patch(
+        "qubes.vm.mix.dvmtemplate.DVMTemplateMixin.remove_preload_excess"
+    )
+    def test_040_dvm_preload_set_template_for_dispvms(
+        self, mock_remove, mock_events
+    ):
+        # Remove preloads when disabling property.
+        mock_events.side_effect = self.mock_coro
+        self.appvm.template_for_dispvms = False
+        mock_events.assert_not_called()
+        mock_remove.assert_called_once_with(0, reason=mock.ANY)
+
+        # Preload when enabling property.
+        self.appvm.features["preload-dispvm-max"] = "1"
+        mock_events.reset_mock()
+        mock_remove.reset_mock()
+        self.appvm.template_for_dispvms = True
+        mock_remove.assert_not_called()
+        mock_events.assert_called_once_with(
+            "domain-preload-dispvm-start", reason=mock.ANY
+        )
+
+        # Try to disable property if it has dependents.
+        mock_events.reset_mock()
+        self.dispvm = self.app.add_new_vm(
+            qubes.vm.dispvm.DispVM,
+            name="test-dispvm",
+            template=self.appvm,
+            label="red",
+            dispid=42,
+        )
+        self.dispvm_alt = self.app.add_new_vm(
+            qubes.vm.dispvm.DispVM,
+            name="test-dispvm-alt",
+            template=self.appvm,
+            label="red",
+            dispid=43,
+        )
+        with self.assertRaises(qubes.exc.QubesVMInUseError):
+            self.appvm.template_for_dispvms = False
+            mock_remove.assert_not_called()
+            mock_events.assert_not_called()
+
+        # Disabling property when not all dependents are preloads
+        self.appvm.features["preload-dispvm-max"] = 1
+        self.appvm.features["preload-dispvm"] = self.dispvm.name
+        mock_events.reset_mock()
+        mock_remove.reset_mock()
+        with self.assertRaises(qubes.exc.QubesVMInUseError):
+            self.appvm.template_for_dispvms = False
+            mock_remove.assert_not_called()
+            mock_events.assert_not_called()
+
+        # Disabling property when all dependents are preloads
+        self.appvm.features["preload-dispvm-max"] = 2
+        mock_events.reset_mock()
+        mock_remove.reset_mock()
+        self.appvm.features["preload-dispvm"] = (
+            self.dispvm.name + " " + self.dispvm_alt.name
+        )
+        mock_events.reset_mock()
+        mock_remove.reset_mock()
+        self.appvm.template_for_dispvms = False
+        mock_remove.assert_called_once_with(0, reason=mock.ANY)
+        mock_events.assert_not_called()
+
     def test_100_get_preload_templates(self):
         print(qubes.vm.dispvm.get_preload_templates(self.app))
         self.appvm.features["supported-rpc.qubes.WaitForRunningSystem"] = True

qubes/tests/vm/mix/net.py

@@ -213,6 +213,19 @@ def test_170_provides_network_netvm(self):
         self.assertPropertyValue(vm2, "netvm", "", None, "")
         self.assertPropertyValue(vm, "provides_network", False, False, "False")
 
+    @patch("qubes.vm.qubesvm.QubesVM.libvirt_domain")
+    @patch("qubes.vm.qubesvm.QubesVM.is_halted", return_value=False)
+    def test_180_shutdown(self, mock_halted, mock_shutdown):
+        # pylint: disable=unused-argument
+        vm = self.get_vm()
+        self.setup_netvms(vm)
+        with patch.object(vm, "is_running", return_value=True):
+            vm.is_preload = False
+            with self.assertRaises(qubes.exc.QubesVMInUseError):
+                self.loop.run_until_complete(vm.netvm.shutdown())
+            vm.is_preload = True
+            self.loop.run_until_complete(vm.netvm.shutdown())
+
     def test_200_vmid_to_ipv4(self):
         testcases = (
             (1, "0.1"),

qubes/vm/dispvm.py

@@ -511,7 +511,9 @@ async def from_appvm(cls, appvm, preload=False, **kwargs):
                     # The gap is filled after the delay set by the
                     # 'domain-shutdown' of its ancestors. Not refilling now to
                     # deliver a disposable faster.
-                    appvm.remove_preload_from_list([qube.name])
+                    appvm.remove_preload_from_list(
+                        [qube.name], reason="of outdated volume(s)"
+                    )
                     # Delay to not  affect this run.
                     asyncio.ensure_future(
                         qube.delay(delay=2, coros=[qube.cleanup()])
@@ -527,7 +529,9 @@ async def from_appvm(cls, appvm, preload=False, **kwargs):
                 # - Another request to this function will not return the same
                 #   qube.
                 dispvm.features["preload-dispvm-in-progress"] = True
-                appvm.remove_preload_from_list([dispvm.name])
+                appvm.remove_preload_from_list(
+                    [dispvm.name], reason="qube was requested"
+                )
                 dispvm.preload_requested = True
                 app.save()
                 timeout = int(dispvm.qrexec_timeout * 1.2)
@@ -603,7 +607,9 @@ def use_preload(self):
             self.log.warning("Using a preloaded qube before requesting it")
             if not appvm.features.get("internal", None):
                 del self.features["internal"]
-            appvm.remove_preload_from_list([self.name])
+            appvm.remove_preload_from_list(
+                [self.name], reason="qube was used without being requested"
+            )
             self.features["preload-dispvm-in-progress"] = False
         self.app.save()
         asyncio.ensure_future(
@@ -620,8 +626,9 @@ async def _bare_cleanup(self):
     def _preload_cleanup(self):
         """Cleanup preload from list"""
         if self.name in self.template.get_feat_preload():
-            self.log.info("Automatic cleanup removes qube from preload list")
-            self.template.remove_preload_from_list([self.name])
+            self.template.remove_preload_from_list(
+                [self.name], reason="automatic cleanup was called"
+            )
 
     async def cleanup(self):
         """Clean up after the DispVM