From 9c7cf340a12213e7cb015cb5cd999b00dbfe5861 Mon Sep 17 00:00:00 2001 From: topjohnwu Date: Wed, 21 Jun 2023 16:47:20 -0700 Subject: [PATCH] Move pattern matching to Rust --- native/src/Android.mk | 1 - native/src/boot/dtb.cpp | 5 +- native/src/boot/lib.rs | 5 +- native/src/boot/magiskboot.hpp | 3 -- native/src/boot/patch.rs | 85 ++++++++++++++++++++++++++++++++++ native/src/boot/pattern.cpp | 66 -------------------------- native/src/boot/ramdisk.rs | 2 +- 7 files changed, 91 insertions(+), 76 deletions(-) delete mode 100644 native/src/boot/pattern.cpp diff --git a/native/src/Android.mk b/native/src/Android.mk index 0b691ba84..c3e557659 100644 --- a/native/src/Android.mk +++ b/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) diff --git a/native/src/boot/dtb.cpp b/native/src/boot/dtb.cpp index b030942c2..793f57f9a 100644 --- a/native/src/boot/dtb.cpp +++ b/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); diff --git a/native/src/boot/lib.rs b/native/src/boot/lib.rs index 6f14d0318..6182b4aa9 100644 --- a/native/src/boot/lib.rs +++ b/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; } } diff --git a/native/src/boot/magiskboot.hpp b/native/src/boot/magiskboot.hpp index bccdd0f1c..22ef9a28d 100644 --- a/native/src/boot/magiskboot.hpp +++ b/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 data); -size_t patch_encryption(rust::Slice data); - static inline bool check_env(const char *name) { using namespace std::string_view_literals; const char *val = getenv(name); diff --git a/native/src/boot/patch.rs b/native/src/boot/patch.rs index 817e07543..bf38f0154 100644 --- a/native/src/boot/patch.rs +++ b/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 { + 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 { + 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 { + match_patterns!(buf, b"forceencrypt", b"forcefdeorfbe", b"fileencryption") + } + + remove_pattern(buf, match_encryption_pattern) +} + fn hex2byte(hex: &[u8]) -> Vec { let mut v = Vec::new(); v.reserve(hex.len() / 2); diff --git a/native/src/boot/pattern.cpp b/native/src/boot/pattern.cpp deleted file mode 100644 index 8eca66aec..000000000 --- a/native/src/boot/pattern.cpp +++ /dev/null @@ -1,66 +0,0 @@ -#include - -#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(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 data) { - return remove_pattern(data, skip_verity_pattern); -} - -size_t patch_encryption(rust::Slice data) { - return remove_pattern(data, skip_encryption_pattern); -} diff --git a/native/src/boot/ramdisk.rs b/native/src/boot/ramdisk.rs index faf6f66fe..9aa7e5efe 100644 --- a/native/src/boot/ramdisk.rs +++ b/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);