Suport parsing hashtree and ecc data from avb

Add support to parse the AvbHashtreeFooter in the avb image. We are in
particular looking for the ecc and hashtree data so that they can laster
be used for error correction.

Bug: 144388532
Test: unittests pass
Change-Id: I3e53e3e84a58b00346274f9946cafc135702ac82

5 files changed, 265 insertions, 1 deletions

diff --git a/libfec/Android.bp b/libfec/Android.bp
index 78f2fedd..b355dfec 100644
--- a/libfec/Android.bp
+++ b/libfec/Android.bp
@@ -25,6 +25,7 @@ cc_library {
     ],
 
     static_libs: [
+        "libavb",
         "libfec_rs",
     ],
 
diff --git a/libfec/fec_open.cpp b/libfec/fec_open.cpp
index 18731f21..8f3fa22d 100644
--- a/libfec/fec_open.cpp
+++ b/libfec/fec_open.cpp
@@ -18,7 +18,9 @@
 #include <sys/ioctl.h>
 #include <sys/stat.h>
 
+#include <android-base/strings.h>
 #include <ext4_utils/ext4_sb.h>
+#include <libavb/libavb.h>
 #include <squashfs_utils.h>
 
 #if defined(__linux__)
@@ -493,6 +495,145 @@ int fec_get_status(struct fec_handle *f, struct fec_status *s)
     return 0;
 }
 
+static int parse_vbmeta_from_footer(fec_handle *f,
+                                    std::vector<uint8_t> *vbmeta) {
+    if (f->size <= AVB_FOOTER_SIZE) {
+        debug("file size not large enough to be avb images:" PRIu64, f->size);
+        return -1;
+    }
+
+    AvbFooter footer_read;
+    if (!raw_pread(f->fd, &footer_read, AVB_FOOTER_SIZE,
+                   f->size - AVB_FOOTER_SIZE)) {
+        error("failed to read footer: %s", strerror(errno));
+        return -1;
+    }
+
+    AvbFooter footer;
+    if (!avb_footer_validate_and_byteswap(&footer_read, &footer)) {
+        debug("invalid avb footer");
+        return -1;
+    }
+    uint64_t vbmeta_offset = footer.vbmeta_offset;
+    uint64_t vbmeta_size = footer.vbmeta_size;
+    check(vbmeta_offset <= f->size - sizeof(footer) - vbmeta_size);
+
+    std::vector<uint8_t> vbmeta_data(vbmeta_size, 0);
+    // TODO(xunchang) handle the sparse image with libsparse.
+    if (!raw_pread(f->fd, vbmeta_data.data(), vbmeta_data.size(),
+                   vbmeta_offset)) {
+        error("failed to read avb vbmeta: %s", strerror(errno));
+        return -1;
+    }
+
+    if (auto status = avb_vbmeta_image_verify(
+            vbmeta_data.data(), vbmeta_data.size(), nullptr, nullptr);
+        status != AVB_VBMETA_VERIFY_RESULT_OK &&
+        status != AVB_VBMETA_VERIFY_RESULT_OK_NOT_SIGNED) {
+        error("failed to verify avb vbmeta, status: %d", status);
+        return -1;
+    }
+    *vbmeta = std::move(vbmeta_data);
+    return 0;
+}
+
+static int parse_avb_image(fec_handle *f, const std::vector<uint8_t> &vbmeta) {
+    // TODO(xunchang) check if avb verification or hashtree is disabled.
+
+    // Look for the hashtree descriptor, we expect exactly one descriptor in
+    // vbmeta.
+    // TODO(xunchang) handle the image with AvbHashDescriptor.
+    auto parse_descriptor = [](const AvbDescriptor *descriptor,
+                               void *user_data) {
+        if (descriptor &&
+            avb_be64toh(descriptor->tag) == AVB_DESCRIPTOR_TAG_HASHTREE) {
+            auto desp = static_cast<const AvbDescriptor **>(user_data);
+            *desp = descriptor;
+            return false;
+        }
+        return true;
+    };
+
+    const AvbHashtreeDescriptor *hashtree_descriptor_ptr = nullptr;
+    avb_descriptor_foreach(vbmeta.data(), vbmeta.size(), parse_descriptor,
+                           &hashtree_descriptor_ptr);
+
+    AvbHashtreeDescriptor hashtree_descriptor;
+    if (!avb_hashtree_descriptor_validate_and_byteswap(hashtree_descriptor_ptr,
+                                                       &hashtree_descriptor)) {
+        error("failed to verify avb hashtree descriptor");
+        return -1;
+    }
+
+    // The partition name, salt, root append right after the hashtree
+    // descriptor.
+    auto read_ptr = reinterpret_cast<const uint8_t *>(hashtree_descriptor_ptr);
+    // Calculate the offset with respect to the vbmeta; and check both the
+    // salt & root are within the range.
+    uint32_t salt_offset =
+        sizeof(AvbHashtreeDescriptor) + hashtree_descriptor.partition_name_len;
+    uint32_t root_offset = salt_offset + hashtree_descriptor.salt_len;
+    check(hashtree_descriptor.salt_len < vbmeta.size());
+    check(salt_offset < vbmeta.size() - hashtree_descriptor.salt_len);
+    check(hashtree_descriptor.root_digest_len < vbmeta.size());
+    check(root_offset < vbmeta.size() - hashtree_descriptor.root_digest_len);
+    std::vector<uint8_t> salt(
+        read_ptr + salt_offset,
+        read_ptr + salt_offset + hashtree_descriptor.salt_len);
+    std::vector<uint8_t> root_hash(
+        read_ptr + root_offset,
+        read_ptr + root_offset + hashtree_descriptor.root_digest_len);
+
+    // Expect the AVB image has the format:
+    // 1. hashtree
+    // 2. ecc data
+    // 3. vbmeta
+    // 4. avb footer
+    check(hashtree_descriptor.fec_offset ==
+          hashtree_descriptor.tree_offset + hashtree_descriptor.tree_size);
+    check(hashtree_descriptor.fec_offset <=
+          f->size - hashtree_descriptor.fec_size);
+
+    f->data_size = hashtree_descriptor.fec_offset;
+
+    f->ecc.blocks = fec_div_round_up(f->data_size, FEC_BLOCKSIZE);
+    f->ecc.rounds = fec_div_round_up(f->ecc.blocks, f->ecc.rsn);
+    f->ecc.size = hashtree_descriptor.fec_size;
+    f->ecc.start = hashtree_descriptor.fec_offset;
+    // TODO(xunchang) verify the integrity of the ecc data.
+    f->ecc.valid = true;
+
+    std::string hash_algorithm =
+        reinterpret_cast<char *>(hashtree_descriptor.hash_algorithm);
+    int nid = -1;
+    if (android::base::EqualsIgnoreCase(hash_algorithm, "sha1")) {
+        nid = NID_sha1;
+    } else if (android::base::EqualsIgnoreCase(hash_algorithm, "sha256")) {
+        nid = NID_sha256;
+    } else {
+        error("unsupported hash algorithm %s", hash_algorithm.c_str());
+    }
+
+    hashtree_info hashtree;
+    hashtree.initialize(hashtree_descriptor.tree_offset,
+                        hashtree_descriptor.tree_offset / FEC_BLOCKSIZE, salt,
+                        nid);
+    if (hashtree.verify_tree(f, root_hash.data()) != 0) {
+        error("failed to verify hashtree");
+        return -1;
+    }
+
+    // We have validate the hashtree,
+    f->data_size = hashtree.hash_start;
+    f->avb = {
+        .valid = true,
+        .vbmeta = vbmeta,
+        .hashtree = std::move(hashtree),
+    };
+
+    return 0;
+}
+
 /* opens `path' using given options and returns a fec_handle in `handle' if
    successful */
 int fec_open(struct fec_handle **handle, const char *path, int mode, int flags,
@@ -546,6 +687,21 @@ int fec_open(struct fec_handle **handle, const char *path, int mode, int flags,
 
     f->data_size = f->size; /* until ecc and/or verity are loaded */
 
+    std::vector<uint8_t> vbmeta;
+    if (parse_vbmeta_from_footer(f.get(), &vbmeta) == 0) {
+        if (parse_avb_image(f.get(), vbmeta) != 0) {
+            error("failed to parse avb image.");
+            return -1;
+        }
+
+        *handle = f.release();
+        return 0;
+    }
+    // TODO(xunchang) For android, handle the case when vbmeta is in a separate
+    // image. We could use avb_slot_verify() && AvbOps from libavb_user.
+
+    // Fall back to use verity format.
+
     if (load_ecc(f.get()) == -1) {
         debug("error-correcting codes not found from '%s'", path);
     }
diff --git a/libfec/fec_private.h b/libfec/fec_private.h
index 62687994..b28c4294 100644
--- a/libfec/fec_private.h
+++ b/libfec/fec_private.h
@@ -124,6 +124,12 @@ struct verity_info {
     verity_header ecc_header;
 };
 
+struct avb_info {
+    bool valid = false;
+    std::vector<uint8_t> vbmeta;
+    hashtree_info hashtree;
+};
+
 struct fec_handle {
     ecc_info ecc;
     int fd;
@@ -134,10 +140,12 @@ struct fec_handle {
     uint64_t data_size;
     uint64_t pos;
     uint64_t size;
+    // TODO(xunchang) switch to std::optional
     verity_info verity;
+    avb_info avb;
 
     hashtree_info hashtree() const {
-        return verity.hashtree;
+        return avb.valid ? avb.hashtree : verity.hashtree;
     }
 };
 
diff --git a/libfec/test/Android.bp b/libfec/test/Android.bp
index 1bbad077..75182a27 100644
--- a/libfec/test/Android.bp
+++ b/libfec/test/Android.bp
@@ -24,6 +24,7 @@ cc_test_host {
     static_libs: [
         "libfec",
         "libfec_rs",
+        "libavb",
         "libcrypto_utils",
         "libcrypto",
         "libext4_utils",
@@ -47,12 +48,14 @@ cc_test_host {
 
     gtest: true,
     required: [
+        "avbtool",
         "fec",
     ],
     static_libs: [
         "libverity_tree",
         "libfec",
         "libfec_rs",
+        "libavb",
         "libcrypto_utils",
         "libext4_utils",
         "libsquashfs_utils",
diff --git a/libfec/test/fec_unittest.cpp b/libfec/test/fec_unittest.cpp
index ca5d91a8..421eb501 100644
--- a/libfec/test/fec_unittest.cpp
+++ b/libfec/test/fec_unittest.cpp
@@ -102,6 +102,20 @@ class FecUnitTest : public ::testing::Test {
         ASSERT_EQ(0, std::system(android::base::Join(cmd, ' ').c_str()));
     }
 
+    void AddAvbHashtreeFooter(const std::string &image_name,
+                              std::string algorithm = "sha256") {
+        salt_ = std::vector<uint8_t>(64, 10);
+        std::vector<std::string> cmd = {
+            "avbtool",          "add_hashtree_footer",
+            "--salt",           HashTreeBuilder::BytesArrayToString(salt_),
+            "--hash_algorithm", algorithm,
+            "--image",          image_name,
+        };
+        ASSERT_EQ(0, std::system(android::base::Join(cmd, ' ').c_str()));
+
+        BuildHashtree(algorithm);
+    }
+
     std::vector<uint8_t> image_;
     std::vector<uint8_t> salt_;
     std::vector<uint8_t> root_hash_;
@@ -202,3 +216,85 @@ TEST_F(FecUnitTest, VerityImage_FecRead) {
     ASSERT_EQ(1024, fec_pread(handle, read_data.data(), 1024, corrupt_offset));
     ASSERT_EQ(std::vector<uint8_t>(1024, 255), read_data);
 }
+
+TEST_F(FecUnitTest, LoadAvbImage_HashtreeFooter) {
+    TemporaryFile avb_image;
+    ASSERT_TRUE(
+        android::base::WriteFully(avb_image.fd, image_.data(), image_.size()));
+    AddAvbHashtreeFooter(avb_image.path);
+
+    struct fec_handle *handle = nullptr;
+    ASSERT_EQ(0, fec_open(&handle, avb_image.path, O_RDWR, FEC_FS_EXT4, 2));
+    std::unique_ptr<fec_handle> guard(handle);
+
+    ASSERT_EQ(1024 * 1024, handle->data_size);  // filesystem size
+
+    ASSERT_TRUE(handle->avb.valid);
+
+    // check the hashtree.
+    ASSERT_EQ(salt_, handle->hashtree().salt);
+    ASSERT_EQ(1024 * 1024, handle->hashtree().hash_start);
+    // the fec hashtree only stores the hash of the lowest level.
+    ASSERT_EQ(std::vector<uint8_t>(hashtree_content_.begin() + 4096,
+                                   hashtree_content_.end()),
+              handle->hashtree().hash_data);
+    uint64_t hash_size =
+        verity_get_size(handle->hashtree().data_blocks * FEC_BLOCKSIZE, nullptr,
+                        nullptr, SHA256_DIGEST_LENGTH);
+    ASSERT_EQ(hashtree_content_.size(), hash_size);
+
+    fec_ecc_metadata ecc_metadata{};
+    ASSERT_EQ(0, fec_ecc_get_metadata(handle, &ecc_metadata));
+    ASSERT_TRUE(ecc_metadata.valid);
+    ASSERT_EQ(1024 * 1024 + hash_size, ecc_metadata.start);
+    ASSERT_EQ(259, ecc_metadata.blocks);
+}
+
+TEST_F(FecUnitTest, LoadAvbImage_CorrectHashtree) {
+    TemporaryFile avb_image;
+    ASSERT_TRUE(
+        android::base::WriteFully(avb_image.fd, image_.data(), image_.size()));
+    AddAvbHashtreeFooter(avb_image.path);
+
+    uint64_t corrupt_offset = 1024 * 1024 + 2 * 4096 + 50;
+    ASSERT_EQ(corrupt_offset, lseek64(avb_image.fd, corrupt_offset, 0));
+    std::vector<uint8_t> corruption(20, 5);
+    ASSERT_TRUE(android::base::WriteFully(avb_image.fd, corruption.data(),
+                                          corruption.size()));
+
+    struct fec_handle *handle = nullptr;
+    ASSERT_EQ(0, fec_open(&handle, avb_image.path, O_RDWR, FEC_FS_EXT4, 2));
+    std::unique_ptr<fec_handle> guard(handle);
+
+    ASSERT_EQ(1024 * 1024, handle->data_size);  // filesystem size
+    fec_ecc_metadata ecc_metadata{};
+    ASSERT_EQ(0, fec_ecc_get_metadata(handle, &ecc_metadata));
+    ASSERT_TRUE(ecc_metadata.valid);
+}
+
+TEST_F(FecUnitTest, AvbImage_FecRead) {
+    TemporaryFile avb_image;
+    ASSERT_TRUE(
+        android::base::WriteFully(avb_image.fd, image_.data(), image_.size()));
+    AddAvbHashtreeFooter(avb_image.path, "sha1");
+
+    uint64_t corrupt_offset = 4096 * 10;
+    ASSERT_EQ(corrupt_offset, lseek64(avb_image.fd, corrupt_offset, 0));
+    std::vector<uint8_t> corruption(50, 99);
+    ASSERT_TRUE(android::base::WriteFully(avb_image.fd, corruption.data(),
+                                          corruption.size()));
+
+    std::vector<uint8_t> read_data(1024, 0);
+    struct fec_handle *handle = nullptr;
+    ASSERT_EQ(0, fec_open(&handle, avb_image.path, O_RDWR, FEC_FS_EXT4, 2));
+    std::unique_ptr<fec_handle> guard(handle);
+
+    // Verify the hashtree has the expected content.
+    ASSERT_EQ(std::vector<uint8_t>(hashtree_content_.begin() + 4096,
+                                   hashtree_content_.end()),
+              handle->hashtree().hash_data);
+
+    // Verify the corruption gets corrected.
+    ASSERT_EQ(1024, fec_pread(handle, read_data.data(), 1024, corrupt_offset));
+    ASSERT_EQ(std::vector<uint8_t>(1024, 10), read_data);
+}