From 6e2ecfc497fd2e77962aeed89d6aa3a6f900c182 Mon Sep 17 00:00:00 2001
From: Dan Bungert <daniel.bungert@canonical.com>
Date: Tue, 19 Jul 2022 11:56:04 -0600
Subject: [PATCH] filesystem: only use logical part renum if forced

---
 subiquity/models/filesystem.py            | 33 +++++++++++++----------
 subiquity/models/tests/test_filesystem.py | 26 ++++++++++++++++++
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/subiquity/models/filesystem.py b/subiquity/models/filesystem.py
index f66097ec..500ad774 100644
--- a/subiquity/models/filesystem.py
+++ b/subiquity/models/filesystem.py
@@ -519,6 +519,9 @@ class _Device(_Formattable, ABC):
     def partitions_by_offset(self):
         return sorted(self._partitions, key=lambda p: p.offset)
 
+    def partitions_by_number(self):
+        return sorted(self._partitions, key=lambda p: p.number)
+
     @property
     def used(self):
         if self._is_entirely_used():
@@ -671,10 +674,10 @@ class Disk(_Device):
             return None
         return id.encode('utf-8').decode('unicode_escape').strip()
 
-    def renumber_logical_partitions(self):
-        parts = [p for p in self.partitions_by_offset() if p.is_logical]
-        primary_limit = self.alignment_data().primary_part_limit
-        next_num = primary_limit + 1
+    def renumber_logical_partitions(self, removed_partition):
+        parts = [p for p in self.partitions_by_number()
+                 if p.is_logical and p.number > removed_partition.number]
+        next_num = removed_partition.number
         for part in parts:
             part.number = next_num
             next_num += 1
@@ -700,17 +703,19 @@ class Partition(_Formattable):
         if self.number is not None:
             return
 
+        used_nums = {p.number for p in self.device._partitions
+                     if p.number is not None
+                     if p.is_logical == self.is_logical}
+        primary_limit = self.device.alignment_data().primary_part_limit
         if self.is_logical:
-            self.device.renumber_logical_partitions()
+            possible_nums = range(primary_limit + 1, 129)
         else:
-            used_nums = {p.number for p in self.device._partitions
-                         if p.number is not None and not p.is_logical}
-            primary_limit = self.device.alignment_data().primary_part_limit
-            for num in range(1, primary_limit + 1):
-                if num not in used_nums:
-                    self.number = num
-                    return
-            raise Exception('Exceeded number of available primary partitions')
+            possible_nums = range(1, primary_limit + 1)
+        for num in possible_nums:
+            if num not in used_nums:
+                self.number = num
+                return
+        raise Exception('Exceeded number of available partitions')
 
     def available(self):
         if self.flag in ['bios_grub', 'prep'] or self.grub_device:
@@ -1466,7 +1471,7 @@ class FilesystemModel(object):
         for p2 in movable_trailing_partitions_and_gap_size(part)[0]:
             p2.offset -= part.size
         self._remove(part)
-        part.device.renumber_logical_partitions()
+        part.device.renumber_logical_partitions(part)
         if len(part.device._partitions) == 0:
             part.device.ptable = None
 
diff --git a/subiquity/models/tests/test_filesystem.py b/subiquity/models/tests/test_filesystem.py
index 80cfbbb2..802eda0d 100644
--- a/subiquity/models/tests/test_filesystem.py
+++ b/subiquity/models/tests/test_filesystem.py
@@ -739,6 +739,32 @@ class TestPartitionNumbering(unittest.TestCase):
         for p in parts:
             self.assert_next(p)
 
+    @parameterized.expand(
+        [[pt, primaries, i]
+         for pt, primaries in (('msdos', 4), ('vtoc', 3))
+         for i in range(3)]
+    )
+    def test_out_of_offset_order(self, ptable, primaries, idx_to_remove):
+        m = make_model(storage_version=2)
+        d1 = make_disk(m, ptable=ptable, size=100 << 20)
+        self.assert_next(make_partition(m, d1, size=10 << 20))
+        self.assert_next(make_partition(m, d1, flag='extended'))
+        self.cur_idx = primaries + 1
+        parts = []
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=30 << 20))
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=20 << 20))
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=40 << 20))
+        for p in parts:
+            self.assert_next(p)
+        to_remove = parts.pop(idx_to_remove)
+        m.remove_partition(to_remove)
+        self.cur_idx = primaries + 1
+        for p in parts:
+            self.assert_next(p)
+
     @parameterized.expand([
         [1, 'msdos', 4],
         [1, 'vtoc', 3],