From a9b3443e5348fb14e78b32b392fe2887584fd28c Mon Sep 17 00:00:00 2001
From: Steve Muckle <smuckle@google.com>
Date: Tue, 12 May 2020 16:21:41 -0700
Subject: [PATCH] fastbootd: copy AVB footer on boot image to end of block
 device

If the flashed boot image is smaller than the block device, the AVB
footer will not be at the end of the partition. Although images are
normally created to match the partition size the GKI boot.img must work
on all devices, and the size of the boot partition will vary.

Copy the AVB footer to the end of the partition before flashing, if it
is not there already.

Bug: 156036850
Change-Id: I11f0c7d32d1b6c74edd4f84f815d175605280cb8
---
 fastboot/Android.bp          |  1 +
 fastboot/device/flashing.cpp | 30 +++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index 3a2deb7dd..6b49fc74e 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -149,6 +149,7 @@ cc_binary {
     ],
 
     header_libs: [
+        "avb_headers",
         "libsnapshot_headers",
     ]
 }
diff --git a/fastboot/device/flashing.cpp b/fastboot/device/flashing.cpp
index fd6ff8ea4..1bf4c9c9c 100644
--- a/fastboot/device/flashing.cpp
+++ b/fastboot/device/flashing.cpp
@@ -31,6 +31,7 @@
 #include <ext4_utils/ext4_utils.h>
 #include <fs_mgr_overlayfs.h>
 #include <fstab/fstab.h>
+#include <libavb/libavb.h>
 #include <liblp/builder.h>
 #include <liblp/liblp.h>
 #include <libsnapshot/snapshot.h>
@@ -122,6 +123,27 @@ int FlashBlockDevice(int fd, std::vector<char>& downloaded_data) {
     }
 }
 
+static void CopyAVBFooter(std::vector<char>* data, const uint64_t block_device_size) {
+    if (data->size() < AVB_FOOTER_SIZE) {
+        return;
+    }
+    std::string footer;
+    uint64_t footer_offset = data->size() - AVB_FOOTER_SIZE;
+    for (int idx = 0; idx < AVB_FOOTER_MAGIC_LEN; idx++) {
+        footer.push_back(data->at(footer_offset + idx));
+    }
+    if (0 != footer.compare(AVB_FOOTER_MAGIC)) {
+        return;
+    }
+
+    // copy AVB footer from end of data to end of block device
+    uint64_t original_data_size = data->size();
+    data->resize(block_device_size, 0);
+    for (int idx = 0; idx < AVB_FOOTER_SIZE; idx++) {
+        data->at(block_device_size - 1 - idx) = data->at(original_data_size - 1 - idx);
+    }
+}
+
 int Flash(FastbootDevice* device, const std::string& partition_name) {
     PartitionHandle handle;
     if (!OpenPartition(device, partition_name, &handle)) {
@@ -131,8 +153,14 @@ int Flash(FastbootDevice* device, const std::string& partition_name) {
     std::vector<char> data = std::move(device->download_data());
     if (data.size() == 0) {
         return -EINVAL;
-    } else if (data.size() > get_block_device_size(handle.fd())) {
+    }
+    uint64_t block_device_size = get_block_device_size(handle.fd());
+    if (data.size() > block_device_size) {
         return -EOVERFLOW;
+    } else if (data.size() < block_device_size &&
+               (partition_name == "boot" || partition_name == "boot_a" ||
+                partition_name == "boot_b")) {
+        CopyAVBFooter(&data, block_device_size);
     }
     WipeOverlayfsForPartition(device, partition_name);
     int result = FlashBlockDevice(handle.fd(), data);