From 0acc39cec0816d60259d36037111123a18073acc Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Mon, 15 Sep 2025 11:24:14 -0700
Subject: [PATCH] Use bitflags to implement BootState

---
 native/src/Cargo.lock         |  1 +
 native/src/core/Cargo.toml    |  1 +
 native/src/core/bootstages.rs | 21 ++++++++++++++++-----
 native/src/core/daemon.rs     | 25 ++-----------------------
 4 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/native/src/Cargo.lock b/native/src/Cargo.lock
index 5e679c0de..3c83618ec 100644
--- a/native/src/Cargo.lock
+++ b/native/src/Cargo.lock
@@ -644,6 +644,7 @@ dependencies = [
  "argh",
  "base",
  "bit-set",
+ "bitflags",
  "bytemuck",
  "cxx",
  "cxx-gen",
diff --git a/native/src/core/Cargo.toml b/native/src/core/Cargo.toml
index b5d5783de..f0fdaa60e 100644
--- a/native/src/core/Cargo.toml
+++ b/native/src/core/Cargo.toml
@@ -27,3 +27,4 @@ thiserror = { workspace = true }
 bit-set = { workspace = true }
 argh = { workspace = true }
 nix = { workspace = true, features = ["fs", "mount", "poll", "signal", "term", "user", "zerocopy"] }
+bitflags = { workspace = true }
diff --git a/native/src/core/bootstages.rs b/native/src/core/bootstages.rs
index 02f8041ff..b60bdf381 100644
--- a/native/src/core/bootstages.rs
+++ b/native/src/core/bootstages.rs
@@ -1,5 +1,5 @@
 use crate::consts::{APP_PACKAGE_NAME, BBPATH, DATABIN, MODULEROOT, SECURE_DIR};
-use crate::daemon::{BootState, MagiskD};
+use crate::daemon::MagiskD;
 use crate::ffi::{
     DbEntryKey, RequestCode, check_key_combo, exec_common_scripts, exec_module_scripts,
     get_magisk_tmp, initialize_denylist,
@@ -11,12 +11,23 @@ use crate::resetprop::get_prop;
 use crate::selinux::restorecon;
 use base::const_format::concatcp;
 use base::{BufReadExt, FsPathBuilder, ResultExt, cstr, error, info};
+use bitflags::bitflags;
 use nix::fcntl::OFlag;
 use std::io::BufReader;
 use std::os::unix::net::UnixStream;
 use std::process::{Command, Stdio};
 use std::sync::atomic::Ordering;
 
+bitflags! {
+    #[derive(Default)]
+    pub struct BootState : u32 {
+        const PostFsDataDone = 1 << 0;
+        const LateStartDone = 1 << 1;
+        const BootComplete = 1 << 2;
+        const SafeMode = 1 << 3;
+    }
+}
+
 impl MagiskD {
     fn setup_magisk_env(&self) -> bool {
         info!("* Initializing Magisk environment");
@@ -185,9 +196,9 @@ impl MagiskD {
             RequestCode::POST_FS_DATA => {
                 if check_data() && !state.contains(BootState::PostFsDataDone) {
                     if self.post_fs_data() {
-                        state.set(BootState::SafeMode);
+                        state.insert(BootState::SafeMode);
                     }
-                    state.set(BootState::PostFsDataDone);
+                    state.insert(BootState::PostFsDataDone);
                 }
             }
             RequestCode::LATE_START => {
@@ -195,13 +206,13 @@ impl MagiskD {
                 if state.contains(BootState::PostFsDataDone) && !state.contains(BootState::SafeMode)
                 {
                     self.late_start();
-                    state.set(BootState::LateStartDone);
+                    state.insert(BootState::LateStartDone);
                 }
             }
             RequestCode::BOOT_COMPLETE => {
                 drop(client);
                 if state.contains(BootState::PostFsDataDone) {
-                    state.set(BootState::BootComplete);
+                    state.insert(BootState::BootComplete);
                     self.boot_complete()
                 }
             }
diff --git a/native/src/core/daemon.rs b/native/src/core/daemon.rs
index 6a58086df..ea58819be 100644
--- a/native/src/core/daemon.rs
+++ b/native/src/core/daemon.rs
@@ -1,3 +1,4 @@
+use crate::bootstages::BootState;
 use crate::consts::{
     MAGISK_FILE_CON, MAGISK_FULL_VER, MAGISK_PROC_CON, MAGISK_VER_CODE, MAGISK_VERSION,
     MAIN_CONFIG, MAIN_SOCKET, ROOTMNT, ROOTOVL,
@@ -38,28 +39,6 @@ use std::sync::{Mutex, OnceLock};
 // Global magiskd singleton
 pub static MAGISKD: OnceLock<MagiskD> = OnceLock::new();
 
-#[repr(u32)]
-pub enum BootState {
-    PostFsDataDone = (1 << 0),
-    LateStartDone = (1 << 1),
-    BootComplete = (1 << 2),
-    SafeMode = (1 << 3),
-}
-
-#[derive(Default)]
-#[repr(transparent)]
-pub struct BootStateFlags(u32);
-
-impl BootStateFlags {
-    pub fn contains(&self, stage: BootState) -> bool {
-        (self.0 & stage as u32) != 0
-    }
-
-    pub fn set(&mut self, stage: BootState) {
-        self.0 |= stage as u32;
-    }
-}
-
 pub const AID_ROOT: i32 = 0;
 pub const AID_SHELL: i32 = 2000;
 pub const AID_APP_START: i32 = 10000;
@@ -78,7 +57,7 @@ pub const fn to_user_id(uid: i32) -> i32 {
 pub struct MagiskD {
     pub sql_connection: Mutex<Option<Sqlite3>>,
     pub manager_info: Mutex<ManagerInfo>,
-    pub boot_stage_lock: Mutex<BootStateFlags>,
+    pub boot_stage_lock: Mutex<BootState>,
     pub module_list: OnceLock<Vec<ModuleInfo>>,
     pub zygisk_enabled: AtomicBool,
     pub zygisk: Mutex<ZygiskState>,