1 files changed, 38 insertions, 8 deletions

diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index a7fc628c4..2fe3b1a9e 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -51,6 +51,7 @@
 #include <utility>
 #include <vector>
 
+#include <android-base/endian.h>
 #include <android-base/file.h>
 #include <android-base/macros.h>
 #include <android-base/parseint.h>
@@ -59,6 +60,7 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <build/version.h>
+#include <libavb/libavb.h>
 #include <liblp/liblp.h>
 #include <platform_tools_version.h>
 #include <sparse/sparse.h>
@@ -919,33 +921,50 @@ static bool load_buf(const char* fname, struct fastboot_buffer* buf) {
     return load_buf_fd(fd.release(), buf);
 }
 
-static void rewrite_vbmeta_buffer(struct fastboot_buffer* buf) {
+static void rewrite_vbmeta_buffer(struct fastboot_buffer* buf, bool vbmeta_in_boot) {
     // Buffer needs to be at least the size of the VBMeta struct which
     // is 256 bytes.
     if (buf->sz < 256) {
         return;
     }
 
-    int fd = make_temporary_fd("vbmeta rewriting");
-
     std::string data;
     if (!android::base::ReadFdToString(buf->fd, &data)) {
         die("Failed reading from vbmeta");
     }
 
+    uint64_t vbmeta_offset = 0;
+    if (vbmeta_in_boot) {
+        // Tries to locate top-level vbmeta from boot.img footer.
+        uint64_t footer_offset = buf->sz - AVB_FOOTER_SIZE;
+        if (0 != data.compare(footer_offset, AVB_FOOTER_MAGIC_LEN, AVB_FOOTER_MAGIC)) {
+            die("Failed to find AVB_FOOTER at offset: %" PRId64, footer_offset);
+        }
+        const AvbFooter* footer = reinterpret_cast<const AvbFooter*>(data.c_str() + footer_offset);
+        vbmeta_offset = be64toh(footer->vbmeta_offset);
+    }
+    // Ensures there is AVB_MAGIC at vbmeta_offset.
+    if (0 != data.compare(vbmeta_offset, AVB_MAGIC_LEN, AVB_MAGIC)) {
+        die("Failed to find AVB_MAGIC at offset: %" PRId64, vbmeta_offset);
+    }
+
+    fprintf(stderr, "Rewriting vbmeta struct at offset: %" PRId64 "\n", vbmeta_offset);
+
     // There's a 32-bit big endian |flags| field at offset 120 where
     // bit 0 corresponds to disable-verity and bit 1 corresponds to
     // disable-verification.
     //
     // See external/avb/libavb/avb_vbmeta_image.h for the layout of
     // the VBMeta struct.
+    uint64_t flags_offset = 123 + vbmeta_offset;
     if (g_disable_verity) {
-        data[123] |= 0x01;
+        data[flags_offset] |= 0x01;
     }
     if (g_disable_verification) {
-        data[123] |= 0x02;
+        data[flags_offset] |= 0x02;
     }
 
+    int fd = make_temporary_fd("vbmeta rewriting");
     if (!android::base::WriteStringToFd(data, fd)) {
         die("Failed writing to modified vbmeta");
     }
@@ -954,14 +973,25 @@ static void rewrite_vbmeta_buffer(struct fastboot_buffer* buf) {
     lseek(fd, 0, SEEK_SET);
 }
 
+static bool has_vbmeta_partition() {
+    std::string partition_type;
+    return fb->GetVar("partition-type:vbmeta", &partition_type) == fastboot::SUCCESS ||
+           fb->GetVar("partition-type:vbmeta_a", &partition_type) == fastboot::SUCCESS ||
+           fb->GetVar("partition-type:vbmeta_b", &partition_type) == fastboot::SUCCESS;
+}
+
 static void flash_buf(const std::string& partition, struct fastboot_buffer *buf)
 {
     sparse_file** s;
 
     // Rewrite vbmeta if that's what we're flashing and modification has been requested.
-    if ((g_disable_verity || g_disable_verification) &&
-        (partition == "vbmeta" || partition == "vbmeta_a" || partition == "vbmeta_b")) {
-        rewrite_vbmeta_buffer(buf);
+    if (g_disable_verity || g_disable_verification) {
+        if (partition == "vbmeta" || partition == "vbmeta_a" || partition == "vbmeta_b") {
+            rewrite_vbmeta_buffer(buf, false /* vbmeta_in_boot */);
+        } else if (!has_vbmeta_partition() &&
+                   (partition == "boot" || partition == "boot_a" || partition == "boot_b")) {
+            rewrite_vbmeta_buffer(buf, true /* vbmeta_in_boot */ );
+        }
     }
 
     switch (buf->type) {