C++/Rust 2 way binding for MagiskD

native/src/core/daemon.rs

@@ -1,23 +1,51 @@
 use std::fs::File;
-use std::io;
 use std::io::BufReader;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Mutex, OnceLock};
+use std::{io, mem};
 
 use base::libc::{O_CLOEXEC, O_RDONLY};
 use base::{
-    cstr, BufReadExt, Directory, FsPathBuf, ResultExt, Utf8CStr, Utf8CStrBuf, Utf8CStrBufArr,
-    Utf8CStrBufRef, WalkResult,
+    cstr, libc, open_fd, BufReadExt, Directory, FsPathBuf, ResultExt, Utf8CStr, Utf8CStrBuf,
+    Utf8CStrBufArr, Utf8CStrBufRef, WalkResult,
 };
 
+use crate::ffi::{CxxMagiskD, RequestCode};
 use crate::logging::magisk_logging;
 use crate::{get_prop, MAIN_CONFIG};
 
 // Global magiskd singleton
 pub static MAGISKD: OnceLock<MagiskD> = OnceLock::new();
 
+#[repr(u32)]
+enum BootState {
+    PostFsDataDone = (1 << 0),
+    LateStartDone = (1 << 1),
+    BootComplete = (1 << 2),
+    SafeMode = (1 << 3),
+}
+
+#[derive(Default)]
+#[repr(transparent)]
+struct BootStateFlags(AtomicU32);
+
+impl BootStateFlags {
+    fn contains(&self, stage: BootState) -> bool {
+        let v = self.0.load(Ordering::Relaxed);
+        (v & 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);
+    }
+}
+
 #[derive(Default)]
 pub struct MagiskD {
     pub logd: Mutex<Option<File>>,
+    boot_stage_lock: Mutex<()>,
+    boot_state_flags: BootStateFlags,
     is_emulator: bool,
     is_recovery: bool,
 }
@@ -30,6 +58,50 @@ impl MagiskD {
     pub fn is_recovery(&self) -> bool {
         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 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);
+                }
+                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)
+                {
+                    self.as_cxx().late_start();
+                    self.boot_state_flags.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);
+                    self.as_cxx().boot_complete()
+                }
+            }
+            _ => {
+                unsafe { libc::close(client) };
+            }
+        }
+    }
+
+    #[inline(always)]
+    fn as_cxx(&self) -> &CxxMagiskD {
+        unsafe { mem::transmute(self) }
+    }
 }
 
 mod cxx_extern {
@@ -69,15 +141,46 @@ pub fn daemon_entry() {
     }
 
     let magiskd = MagiskD {
-        logd: Default::default(),
         is_emulator,
         is_recovery,
+        ..Default::default()
     };
     magiskd.start_log_daemon();
     MAGISKD.set(magiskd).ok();
     magisk_logging();
 }
 
+fn check_data() -> bool {
+    if let Ok(fd) = open_fd!(cstr!("/proc/mounts"), O_RDONLY | O_CLOEXEC) {
+        let file = File::from(fd);
+        let mut mnt = false;
+        BufReader::new(file).foreach_lines(|line| {
+            if line.contains(" /data ") && !line.contains("tmpfs") {
+                mnt = true;
+                return false;
+            }
+            true
+        });
+        if !mnt {
+            return false;
+        }
+        let crypto = get_prop(cstr!("ro.crypto.state"), false);
+        return if !crypto.is_empty() {
+            if crypto != "encrypted" {
+                // Unencrypted, we can directly access data
+                true
+            } else {
+                // Encrypted, check whether vold is started
+                !get_prop(cstr!("init.svc.vold"), false).is_empty()
+            }
+        } else {
+            // ro.crypto.state is not set, assume it's unencrypted
+            true
+        };
+    }
+    false
+}
+
 pub fn get_magiskd() -> &'static MagiskD {
     unsafe { MAGISKD.get().unwrap_unchecked() }
 }