Use rust to implement collect/reset overlay context

Author: yujincheng08

native/src/external/crt0

@@ -1,8 +1,9 @@
 #![feature(format_args_nl)]
+#![feature(once_cell_try)]
 
 use logging::setup_klog;
 use mount::{is_device_mounted, switch_root};
-use rootdir::inject_magisk_rc;
+use rootdir::{inject_magisk_rc, collect_overlay_contexts, reset_overlay_contexts};
 // Has to be pub so all symbols in that crate is included
 pub use magiskpolicy;
 
@@ -18,6 +19,8 @@ pub mod ffi {
         fn inject_magisk_rc(fd: i32, tmp_dir: Utf8CStrRef);
         fn switch_root(path: Utf8CStrRef);
         fn is_device_mounted(dev: u64, target: Pin<&mut CxxString>) -> bool;
+        fn collect_overlay_contexts(src: Utf8CStrRef) -> bool;
+        fn reset_overlay_contexts() -> bool;
     }
 
     unsafe extern "C++" {

native/src/init/lib.rs

@@ -1,8 +1,6 @@
 #include <sys/mount.h>
 #include <libgen.h>
 #include <sys/sysmacros.h>
-#include <syscall.h>
-#include <sys/xattr.h>
 
 #include <sepolicy.hpp>
 #include <consts.hpp>
@@ -16,13 +14,6 @@ using namespace std;
 static vector<string> rc_list;
 static string magic_mount_list;
 
-struct FileContext {
-    std::string path;
-    std::string con;
-};
-
-static std::vector<FileContext> mount_contexts;
-
 #define NEW_INITRC_DIR  "/system/etc/init/hw"
 #define INIT_RC         "init.rc"
 
@@ -50,46 +41,6 @@ static void magic_mount(const string &sdir, const string &ddir = "") {
     }
 }
 
-static int setfilecon(const char* path, const char* con) {
-    int ret = syscall(__NR_setxattr, path, XATTR_NAME_SELINUX, con, strlen(con) + 1, 0);
-    if (ret == -1) PLOGE("setfilecon %s %s", path, con);
-    return ret;
-}
-
-static std::string getfilecon(const char* path) {
-    char buf[1024];
-    ssize_t sz = syscall(__NR_getxattr, path, XATTR_NAME_SELINUX, buf, sizeof(buf));
-    if (sz == -1) {
-        PLOGE("getfilecon %s", path);
-        return "";
-    }
-    return buf;
-}
-
-static void collect_overlay_contexts(const string &sdir, const string &ddir = "") {
-    auto dir = xopen_dir(sdir.data());
-    if (!dir) return;
-    for (dirent *entry; (entry = xreaddir(dir.get()));) {
-        string src = sdir + "/" + entry->d_name;
-        string dest = ddir + "/" + entry->d_name;
-        if (access(dest.data(), F_OK) == 0) {
-            if (entry->d_type == DT_DIR) {
-                // Recursive
-                collect_overlay_contexts(src, dest);
-            } else {
-                mount_contexts.emplace_back(dest, getfilecon(dest.data()));
-            }
-        }
-    }
-}
-
-void reset_overlay_contexts() {
-    for (auto &attr: mount_contexts) {
-        LOGD("set %s -> %s", attr.path.c_str(), attr.con.c_str());
-        setfilecon(attr.path.c_str(), attr.con.c_str());
-    }
-}
-
 static void patch_rc_scripts(const char *src_path, const char *tmp_path, bool writable) {
     auto src_dir = xopen_dir(src_path);
     if (!src_dir) return;
@@ -380,7 +331,7 @@ void MagiskInit::patch_ro_root() {
     // Extract overlay archives
     extract_files(false);
 
-    collect_overlay_contexts(ROOTOVL);
+    rust::collect_overlay_contexts(ROOTOVL);
 
     // Oculus Go will use a special sepolicy if unlocked
     if (access("/sepolicy.unlocked", F_OK) == 0) {

native/src/init/rootdir.cpp

@@ -1,9 +1,57 @@
+use base::{debug, libc, Directory, LibcReturn, ResultExt, Utf8CStr, Utf8CStrBuf, Utf8CStrBufArr, WalkResult};
 use std::fs::File;
 use std::io::Write;
 use std::mem;
 use std::os::fd::{FromRawFd, RawFd};
+use std::sync::OnceLock;
 
-use base::{debug, Utf8CStr};
+pub static OVERLAY_ATTRS: OnceLock<Vec<(String, String)>> = OnceLock::new();
+
+const XATTR_NAME_SELINUX: &[u8] = b"security.selinux\0";
+
+pub fn collect_overlay_contexts(src: &Utf8CStr) -> bool {
+    OVERLAY_ATTRS.get_or_try_init(|| -> Result<_, std::io::Error> {
+        let mut contexts = vec![];
+        let mut con = Utf8CStrBufArr::default();
+        let mut path = Utf8CStrBufArr::default();
+        Directory::open(src)?.post_order_walk(|f| {
+            f.path(&mut path)?;
+
+            unsafe {
+                let sz = libc::lgetxattr(
+                    path.as_ptr(),
+                    XATTR_NAME_SELINUX.as_ptr().cast(),
+                    con.as_mut_ptr().cast(),
+                    con.capacity(),
+                );
+                con.set_len((sz - 1) as usize);
+            }
+
+            debug!("collect context: {:?} -> {:?}", path, con);
+            contexts.push((path.to_string(), con.to_string()));
+            Ok(WalkResult::Continue)
+        })?;
+        Ok(contexts)
+    }).is_ok()
+}
+
+pub fn reset_overlay_contexts() -> bool {
+    OVERLAY_ATTRS.get().map(|attrs| {
+        for (dest, con) in attrs.iter() {
+            debug!("set context: {} -> {}", dest, con);
+            unsafe {
+                libc::lsetxattr(
+                    dest.as_ptr().cast(),
+                    XATTR_NAME_SELINUX.as_ptr().cast(),
+                    con.as_ptr().cast(),
+                    con.len() + 1,
+                    0,
+                ).as_os_err().log().ok()?;
+            }
+        }
+        Some(())
+    }).is_some()
+}
 
 pub fn inject_magisk_rc(fd: RawFd, tmp_dir: &Utf8CStr) {
     debug!("Injecting magisk rc");

native/src/init/rootdir.rs

@@ -132,8 +132,7 @@ bool MagiskInit::hijack_sepolicy() {
     sepol->to_file(SELINUX_LOAD);
 
     // restore mounted files' context after sepolicy loaded
-    void reset_overlay_contexts();
-    reset_overlay_contexts();
+    rust::reset_overlay_contexts();
 
     // Write to the enforce node ONLY after sepolicy is loaded. We need to make sure
     // the actual init process is blocked until sepolicy is loaded, or else