Use rust to implement collect/reset overlay context

native/src/external/crt0

@@ -1 +1 @@
-Subproject commit ee73137c106cb01a506a98a460457aecef5a1217
+Subproject commit 0ca1dea7e4e741b48fe94697b563cff712322591

native/src/init/lib.rs

@@ -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);
+        fn reset_overlay_contexts();
     }
 
     unsafe extern "C++" {

native/src/init/rootdir.cpp

@@ -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.rs

@@ -1,9 +1,83 @@
+use base::{
+    debug, libc, Directory, LibcReturn, LoggedResult, 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";
+
+fn get_context<const N: usize>(path: &str, con: &mut Utf8CStrBufArr<N>) -> std::io::Result<()> {
+    unsafe {
+        let sz = libc::lgetxattr(
+            path.as_ptr().cast(),
+            XATTR_NAME_SELINUX.as_ptr().cast(),
+            con.as_mut_ptr().cast(),
+            con.capacity(),
+        )
+        .check_os_err()?;
+        con.set_len((sz - 1) as usize);
+    }
+    Ok(())
+}
+
+fn set_context(path: &str, con: &str) -> std::io::Result<()> {
+    unsafe {
+        libc::lsetxattr(
+            path.as_ptr().cast(),
+            XATTR_NAME_SELINUX.as_ptr().cast(),
+            con.as_ptr().cast(),
+            con.len() + 1,
+            0,
+        )
+        .as_os_err()
+    }
+}
+
+pub fn collect_overlay_contexts(src: &Utf8CStr) {
+    OVERLAY_ATTRS
+        .get_or_try_init(|| -> LoggedResult<_> {
+            let mut contexts = vec![];
+            let mut con = Utf8CStrBufArr::default();
+            let mut path = Utf8CStrBufArr::default();
+            let mut src = Directory::open(src)?;
+            src.path(&mut path)?;
+            let src_len = path.len();
+            src.post_order_walk(|f| {
+                f.path(&mut path)?;
+
+                let path = &path[src_len..];
+                if get_context(path, &mut con)
+                    .log_with_msg(|w| w.write_fmt(format_args!("collect context {}", path)))
+                    .is_ok()
+                {
+                    debug!("collect context: {:?} -> {:?}", path, con);
+                    contexts.push((path.to_string(), con.to_string()));
+                }
+
+                Ok(WalkResult::Continue)
+            })?;
+            Ok(contexts)
+        })
+        .ok();
+}
+
+pub fn reset_overlay_contexts() {
+    OVERLAY_ATTRS.get().map(|attrs| {
+        for (path, con) in attrs.iter() {
+            debug!("set context: {} -> {}", path, con);
+            set_context(path, con)
+                .log_with_msg(|w| w.write_fmt(format_args!("reset context {}", path)))
+                .ok();
+        }
+        Some(())
+    });
+}
 
 pub fn inject_magisk_rc(fd: RawFd, tmp_dir: &Utf8CStr) {
     debug!("Injecting magisk rc");

native/src/init/selinux.cpp

@@ -132,8 +132,7 @@ bool MagiskInit::hijack_sepolicy() {
     sepol->to_file(SELINUX_LOAD);
 
     // restore mounted files' context after sepolicy loaded
-    void reset_overlay_contexts();
+    rust::reset_overlay_contexts();
-    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