Merge ~waveform/ubuntu-release-upgrader:lp2065051-noble into ubuntu-release-upgrader:ubuntu/noble

[Dave Jones]

[Nick Rosbrook (enr0n)]

review: Approve

diff --git a/DistUpgrade/DistUpgradeQuirks.py b/DistUpgrade/DistUpgradeQuirks.py
index d3a0bbd..292394b 100644
--- a/DistUpgrade/DistUpgradeQuirks.py
+++ b/DistUpgrade/DistUpgradeQuirks.py
@@ -156,6 +156,9 @@ class DistUpgradeQuirks(object):
         if 'ubuntu-desktop-raspi' in cache:
             if cache['ubuntu-desktop-raspi'].is_installed:
                 self._replace_fkms_overlay()
+        if 'ubuntu-server-raspi' in cache:
+            if cache['ubuntu-server-raspi'].is_installed:
+                self._add_kms_overlay()
         if 'linux-firmware-raspi2' in cache:
             if cache['linux-firmware-raspi2'].is_installed:
                 self._remove_uboot_on_rpi()
@@ -1284,6 +1287,23 @@ class DistUpgradeQuirks(object):
                 self._snap_list[snap] = snap_object
         return self._snap_list
 
+    def _replace_pi_boot_config(self, old_config, new_config,
+                                boot_config_filename, failure_action):
+        try:
+            boot_backup_filename = boot_config_filename + '.distUpgrade'
+            with open(boot_backup_filename, 'w', encoding='utf-8') as f:
+                f.write(old_config)
+        except IOError as exc:
+            logging.error("unable to write boot config backup to %s: %s; %s",
+                          boot_backup_filename, exc, failure_action)
+            return
+        try:
+            with open(boot_config_filename, 'w', encoding='utf-8') as f:
+                f.write(new_config)
+        except IOError as exc:
+            logging.error("unable to write new boot config to %s: %s; %s",
+                          boot_config_filename, exc, failure_action)
+
     def _replace_fkms_overlay(self, boot_dir='/boot/firmware'):
         failure_action = (
             "You may need to replace the vc4-fkms-v3d overlay with "
@@ -1318,22 +1338,80 @@ class DistUpgradeQuirks(object):
             logging.warning("no fkms overlay or camera firmware line found "
                             "in %s", boot_config_filename)
             return
+        self._replace_pi_boot_config(
+            boot_config, new_config, boot_config_filename, failure_action)
 
+    def _add_kms_overlay(self, boot_dir='/boot/firmware'):
+        failure_action = (
+            "You may need to add dtoverlay=vc4-kms-v3d to an [all] section "
+            "in config.txt on your boot partition")
+        added_lines = [
+            '# added by do-release-upgrade (LP: #2065051)',
+            'dtoverlay=vc4-kms-v3d',
+            'disable_fw_kms_setup=1',
+            '',
+            '[pi3+]',
+            'dtoverlay=vc4-kms-v3d,cma-128',
+            '',
+            '[pi02]',
+            'dtoverlay=vc4-kms-v3d,cma-128',
+            '',
+            '[all]',
+        ]
         try:
-            boot_backup_filename = os.path.join(
-                boot_dir, 'config.txt.distUpgrade')
-            with open(boot_backup_filename, 'w', encoding='utf-8') as f:
-                f.write(boot_config)
-        except IOError as exc:
-            logging.error("unable to write boot config backup to %s: %s; %s",
-                          boot_backup_filename, exc, failure_action)
+            boot_config_filename = os.path.join(boot_dir, 'config.txt')
+            with open(boot_config_filename, 'r', encoding='utf-8') as f:
+                boot_config = f.read()
+        except FileNotFoundError:
+            logging.error("failed to open boot configuration in %s; %s",
+                          boot_config_filename, failure_action)
             return
-        try:
-            with open(boot_config_filename, 'w', encoding='utf-8') as f:
-                f.write(new_config)
-        except IOError as exc:
-            logging.error("unable to write new boot config to %s: %s; %s",
-                          boot_config_filename, exc, failure_action)
+
+        def find_insertion_point(lines):
+            # Returns the zero-based index of the dtoverlay=vc4-kms-v3d line in
+            # an [all] section, if one exists, or the last line of the last
+            # [all] section of the file, if one does not exist
+            in_all = True
+            last = 0
+            for index, line in enumerate(lines):
+                line = line.rstrip()
+                if in_all:
+                    last = index
+                    # startswith used to cope with any trailing dtparams
+                    if line.startswith('dtoverlay=vc4-kms-v3d'):
+                        return last
+                    elif line.startswith('[') and line.endswith(']'):
+                        in_all = line == '[all]'
+                    elif line.startswith('include '):
+                        # [sections] are included from includes verbatim, hence
+                        # (without reading the included file) we must assume
+                        # we're no longer in an [all] section
+                        in_all = False
+                else:
+                    in_all = line == '[all]'
+            return last
+
+        def add_kms_overlay(lines):
+            insert_point = find_insertion_point(lines)
+            try:
+                if lines[insert_point].startswith('dtoverlay=vc4-kms-v3d'):
+                    return lines
+            except IndexError:
+                # Empty config, apparently!
+                pass
+            lines[insert_point:insert_point] = added_lines
+            return lines
+
+        lines = [line.rstrip() for line in boot_config.splitlines()]
+        lines = add_kms_overlay(lines)
+        new_config = ''.join(line + '\n' for line in lines)
+
+        if new_config == boot_config:
+            logging.warning("no addition of KMS overlay required in %s",
+                            boot_config_filename)
+            return
+        self._replace_pi_boot_config(
+            boot_config, new_config, boot_config_filename, failure_action)
 
     def _remove_uboot_on_rpi(self, boot_dir='/boot/firmware'):
         kernel_line = 'kernel=vmlinuz'
@@ -1452,21 +1530,8 @@ class DistUpgradeQuirks(object):
             logging.warning("no u-boot removal performed in %s",
                             boot_config_filename)
             return
-
-        try:
-            boot_backup_filename = boot_config_filename + '.distUpgrade'
-            with open(boot_backup_filename, 'w', encoding='utf-8') as f:
-                f.write(boot_config)
-        except IOError as exc:
-            logging.error("unable to write boot config backup to %s: %s; %s",
-                          boot_backup_filename, exc, failure_action)
-            return
-        try:
-            with open(boot_config_filename, 'w', encoding='utf-8') as f:
-                f.write(new_config)
-        except IOError as exc:
-            logging.error("unable to write new boot config to %s: %s; %s",
-                          boot_config_filename, exc, failure_action)
+        self._replace_pi_boot_config(
+            boot_config, new_config, boot_config_filename, failure_action)
 
     def _set_generic_font(self):
         """ Due to changes to the Ubuntu font we enable a generic font
diff --git a/debian/changelog b/debian/changelog
index 862afb2..825a2d7 100644
--- a/debian/changelog
+++ b/debian/changelog
@@ -1,3 +1,10 @@
+ubuntu-release-upgrader (1:24.04.18) noble; urgency=medium
+
+  * DistUpgradeQuirks: add quirk to migrate Pi server users to the KMS
+    overlay (LP: #2065051)
+
+ -- Dave Jones <[e-mail]>  Tue, 07 May 2024 18:54:12 +0100
+
 ubuntu-release-upgrader (1:24.04.17) noble; urgency=medium
 
   [ Nick Rosbrook ]
diff --git a/tests/test_quirks.py b/tests/test_quirks.py
index f40d502..a3fabec 100644
--- a/tests/test_quirks.py
+++ b/tests/test_quirks.py
@@ -522,6 +522,85 @@ dtoverlay=vc4-kms-v3d,cma-256
             with open(os.path.join(boot_dir, 'config.txt')) as f:
                 self.assertTrue(f.read() == expected_config)
 
+    def test_add_kms_overlay_no_config(self):
+        with tempfile.TemporaryDirectory() as boot_dir:
+            mock_controller = mock.Mock()
+
+            q = DistUpgradeQuirks(mock_controller, mock.Mock())
+
+            q._add_kms_overlay(boot_dir)
+            self.assertFalse(os.path.exists(os.path.join(
+                boot_dir, 'config.txt.distUpgrade')))
+
+    def test_add_kms_overlay_no_changes(self):
+        with tempfile.TemporaryDirectory() as boot_dir:
+            boot_config = """\
+arm_64bit=1
+kernel=vmlinuz
+initramfs initrd.img followkernel
+
+# This line is implicitly in an [all] section and
+# should prevent the quirk from doing anything
+dtoverlay=vc4-kms-v3d,cma-128
+
+[pi4]
+max_framebuffers=2
+"""
+            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
+                f.write(boot_config)
+
+            mock_controller = mock.Mock()
+            q = DistUpgradeQuirks(mock_controller, mock.Mock())
+            q._add_kms_overlay(boot_dir)
+
+            self.assertFalse(os.path.exists(os.path.join(
+                boot_dir, 'config.txt.distUpgrade')))
+            with open(os.path.join(boot_dir, 'config.txt')) as f:
+                self.assertTrue(f.read() == boot_config)
+
+    def test_add_kms_overlay_with_changes(self):
+        with tempfile.TemporaryDirectory() as boot_dir:
+            config_txt = """\
+arm_64bit=1
+kernel=vmlinuz
+initramfs initrd.img followkernel
+
+[pi4]
+max_framebuffers=2
+"""
+            expected_config_txt = """\
+arm_64bit=1
+kernel=vmlinuz
+initramfs initrd.img followkernel
+
+# added by do-release-upgrade (LP: #2065051)
+dtoverlay=vc4-kms-v3d
+disable_fw_kms_setup=1
+
+[pi3+]
+dtoverlay=vc4-kms-v3d,cma-128
+
+[pi02]
+dtoverlay=vc4-kms-v3d,cma-128
+
+[all]
+[pi4]
+max_framebuffers=2
+"""
+            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
+                f.write(config_txt)
+
+            mock_controller = mock.Mock()
+            q = DistUpgradeQuirks(mock_controller, mock.Mock())
+            q._add_kms_overlay(boot_dir)
+
+            self.assertTrue(os.path.exists(os.path.join(
+                boot_dir, 'config.txt.distUpgrade')))
+            self.assertTrue(os.path.exists(os.path.join(
+                boot_dir, 'config.txt')))
+            with open(os.path.join(boot_dir, 'config.txt')) as f:
+                self.assertTrue(f.read() == expected_config_txt)
+
     def test_remove_uboot_no_config(self):
         with tempfile.TemporaryDirectory() as boot_dir:
             mock_controller = mock.Mock()