From 30c436ecde463bbd391ebb6e855604d7c5164fa4 Mon Sep 17 00:00:00 2001 From: catborise Date: Mon, 28 Sep 2026 14:55:24 +0300 Subject: [PATCH] fix(snapshot): preserve loader XML during snapshot operations --- vrtManager/instance.py | 52 +++++++++++++------- vrtManager/test_xml_safety.py | 90 ++++++++++++++++++++++++++++++++++- 2 files changed, 125 insertions(+), 17 deletions(-) diff --git a/vrtManager/instance.py b/vrtManager/instance.py index 2aa87e98..a4be0d3d 100644 --- a/vrtManager/instance.py +++ b/vrtManager/instance.py @@ -1,8 +1,11 @@ import contextlib import json +import logging import os.path import time +logger = logging.getLogger(__name__) + try: from libvirt import ( VIR_DOMAIN_AFFECT_CONFIG, @@ -1277,28 +1280,43 @@ class wvmInstance(wvmConnect): state, time.time(), ) - self.change_snapshot_xml() + changed = self.change_snapshot_xml() try: xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE) xml += """0 """ self._snapshotCreateXML(xml, 0) finally: - self.recover_snapshot_xml() + if changed: + self.recover_snapshot_xml() def change_snapshot_xml(self): - xml_temp = self._XMLDesc(VIR_DOMAIN_XML_SECURE).replace( - "", - "", - ) - self._defineXML(xml_temp) + try: + raw_xml = self._XMLDesc(VIR_DOMAIN_XML_SECURE) + tree = ElementTree.fromstring(raw_xml) + loader = tree.find("./os/loader") + if loader is not None and loader.get("type") == "pflash": + loader.set("type", "rom") + self._defineXML(ElementTree.tostring(tree, encoding="unicode")) + return True + return False + except Exception as e: + logger.error("Failed to change snapshot XML loader: %s", e) + raise def recover_snapshot_xml(self): - xml_temp = self._XMLDesc(VIR_DOMAIN_XML_SECURE).replace( - "", - "", - ) - self._defineXML(xml_temp) + try: + raw_xml = self._XMLDesc(VIR_DOMAIN_XML_SECURE) + tree = ElementTree.fromstring(raw_xml) + loader = tree.find("./os/loader") + if loader is not None and loader.get("type") == "rom": + loader.set("type", "pflash") + self._defineXML(ElementTree.tostring(tree, encoding="unicode")) + return True + return False + except Exception as e: + logger.error("Failed to recover snapshot XML loader: %s", e) + raise def create_external_snapshot(self, name, date=None, desc=None): creation_time = time.time() @@ -1318,7 +1336,7 @@ class wvmInstance(wvmConnect): creation_time, ) - self.change_snapshot_xml() + changed = self.change_snapshot_xml() try: xml += self._XMLDesc(VIR_DOMAIN_XML_SECURE) xml += """0 @@ -1326,7 +1344,8 @@ class wvmInstance(wvmConnect): self._snapshotCreateXML(xml, VIR_DOMAIN_SNAPSHOT_CREATE_DISK_ONLY) finally: - self.recover_snapshot_xml() + if changed: + self.recover_snapshot_xml() self.refresh_instance_pools() def get_external_snapshots(self): @@ -1402,12 +1421,13 @@ class wvmInstance(wvmConnect): snap.delete(0) def snapshot_revert(self, snapshot): - self.change_snapshot_xml() + changed = self.change_snapshot_xml() try: snap = self.instance.snapshotLookupByName(snapshot, 0) self.instance.revertToSnapshot(snap, 0) finally: - self.recover_snapshot_xml() + if changed: + self.recover_snapshot_xml() def get_managed_save_image(self): return self.instance.hasManagedSaveImage(0) diff --git a/vrtManager/test_xml_safety.py b/vrtManager/test_xml_safety.py index 70d9ed9a..70bef485 100644 --- a/vrtManager/test_xml_safety.py +++ b/vrtManager/test_xml_safety.py @@ -69,9 +69,28 @@ class TestXmlSafety(unittest.TestCase): inst.create_snapshot(bad_name, bad_desc) inst.change_snapshot_xml.assert_called_once() - # Verify recover_snapshot_xml was called in finally despite error inst.recover_snapshot_xml.assert_called_once() + def test_snapshot_xml_payload_is_escaped(self): + inst = wvmInstance.__new__(wvmInstance) + inst.get_status = MagicMock(return_value=5) # shutoff + inst.change_snapshot_xml = MagicMock() + inst.recover_snapshot_xml = MagicMock() + inst._XMLDesc = MagicMock(return_value="") + inst._snapshotCreateXML = MagicMock(return_value=None) + + bad_name = "snap"es" + bad_desc = "desc&and" + + inst.create_snapshot(bad_name, bad_desc) + + inst._snapshotCreateXML.assert_called_once() + created_xml = inst._snapshotCreateXML.call_args[0][0] + self.assertNotIn("", created_xml) + self.assertNotIn("", created_xml) + self.assertIn("<inject>&quotes", created_xml) + self.assertIn("<tag>&and", created_xml) + def test_external_snapshot_recovery(self): inst = wvmInstance.__new__(wvmInstance) inst.get_status = MagicMock(return_value=1) @@ -87,6 +106,75 @@ class TestXmlSafety(unittest.TestCase): inst.recover_snapshot_xml.assert_called_once() inst.refresh_instance_pools.assert_called_once() + def test_secure_boot_loader_change_and_recovery(self): + inst = wvmInstance.__new__(wvmInstance) + inst.wvm = MagicMock() + uefi_xml = "/usr/share/OVMF/OVMF.fd" + inst._XMLDesc = MagicMock(return_value=uefi_xml) + + inst.change_snapshot_xml() + inst.wvm.defineXML.assert_called_once() + changed_xml = inst.wvm.defineXML.call_args[0][0] + self.assertIn("type=\"rom\"", changed_xml) + self.assertIn("secure=\"yes\"", changed_xml) + + # Test recovery from rom back to pflash + inst._XMLDesc = MagicMock(return_value=changed_xml) + inst.wvm.defineXML.reset_mock() + inst.recover_snapshot_xml() + inst.wvm.defineXML.assert_called_once() + recovered_xml = inst.wvm.defineXML.call_args[0][0] + self.assertIn("type=\"pflash\"", recovered_xml) + self.assertIn("secure=\"yes\"", recovered_xml) + + def test_original_rom_loader_preserved(self): + inst = wvmInstance.__new__(wvmInstance) + inst._XMLDesc = MagicMock(return_value="/usr/share/rom.bin") + inst._defineXML = MagicMock() + inst._snapshotCreateXML = MagicMock(return_value=None) + inst.get_status = MagicMock(return_value=5) + + # change_snapshot_xml should return False and not define XML + self.assertFalse(inst.change_snapshot_xml()) + inst._defineXML.assert_not_called() + + # Create snapshot should not invoke recover_snapshot_xml if not changed + inst.change_snapshot_xml = MagicMock(return_value=False) + inst.recover_snapshot_xml = MagicMock() + inst.create_snapshot("snap_rom", "testing rom preservation") + inst.change_snapshot_xml.assert_called_once() + inst.recover_snapshot_xml.assert_not_called() + + def test_malformed_xml_raises_in_change(self): + inst = wvmInstance.__new__(wvmInstance) + inst._XMLDesc = MagicMock(return_value="/path") + inst._defineXML = MagicMock(side_effect=RuntimeError("libvirt defineXML failed")) + with self.assertRaises(RuntimeError): + inst.recover_snapshot_xml() + + + def test_snapshot_revert_rom_loader_preserved(self): + inst = wvmInstance.__new__(wvmInstance) + inst.instance = MagicMock() + mock_snap = MagicMock() + inst.instance.snapshotLookupByName = MagicMock(return_value=mock_snap) + inst.instance.revertToSnapshot = MagicMock() + inst.change_snapshot_xml = MagicMock(return_value=False) + inst.recover_snapshot_xml = MagicMock() + + inst.snapshot_revert("snap_rom") + + inst.change_snapshot_xml.assert_called_once() + inst.instance.snapshotLookupByName.assert_called_once_with("snap_rom", 0) + inst.instance.revertToSnapshot.assert_called_once_with(mock_snap, 0) + inst.recover_snapshot_xml.assert_not_called() + if __name__ == "__main__": unittest.main()