From fa1711f84f07f05cc566ee6d56d575ca812a2058 Mon Sep 17 00:00:00 2001
From: Sunil Mohan Adapa <sunil@medhas.org>
Date: Sun, 26 Feb 2023 06:35:41 +0530
Subject: [PATCH] container: Fix resizing disk image containing multiple
 partitions

Closes: #2324.

The arm64 image contains two partitions in a GPT partition table. First is the
EFI partition and the second is the root partition. The container script
currently assumes that there will only be one partition in the image file. Fix
this by picking up the partition number of the last partition and resizing that.
GPT partition table also requires relocating the second copy of the partition
table to the end before partition resize can succeed.

Tests:

- Create testing containers in arm64 and amd64 architectures.

Signed-off-by: Sunil Mohan Adapa <sunil@medhas.org>
Reviewed-by: James Valleroy <jvalleroy@mailbox.org>
---
 container | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/container b/container
index cbb1fadc1..eab4dae2f 100755
--- a/container
+++ b/container
@@ -368,6 +368,7 @@ def _verify_dependencies():
         'gpg': 'gpg',
         'xz': 'xz-utils',
         'parted': 'parted',
+        'sgdisk': 'gdisk',
         'btrfs': 'btrfs-progs',
         'nmcli': 'network-manager',
         'ssh': 'openssh-client',
@@ -522,6 +523,21 @@ def _download_disk_image(distribution, force=False):
     return _extract_image(target_file)
 
 
+def _get_partition_info(image_file):
+    """Return the number of the final partition in the image file."""
+    process = subprocess.run(
+        ['sudo', 'parted', '--script',
+         str(image_file), 'print'], stdout=subprocess.PIPE, check=True)
+    lines = process.stdout.decode().splitlines()
+    last_partition_number = lines[-2].split()[0]
+    partition_table_type = 'msdos'
+    for line in lines:
+        if line.startswith('Partition Table:'):
+            partition_table_type = line.partition(': ')[2]
+
+    return partition_table_type, last_partition_number
+
+
 def _resize_disk_image(image_file, new_size):
     """Resize the disk image if has not already been."""
     if new_size[-1] != 'G':
@@ -542,14 +558,22 @@ def _resize_disk_image(image_file, new_size):
         ['truncate', '--size',
          str(new_size_bytes),
          str(image_file)], check=True)
+    partition_table_type, last_partition_number = _get_partition_info(
+        image_file)
+    if partition_table_type == 'gpt':
+        subprocess.run(
+            ['sudo', 'sgdisk', '--move-second-header',
+             str(image_file)], check=True)
+
     subprocess.run([
         'sudo', 'parted', '--align=optimal', '--script',
-        str(image_file), 'resizepart', '1', '100%'
+        str(image_file), 'resizepart', last_partition_number, '100%'
     ], check=True)
     process = subprocess.run(
         ['sudo', 'kpartx', '-avs', str(image_file)], stdout=subprocess.PIPE,
         check=True)
-    partition = '/dev/mapper/' + process.stdout.decode().split()[2]
+    last_line = process.stdout.decode().splitlines()[-1]
+    partition = '/dev/mapper/' + last_line.split()[2]
     subprocess.run(['sudo', 'btrfstune', '-uf', partition], check=True)
 
     with tempfile.TemporaryDirectory(