bootloader: stop MenuEntry.contents getting clobbered by setNextBoot

tests/test_bootloader.py

@@ -52,12 +52,35 @@ def test_no_hypervisor(self):
         with self.assertRaises(RuntimeError):
             Bootloader.readGrub2("tests/data/grub-no-hypervisor.cfg")
 
+    def test_set_grub_variable(self):
+        tmpdir = mkdtemp(prefix="testbl")
+        env = os.path.join(tmpdir, 'grubenv')
+        bl = Bootloader("", "", env_block=env)
+
+        self.assertFalse(os.path.isfile(env))
+
+        self.assertTrue(bl.setGrubVariable("waffles=true"))
+
+        self.assertTrue(os.path.isfile(env))
+        self.assertGreater(os.path.getsize(env), 0)
+
+    def test_set_variable_no_envfile(self):
+        """
+        Test that calling setGrubVariable() without setting an envfile first
+        will throw an exception.
+        """
+        bl = Bootloader("", "", env_block=None)
+
+        with self.assertRaises(AssertionError):
+            bl.setGrubVariable("waffles=true")
+
 
 class TestMenuEntry(unittest.TestCase):
     def setUp(self):
         self.tmpdir = mkdtemp(prefix="testbl")
         self.fn = os.path.join(self.tmpdir, 'grub.cfg')
         self.bl = Bootloader('grub2', self.fn)
+        self.env = os.path.join(self.tmpdir, 'grubenv')
 
     def tearDown(self):
         shutil.rmtree(self.tmpdir)
@@ -137,6 +160,136 @@ def test_new_linux(self):
 }
 ''')
 
+    def test_arbitrary_contents(self):
+        """ Test that arbitrary data can be injected into the MenuEntry.contents field. """
+        e = MenuEntry(hypervisor='xen.efi', hypervisor_args='a',
+                      kernel='vmlinuz', kernel_args='b',
+                      initrd='initrd.img',
+                      title='menu_name')
+
+        e.contents.append("\textra data line 1")
+        e.contents.append("\textra data line 2")
+
+        e.entry_format = Grub2Format.XEN_BOOT
+
+        self.bl.append('menu_name', e)
+        self.bl.commit()
+
+        with open_with_codec_handling(self.fn, 'r') as f:
+            content = f.read()
+
+        self.assertEqual(content,
+'''menuentry 'menu_name' {
+	extra data line 1
+	extra data line 2
+	xen_hypervisor xen.efi a
+	xen_module vmlinuz b
+	xen_module initrd.img
+}
+''')
+
+    def test_chainloader(self):
+        e = MenuEntry(hypervisor='xen.efi', hypervisor_args='a',
+                      kernel='vmlinuz', kernel_args='b',
+                      initrd='initrd.img',
+                      title='menu_name')
+
+        e.contents.append("\textra data line 1")
+        e.entry_format = Grub2Format.XEN_BOOT
+        e.setRpuChainloader("/EFI/installer/shimx64.efi", "GUARD_VAR", "ESP_LABEL")
+
+        self.bl.append('menu_name', e)
+        self.bl.commit()
+
+        with open_with_codec_handling(self.fn, 'r') as f:
+            content = f.read()
+
+        self.assertEqual(content,
+'''menuentry 'menu_name' {
+	if [ "${GUARD_VAR}" = "1" ]; then
+		unset GUARD_VAR
+		save_env GUARD_VAR
+		search --label --set root ESP_LABEL
+		chainloader /EFI/installer/shimx64.efi
+	else
+	extra data line 1
+	xen_hypervisor xen.efi a
+	xen_module vmlinuz b
+	xen_module initrd.img
+	fi
+}
+''')
+
+    def test_contents_not_clobbered(self):
+        """
+        Test that MenuEntry.contents is not clobbered by setNextBoot
+        """
+
+        self.assertIsNone(self.bl.env_block)
+        self.bl.env_block = self.env
+
+        e = MenuEntry(hypervisor='xen.efi', hypervisor_args='a',
+                      kernel='vmlinuz', kernel_args='b',
+                      initrd='initrd.img',
+                      title='menu_title')
+
+        e.contents.append("\textra data line 1")
+        e.contents.append("\textra data line 2")
+
+        e.entry_format = Grub2Format.XEN_BOOT
+
+        self.bl.append('menu_title', e)
+        self.assertTrue(self.bl.setNextBoot('menu_title'))
+        self.bl.commit()
+
+
+        with open_with_codec_handling(self.fn, 'r') as f:
+            content = f.read()
+
+        self.assertEqual(content,
+'''menuentry 'menu_title' {
+	unset override_entry
+	save_env override_entry
+	extra data line 1
+	extra data line 2
+	xen_hypervisor xen.efi a
+	xen_module vmlinuz b
+	xen_module initrd.img
+}
+''')
+
+    def test_setnextboot_is_indempotent(self):
+        self.bl.env_block = self.env
+
+        e = MenuEntry(hypervisor='xen.efi', hypervisor_args='a',
+                      kernel='vmlinuz', kernel_args='b',
+                      initrd='initrd.img',
+                      title='menu_title')
+
+        e.entry_format = Grub2Format.XEN_BOOT
+
+        self.bl.append('menu_title', e)
+
+        # Calling twice should have thte same effect as calling once
+        self.assertTrue(self.bl.setNextBoot('menu_title'))
+        self.assertTrue(self.bl.setNextBoot('menu_title'))
+
+        self.bl.commit()
+
+
+        with open_with_codec_handling(self.fn, 'r') as f:
+            content = f.read()
+
+        self.assertEqual(content,
+'''menuentry 'menu_title' {
+	unset override_entry
+	save_env override_entry
+	xen_hypervisor xen.efi a
+	xen_module vmlinuz b
+	xen_module initrd.img
+}
+''')
+
 
 class TestLinuxBootloader(unittest.TestCase):
     def setUp(self):

xcp/bootloader.py

@@ -40,7 +40,7 @@
 
 from .compat import open_textfile
 
-COUNTER = 0
+_counter = 0
 
 class Grub2Format(Enum):
     MULTIBOOT2 = 0
@@ -61,6 +61,14 @@ def __init__(self, hypervisor, hypervisor_args, kernel, kernel_args,
         self.title = title
         self.root = root
         self.entry_format = None  # type: Grub2Format | None
+        self.chainloader = None
+        self.guard_var = None
+        self.esp_label = None
+
+    def setRpuChainloader(self, chainloader, guard_var, esp_label):
+        self.chainloader = chainloader
+        self.guard_var = guard_var
+        self.esp_label = esp_label
 
     def getHypervisorArgs(self):
         return re.findall(r'\S[^ "]*(?:"[^"]*")?\S*', self.hypervisor_args)
@@ -127,7 +135,7 @@ def readGrub2(cls, src_file):
         entry_format = Grub2Format.MULTIBOOT2
 
         def create_label(title):
-            global COUNTER
+            global _counter
 
             if title == branding.PRODUCT_BRAND:
                 return 'xe'
@@ -142,8 +150,8 @@ def create_label(title):
                     return 'fallback-serial'
                 else:
                     return 'fallback'
-            COUNTER += 1
-            return "label%d" % COUNTER
+            _counter += 1
+            return "label%d" % _counter
 
         def parse_boot_entry(line):
             parts = line.split(None, 2)  # Split into at most 3 parts
@@ -318,6 +326,14 @@ def writeGrub2(self, dst_file = None):
             extra = m.extra if m.extra else ' '
             print("menuentry '%s'%s{" % (m.title, extra), file=fh)
 
+            if m.chainloader and m.guard_var:
+                print(f"\tif [ \"${{{m.guard_var}}}\" = \"1\" ]; then",  file=fh)
+                print(f"\t\tunset {m.guard_var}", file=fh)
+                print(f"\t\tsave_env {m.guard_var}", file=fh)
+                print(f"\t\tsearch --label --set root {m.esp_label}", file=fh)
+                print(f"\t\tchainloader {m.chainloader}", file=fh)
+                print("\telse", file=fh)
+
             try:
                 contents = "\n".join(m.contents)
                 if contents:
@@ -349,6 +365,9 @@ def writeGrub2(self, dst_file = None):
             else:
                 raise AssertionError("Unreachable")
 
+            if m.chainloader and m.guard_var:
+                print("\tfi", file=fh)
+
             print("}", file=fh)
         if not hasattr(dst_file, 'name'):
             fh.close()
@@ -374,15 +393,21 @@ def setNextBoot(self, entry):
             return False
 
         clear_default = ['\tunset override_entry', '\tsave_env override_entry']
-        self.menu[entry].contents = clear_default
+        if clear_default[0] not in self.menu[entry].contents:
+            self.menu[entry].contents = clear_default + self.menu[entry].contents
 
         for i in range(len(self.menu_order)):
             if self.menu_order[i] == entry:
-                cmd = ['grub-editenv', self.env_block, 'set', 'override_entry=%d' % i]
-                return xcp.cmd.runCmd(cmd) == 0
+                return self.setGrubVariable('override_entry=%d' % i)
 
         return False
 
+    def setGrubVariable(self, var):
+        if self.env_block is None:
+            raise AssertionError("No grubenv file")
+        cmd = ['grub-editenv', self.env_block, 'set', var]
+        return xcp.cmd.runCmd(cmd) == 0
+
     @classmethod
     def newDefault(cls, kernel_link_name, initrd_link_name, root = '/'):
         b = cls.loadExisting(root)