From c7083659aacb1d44cc83adcbb07ab0dbd4bd1b49 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Mon, 27 Nov 2023 17:40:58 +0800
Subject: [PATCH] Directly guard boot state with mutex

---
 native/src/core/bootstages.cpp   |  7 ++++--
 native/src/core/daemon.rs        | 37 +++++++++++++-------------------
 native/src/core/include/core.hpp |  2 +-
 native/src/core/lib.rs           |  3 +--
 4 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/native/src/core/bootstages.cpp b/native/src/core/bootstages.cpp
index e876de0cf..ea396e77b 100644
--- a/native/src/core/bootstages.cpp
+++ b/native/src/core/bootstages.cpp
@@ -304,7 +304,7 @@ static bool check_key_combo() {
 
 extern int disable_deny();
 
-void MagiskD::post_fs_data() const {
+bool MagiskD::post_fs_data() const {
     as_rust().setup_logfile();
 
     LOGI("** post-fs-data mode running\n");
@@ -313,6 +313,8 @@ void MagiskD::post_fs_data() const {
     mount_mirrors();
     prune_su_access();
 
+    bool safe_mode = false;
+
     if (access(SECURE_DIR, F_OK) != 0) {
         LOGE(SECURE_DIR " is not present, abort\n");
         goto early_abort;
@@ -325,7 +327,7 @@ void MagiskD::post_fs_data() const {
 
     if (get_prop("persist.sys.safemode", true) == "1" ||
         get_prop("ro.sys.safemode") == "1" || check_key_combo()) {
-        as_rust().enable_safe_mode();
+        safe_mode = true;
         // Disable all modules and denylist so next boot will be clean
         disable_modules();
         disable_deny();
@@ -341,6 +343,7 @@ void MagiskD::post_fs_data() const {
 early_abort:
     // We still do magic mount because root itself might need it
     load_modules();
+    return safe_mode;
 }
 
 void MagiskD::late_start() const {
diff --git a/native/src/core/daemon.rs b/native/src/core/daemon.rs
index 1986203aa..25e3858db 100644
--- a/native/src/core/daemon.rs
+++ b/native/src/core/daemon.rs
@@ -1,6 +1,5 @@
 use std::fs::File;
 use std::io::BufReader;
-use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Mutex, OnceLock};
 use std::{io, mem};
 
@@ -27,25 +26,22 @@ enum BootState {
 
 #[derive(Default)]
 #[repr(transparent)]
-struct BootStateFlags(AtomicU32);
+struct BootStateFlags(u32);
 
 impl BootStateFlags {
     fn contains(&self, stage: BootState) -> bool {
-        let v = self.0.load(Ordering::Relaxed);
-        (v & stage as u32) != 0
+        (self.0 & stage as u32) != 0
     }
 
-    fn set(&self, stage: BootState) {
-        let v = self.0.load(Ordering::Relaxed);
-        self.0.store(v | stage as u32, Ordering::Relaxed);
+    fn set(&mut self, stage: BootState) {
+        self.0 |= stage as u32;
     }
 }
 
 #[derive(Default)]
 pub struct MagiskD {
     pub logd: Mutex<Option<File>>,
-    boot_stage_lock: Mutex<()>,
-    boot_state_flags: BootStateFlags,
+    boot_stage_lock: Mutex<BootStateFlags>,
     is_emulator: bool,
     is_recovery: bool,
 }
@@ -59,36 +55,33 @@ impl MagiskD {
         self.is_recovery
     }
 
-    pub fn enable_safe_mode(&self) {
-        self.boot_state_flags.set(BootState::SafeMode)
-    }
-
     pub fn boot_stage_handler(&self, client: i32, code: i32) {
         // Make sure boot stage execution is always serialized
-        let _guard = self.boot_stage_lock.lock().unwrap();
+        let mut state = self.boot_stage_lock.lock().unwrap();
 
         let code = RequestCode { repr: code };
         match code {
             RequestCode::POST_FS_DATA => {
-                if check_data() && !self.boot_state_flags.contains(BootState::PostFsDataDone) {
-                    self.as_cxx().post_fs_data();
-                    self.boot_state_flags.set(BootState::PostFsDataDone);
+                if check_data() && !state.contains(BootState::PostFsDataDone) {
+                    if self.as_cxx().post_fs_data() {
+                        state.set(BootState::SafeMode);
+                    }
+                    state.set(BootState::PostFsDataDone);
                 }
                 unsafe { libc::close(client) };
             }
             RequestCode::LATE_START => {
                 unsafe { libc::close(client) };
-                if self.boot_state_flags.contains(BootState::PostFsDataDone)
-                    && !self.boot_state_flags.contains(BootState::SafeMode)
+                if state.contains(BootState::PostFsDataDone) && !state.contains(BootState::SafeMode)
                 {
                     self.as_cxx().late_start();
-                    self.boot_state_flags.set(BootState::LateStartDone);
+                    state.set(BootState::LateStartDone);
                 }
             }
             RequestCode::BOOT_COMPLETE => {
                 unsafe { libc::close(client) };
-                if !self.boot_state_flags.contains(BootState::SafeMode) {
-                    self.boot_state_flags.set(BootState::BootComplete);
+                if !state.contains(BootState::SafeMode) {
+                    state.set(BootState::BootComplete);
                     self.as_cxx().boot_complete()
                 }
             }
diff --git a/native/src/core/include/core.hpp b/native/src/core/include/core.hpp
index c098fcd16..74a53ef12 100644
--- a/native/src/core/include/core.hpp
+++ b/native/src/core/include/core.hpp
@@ -32,7 +32,7 @@ struct MagiskD {
 
     // C++ implementation
     void reboot() const;
-    void post_fs_data() const;
+    bool post_fs_data() const;
     void late_start() const;
     void boot_complete() const;
 
diff --git a/native/src/core/lib.rs b/native/src/core/lib.rs
index 3cdb96941..bd55eb65b 100644
--- a/native/src/core/lib.rs
+++ b/native/src/core/lib.rs
@@ -57,7 +57,7 @@ pub mod ffi {
 
         #[cxx_name = "MagiskD"]
         type CxxMagiskD;
-        fn post_fs_data(self: &CxxMagiskD);
+        fn post_fs_data(self: &CxxMagiskD) -> bool;
         fn late_start(self: &CxxMagiskD);
         fn boot_complete(self: &CxxMagiskD);
     }
@@ -87,7 +87,6 @@ pub mod ffi {
         fn is_emulator(self: &MagiskD) -> bool;
         fn is_recovery(self: &MagiskD) -> bool;
         fn boot_stage_handler(self: &MagiskD, client: i32, code: i32);
-        fn enable_safe_mode(self: &MagiskD);
     }
 }