Merge pull request #1720 from mwhudson/configure-rp-boot

configure reset partition boot

apt-deps.txt

@@ -2,6 +2,7 @@ build-essential
 cloud-init
 curl
 dctrl-tools
+efibootmgr
 fuseiso
 gettext
 gir1.2-umockdev-1.0

subiquity/server/controllers/install.py

@@ -26,6 +26,10 @@ import tempfile
 from typing import Any, Dict, List, Optional
 
 from curtin.config import merge_config
+from curtin.util import (
+    get_efibootmgr,
+    is_uefi_bootable,
+    )
 import yaml
 
 from subiquitycore.async_helpers import (
@@ -33,12 +37,17 @@ from subiquitycore.async_helpers import (
     run_in_thread,
     )
 from subiquitycore.context import with_context
-from subiquitycore.file_util import write_file, generate_config_yaml
+from subiquitycore.file_util import (
+    write_file,
+    generate_config_yaml,
+    generate_timestamped_header,
+    )
 from subiquitycore.utils import arun_command, log_process_streams
 
 from subiquity.common.errorreport import ErrorReportKind
 from subiquity.common.types import (
     ApplicationState,
+    PackageInstallState,
     )
 from subiquity.journald import (
     journald_listen,
@@ -380,6 +389,63 @@ class InstallController(SubiquityController):
                     step_config=self.rp_config(logs_dir, mp.p()),
                     source='cp:///cdrom',
                     )
+            await self.create_rp_boot_entry(context=context, rp=rp)
+
+    @with_context(description="creating boot entry for reset partition")
+    async def create_rp_boot_entry(self, context, rp):
+        fs_controller = self.app.controllers.Filesystem
+        if not fs_controller.reset_partition_only:
+            cp = await self.app.command_runner.run(
+                ['lsblk', '-n', '-o', 'UUID', rp.path],
+                capture=True)
+            uuid = cp.stdout.decode('ascii').strip()
+            conf = grub_reset_conf.format(
+                HEADER=generate_timestamped_header(),
+                PARTITION=rp.number,
+                UUID=uuid)
+            with open(self.tpath('etc/grub.d/99_reset'), 'w') as fp:
+                os.chmod(fp.fileno(), 0o755)
+                fp.write(conf)
+            await run_curtin_command(
+                self.app, context, "in-target", "-t", self.tpath(), "--",
+                "update-grub", private_mounts=False)
+        if self.app.opts.dry_run and not is_uefi_bootable():
+            # Can't even run efibootmgr in this case.
+            return
+        state = await self.app.package_installer.install_pkg('efibootmgr')
+        if state != PackageInstallState.DONE:
+            raise RuntimeError('could not install efibootmgr')
+        efi_state_before = get_efibootmgr('/')
+        cmd = [
+            'efibootmgr', '--create',
+            '--loader', '\\EFI\\boot\\shimx64.efi',
+            '--disk', rp.device.path,
+            '--part', str(rp.number),
+            '--label', "Restore Ubuntu to factory state",
+            ]
+        await self.app.command_runner.run(cmd)
+        efi_state_after = get_efibootmgr('/')
+        new_bootnums = (
+            set(efi_state_after.entries) - set(efi_state_before.entries))
+        if not new_bootnums:
+            return
+        new_bootnum = new_bootnums.pop()
+        new_entry = efi_state_after.entries[new_bootnum]
+        was_dup = False
+        for entry in efi_state_before.entries.values():
+            if entry.path == new_entry.path and entry.name == new_entry.name:
+                was_dup = True
+        if was_dup:
+            cmd = [
+                'efibootmgr', '--delete-bootnum',
+                '--bootnum', new_bootnum,
+                ]
+        else:
+            cmd = [
+                'efibootmgr',
+                '--bootorder', ','.join(efi_state_before.order),
+                ]
+        await self.app.command_runner.run(cmd)
 
     @with_context(description="creating fstab")
     async def create_core_boot_classic_fstab(self, *, context):
@@ -417,6 +483,9 @@ class InstallController(SubiquityController):
             else:
                 for_install_path = ''
 
+            if self.app.controllers.Filesystem.reset_partition:
+                self.app.package_installer.start_installing_pkg('efibootmgr')
+
             await self.curtin_install(
                 context=context, source='cp://' + for_install_path)
 
@@ -590,3 +659,18 @@ Unattended-Upgrade::Allowed-Origins {
         "${distro_id}ESM:${distro_codename}-infra-security";
 };
 """
+
+grub_reset_conf = """\
+#!/bin/sh
+{HEADER}
+
+set -e
+
+cat << EOF
+menuentry "Restore Ubuntu to factory state" {
+      search --no-floppy --hint '(hd0,{PARTITION})' --set --fs-uuid {UUID}
+      linux  /casper/vmlinuz uuid={UUID} nopersistent
+      initrd /casper/initrd
+}
+EOF
+"""

subiquity/server/controllers/tests/test_install.py

@@ -16,8 +16,19 @@
 from pathlib import Path
 import subprocess
 import unittest
-from unittest.mock import ANY, Mock, mock_open, patch
+from unittest.mock import (
+    ANY,
+    AsyncMock,
+    call,
+    Mock,
+    mock_open,
+    patch,
+    )
 
+from curtin.util import EFIBootEntry, EFIBootState
+
+from subiquity.common.types import PackageInstallState
+from subiquity.models.tests.test_filesystem import make_model_and_partition
 from subiquity.server.controllers.install import (
     InstallController,
     )
@@ -115,6 +126,64 @@ class TestWriteConfig(unittest.IsolatedAsyncioTestCase):
             })
 
 
+efi_state_no_rp = EFIBootState(
+    current='0000',
+    timeout='0 seconds',
+    order=['0000', '0002'],
+    entries={
+        '0000': EFIBootEntry(
+            name='ubuntu',
+            path='HD(1,GPT,...)/File(\\EFI\\ubuntu\\shimx64.efi)'),
+        '0001': EFIBootEntry(
+            name='Windows Boot Manager',
+            path='HD(1,GPT,...,0x82000)/File(\\EFI\\bootmgfw.efi'),
+        '0002': EFIBootEntry(
+            name='Linux-Firmware-Updater',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)\\.fwupd'),
+        })
+
+efi_state_with_rp = EFIBootState(
+    current='0000',
+    timeout='0 seconds',
+    order=['0000', '0002', '0003'],
+    entries={
+        '0000': EFIBootEntry(
+            name='ubuntu',
+            path='HD(1,GPT,...)/File(\\EFI\\ubuntu\\shimx64.efi)'),
+        '0001': EFIBootEntry(
+            name='Windows Boot Manager',
+            path='HD(1,GPT,...,0x82000)/File(\\EFI\\bootmgfw.efi'),
+        '0002': EFIBootEntry(
+            name='Linux-Firmware-Updater',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)\\.fwupd'),
+        '0003': EFIBootEntry(
+            name='Restore Ubuntu to factory state',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)'),
+        })
+
+efi_state_with_dup_rp = EFIBootState(
+    current='0000',
+    timeout='0 seconds',
+    order=['0000', '0002', '0004'],
+    entries={
+        '0000': EFIBootEntry(
+            name='ubuntu',
+            path='HD(1,GPT,...)/File(\\EFI\\ubuntu\\shimx64.efi)'),
+        '0001': EFIBootEntry(
+            name='Windows Boot Manager',
+            path='HD(1,GPT,...,0x82000)/File(\\EFI\\bootmgfw.efi'),
+        '0002': EFIBootEntry(
+            name='Linux-Firmware-Updater',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)\\.fwupd'),
+        '0003': EFIBootEntry(
+            name='Restore Ubuntu to factory state',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)'),
+        '0004': EFIBootEntry(
+            name='Restore Ubuntu to factory state',
+            path='HD(1,GPT,...,0x800,0x100000)/File(\\shimx64.efi)'),
+        })
+
+
 class TestInstallController(unittest.IsolatedAsyncioTestCase):
     def setUp(self):
         self.controller = InstallController(make_app())
@@ -141,3 +210,56 @@ class TestInstallController(unittest.IsolatedAsyncioTestCase):
         with patch(run_curtin, side_effect=(error, error, error, error)):
             with self.assertRaises(subprocess.CalledProcessError):
                 await self.controller.install_package(package="git")
+
+    def setup_rp_test(self):
+        app = self.controller.app
+        app.opts.dry_run = False
+        fsc = app.controllers.Filesystem
+        fsc.reset_partition_only = True
+        app.package_installer = Mock()
+        app.command_runner = Mock()
+        self.run = app.command_runner.run = AsyncMock()
+        app.package_installer.install_pkg = AsyncMock()
+        app.package_installer.install_pkg.return_value = \
+            PackageInstallState.DONE
+        fsm, self.part = make_model_and_partition()
+
+    @patch("subiquity.server.controllers.install.get_efibootmgr")
+    async def test_create_rp_boot_entry_add(self, m_get_efibootmgr):
+        m_get_efibootmgr.side_effect = iter([
+            efi_state_no_rp, efi_state_with_rp])
+        self.setup_rp_test()
+        await self.controller.create_rp_boot_entry(rp=self.part)
+        calls = [
+            call([
+                'efibootmgr', '--create',
+                '--loader', '\\EFI\\boot\\shimx64.efi',
+                '--disk', self.part.device.path,
+                '--part', str(self.part.number),
+                '--label', "Restore Ubuntu to factory state",
+                ]),
+            call([
+                'efibootmgr', '--bootorder', '0000,0002',
+                ]),
+            ]
+        self.run.assert_has_awaits(calls)
+
+    @patch("subiquity.server.controllers.install.get_efibootmgr")
+    async def test_create_rp_boot_entry_dup(self, m_get_efibootmgr):
+        m_get_efibootmgr.side_effect = iter([
+            efi_state_with_rp, efi_state_with_dup_rp])
+        self.setup_rp_test()
+        await self.controller.create_rp_boot_entry(rp=self.part)
+        calls = [
+            call([
+                'efibootmgr', '--create',
+                '--loader', '\\EFI\\boot\\shimx64.efi',
+                '--disk', self.part.device.path,
+                '--part', str(self.part.number),
+                '--label', "Restore Ubuntu to factory state",
+                ]),
+            call([
+                'efibootmgr', '--delete-bootnum', '--bootnum', '0004',
+                ]),
+            ]
+        self.run.assert_has_awaits(calls)