From 3a39dd404917a8f17aaa4deafed1241884d78b99 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Wed, 8 Jun 2022 23:23:39 -0700
Subject: [PATCH] Update ramdisk restore implementation

---
 native/jni/boot/cpio.cpp    |  6 ++--
 native/jni/boot/cpio.hpp    |  4 +--
 native/jni/boot/ramdisk.cpp | 57 +++++++++++++++++++++----------------
 3 files changed, 38 insertions(+), 29 deletions(-)

diff --git a/native/jni/boot/cpio.cpp b/native/jni/boot/cpio.cpp
index d0aa4bdd8..5c63fd61f 100644
--- a/native/jni/boot/cpio.cpp
+++ b/native/jni/boot/cpio.cpp
@@ -53,7 +53,9 @@ void cpio::dump(const char *file) {
     dump(xfopen(file, "we"));
 }
 
-void cpio::rm(entry_map::iterator &it) {
+void cpio::rm(entry_map::iterator it) {
+    if (it == entries.end())
+        return;
     fprintf(stderr, "Remove [%s]\n", it->first.data());
     entries.erase(it);
 }
@@ -188,7 +190,7 @@ void cpio::ln(const char *target, const char *name) {
     fprintf(stderr, "Create symlink [%s] -> [%s]\n", name, target);
 }
 
-void cpio::mv(entry_map::iterator &it, const char *name) {
+void cpio::mv(entry_map::iterator it, const char *name) {
     fprintf(stderr, "Move [%s] -> [%s]\n", it->first.data(), name);
     auto e = it->second.release();
     entries.erase(it);
diff --git a/native/jni/boot/cpio.hpp b/native/jni/boot/cpio.hpp
index 8a1b9ab64..3feb9fb46 100644
--- a/native/jni/boot/cpio.hpp
+++ b/native/jni/boot/cpio.hpp
@@ -45,8 +45,8 @@ protected:
     entry_map entries;
 
     static void extract_entry(const entry_map::value_type &e, const char *file);
-    void rm(entry_map::iterator &it);
-    void mv(entry_map::iterator &it, const char *name);
+    void rm(entry_map::iterator it);
+    void mv(entry_map::iterator it, const char *name);
 
 private:
     void dump(FILE *out);
diff --git a/native/jni/boot/ramdisk.cpp b/native/jni/boot/ramdisk.cpp
index 5fd61c5fc..e5f58d8be 100644
--- a/native/jni/boot/ramdisk.cpp
+++ b/native/jni/boot/ramdisk.cpp
@@ -82,11 +82,10 @@ int magisk_cpio::test() {
 }
 
 #define for_each_line(line, buf, size) \
-for (line = (char *) buf; line < (char *) buf + size && line[0]; line = strchr(line + 1, '\n') + 1)
+for (char *line = (char *) buf; line < (char *) buf + size && line[0]; line = strchr(line + 1, '\n') + 1)
 
 char *magisk_cpio::sha1() {
     char sha1[41];
-    char *line;
     for (auto &e : entries) {
         if (e.first == "init.magisk.rc" || e.first == "overlay/init.magisk.rc") {
             for_each_line(line, e.second->data, e.second->filesize) {
@@ -112,40 +111,48 @@ char *magisk_cpio::sha1() {
 }
 
 #define for_each_str(str, buf, size) \
-for (str = (char *) buf; str < (char *) buf + size; str = str += strlen(str) + 1)
+for (char *str = (char *) buf; str < (char *) buf + size; str += strlen(str) + 1)
 
 void magisk_cpio::restore() {
-    // If the backup init is missing, this means that the boot ramdisk
-    // was created from scratch by us. Nothing needs to be restored.
-    if (exists(".backup") && !exists(".backup/init")) {
+    // Collect files
+    auto bk = entries.end();
+    auto rl = entries.end();
+    auto mg = entries.end();
+    vector<entry_map::iterator> backups;
+    for (auto it = entries.begin(); it != entries.end(); ++it) {
+        if (it->first == ".backup") {
+            bk = it;
+        } else if (it->first == ".backup/.rmlist") {
+            rl = it;
+        } else if (it->first == ".backup/.magisk") {
+            mg = it;
+        } else if (str_starts(it->first, ".backup/")) {
+            backups.emplace_back(it);
+        }
+    }
+
+    // If the .backup folder is effectively empty, this means that the boot ramdisk was
+    // created from scratch by an old broken magiskboot. This is just a hacky workaround.
+    if (bk != entries.end() && mg != entries.end() && rl == entries.end() && backups.empty()) {
         fprintf(stderr, "Remove all in ramdisk\n");
         entries.clear();
         return;
     }
 
-    if (auto it = entries.find(".backup/.rmlist"); it != entries.end()) {
-        char *file;
-        for_each_str(file, it->second->data, it->second->filesize) {
+    // Remove files
+    rm(bk);
+    rm(mg);
+    if (rl != entries.end()) {
+        for_each_str(file, rl->second->data, rl->second->filesize) {
             rm(file);
         }
-        rm(it);
+        rm(rl);
     }
 
-    for (auto it = entries.begin(); it != entries.end();) {
-        auto cur = it++;
-        if (str_starts(cur->first, ".backup")) {
-            if (cur->first.length() == 7 || &cur->first[8] == ".magisk"sv) {
-                rm(cur);
-            } else {
-                mv(cur, &cur->first[8]);
-            }
-        } else if (str_starts(cur->first, "magisk") ||
-                cur->first == "overlay/init.magisk.rc" ||
-                cur->first == "sbin/magic_mask.sh" ||
-                cur->first == "init.magisk.rc") {
-            // Some known stuff we can remove
-            rm(cur);
-        }
+    // Restore files
+    for (auto it : backups) {
+        const char *name = &it->first[8];
+        mv(it, name);
     }
 }