Guard boot stages more precisely

Close #6468, fix #6148

native/src/core/bootstages.cpp

@@ -17,7 +17,17 @@
 
 using namespace std;
 
-static bool safe_mode = false;
+// Boot stage state
+enum : int {
+    FLAG_NONE = 0,
+    FLAG_POST_FS_DATA_DONE = (1 << 0),
+    FLAG_LATE_START_DONE = (1 << 1),
+    FLAG_BOOT_COMPLETE = (1 << 2),
+    FLAG_SAFE_MODE = (1 << 3),
+};
+
+static int boot_state = FLAG_NONE;
+
 bool zygisk_enabled = false;
 
 /*********
@@ -272,21 +282,15 @@ static bool check_key_combo() {
  * Boot Stage Handlers *
  ***********************/
 
-static pthread_mutex_t stage_lock = PTHREAD_MUTEX_INITIALIZER;
 extern int disable_deny();
 
-void post_fs_data(int client) {
-    close(client);
-
-    mutex_guard lock(stage_lock);
-
+static void post_fs_data() {
     if (getenv("REMOUNT_ROOT"))
         xmount(nullptr, "/", nullptr, MS_REMOUNT | MS_RDONLY, nullptr);
 
     if (!check_data())
-        goto unblock_init;
+        return;
 
-    DAEMON_STATE = STATE_POST_FS_DATA;
     setup_logfile(true);
 
     LOGI("** post-fs-data mode running\n");
@@ -315,7 +319,7 @@ void post_fs_data(int client) {
 
     if (getprop("persist.sys.safemode", true) == "1" ||
         getprop("ro.sys.safemode") == "1" || check_key_combo()) {
-        safe_mode = true;
+        boot_state |= FLAG_SAFE_MODE;
         // Disable all modules and denylist so next boot will be clean
         disable_modules();
         disable_deny();
@@ -331,40 +335,26 @@ void post_fs_data(int client) {
 early_abort:
     // We still do magic mount because root itself might need it
     magic_mount();
-    DAEMON_STATE = STATE_POST_FS_DATA_DONE;
-
-unblock_init:
-    close(xopen(UNBLOCKFILE, O_RDONLY | O_CREAT, 0));
+    boot_state |= FLAG_POST_FS_DATA_DONE;
 }
 
-void late_start(int client) {
-    close(client);
-
-    mutex_guard lock(stage_lock);
-    run_finally fin([]{ DAEMON_STATE = STATE_LATE_START_DONE; });
+static void late_start() {
     setup_logfile(false);
 
     LOGI("** late_start service mode running\n");
 
-    if (DAEMON_STATE < STATE_POST_FS_DATA_DONE || safe_mode)
-        return;
-
     exec_common_scripts("service");
     exec_module_scripts("service");
+
+    boot_state |= FLAG_LATE_START_DONE;
 }
 
-void boot_complete(int client) {
-    close(client);
-
-    mutex_guard lock(stage_lock);
-    DAEMON_STATE = STATE_BOOT_COMPLETE;
+static void boot_complete() {
+    boot_state |= FLAG_BOOT_COMPLETE;
     setup_logfile(false);
 
     LOGI("** boot-complete triggered\n");
 
-    if (safe_mode)
-        return;
-
     // At this point it's safe to create the folder
     if (access(SECURE_DIR, F_OK) != 0)
         xmkdir(SECURE_DIR, 0700);
@@ -374,10 +364,26 @@ void boot_complete(int client) {
     get_manager(0, nullptr, true);
 }
 
-void zygote_restart(int client) {
-    close(client);
+void boot_stage_handler(int code) {
+    // Make sure boot stage execution is always serialized
+    static pthread_mutex_t stage_lock = PTHREAD_MUTEX_INITIALIZER;
+    mutex_guard lock(stage_lock);
 
-    LOGI("** zygote restarted\n");
-    pkg_xml_ino = 0;
-    prune_su_access();
+    switch (code) {
+    case MainRequest::POST_FS_DATA:
+        if ((boot_state & FLAG_POST_FS_DATA_DONE) == 0)
+            post_fs_data();
+        close(xopen(UNBLOCKFILE, O_RDONLY | O_CREAT, 0));
+        break;
+    case MainRequest::LATE_START:
+        if ((boot_state & FLAG_POST_FS_DATA_DONE) && (boot_state & FLAG_SAFE_MODE) == 0)
+            late_start();
+        break;
+    case MainRequest::BOOT_COMPLETE:
+        if ((boot_state & FLAG_SAFE_MODE) == 0)
+            boot_complete();
+        break;
+    default:
+        __builtin_unreachable();
+    }
 }

native/src/core/core.hpp

@@ -4,18 +4,8 @@
 #include <vector>
 
 extern bool RECOVERY_MODE;
-extern int DAEMON_STATE;
 extern std::atomic<ino_t> pkg_xml_ino;
 
-// Daemon state
-enum : int {
-    STATE_NONE,
-    STATE_POST_FS_DATA,
-    STATE_POST_FS_DATA_DONE,
-    STATE_LATE_START_DONE,
-    STATE_BOOT_COMPLETE
-};
-
 void unlock_blocks();
 void reboot();
 void start_log_daemon();

native/src/core/daemon.cpp

@@ -21,7 +21,6 @@ int SDK_INT = -1;
 string MAGISKTMP;
 
 bool RECOVERY_MODE = false;
-int DAEMON_STATE = STATE_NONE;
 
 static struct stat self_st;
 
@@ -143,17 +142,11 @@ static void handle_request_async(int client, int code, const sock_cred &cred) {
     case MainRequest::SUPERUSER:
         su_daemon_handler(client, &cred);
         break;
-    case MainRequest::POST_FS_DATA:
-        post_fs_data(client);
-        break;
-    case MainRequest::LATE_START:
-        late_start(client);
-        break;
-    case MainRequest::BOOT_COMPLETE:
-        boot_complete(client);
-        break;
     case MainRequest::ZYGOTE_RESTART:
-        zygote_restart(client);
+        close(client);
+        LOGI("** zygote restarted\n");
+        pkg_xml_ino = 0;
+        prune_su_access();
         break;
     case MainRequest::SQLITE_CMD:
         exec_sql(client);
@@ -232,7 +225,9 @@ static void handle_request(pollfd *pfd) {
     }
 
     code = read_int(client);
-    if (code < 0 || code >= MainRequest::END || code == MainRequest::_SYNC_BARRIER_) {
+    if (code < 0 || code >= MainRequest::END ||
+        code == MainRequest::_SYNC_BARRIER_ ||
+        code == MainRequest::_STAGE_BARRIER_) {
         // Unknown request code
         goto done;
     }
@@ -274,10 +269,12 @@ static void handle_request(pollfd *pfd) {
     if (code < MainRequest::_SYNC_BARRIER_) {
         handle_request_sync(client, code);
         goto done;
-    }
-
-    // Handle async requests in another thread
+    } else if (code < MainRequest::_STAGE_BARRIER_) {
         exec_task([=] { handle_request_async(client, code, cred); });
+    } else {
+        close(client);
+        exec_task([=] { boot_stage_handler(code); });
+    }
     return;
 
 done:

native/src/include/daemon.hpp

@@ -31,15 +31,19 @@ enum : int {
     _SYNC_BARRIER_,
 
     SUPERUSER,
-    POST_FS_DATA,
-    LATE_START,
-    BOOT_COMPLETE,
     ZYGOTE_RESTART,
     DENYLIST,
     SQLITE_CMD,
     REMOVE_MODULES,
     ZYGISK,
     ZYGISK_PASSTHROUGH,
+
+    _STAGE_BARRIER_,
+
+    POST_FS_DATA,
+    LATE_START,
+    BOOT_COMPLETE,
+
     END,
 };
 }
@@ -84,10 +88,7 @@ extern std::atomic<int> logd_fd;
 extern "C" void magisk_log_write(int prio, const char *msg, int len);
 
 // Daemon handlers
-void post_fs_data(int client);
-void late_start(int client);
-void boot_complete(int client);
-void zygote_restart(int client);
+void boot_stage_handler(int code);
 void denylist_handler(int client, const sock_cred *cred);
 void su_daemon_handler(int client, const sock_cred *cred);
 void zygisk_handler(int client, const sock_cred *cred);