From 228570640eb42333d42b79d17e6301fc0513adb6 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Tue, 23 Nov 2021 22:14:12 -0800
Subject: [PATCH] Introduce KEEPVBMETAFLAG env variable

Close #4447, close #4906, close #4901, close #4964
---
 native/jni/magiskboot/bootimg.cpp |  4 +++-
 native/jni/magiskboot/main.cpp    | 22 ++++++++++++----------
 scripts/boot_patch.sh             |  3 +++
 scripts/util_functions.sh         |  2 ++
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/native/jni/magiskboot/bootimg.cpp b/native/jni/magiskboot/bootimg.cpp
index 7aac9afeb..2e2ff28a7 100644
--- a/native/jni/magiskboot/bootimg.cpp
+++ b/native/jni/magiskboot/bootimg.cpp
@@ -768,7 +768,9 @@ void repack(const char *src_img, const char *out_img, bool skip_comp) {
         memcpy(footer, boot.avb_footer, sizeof(AvbFooter));
         footer->original_image_size = __builtin_bswap64(off.total);
         footer->vbmeta_offset = __builtin_bswap64(off.vbmeta);
-        vbmeta->flags = __builtin_bswap32(3);
+        if (!check_env("KEEPVBMETAFLAG")) {
+            vbmeta->flags = __builtin_bswap32(3);
+        }
     }
 
     if (boot.flags[DHTB_FLAG]) {
diff --git a/native/jni/magiskboot/main.cpp b/native/jni/magiskboot/main.cpp
index b8cfa6e0b..5cc88caf8 100644
--- a/native/jni/magiskboot/main.cpp
+++ b/native/jni/magiskboot/main.cpp
@@ -6,7 +6,7 @@
 
 using namespace std;
 
-static void print_methods() {
+static void print_formats() {
     for (int fmt = GZIP; fmt < LZOP; ++fmt) {
         fprintf(stderr, "%s ", fmt2name[(format_t) fmt]);
     }
@@ -33,8 +33,10 @@ Supported actions:
     Repack boot image components from current directory
     to [outbootimg], or new-boot.img if not specified.
     If '-n' is provided, it will not attempt to recompress ramdisk.cpio,
-    otherwise it will compress ramdisk.cpio and kernel with the same method
-    in <origbootimg> if the file provided is not already compressed.
+    otherwise it will compress ramdisk.cpio and kernel with the same format
+    as in <origbootimg> if the file provided is not already compressed.
+    All disable flags will be set in the vbmeta header, unless env variable
+    KEEPVBMETAFLAG is set to true.
 
   hexpatch <file> <hexpattern1> <hexpattern2>
     Search <hexpattern1> in <file>, and replace with <hexpattern2>
@@ -91,21 +93,21 @@ Supported actions:
   cleanup
     Cleanup the current working directory
 
-  compress[=method] <infile> [outfile]
-    Compress <infile> with [method] (default: gzip), optionally to [outfile]
+  compress[=format] <infile> [outfile]
+    Compress <infile> with [format] (default: gzip), optionally to [outfile]
     <infile>/[outfile] can be '-' to be STDIN/STDOUT
-    Supported methods: )EOF", arg0);
+    Supported formats: )EOF", arg0);
 
-    print_methods();
+    print_formats();
 
     fprintf(stderr, R"EOF(
 
   decompress <infile> [outfile]
-    Detect method and decompress <infile>, optionally to [outfile]
+    Detect format and decompress <infile>, optionally to [outfile]
     <infile>/[outfile] can be '-' to be STDIN/STDOUT
-    Supported methods: )EOF");
+    Supported formats: )EOF");
 
-    print_methods();
+    print_formats();
 
     fprintf(stderr, "\n\n");
     exit(1);
diff --git a/scripts/boot_patch.sh b/scripts/boot_patch.sh
index 6084b7df1..721bd9511 100644
--- a/scripts/boot_patch.sh
+++ b/scripts/boot_patch.sh
@@ -70,9 +70,11 @@ fi
 # Flags
 [ -z $KEEPVERITY ] && KEEPVERITY=false
 [ -z $KEEPFORCEENCRYPT ] && KEEPFORCEENCRYPT=false
+[ -z $KEEPVBMETAFLAG ] && KEEPVBMETAFLAG=false
 [ -z $RECOVERYMODE ] && RECOVERYMODE=false
 export KEEPVERITY
 export KEEPFORCEENCRYPT
+export KEEPVBMETAFLAG
 
 chmod -R 755 .
 
@@ -149,6 +151,7 @@ ui_print "- Patching ramdisk"
 
 echo "KEEPVERITY=$KEEPVERITY" > config
 echo "KEEPFORCEENCRYPT=$KEEPFORCEENCRYPT" >> config
+echo "KEEPVBMETAFLAG=$KEEPVBMETAFLAG" >> config
 echo "RECOVERYMODE=$RECOVERYMODE" >> config
 [ ! -z $SHA1 ] && echo "SHA1=$SHA1" >> config
 
diff --git a/scripts/util_functions.sh b/scripts/util_functions.sh
index d7b9a4d39..b1022189b 100644
--- a/scripts/util_functions.sh
+++ b/scripts/util_functions.sh
@@ -382,6 +382,7 @@ get_flags() {
   getvar KEEPVERITY
   getvar KEEPFORCEENCRYPT
   getvar RECOVERYMODE
+  getvar KEEPVBMETAFLAG
   if [ -z $KEEPVERITY ]; then
     if $SYSTEM_ROOT; then
       KEEPVERITY=true
@@ -402,6 +403,7 @@ get_flags() {
       KEEPFORCEENCRYPT=false
     fi
   fi
+  [ -z $KEEPVBMETAFLAG ] && KEEPVBMETAFLAG=false
   [ -z $RECOVERYMODE ] && RECOVERYMODE=false
 }