1 files changed, 96 insertions, 0 deletions
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);
+}