From 1f974cb220efa63951226a6359b1f346eff04abc Mon Sep 17 00:00:00 2001
From: ChsBuffer <33744752+chsbuffer@users.noreply.github.com>
Date: Wed, 7 Aug 2024 12:13:12 +0800
Subject: [PATCH] Use `androidboot.partition_map` as a fallback for matching
 partition names in the preinit finding.

---
 native/src/init/getinfo.cpp | 36 ++++++++++++++++++++++++++++++++++++
 native/src/init/init.hpp    |  2 ++
 native/src/init/mount.cpp   |  8 ++++++++
 3 files changed, 46 insertions(+)

diff --git a/native/src/init/getinfo.cpp b/native/src/init/getinfo.cpp
index 7832d4815..5d5bf1e32 100644
--- a/native/src/init/getinfo.cpp
+++ b/native/src/init/getinfo.cpp
@@ -207,6 +207,42 @@ void load_kernel_info(BootConfig *config) {
     config->print();
 }
 
+// `androidboot.partition_map` allows associating a partition name for a raw block device
+// through a comma separated and semicolon deliminated list. For example,
+// `androidboot.partition_map=vdb,metadata;vdc,userdata` maps `vdb` to `metadata` and `vdc` to
+// `userdata`.
+// https://android.googlesource.com/platform/system/core/+/refs/heads/android13-release/init/devices.cpp#191
+
+kv_pairs load_partition_map() {
+    kv_pairs partition_map;
+    auto fn = [&partition_map](const kv_pairs &kv) {
+        for (const auto &[key, value] : kv) {
+            if (key != "androidboot.partition_map")
+                continue; 
+            for (const auto &map : split(value, ";")) {
+                auto map_pieces = split(map, ",");
+                if (map_pieces.size() != 2) {
+                    LOGE("Expected a comma separated device,partition mapping, but found '%s'", map.c_str());
+                    continue;
+                }
+                partition_map.emplace_back(map_pieces[0], map_pieces[1]);
+            }
+        }
+    };
+
+    fn(parse_cmdline(full_read("/proc/cmdline")));
+    fn(parse_bootconfig(full_read("/proc/bootconfig")));
+    return partition_map;
+}
+
+std::string get_partition_name_for_device(const kv_pairs &partition_map, const std::string &query_device){
+    for (const auto &[device, partition] : partition_map) {
+        if (query_device == device)
+            return partition;
+    }
+    return {};
+}
+
 bool check_two_stage() {
     if (access("/apex", F_OK) == 0)
         return true;
diff --git a/native/src/init/init.hpp b/native/src/init/init.hpp
index e5c5e2d01..16e90c37b 100644
--- a/native/src/init/init.hpp
+++ b/native/src/init/init.hpp
@@ -29,6 +29,8 @@ extern std::vector<std::string> mount_list;
 int magisk_proxy_main(int argc, char *argv[]);
 bool unxz(out_stream &strm, rust::Slice<const uint8_t> bytes);
 void load_kernel_info(BootConfig *config);
+kv_pairs load_partition_map();
+std::string get_partition_name_for_device(const kv_pairs &partition_map, const std::string &query_device);
 bool check_two_stage();
 const char *backup_init();
 void restore_ramdisk_init();
diff --git a/native/src/init/mount.cpp b/native/src/init/mount.cpp
index 93d437774..4b462ea92 100644
--- a/native/src/init/mount.cpp
+++ b/native/src/init/mount.cpp
@@ -56,6 +56,8 @@ static void collect_devices() {
             if (access(path, F_OK) == 0) {
                 auto name = rtrim(full_read(path));
                 strcpy(dev.dmname, name.data());
+            } else {
+                dev.dmname[0] = '\0';
             }
             sprintf(path, "/sys/dev/block/%s", entry->d_name);
             xrealpath(path, dev.devpath, sizeof(dev.devpath));
@@ -70,11 +72,17 @@ static struct {
 } blk_info;
 
 static dev_t setup_block() {
+    const auto partition_map = load_partition_map();
     if (dev_list.empty())
         collect_devices();
 
     for (int tries = 0; tries < 3; ++tries) {
         for (auto &dev : dev_list) {
+            // use androidboot.partition_map as partname fallback.
+            if (dev.partname[0] == '\0') {
+                auto partname = get_partition_name_for_device(partition_map, dev.devname);
+                strcpy(dev.partname, partname.c_str());
+            }
             if (strcasecmp(dev.partname, blk_info.partname) == 0)
                 LOGD("Setup %s: [%s] (%d, %d)\n", dev.partname, dev.devname, dev.major, dev.minor);
             else if (strcasecmp(dev.dmname, blk_info.partname) == 0)