Preload disposables

For: QubesOS/qubes-issues#1512

Author: ben-grande

linux/aux-tools/preload-dispvm

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+
+#import asyncio
+import qubesadmin
+
+def get_apps():
+    domains = qubesadmin.Qubes().domains
+    return [qube for qube in domains
+    if int(qube.features.get("dispvm-preload-max", 0)) > 0
+    and qube.klass == "AppVM"
+    and getattr(qube, "template_for_dispvms", False)
+    ]
+
+#async def main():
+#    appvms = get_apps()
+#    ## TODO: How to fire events outside of qubesd? How to load DispVM by hand.
+#    #o = MyClass()
+#    #o.events_enabled = True
+#    #effect = o.fire_event('event1')
+#    event = "domain-preloaded-dispvm-autostart"
+#    tasks = [qube.fire_event_async(event) for qube in appvms]
+#    await asyncio.gather(*tasks)
+
+def main():
+    appvms = get_apps()
+    method = "admin.vm.CreateDisposable"
+    for qube in appvms:
+        qube.qubesd_call("dom0", method, "preload-autostart")
+
+if __name__ == "__main__":
+    #asyncio.run(main())
+    main()

linux/systemd/Makefile

@@ -9,6 +9,7 @@ install:
 	cp qubes-vm@.service $(DESTDIR)$(UNITDIR)
 	cp qubes-qmemman.service $(DESTDIR)$(UNITDIR)
 	cp qubesd.service $(DESTDIR)$(UNITDIR)
+	cp qubes-preload-dispvm.service $(DESTDIR)$(UNITDIR)
 	install -d $(DESTDIR)$(UNITDIR)/lvm2-pvscan@.service.d
 	install -m 0644 lvm2-pvscan@.service.d_30_qubes.conf \
 		$(DESTDIR)$(UNITDIR)/lvm2-pvscan@.service.d/30_qubes.conf

linux/systemd/qubes-preload-dispvm.service

@@ -0,0 +1,17 @@
+[Unit]
+Description=Preload Qubes DispVMs
+After=qubes-vm@.service
+# TODO: or 'Requires='? Should a failure to autostart sys-(usb|net) make this
+# unit be skipped?
+Wants=qubes-vm@.service
+ConditionKernelCommandLine=!qubes.skip_autostart
+
+[Service]
+Type=oneshot
+Environment=DISPLAY=:0
+ExecStart=/usr/lib/qubes/preload-dispvm
+Group=qubes
+RemainAfterExit=yes
+
+[Install]
+WantedBy=multi-user.target

qubes/api/admin.py

@@ -1272,7 +1272,13 @@ async def _vm_create(
 
     @qubes.api.method("admin.vm.CreateDisposable", scope="global", write=True)
     async def create_disposable(self, untrusted_payload):
-        self.enforce(not self.arg)
+        self.enforce(self.arg in [None, "preload", "preload-autostart"])
+        preload = False
+        preload_autostart = False
+        if self.arg == "preload":
+            preload = True
+        if self.arg == "preload-autostart":
+            preload_autostart = True
         if untrusted_payload not in (b"", b"uuid"):
             raise qubes.exc.QubesValueError(
                 "Invalid payload for admin.vm.CreateDisposable: "
@@ -1286,7 +1292,15 @@ async def create_disposable(self, untrusted_payload):
 
         self.fire_event_for_permission(dispvm_template=dispvm_template)
 
-        dispvm = await qubes.vm.dispvm.DispVM.from_appvm(dispvm_template)
+        if preload_autostart:
+            await (
+                self.dest.fire_event_async("domain-preloaded-dispvm-autostart")
+            )
+            return
+
+        dispvm = await qubes.vm.dispvm.DispVM.from_appvm(
+            dispvm_template, preload=preload
+        )
         # TODO: move this to extension (in race-free fashion, better than here)
         dispvm.tags.add("created-by-" + str(self.src))
         dispvm.tags.add("disp-created-by-" + str(self.src))

qubes/tests/api_admin.py

@@ -3716,6 +3716,35 @@ def test_642_vm_create_disposable_not_allowed(self, storage_mock):
             self.call_mgmt_func(b"admin.vm.CreateDisposable", b"test-vm1")
         self.assertFalse(self.app.save.called)
 
+    @unittest.mock.patch("qubes.storage.Storage.create")
+    def test_643_vm_create_disposable_preload(self, mock_storage):
+        mock_storage.side_effect = self.dummy_coro
+        self.vm.template_for_dispvms = True
+        self.vm.features["preload-dispvm-max"] = 1
+        self.app.default_dispvm = self.vm
+        retval = self.call_mgmt_func(
+            b"admin.vm.CreateDisposable", b"dom0", arg="preload"
+        )
+        dispvm_preload = self.vm.features.get("dispvm-preload", "").split(" ")
+        self.assertIn(retval, dispvm_preload)
+        mock_storage.assert_called_once_with()
+        self.assertTrue(self.app.save.called)
+
+    @unittest.mock.patch("qubes.storage.Storage.create")
+    def test_643_vm_create_disposable_preload_autostart(self, mock_storage):
+        mock_storage.side_effect = self.dummy_coro
+        self.vm.template_for_dispvms = True
+        self.vm.features["preload-dispvm-max"] = 1
+        self.app.default_dispvm = self.vm
+        retval = self.call_mgmt_func(
+            b"admin.vm.CreateDisposable", b"dom0", arg="preload-autostart"
+        )
+        # TODO: doesn't return any value, so how to check if it was preloaded?
+        #dispvm_preload = self.vm.features.get("preload-dispvm", "").split(" ")
+        self.assertIsNone(retval)
+        mock_storage.assert_called_once_with()
+        self.assertTrue(self.app.save.called)
+
     def test_650_vm_device_set_mode_required(self):
         assignment = DeviceAssignment(
             VirtualDevice(Port(self.vm, "1234", "testclass"), device_id="bee"),

qubes/tests/integ/dispvm.py

@@ -206,6 +206,7 @@ def tearDown(self):
         self.app.default_dispvm = None
         super(TC_20_DispVMMixin, self).tearDown()
 
+    # TODO: Test if run_service() marks the prelaoded DispVM as used.
     def test_010_simple_dvm_run(self):
         dispvm = self.loop.run_until_complete(
             qubes.vm.dispvm.DispVM.from_appvm(self.disp_base)

qubes/tests/vm/dispvm.py

@@ -84,6 +84,7 @@ def cleanup_dispvm(self):
     async def mock_coro(self, *args, **kwargs):
         pass
 
+    # TODO: Test creating preloaded disposable and if features are correct.
     @mock.patch("os.symlink")
     @mock.patch("os.makedirs")
     @mock.patch("qubes.storage.Storage")

qubes/vm/dispvm.py

@@ -20,7 +20,10 @@
 
 """ A disposable vm implementation """
 
+import asyncio
 import copy
+import psutil
+import subprocess
 
 import qubes.vm.qubesvm
 import qubes.vm.appvm
@@ -187,11 +190,170 @@ def __init__(self, app, xml, *args, **kwargs):
             self.features.update(template.features)
             self.tags.update(template.tags)
 
+    def get_feat_preload(self, feature):
+        if feature not in ["preload-dispvm", "preload-dispvm-max"]:
+            raise qubes.exc.QubesException("Invalid feature provided")
+
+        if feature == "preload-dispvm":
+            default = ""
+        elif feature == "preload-dispvm-max":
+            default = 0
+
+        value = self.features.check_with_template(feature, default)
+
+        if feature == "preload-dispvm":
+            return value.split(" ")
+        if feature == "preload-dispvm-max":
+            return int(value)
+        return None
+
+    def is_preloaded(self):
+        preload_dispvm = self.get_feat_preload("preload-dispvm")
+        if not preload_dispvm:
+            return False
+        if self.name not in preload_dispvm:
+            return False
+        return True
+
+    async def mark_preloaded(self):
+        """
+        Create preloaded DispVM.
+
+        Template from which the VM should be created.
+
+        :return:
+        """
+        preload_dispvm = self.get_feat_preload("preload-dispvm")
+        if preload_dispvm:
+            preload_dispvm.append(self.name)
+        else:
+            preload_dispvm = [self.name]
+
+        appvm = getattr(self, "template")
+        appvm.features["preload-dispvm"] = " ".join(preload_dispvm)
+        self.features["internal"] = True
+
+    async def use_preloaded(self):
+        """
+        Mark preloaded DispVM as used.
+
+        :return:
+        """
+        appvm = getattr(self, "template")
+
+        preload_dispvm = self.get_feat_preload("preload-dispvm")
+        if self.name not in preload_dispvm:
+            raise qubes.exc.QubesException("DispVM is not preloaded")
+
+        preload_dispvm = " ".join(preload_dispvm.remove(self.name))
+        appvm.features["preload-dispvm"] = preload_dispvm
+        self.features["internal"] = False
+        await appvm.fire_event_async(
+            "domain-preloaded-dispvm-used", dispvm=self
+        )
+
+    @qubes.events.handler(
+        "domain-preloaded-dispvm-used", "domain-preloaded-dispvm-autostart"
+    )
+    async def on_domain_preloaded_dispvm_used(self, event, delay=5, **kwargs):  # pylint: disable=unused-argument
+        """When preloaded DispVM is used or after boot, preload another one.
+
+        :param event: event which was fired
+        :param delay: delay between trials
+        :returns:
+        """
+        await asyncio.sleep(delay)
+        while True:
+            # TODO: Is there existing Qubes code that checks available memory
+            # before starting a qube?
+            memory = getattr(self, "memory", 0)
+            available_memory = (
+                psutil.virtual_memory().available / (1024 * 1024)
+            )
+            threshold = 1024 * 5
+            if memory >= (available_memory - threshold):
+                ## TODO: how to pass arg?
+                await qubes.vm.dispvm.DispVM.from_appvm(
+                    self, preload=True
+                ).start()
+                #await qubes.api.admin.QubesAdminAPI.create_disposable(
+                #    self.app, b"dom0", "admin.vm.CreateDisposable", b"dom0", b"preload"
+                #)
+                # TODO: what to do if the maximum is never reached on autostart
+                # as there is not enough memory, and then a preloaded DispVM is
+                # used, calling for the creation of another one, while the
+                # autostart will also try to create one. Is this a race
+                # condition?
+                # TODO: fire event after start of all qubes that are set to
+                # autostart.
+                if event == "domain-preloaded-dispvm-autostart":
+                    preload_dispvm_max = self.get_feat_preload(
+                        "preload-dispvm-max"
+                    )
+                    preload_dispvm = self.get_feat_preload("preload-dispvm")
+                    if (
+                        preload_dispvm
+                        and len(preload_dispvm) < preload_dispvm_max
+                    ):
+                        continue
+                break
+            await asyncio.sleep(delay)
+
     @qubes.events.handler("domain-load")
     def on_domain_loaded(self, event):
         """When domain is loaded assert that this vm has a template."""  # pylint: disable=unused-argument
         assert self.template
 
+    @qubes.events.handler("domain-start")
+    # W0236 (invalid-overridden-method) Method 'on_domain_started' was expected
+    # to be 'non-async', found it instead as 'async'
+    # TODO: Seems to conflict with qubes.vm.mix.net, which is pretty strange.
+    # Larger bug? qubes.vm.qubesvm.QubesVM has NetVMMixin... which conflicts...
+    async def on_domain_started(self, event, **kwargs):
+        """Pause preloaded domains as soon as they start."""
+        # TODO:
+        # Marek: Test if pause isn't too early. Some services (especially:
+        #   gui-agent) may still be starting.  qubes.WaitForSession service may
+        #   help (ensure to use async handler to not block qubesd while waiting
+        #   on it).
+        no_gui_sleep = 15
+        gui_timeout = 30
+        if self.is_preloaded():
+            gui = self.features.get("gui", None)
+            if not gui:
+                asyncio.sleep(no_gui_sleep)
+                self.pause()
+                return
+
+            proc = None
+            try:
+                proc = await asyncio.wait_for(
+                    self.run_service(
+                        "qubes.WaitForSession",
+                        user=self.default_user,
+                        stdout=subprocess.DEVNULL,
+                        stderr=subprocess.DEVNULL,
+                    ),
+                    timeout=gui_timeout,
+                )
+            except asyncio.TimeoutError:
+                ## TODO: should timeout be treated as an error/qubes.exc?
+                return
+            except (subprocess.CalledProcessError,qubes.exc.QubesException):
+                raise qubes.exc.QubesException(
+                    "Failed to run QUBESRPC qubes.WaitForSession"
+                )
+            finally:
+                if proc is not None:
+                    proc.terminate()
+                self.pause()
+
+    @qubes.events.handler("domain-unpaused")
+    async def on_domain_unpaused(self):
+        """Mark unpaused preloaded domains as used."""
+        if self.is_preloaded():
+            await self.use_preloaded()
+
     @qubes.events.handler("property-pre-reset:template")
     def on_property_pre_reset_template(self, event, name, oldvalue=None):
         """Forbid deleting template of VM"""  # pylint: disable=unused-argument
@@ -228,11 +390,12 @@ async def _auto_cleanup(self):
             self.app.save()
 
     @classmethod
-    async def from_appvm(cls, appvm, **kwargs):
+    async def from_appvm(cls, appvm, preload=False, **kwargs):
         """Create a new instance from given AppVM
 
         :param qubes.vm.appvm.AppVM appvm: template from which the VM should \
             be created
+        :param bool preload: Whether to preload a disposable
         :returns: new disposable vm
 
         *kwargs* are passed to the newly created VM
@@ -251,10 +414,32 @@ async def from_appvm(cls, appvm, **kwargs):
                 "template_for_dispvms=False"
             )
         app = appvm.app
+
+        if preload:
+            preload_dispvm_max = appvm.get_feat_preload("preload-dispvm-max")
+            if preload_dispvm_max == 0:
+                return
+            preload_dispvm = appvm.get_feat_preload("preload-dispvm")
+            if preload_dispvm and len(preload_dispvm) >= preload_dispvm_max:
+                raise qubes.exc.QubesException(
+                    "Failed to create preloaded disposable, limit of "
+                    "preloaded DispVMs reached"
+                )
+        else:
+            preload_dispvm = appvm.get_feat_preload("preload-dispvm")
+            if preload_dispvm:
+                dispvm = app.domains[preload_dispvm[0]]
+                await dispvm.use_preloaded()
+                return dispvm
+
         dispvm = app.add_new_vm(
             cls, template=appvm, auto_cleanup=True, **kwargs
         )
         await dispvm.create_on_disk()
+
+        if preload:
+            await dispvm.mark_preloaded()
+
         app.save()
         return dispvm
 

rpm_spec/core-dom0.spec.in

@@ -546,6 +546,7 @@ done
 %{python3_sitelib}/qubes/qmemman/domainstate.py
 %{python3_sitelib}/qubes/qmemman/systemstate.py
 
+/usr/lib/qubes/preload-dispvm
 /usr/lib/qubes/cleanup-dispvms
 /usr/lib/qubes/fix-dir-perms.sh
 /usr/lib/qubes/startup-misc.sh
@@ -556,6 +557,7 @@ done
 %{_unitdir}/qubes-qmemman.service
 %{_unitdir}/qubes-vm@.service
 %{_unitdir}/qubesd.service
+%{_unitdir}/qubes-preload-dispvm.service
 %attr(2770,root,qubes) %dir /var/lib/qubes
 %attr(2770,root,qubes) %dir /var/lib/qubes/vm-templates
 %attr(2770,root,qubes) %dir /var/lib/qubes/appvms