From 2d64eaeff8269d3e0c0a9de26a33d4a644a25747 Mon Sep 17 00:00:00 2001
From: Guillaume Chinal <guiiix@invisiblethingslab.com>
Date: Mon, 20 Jan 2025 21:19:31 +0100
Subject: [PATCH] add custom_persist extension

---
 doc/qubes-ext.rst           |   1 +
 qubes/ext/custom_persist.py | 102 ++++++++++++++++++++++++++++++++++++
 rpm_spec/core-dom0.spec.in  |   1 +
 setup.py                    |   1 +
 4 files changed, 105 insertions(+)
 create mode 100644 qubes/ext/custom_persist.py

diff --git a/doc/qubes-ext.rst b/doc/qubes-ext.rst
index 5eea811d5..6804edb55 100644
--- a/doc/qubes-ext.rst
+++ b/doc/qubes-ext.rst
@@ -14,6 +14,7 @@ Extensions defined here
 .. autoclass:: qubes.ext.admin.AdminExtension
 .. autoclass:: qubes.ext.block.BlockDeviceExtension
 .. autoclass:: qubes.ext.core_features.CoreFeatures
+.. autoclass:: qubes.ext.custom_persist.CustomPersist
 .. autoclass:: qubes.ext.gui.GUI
 .. autoclass:: qubes.ext.pci.PCIDeviceExtension
 .. autoclass:: qubes.ext.r3compatibility.R3Compatibility
diff --git a/qubes/ext/custom_persist.py b/qubes/ext/custom_persist.py
new file mode 100644
index 000000000..5f423cfc1
--- /dev/null
+++ b/qubes/ext/custom_persist.py
@@ -0,0 +1,102 @@
+# -*- encoding: utf-8 -*-
+#
+# The Qubes OS Project, http://www.qubes-os.org
+#
+# Copyright (C) 2024 Guillaume Chinal <guiiix@invisiblethingslab.com>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this library; if not, see <https://www.gnu.org/licenses/>.
+
+import qubes.ext
+import qubes.config
+
+FEATURE_PREFIX = "custom-persist."
+QDB_PREFIX = "/persist/"
+QDB_KEY_LIMIT = 63
+
+
+class CustomPersist(qubes.ext.Extension):
+    """This extension allows to create minimal-state APP with by configuring an
+    exhaustive list of bind dirs(and files)
+    """
+
+    @staticmethod
+    def _extract_key_from_feature(feature) -> str:
+        return feature[len(FEATURE_PREFIX) :]
+
+    @staticmethod
+    def _is_expected_feature(feature) -> bool:
+        return feature.startswith(FEATURE_PREFIX)
+
+    @staticmethod
+    def _is_valid_key(key, vm) -> bool:
+        if not key:
+            vm.log.warning(
+                "Got empty custom-persist key, ignoring: {}".format(key)
+            )
+            return False
+
+        # QubesDB key length limit
+        key_maxlen = QDB_KEY_LIMIT - len(QDB_PREFIX)
+        if len(key) >= key_maxlen:
+            vm.log.warning(
+                "custom-persist key is too long (max {}), ignoring: "
+                "{}".format(key_maxlen, key)
+            )
+            return False
+        return True
+
+    def _write_db_value(self, feature, value, vm):
+        vm.untrusted_qdb.write(
+            "{}{}".format(QDB_PREFIX, self._extract_key_from_feature(feature)),
+            str(value),
+        )
+
+    @qubes.ext.handler("domain-qdb-create")
+    def on_domain_qdb_create(self, vm, event):
+        """Actually export features"""
+        # pylint: disable=unused-argument
+        for feature, value in vm.features.items():
+            if self._is_expected_feature(feature) and self._is_valid_key(
+                self._extract_key_from_feature(feature), vm
+            ):
+                self._write_db_value(feature, value, vm)
+
+    @qubes.ext.handler("domain-feature-set:*")
+    def on_domain_feature_set(self, vm, event, feature, value, oldvalue=None):
+        """Inject persist keys in QubesDB in runtime"""
+        # pylint: disable=unused-argument
+
+        if not self._is_expected_feature(feature):
+            return
+
+        if not self._is_valid_key(self._extract_key_from_feature(feature), vm):
+            return
+
+        if not vm.is_running():
+            return
+
+        self._write_db_value(feature, value, vm)
+
+    @qubes.ext.handler("domain-feature-delete:*")
+    def on_domain_feature_delete(self, vm, event, feature):
+        """Update /vm-config/ QubesDB tree in runtime"""
+        # pylint: disable=unused-argument
+        if not vm.is_running():
+            return
+        if not feature.startswith(FEATURE_PREFIX):
+            return
+
+        vm.untrusted_qdb.rm(
+            "{}{}".format(QDB_PREFIX, self._extract_key_from_feature(feature))
+        )
diff --git a/rpm_spec/core-dom0.spec.in b/rpm_spec/core-dom0.spec.in
index 2ea3d2eb1..efa36e831 100644
--- a/rpm_spec/core-dom0.spec.in
+++ b/rpm_spec/core-dom0.spec.in
@@ -443,6 +443,7 @@ done
 %{python3_sitelib}/qubes/ext/backup_restore.py
 %{python3_sitelib}/qubes/ext/block.py
 %{python3_sitelib}/qubes/ext/core_features.py
+%{python3_sitelib}/qubes/ext/custom_persist.py
 %{python3_sitelib}/qubes/ext/gui.py
 %{python3_sitelib}/qubes/ext/audio.py
 %{python3_sitelib}/qubes/ext/pci.py
diff --git a/setup.py b/setup.py
index 69bace60f..bc2d4f697 100644
--- a/setup.py
+++ b/setup.py
@@ -64,6 +64,7 @@ def run(self):
                 'qubes.ext.backup_restore = '
                 'qubes.ext.backup_restore:BackupRestoreExtension',
                 'qubes.ext.core_features = qubes.ext.core_features:CoreFeatures',
+                'qubes.ext.custom_persist = qubes.ext.custom_persist:CustomPersist',
                 'qubes.ext.gui = qubes.ext.gui:GUI',
                 'qubes.ext.audio = qubes.ext.audio:AUDIO',
                 'qubes.ext.r3compatibility = qubes.ext.r3compatibility:R3Compatibility',