Merge pull request #1761 from mwhudson/rp-hacky-fixes

several fixes around reset partition handling

subiquity/models/filesystem.py

@@ -1820,6 +1820,9 @@ def all_raids(self):
     def all_volgroups(self):
         return self._all(type="lvm_volgroup")
 
+    def partition_by_partuuid(self, partuuid: str) -> Optional[Partition]:
+        return self._one(type="partition", uuid=partuuid)
+
     def _remove(self, obj):
         _remove_backlinks(obj)
         self._actions.remove(obj)

subiquity/server/controllers/install.py

@@ -15,13 +15,16 @@
 
 import asyncio
 import copy
+import glob
 import json
 import logging
 import os
 import re
+import shlex
 import shutil
 import subprocess
 import tempfile
+import uuid
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
@@ -32,11 +35,11 @@
 from subiquity.common.errorreport import ErrorReportKind
 from subiquity.common.types import ApplicationState, PackageInstallState
 from subiquity.journald import journald_listen
-from subiquity.models.filesystem import ActionRenderMode
+from subiquity.models.filesystem import ActionRenderMode, Partition
 from subiquity.server.controller import SubiquityController
 from subiquity.server.curtin import run_curtin_command, start_curtin_command
 from subiquity.server.kernel import list_installed_kernels
-from subiquity.server.mounter import Mounter
+from subiquity.server.mounter import Mounter, Mountpoint
 from subiquity.server.types import InstallerChannels
 from subiquitycore.async_helpers import run_bg_task, run_in_thread
 from subiquitycore.context import with_context
@@ -421,35 +424,89 @@ async def run_curtin_step(name, stages, step_config, source=None):
                 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
+            new_casper_uuid = await self.adjust_rp(rp, mp)
+            await self.configure_rp_boot(
+                context=context, rp=rp, casper_uuid=new_casper_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.
+        else:
+            await self.maybe_configure_exiting_rp_boot(context=context)
+
+    async def adjust_rp(self, rp: Partition, mp: Mountpoint) -> str:
+        if self.app.opts.dry_run:
             return
+        # Once the installer has been copied to the RP, we need to make two
+        # adjustments:
+        #
+        # 1. set a new "casper uuid" so that booting from the install
+        #    media again, or booting from the RP but with the install
+        #    media still attached, does not get confused about which
+        #    device to use as /cdrom.
+        #
+        # 2. add "rp-partuuid" to the kernel command line in grub.cfg
+        #    so that subiquity can identify when it is running from
+        #    the recovery partition and add a reference to it to
+        #    grub.cfg on the target system in that case.
+        grub_cfg_path = mp.p("boot/grub/grub.cfg")
+        new_cfg = []
+        new_casper_uuid = str(uuid.uuid4())
+        cp = await self.app.command_runner.run(
+            ["lsblk", "-n", "-o", "PARTUUID", rp.path], capture=True
+        )
+        rp_uuid = cp.stdout.decode("ascii").strip()
+        with open(grub_cfg_path) as fp:
+            for line in fp:
+                words = shlex.split(line)
+                if words and words[0] == "linux" and "---" in words:
+                    index = words.index("---")
+                    words[index - 1 : index - 1] = [
+                        "uuid=" + new_casper_uuid,
+                        "rp-partuuid=" + rp_uuid,
+                    ]
+                    new_cfg.append(shlex.join(words) + "\n")
+                else:
+                    new_cfg.append(line)
+        with open(grub_cfg_path, "w") as fp:
+            fp.write("".join(new_cfg))
+        for casper_uuid_file in glob.glob(mp.p(".disk/casper-uuid-*")):
+            with open(casper_uuid_file, "w") as fp:
+                fp.write(new_casper_uuid + "\n")
+        return new_casper_uuid
+
+    @with_context(description="configuring grub menu entry for factory reset")
+    async def configure_rp_boot_grub(self, context, rp: Partition, casper_uuid: str):
+        # Add a grub menu entry to boot from the RP
+        cp = await self.app.command_runner.run(
+            ["lsblk", "-n", "-o", "UUID,PARTUUID", rp.path], capture=True
+        )
+        fs_uuid, rp_uuid = cp.stdout.decode("ascii").strip().split()
+        conf = grub_reset_conf.format(
+            HEADER=generate_timestamped_header(),
+            PARTITION=rp.number,
+            FS_UUID=fs_uuid,
+            CASPER_UUID=casper_uuid,
+            RP_UUID=rp_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,
+        )
+
+    @with_context(description="configuring UEFI menu entry for factory reset")
+    async def configure_rp_boot_uefi(self, context, rp: Partition):
+        # Add an UEFI boot entry to point at the RP
+        # Details:
+        # 1. Do not leave duplicate entries
+        # 2. Do not leave the boot entry in BootOrder
+        # 3. Set BootNext to boot from RP in the reset-partition-only case.
         state = await self.app.package_installer.install_pkg("efibootmgr")
         if state != PackageInstallState.DONE:
             raise RuntimeError("could not install efibootmgr")
@@ -458,7 +515,7 @@ async def create_rp_boot_entry(self, context, rp):
             "efibootmgr",
             "--create",
             "--loader",
-            "\\EFI\\boot\\shimx64.efi",
+            "\\EFI\\boot\\bootx64.efi",
             "--disk",
             rp.device.path,
             "--part",
@@ -470,6 +527,7 @@ async def create_rp_boot_entry(self, context, rp):
         efi_state_after = get_efibootmgr("/")
         new_bootnums = set(efi_state_after.entries) - set(efi_state_before.entries)
         if not new_bootnums:
+            # Will probably only happen in dry-run mode.
             return
         new_bootnum = new_bootnums.pop()
         new_entry = efi_state_after.entries[new_bootnum]
@@ -489,18 +547,51 @@ async def create_rp_boot_entry(self, context, rp):
             cmd = [
                 "efibootmgr",
                 "--bootorder",
-                ",".join(efi_state_before.order),
+                ",".join(efi_state_before.order + [new_bootnum]),
             ]
             rp_bootnum = new_bootnum
         await self.app.command_runner.run(cmd)
-        if not fs_controller.reset_partition_only:
+        if self.model.target is None:
             cmd = [
                 "efibootmgr",
                 "--bootnext",
                 rp_bootnum,
             ]
             await self.app.command_runner.run(cmd)
 
+    async def configure_rp_boot(self, context, rp: Partition, casper_uuid: str):
+        if self.model.target is not None and not self.opts.dry_run:
+            await self.configure_rp_boot_grub(
+                context=context, rp=rp, casper_uuid=casper_uuid
+            )
+        if self.app.opts.dry_run and not is_uefi_bootable():
+            # Can't even run efibootmgr in this case.
+            return
+        await self.configure_rp_boot_uefi(context=context, rp=rp)
+
+    async def maybe_configure_exiting_rp_boot(self, context):
+        # We are not creating a reset partition here if we are running
+        # from one we still want to configure booting from it.
+
+        # Look for the command line argument added in adjust_rp)
+        # above.
+        rp_partuuid = self.app.kernel_cmdline.get("rp-partuuid")
+        if rp_partuuid is None:
+            # Most likely case: we are not running from an reset partition
+            return
+        rp = self.app.base_model.filesystem.partition_by_partuuid(rp_partuuid)
+        if rp is None:
+            # This shouldn't happen, but don't crash.
+            return
+        casper_uuid = None
+        for casper_uuid_file in glob.glob("/cdrom/.disk/casper-uuid-*"):
+            with open(casper_uuid_file) as fp:
+                casper_uuid = fp.read().strip()
+        if casper_uuid is None:
+            # This also shouldn't happen, but, again, don't crash.
+            return
+        await self.configure_rp_boot(context=context, rp=rp, casper_uuid=casper_uuid)
+
     @with_context(description="creating fstab")
     async def create_core_boot_classic_fstab(self, *, context):
         with open(self.tpath("etc/fstab"), "w") as fp:
@@ -739,10 +830,10 @@ async def stop_unattended_upgrades(self):
 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
+menuentry "Restore Ubuntu to factory state" {{
+      search --no-floppy --hint '(hd0,{PARTITION})' --set --fs-uuid {FS_UUID}
+      linux  /casper/vmlinuz uuid={CASPER_UUID} rp-partuuid={RP_UUID} nopersistent
       initrd /casper/initrd
-}
+}}
 EOF
 """

subiquity/server/controllers/tests/test_install.py

@@ -13,7 +13,10 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
+import shutil
 import subprocess
+import tempfile
 import unittest
 from pathlib import Path
 from unittest.mock import ANY, AsyncMock, Mock, call, mock_open, patch
@@ -198,7 +201,9 @@ def setUp(self):
         self.controller = InstallController(make_app())
         self.controller.app.report_start_event = Mock()
         self.controller.app.report_finish_event = Mock()
-        self.controller.model.target = "/target"
+        self.controller.model.target = tempfile.mkdtemp()
+        os.makedirs(os.path.join(self.controller.model.target, "etc/grub.d"))
+        self.addCleanup(shutil.rmtree, self.controller.model.target)
 
     @patch("asyncio.sleep")
     async def test_install_package(self, m_sleep):
@@ -221,30 +226,32 @@ async def test_install_package(self, m_sleep):
             with self.assertRaises(subprocess.CalledProcessError):
                 await self.controller.install_package(package="git")
 
-    def setup_rp_test(self):
+    def setup_rp_test(self, lsblk_output=b"lsblk_output"):
         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.command_runner = AsyncMock()
+        self.run = app.command_runner.run = AsyncMock(
+            return_value=subprocess.CompletedProcess((), 0, stdout=lsblk_output)
+        )
         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):
+    async def test_configure_rp_boot_uefi_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)
+        await self.controller.configure_rp_boot_uefi(rp=self.part)
         calls = [
             call(
                 [
                     "efibootmgr",
                     "--create",
                     "--loader",
-                    "\\EFI\\boot\\shimx64.efi",
+                    "\\EFI\\boot\\bootx64.efi",
                     "--disk",
                     self.part.device.path,
                     "--part",
@@ -257,24 +264,62 @@ async def test_create_rp_boot_entry_add(self, m_get_efibootmgr):
                 [
                     "efibootmgr",
                     "--bootorder",
-                    "0000,0002",
+                    "0000,0002,0003",
                 ]
             ),
         ]
         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):
+    async def test_configure_rp_boot_uefi_bootnext(self, m_get_efibootmgr):
+        m_get_efibootmgr.side_effect = iter([efi_state_no_rp, efi_state_with_rp])
+        self.setup_rp_test()
+        self.controller.app.base_model.target = None
+        await self.controller.configure_rp_boot_uefi(rp=self.part)
+        calls = [
+            call(
+                [
+                    "efibootmgr",
+                    "--create",
+                    "--loader",
+                    "\\EFI\\boot\\bootx64.efi",
+                    "--disk",
+                    self.part.device.path,
+                    "--part",
+                    str(self.part.number),
+                    "--label",
+                    "Restore Ubuntu to factory state",
+                ]
+            ),
+            call(
+                [
+                    "efibootmgr",
+                    "--bootorder",
+                    "0000,0002,0003",
+                ]
+            ),
+            call(
+                [
+                    "efibootmgr",
+                    "--bootnext",
+                    "0003",
+                ]
+            ),
+        ]
+        self.run.assert_has_awaits(calls)
+
+    @patch("subiquity.server.controllers.install.get_efibootmgr")
+    async def test_configure_rp_boot_uefi_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)
+        await self.controller.configure_rp_boot_uefi(rp=self.part)
         calls = [
             call(
                 [
                     "efibootmgr",
                     "--create",
                     "--loader",
-                    "\\EFI\\boot\\shimx64.efi",
+                    "\\EFI\\boot\\bootx64.efi",
                     "--disk",
                     self.part.device.path,
                     "--part",
@@ -293,3 +338,15 @@ async def test_create_rp_boot_entry_dup(self, m_get_efibootmgr):
             ),
         ]
         self.run.assert_has_awaits(calls)
+
+    async def test_configure_rp_boot_grub(self):
+        fsuuid, partuuid = "fsuuid", "partuuid"
+        self.setup_rp_test(f"{fsuuid}\t{partuuid}".encode("ascii"))
+        await self.controller.configure_rp_boot_grub(
+            rp=self.part, casper_uuid="casper-uuid"
+        )
+        with open(self.controller.tpath("etc/grub.d/99_reset")) as fp:
+            cfg = fp.read()
+        self.assertIn("--fs-uuid fsuuid", cfg)
+        self.assertIn("rp-partuuid=partuuid", cfg)
+        self.assertIn("uuid=casper-uuid", cfg)