libsnapshot: Add test for accounting for hash tree

Test: libsnapshot_Test
Bug: 145180464
Change-Id: If8546dea89fdd7ec7499522a232a777699c52d82
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 7d6e78f..51b969b 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -716,6 +716,45 @@
         return AssertionSuccess();
     }
 
+    AssertionResult MapUpdateSnapshot(const std::string& name, std::string* path = nullptr) {
+        std::string real_path;
+        if (!sm->MapUpdateSnapshot(
+                    CreateLogicalPartitionParams{
+                            .block_device = fake_super,
+                            .metadata_slot = 1,
+                            .partition_name = name,
+                            .timeout_ms = 10s,
+                            .partition_opener = opener_.get(),
+                    },
+                    &real_path)) {
+            return AssertionFailure() << "Unable to map snapshot " << name;
+        }
+        if (path) {
+            *path = real_path;
+        }
+        return AssertionSuccess() << "Mapped snapshot " << name << " to " << real_path;
+    }
+
+    AssertionResult WriteSnapshotAndHash(const std::string& name,
+                                         std::optional<size_t> size = std::nullopt) {
+        std::string path;
+        auto res = MapUpdateSnapshot(name, &path);
+        if (!res) {
+            return res;
+        }
+
+        std::string size_string = size ? (std::to_string(*size) + " bytes") : "";
+
+        if (!WriteRandomData(path, size, &hashes_[name])) {
+            return AssertionFailure() << "Unable to write " << size_string << " to " << path
+                                      << " for partition " << name;
+        }
+
+        return AssertionSuccess() << "Written " << size_string << " to " << path
+                                  << " for snapshot partition " << name
+                                  << ", hash: " << hashes_[name];
+    }
+
     std::unique_ptr<TestPartitionOpener> opener_;
     DeltaArchiveManifest manifest_;
     std::unique_ptr<MetadataBuilder> src_;
@@ -1313,6 +1352,56 @@
     EXPECT_FALSE(test_device->IsSlotUnbootable(0));
 }
 
+TEST_F(SnapshotUpdateTest, Hashtree) {
+    constexpr auto partition_size = 4_MiB;
+    constexpr auto data_size = 3_MiB;
+    constexpr auto hashtree_size = 512_KiB;
+    constexpr auto fec_size = partition_size - data_size - hashtree_size;
+
+    const auto block_size = manifest_.block_size();
+    SetSize(sys_, partition_size);
+
+    auto e = sys_->add_operations()->add_dst_extents();
+    e->set_start_block(0);
+    e->set_num_blocks(data_size / block_size);
+
+    // Set hastree extents.
+    sys_->mutable_hash_tree_data_extent()->set_start_block(0);
+    sys_->mutable_hash_tree_data_extent()->set_num_blocks(data_size / block_size);
+
+    sys_->mutable_hash_tree_extent()->set_start_block(data_size / block_size);
+    sys_->mutable_hash_tree_extent()->set_num_blocks(hashtree_size / block_size);
+
+    // Set FEC extents.
+    sys_->mutable_fec_data_extent()->set_start_block(0);
+    sys_->mutable_fec_data_extent()->set_num_blocks((data_size + hashtree_size) / block_size);
+
+    sys_->mutable_fec_extent()->set_start_block((data_size + hashtree_size) / block_size);
+    sys_->mutable_fec_extent()->set_num_blocks(fec_size / block_size);
+
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Write some data to target partition.
+    ASSERT_TRUE(WriteSnapshotAndHash("sys_b", partition_size));
+
+    // Finish update.
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+
+    // Check that the target partition have the same content. Hashtree and FEC extents
+    // should be accounted for.
+    ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
+}
+
 class FlashAfterUpdateTest : public SnapshotUpdateTest,
                              public WithParamInterface<std::tuple<uint32_t, bool>> {
   public:
diff --git a/fs_mgr/libsnapshot/test_helpers.cpp b/fs_mgr/libsnapshot/test_helpers.cpp
index 312fa3e..2d62347 100644
--- a/fs_mgr/libsnapshot/test_helpers.cpp
+++ b/fs_mgr/libsnapshot/test_helpers.cpp
@@ -62,24 +62,6 @@
     return PartitionOpener::GetDeviceString(partition_name);
 }
 
-bool WriteRandomData(const std::string& path) {
-    unique_fd rand(open("/dev/urandom", O_RDONLY));
-    unique_fd fd(open(path.c_str(), O_WRONLY));
-
-    char buf[4096];
-    while (true) {
-        ssize_t n = TEMP_FAILURE_RETRY(read(rand.get(), buf, sizeof(buf)));
-        if (n <= 0) return false;
-        if (!WriteFully(fd.get(), buf, n)) {
-            if (errno == ENOSPC) {
-                return true;
-            }
-            PLOG(ERROR) << "Cannot write " << path;
-            return false;
-        }
-    }
-}
-
 std::string ToHexString(const uint8_t* buf, size_t len) {
     char lookup[] = "0123456789abcdef";
     std::string out(len * 2 + 1, '\0');
@@ -91,6 +73,47 @@
     return out;
 }
 
+bool WriteRandomData(const std::string& path, std::optional<size_t> expect_size,
+                     std::string* hash) {
+    unique_fd rand(open("/dev/urandom", O_RDONLY));
+    unique_fd fd(open(path.c_str(), O_WRONLY));
+
+    SHA256_CTX ctx;
+    if (hash) {
+        SHA256_Init(&ctx);
+    }
+
+    char buf[4096];
+    size_t total_written = 0;
+    while (!expect_size || total_written < *expect_size) {
+        ssize_t n = TEMP_FAILURE_RETRY(read(rand.get(), buf, sizeof(buf)));
+        if (n <= 0) return false;
+        if (!WriteFully(fd.get(), buf, n)) {
+            if (errno == ENOSPC) {
+                break;
+            }
+            PLOG(ERROR) << "Cannot write " << path;
+            return false;
+        }
+        total_written += n;
+        if (hash) {
+            SHA256_Update(&ctx, buf, n);
+        }
+    }
+
+    if (expect_size && total_written != *expect_size) {
+        PLOG(ERROR) << "Written " << total_written << " bytes, expected " << *expect_size;
+        return false;
+    }
+
+    if (hash) {
+        uint8_t out[32];
+        SHA256_Final(out, &ctx);
+        *hash = ToHexString(out, sizeof(out));
+    }
+    return true;
+}
+
 std::optional<std::string> GetHash(const std::string& path) {
     std::string content;
     if (!android::base::ReadFileToString(path, &content, true)) {
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index 9083843..2bf1b57 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -137,8 +137,11 @@
 // Helper for error-spam-free cleanup.
 void DeleteBackingImage(android::fiemap::IImageManager* manager, const std::string& name);
 
-// Write some random data to the given device. Will write until reaching end of the device.
-bool WriteRandomData(const std::string& device);
+// Write some random data to the given device.
+// If expect_size is not specified, will write until reaching end of the device.
+// Expect space of |path| is multiple of 4K.
+bool WriteRandomData(const std::string& path, std::optional<size_t> expect_size = std::nullopt,
+                     std::string* hash = nullptr);
 
 std::optional<std::string> GetHash(const std::string& path);