From 7abb281296044ddd0d9dbf6c71f07ce53713bd42 Mon Sep 17 00:00:00 2001
From: Oliver Smith <ollieparanoid@postmarketos.org>
Date: Tue, 12 Apr 2022 09:42:58 +0200
Subject: [PATCH] partitions_mount: fix finding the partition (MR 2178)

Instead of trying both "${dev}1" and "${dev}p1" as partition paths, only
try one of them depending on if "${dev}" ends in a number or not. This
fixes getting a wrong /dev/loop11 partition by accident if there are
many loop devices, instead of only looking for the correct path
/dev/loop1p1. People reported this happening with snaps on ubuntu.
---
 pmb/install/partition.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/pmb/install/partition.py b/pmb/install/partition.py
index 57230a8f..761b1af9 100644
--- a/pmb/install/partition.py
+++ b/pmb/install/partition.py
@@ -19,23 +19,27 @@ def partitions_mount(args, layout, sdcard):
         img_path = "/home/pmos/rootfs/" + args.device + ".img"
         prefix = pmb.install.losetup.device_by_back_file(args, img_path)
 
-    partition_prefix = None
     tries = 20
+
+    # Devices ending with a number have a "p" before the partition number,
+    # /dev/sda1 has no "p", but /dev/mmcblk0p1 has. See add_partition() in
+    # block/partitions/core.c of linux.git.
+    partition_prefix = prefix
+    if str.isdigit(prefix[-1:]):
+        partition_prefix = f"{prefix}p"
+
+    found = False
     for i in range(tries):
-        for symbol in ["p", ""]:
-            if os.path.exists(prefix + symbol + "1"):
-                partition_prefix = symbol
-        if partition_prefix is not None:
+        if os.path.exists(f"{partition_prefix}1"):
+            found = True
             break
         logging.debug(f"NOTE: ({i + 1}/{tries}) failed to find the install "
                       "partition. Retrying...")
         time.sleep(0.1)
 
-    if partition_prefix is None:
-        raise RuntimeError("Unable to find the partition prefix,"
-                           " expected the first partition of " +
-                           prefix + " to be located at " + prefix +
-                           "1 or " + prefix + "p1!")
+    if not found:
+        raise RuntimeError(f"Unable to find the first partition of {prefix}, "
+                           f"expected it to be at {partition_prefix}1!")
 
     partitions = [layout["boot"], layout["root"]]
 
@@ -43,7 +47,7 @@ def partitions_mount(args, layout, sdcard):
         partitions += [layout["kernel"]]
 
     for i in partitions:
-        source = prefix + partition_prefix + str(i)
+        source = f"{partition_prefix}{i}"
         target = args.work + "/chroot_native/dev/installp" + str(i)
         pmb.helpers.mount.bind_file(args, source, target)