Consistent preload template gathering

- Easier copy-paste when code duplication is necessary;
- Only fires event if the qube should be really preloaded, which is more
  difficult to do considering the global feature. The prior behavior
  didn't error out, just a warning log;
- Add tests

For: QubesOS/qubes-issues#1512

Author: ben-grande

linux/aux-tools/preload-dispvm

@@ -12,39 +12,43 @@ import concurrent.futures
 import qubesadmin
 
 
-def get_max(qube):
-    return int(qube.features.get("preload-dispvm-max", 0) or 0)
+def get_preload_max(qube) -> int | None:
+    value = qube.features.get("preload-dispvm-max", None)
+    return int(value) if value else value
 
 
 async def main():
     app = qubesadmin.Qubes()
     domains = app.domains
     default_dispvm = getattr(app, "default_dispvm", None)
+    global_max = get_preload_max(domains["dom0"])
     appvms = [
         qube
         for qube in domains
-        if get_max(qube) > 0
-        and (
-            (
-                qube.klass == "AppVM"
-                and getattr(qube, "template_for_dispvms", False)
+        if (
+            qube.klass == "AppVM"
+            and getattr(qube, "template_for_dispvms", False)
+            and (
+                (qube != default_dispvm and get_preload_max(qube))
+                or (
+                    (qube == default_dispvm and global_max)
+                    or (global_max is None and get_preload_max(qube))
+                )
             )
-            or (qube.name == "dom0" and default_dispvm)
         )
     ]
     method = "admin.vm.CreateDisposable"
     loop = asyncio.get_running_loop()
     tasks = []
-    if "dom0" in appvms and default_dispvm in appvms:
-        appvms.remove(default_dispvm)
     with concurrent.futures.ThreadPoolExecutor() as executor:
         for qube in appvms:
-            maximum = get_max(qube)
-            msg = f"{qube}:{maximum}"
-            if qube.name == "dom0":
-                qube = default_dispvm
-                msg = "global:" + msg
-            print(msg)
+            if qube == default_dispvm and global_max is not None:
+                maximum = global_max
+                msg = f"global:{qube}:{maximum}"
+            else:
+                maximum = get_preload_max(qube)
+                msg = f"{qube}:{maximum}"
+            print(repr(msg))
             exec_args = qube.qubesd_call, qube.name, method, "preload-autostart"
             future = loop.run_in_executor(executor, *exec_args)
             tasks.append(future)

qubes/tests/integ/dispvm.py

@@ -607,6 +607,9 @@ def test_017_preload_autostart(self):
         'admin.vm.CreateDisposable+preload-autostart' which fires the event
         'domain-preload-dispvm-autostart', clearing the current preload list
         and filling with new ones."""
+        # This function doesn't patch qubesadmin.Qubes() of the autostart
+        # script, therefore, it will also refresh preloaded disposables that
+        # are not part of the test.
         logger.info("start")
         self.app.default_dispvm = self.disp_base
 
@@ -615,7 +618,7 @@ def test_017_preload_autostart(self):
             asyncio.create_subprocess_exec("/usr/lib/qubes/preload-dispvm")
         )
         self.loop.run_until_complete(
-            asyncio.wait_for(proc.communicate(), timeout=10)
+            asyncio.wait_for(proc.communicate(), timeout=60)
         )
         self.assertEqual(self.disp_base.get_feat_preload(), [])
 
@@ -625,7 +628,7 @@ def test_017_preload_autostart(self):
         proc = self.loop.run_until_complete(
             asyncio.create_subprocess_exec("/usr/lib/qubes/preload-dispvm")
         )
-        self.loop.run_until_complete(asyncio.wait_for(proc.wait(), timeout=30))
+        self.loop.run_until_complete(asyncio.wait_for(proc.wait(), timeout=60))
         preload_dispvm = self.disp_base.get_feat_preload()
         self.assertEqual(len(old_preload), preload_max)
         self.assertEqual(len(preload_dispvm), preload_max)
@@ -642,7 +645,7 @@ def test_017_preload_autostart(self):
         proc = self.loop.run_until_complete(
             asyncio.create_subprocess_exec("/usr/lib/qubes/preload-dispvm")
         )
-        self.loop.run_until_complete(asyncio.wait_for(proc.wait(), timeout=30))
+        self.loop.run_until_complete(asyncio.wait_for(proc.wait(), timeout=60))
         preload_dispvm = self.disp_base.get_feat_preload()
         self.assertEqual(len(old_preload), preload_max)
         self.assertEqual(len(preload_dispvm), preload_max)

qubes/tests/vm/dispvm.py

@@ -35,12 +35,16 @@
 class TestApp(qubes.tests.vm.TestApp):
     def __init__(self):
         super(TestApp, self).__init__()
-        self.qid_counter = 1
+        self.qid_counter = 0
 
     def add_new_vm(self, cls, **kwargs):
         qid = self.qid_counter
-        self.qid_counter += 1
-        vm = cls(self, None, qid=qid, **kwargs)
+        if self.qid_counter == 0:
+            self.qid_counter += 1
+            vm = cls(self, None, **kwargs)
+        else:
+            self.qid_counter += 1
+            vm = cls(self, None, qid=qid, **kwargs)
         self.domains[vm.name] = vm
         self.domains[vm] = vm
         return vm
@@ -58,6 +62,8 @@ def setUp(self):
             name="linux-kernel"
         )
         self.app.vmm.offline_mode = True
+        self.app.default_dispvm = None
+        self.adminvm = self.app.add_new_vm(qubes.vm.adminvm.AdminVM)
         self.template = self.app.add_new_vm(
             qubes.vm.templatevm.TemplateVM, name="test-template", label="red"
         )
@@ -68,8 +74,12 @@ def setUp(self):
             template=self.template,
             label="red",
         )
-        self.app.domains[self.appvm.name] = self.appvm
-        self.app.domains[self.appvm] = self.appvm
+        self.appvm_alt = self.app.add_new_vm(
+            qubes.vm.appvm.AppVM,
+            name="test-vm-alt",
+            template=self.template,
+            label="red",
+        )
         self.addCleanup(self.cleanup_dispvm)
         self.emitter = qubes.tests.TestEmitter()
 
@@ -83,10 +93,15 @@ def cleanup_dispvm(self):
             del self.dispvm
         self.template.close()
         self.appvm.close()
-        del self.template
+        self.appvm_alt.close()
         del self.appvm
+        del self.appvm_alt
+        del self.template
+        del self.adminvm
+        self.app.close()
         self.app.domains.clear()
         self.app.pools.clear()
+        del self.app
 
     async def mock_coro(self, *args, **kwargs):
         pass
@@ -275,6 +290,49 @@ def test_000_from_appvm_preload_fill_gap(
         mock_symlink.assert_not_called()
         mock_makedirs.assert_called_once()
 
+    def test_000_get_preload_max(self):
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.appvm), None)
+        self.appvm.features["supported-rpc.qubes.WaitForRunningSystem"] = True
+        self.appvm.features["preload-dispvm-max"] = 1
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.appvm), 1)
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.adminvm), None)
+        self.adminvm.features["preload-dispvm-max"] = ""
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.adminvm), "")
+        self.adminvm.features["preload-dispvm-max"] = 2
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.adminvm), 2)
+
+    def test_000_get_preload_templates(self):
+        get_preload_templates = qubes.vm.dispvm.get_preload_templates
+        self.assertEqual(get_preload_templates(self.app), [])
+        self.appvm.template_for_dispvms = True
+        self.appvm_alt.template_for_dispvms = True
+        self.assertEqual(get_preload_templates(self.app), [])
+
+        self.appvm.features["supported-rpc.qubes.WaitForRunningSystem"] = True
+        self.appvm_alt.features["supported-rpc.qubes.WaitForRunningSystem"] = (
+            True
+        )
+        self.appvm.features["preload-dispvm-max"] = 1
+        self.appvm_alt.features["preload-dispvm-max"] = 0
+        self.assertEqual(get_preload_templates(self.app), [self.appvm])
+
+        self.adminvm.features["preload-dispvm-max"] = ""
+        # Still not default_dispvm
+        self.appvm_alt.features["preload-dispvm-max"] = 1
+        self.assertEqual(
+            get_preload_templates(self.app), [self.appvm, self.appvm_alt]
+        )
+
+        with mock.patch.object(self.appvm, "fire_event_async"):
+            self.app.default_dispvm = self.appvm
+        self.assertEqual(get_preload_templates(self.app), [self.appvm_alt])
+
+        self.app.default_dispvm = None
+        self.adminvm.features["preload-dispvm-max"] = 1
+        self.assertEqual(
+            get_preload_templates(self.app), [self.appvm, self.appvm_alt]
+        )
+
     def test_001_from_appvm_reject_not_allowed(self):
         with self.assertRaises(qubes.exc.QubesException):
             dispvm = self.loop.run_until_complete(

qubes/tests/vm/mix/dvmtemplate.py

@@ -259,3 +259,9 @@ def test_013_dvm_preload_get_treshold(self):
                 self.adminvm.features["preload-dispvm-threshold"] = value
                 threshold = self.appvm.get_feat_preload_threshold()
                 self.assertEqual(threshold, int(value or 0) * 1024**2)
+
+    def test_100_get_preload_templates(self):
+        print(qubes.vm.dispvm.get_preload_templates(self.app))
+        self.appvm.features["supported-rpc.qubes.WaitForRunningSystem"] = True
+        self.appvm.features["preload-dispvm-max"] = 1
+        self.assertEqual(qubes.vm.dispvm.get_preload_max(self.appvm), 1)

qubes/vm/dispvm.py

@@ -40,26 +40,32 @@ def _setter_template(self, prop, value):
     return value
 
 
+# Keep in sync with linux/aux-tools/preload-dispvm
+def get_preload_max(qube) -> int | None:
+    value = qube.features.get("preload-dispvm-max", None)
+    return int(value) if value else value
+
+
+# Keep in sync with linux/aux-tools/preload-dispvm
 def get_preload_templates(app) -> list:
     domains = app.domains
-    appvms = []
     default_dispvm = getattr(app, "default_dispvm", None)
-    # Only add default_dispvm now if it will not be added by the other clause.
-    if (
-        int(domains["dom0"].features.get("preload-dispvm-max", 0) or 0) > 0
-        and default_dispvm
-        and int(default_dispvm.features.get("preload-dispvm-max", 0) or 0) == 0
-    ):
-        appvms.append(default_dispvm)
-    appvms.extend(
-        [
-            qube
-            for qube in domains
-            if int(qube.features.get("preload-dispvm-max", 0) or 0) > 0
-            and qube.klass == "AppVM"
+    global_max = get_preload_max(domains["dom0"])
+    appvms = [
+        qube
+        for qube in domains
+        if (
+            qube.klass == "AppVM"
             and getattr(qube, "template_for_dispvms", False)
-        ]
-    )
+            and (
+                (qube != default_dispvm and get_preload_max(qube))
+                or (
+                    (qube == default_dispvm and global_max)
+                    or (global_max is None and get_preload_max(qube))
+                )
+            )
+        )
+    ]
     return appvms