Add global preload to template gathering

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/api/internal.py

@@ -265,9 +265,7 @@ async def suspend_pre(self):
         :return:
         """
 
-        preload_templates = qubes.vm.dispvm.get_preload_templates(
-            self.app.domains
-        )
+        preload_templates = qubes.vm.dispvm.get_preload_templates(self.app)
         for qube in preload_templates:
             qube.remove_preload_excess(0)
 
@@ -416,9 +414,7 @@ async def suspend_post(self):
                         "qubes.SuspendPostAll",
                     )
 
-        preload_templates = qubes.vm.dispvm.get_preload_templates(
-            self.app.domains
-        )
+        preload_templates = qubes.vm.dispvm.get_preload_templates(self.app)
         for qube in preload_templates:
             asyncio.ensure_future(
                 qube.fire_event_async("domain-preload-dispvm-autostart")

qubes/tests/integ/dispvm.py

@@ -298,25 +298,23 @@ async def cleanup_preload_run(self, qube):
 
     def cleanup_preload(self):
         logger.info("start")
-        if "preload-dispvm-max" in self.app.domains[self.disp_base].features:
-            logger.info("has local preload")
-            self.loop.run_until_complete(
-                self.cleanup_preload_run(self.disp_base)
-            )
-            logger.info("deleting local feature")
-            del self.disp_base.features["preload-dispvm-max"]
-        if "preload-dispvm-max" in self.app.domains["dom0"].features:
-            logger.info("has global preload")
-            default_dispvm = self.app.default_dispvm
-            if default_dispvm:
-                self.loop.run_until_complete(
-                    self.cleanup_preload_run(default_dispvm)
-                )
-            logger.info("deleting global feature")
-            del self.app.domains["dom0"].features["preload-dispvm-max"]
+        default_dispvm = self.app.default_dispvm
+        # Clean features from all qubes to avoid them being considered by
+        # tests that target preloads on the whole system, such as
+        # `/usr/lib/qubes/preload-dispvm`.
         if "preload-dispvm-threshold" in self.app.domains["dom0"].features:
             logger.info("deleting global threshold feature")
             del self.app.domains["dom0"].features["preload-dispvm-threshold"]
+        for qube in self.app.domains:
+            if "preload-dispvm-max" not in qube.features:
+                continue
+            logger.info("removing preloaded disposables: '%s'", qube.name)
+            if qube == default_dispvm:
+                self.loop.run_until_complete(
+                    self.cleanup_preload_run(default_dispvm)
+                )
+            logger.info("deleting max preload feature")
+            del qube.features["preload-dispvm-max"]
         logger.info("end")
 
     async def no_preload(self):
@@ -616,6 +614,7 @@ def test_017_preload_autostart(self):
         self.app.default_dispvm = self.disp_base
 
         preload_max = 1
+        logger.info("no refresh to be made")
         proc = self.loop.run_until_complete(
             asyncio.create_subprocess_exec("/usr/lib/qubes/preload-dispvm")
         )
@@ -624,13 +623,14 @@ def test_017_preload_autostart(self):
         )
         self.assertEqual(self.disp_base.get_feat_preload(), [])
 
+        logger.info("refresh to be made")
         self.disp_base.features["preload-dispvm-max"] = str(preload_max)
         self.loop.run_until_complete(self.wait_preload(preload_max))
         old_preload = self.disp_base.get_feat_preload()
         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=40))
         preload_dispvm = self.disp_base.get_feat_preload()
         self.assertEqual(len(old_preload), preload_max)
         self.assertEqual(len(preload_dispvm), preload_max)
@@ -639,6 +639,7 @@ def test_017_preload_autostart(self):
             f"old_preload={old_preload} preload_dispvm={preload_dispvm}",
         )
 
+        logger.info("global refresh to be made")
         preload_max += 1
         self.adminvm.features["preload-dispvm-max"] = str(preload_max)
         self.loop.run_until_complete(self.wait_preload(preload_max))
@@ -647,7 +648,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=40))
         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,14 +40,33 @@ def _setter_template(self, prop, value):
     return value
 
 
-def get_preload_templates(domains) -> list:
-    return [
+# 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
+    default_dispvm = getattr(app, "default_dispvm", None)
+    global_max = get_preload_max(domains["dom0"])
+    appvms = [
         qube
         for qube in domains
-        if int(qube.features.get("preload-dispvm-max", 0) or 0) > 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))
+                )
+            )
+        )
     ]
+    return appvms
 
 
 class DispVM(qubes.vm.qubesvm.QubesVM):