diff --git a/tests/test_bootloader.py b/tests/test_bootloader.py index c67dbfa2..afc08faf 100644 --- a/tests/test_bootloader.py +++ b/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): diff --git a/xcp/bootloader.py b/xcp/bootloader.py index 6bb17459..4454ba12 100644 --- a/xcp/bootloader.py +++ b/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)