Move pattern matching to Rust

native/src/Android.mk

@@ -114,7 +114,6 @@ LOCAL_SRC_FILES := \
     boot/compress.cpp \
     boot/format.cpp \
     boot/dtb.cpp \
-    boot/pattern.cpp \
     boot/boot-rs.cpp
 
 include $(BUILD_EXECUTABLE)

native/src/boot/dtb.cpp

@@ -8,6 +8,7 @@
 #include "magiskboot.hpp"
 #include "dtb.hpp"
 #include "format.hpp"
+#include "boot-rs.hpp"
 
 using namespace std;
 
@@ -151,7 +152,7 @@ static bool dtb_patch(const char *file) {
                     int len;
                     char *value = (char *) fdt_getprop(fdt, node, "fsmgr_flags", &len);
                     byte_data data(value, len);
-                    patched |= (patch_verity(data) != len);
+                    patched |= (rust::patch_verity(data) != len);
                 }
             }
         }
@@ -224,7 +225,7 @@ static bool fdt_patch(void *fdt) {
         int len;
         const void *value = fdt_getprop(fdt, node, "fsmgr_flags", &len);
         heap_data copy = byte_view(value, len).clone();
-        auto patched_sz = patch_verity(copy);
+        auto patched_sz = rust::patch_verity(copy);
         if (patched_sz != len) {
             modified = true;
             fdt_setprop(fdt, node, "fsmgr_flags", copy.buf(), patched_sz);

native/src/boot/lib.rs

@@ -18,10 +18,7 @@ mod ramdisk;
 pub mod ffi {
     unsafe extern "C++" {
         include!("compress.hpp");
-        include!("magiskboot.hpp");
         fn decompress(buf: &[u8], fd: i32) -> bool;
-        fn patch_encryption(buf: &mut [u8]) -> usize;
-        fn patch_verity(buf: &mut [u8]) -> usize;
     }
 
     #[namespace = "rust"]
@@ -34,5 +31,7 @@ pub mod ffi {
 
         unsafe fn cpio_commands(argc: i32, argv: *const *const c_char) -> bool;
         fn hexpatch(file: &[u8], from: &[u8], to: &[u8]) -> bool;
+        fn patch_encryption(buf: &mut [u8]) -> usize;
+        fn patch_verity(buf: &mut [u8]) -> usize;
     }
 }

native/src/boot/magiskboot.hpp

@@ -19,9 +19,6 @@ void repack(const char *src_img, const char *out_img, bool skip_comp = false);
 int split_image_dtb(const char *filename);
 int dtb_commands(int argc, char *argv[]);
 
-size_t patch_verity(rust::Slice<uint8_t> data);
-size_t patch_encryption(rust::Slice<uint8_t> data);
-
 static inline bool check_env(const char *name) {
     using namespace std::string_view_literals;
     const char *val = getenv(name);

native/src/boot/patch.rs

@@ -1,5 +1,90 @@
 use base::{MappedFile, MutBytesExt, ResultExt, Utf8CStr};
 
+// SAFETY: assert(buf.len() >= 1) && assert(len <= buf.len())
+macro_rules! match_patterns {
+    ($buf:ident, $($str:literal), *) => {{
+        let mut len = if *$buf.get_unchecked(0) == b',' { 1 } else { 0 };
+        let b = $buf.get_unchecked(len..);
+        let found = if b.is_empty() {
+            false
+        }
+        $(
+        else if b.starts_with($str) {
+            len += $str.len();
+            true
+        }
+        )*
+        else {
+            false
+        };
+        if found {
+            let b = $buf.get_unchecked(len..);
+            if !b.is_empty() && b[0] == b'=' {
+                for c in b.iter() {
+                    if b" \n\0".contains(c) {
+                        break;
+                    }
+                    len += 1;
+                }
+            }
+            Some(len)
+        } else {
+            None
+        }
+    }};
+}
+
+fn remove_pattern(buf: &mut [u8], pattern_matcher: unsafe fn(&[u8]) -> Option<usize>) -> usize {
+    let mut write = 0_usize;
+    let mut read = 0_usize;
+    let mut sz = buf.len();
+    // SAFETY: assert(write <= read) && assert(read <= buf.len())
+    unsafe {
+        while read < buf.len() {
+            if let Some(len) = pattern_matcher(buf.get_unchecked(read..)) {
+                let skipped = buf.get_unchecked(read..(read + len));
+                // SAFETY: all matching patterns are ASCII bytes
+                let skipped = std::str::from_utf8_unchecked(skipped);
+                eprintln!("Remove pattern [{}]", skipped);
+                sz -= len;
+                read += len;
+            } else {
+                *buf.get_unchecked_mut(write) = *buf.get_unchecked(read);
+                write += 1;
+                read += 1;
+            }
+        }
+    }
+    if let Some(buf) = buf.get_mut(write..) {
+        buf.fill(0);
+    }
+    sz
+}
+
+pub fn patch_verity(buf: &mut [u8]) -> usize {
+    unsafe fn match_verity_pattern(buf: &[u8]) -> Option<usize> {
+        match_patterns!(
+            buf,
+            b"verifyatboot",
+            b"verify",
+            b"avb_keys",
+            b"avb",
+            b"support_scfs",
+            b"fsverity"
+        )
+    }
+
+    remove_pattern(buf, match_verity_pattern)
+}
+
+pub fn patch_encryption(buf: &mut [u8]) -> usize {
+    unsafe fn match_encryption_pattern(buf: &[u8]) -> Option<usize> {
+        match_patterns!(buf, b"forceencrypt", b"forcefdeorfbe", b"fileencryption")
+    }
+
+    remove_pattern(buf, match_encryption_pattern)
+}
+
 fn hex2byte(hex: &[u8]) -> Vec<u8> {
     let mut v = Vec::new();
     v.reserve(hex.len() / 2);

native/src/boot/pattern.cpp

@@ -1,66 +0,0 @@
-#include <base.hpp>
-
-#include "magiskboot.hpp"
-
-#define MATCH(p) else if (strncmp(s + skip, p, sizeof(p) - 1) == 0) skip += (sizeof(p) - 1)
-
-static int skip_verity_pattern(const char *s) {
-    int skip = s[0] == ',';
-
-    if (0) {}
-    MATCH("verifyatboot");
-    MATCH("verify");
-    MATCH("avb_keys");
-    MATCH("avb");
-    MATCH("support_scfs");
-    MATCH("fsverity");
-    else return -1;
-
-    if (s[skip] == '=') {
-        while (!strchr(" \n,", s[skip]))
-            ++skip;
-    }
-    return skip;
-}
-
-static int skip_encryption_pattern(const char *s) {
-    int skip = s[0] == ',';
-
-    if (0) {}
-    MATCH("forceencrypt");
-    MATCH("forcefdeorfbe");
-    MATCH("fileencryption");
-    else return -1;
-
-    if (s[skip] == '=') {
-        while (!strchr(" \n,", s[skip]))
-            ++skip;
-    }
-    return skip;
-}
-
-static size_t remove_pattern(byte_data data, int(*pattern_skip)(const char *)) {
-    char *src = reinterpret_cast<char *>(data.buf());
-    size_t sz = data.sz();
-    int write = 0;
-    int read = 0;
-    while (read < data.sz()) {
-        if (int skip = pattern_skip(src + read); skip > 0) {
-            fprintf(stderr, "Remove pattern [%.*s]\n", skip, src + read);
-            sz -= skip;
-            read += skip;
-        } else {
-            src[write++] = src[read++];
-        }
-    }
-    memset(src + write, 0, data.sz() - write);
-    return sz;
-}
-
-size_t patch_verity(rust::Slice<uint8_t> data) {
-    return remove_pattern(data, skip_verity_pattern);
-}
-
-size_t patch_encryption(rust::Slice<uint8_t> data) {
-    return remove_pattern(data, skip_encryption_pattern);
-}

native/src/boot/ramdisk.rs

@@ -7,7 +7,7 @@ use base::libc::{S_IFDIR, S_IFMT, S_IFREG};
 use base::Utf8CStr;
 
 use crate::cpio::{Cpio, CpioEntry};
-use crate::ffi::{patch_encryption, patch_verity};
+use crate::patch::{patch_encryption, patch_verity};
 
 pub trait MagiskCpio {
     fn patch(&mut self);