Merge changes I8c5ab552,If8546dea

* changes:
  libsnapshot: tests uses common MapUpdateSnapshot/WriteSnapshotAndHash
  libsnapshot: Add test for accounting for hash tree

fs_mgr/libsnapshot/snapshot_test.cpp

@@ -716,6 +716,45 @@ class SnapshotUpdateTest : public SnapshotTest {
         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_;
@@ -762,21 +801,7 @@ TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
 
     // Write some data to target partitions.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        std::string path;
+        ASSERT_TRUE(WriteSnapshotAndHash(name, partition_size));
-        ASSERT_TRUE(sm->MapUpdateSnapshot(
-                CreateLogicalPartitionParams{
-                        .block_device = fake_super,
-                        .metadata_slot = 1,
-                        .partition_name = name,
-                        .timeout_ms = 10s,
-                        .partition_opener = opener_.get(),
-                },
-                &path))
-                << name;
-        ASSERT_TRUE(WriteRandomData(path));
-        auto hash = GetHash(path);
-        ASSERT_TRUE(hash.has_value());
-        hashes_[name] = *hash;
     }
 
     // Assert that source partitions aren't affected.
@@ -890,17 +915,7 @@ TEST_F(SnapshotUpdateTest, SnapshotStatusFileWithoutCow) {
 
     // Check that target partitions can be mapped.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        std::string path;
+        EXPECT_TRUE(MapUpdateSnapshot(name));
-        EXPECT_TRUE(sm->MapUpdateSnapshot(
-                CreateLogicalPartitionParams{
-                        .block_device = fake_super,
-                        .metadata_slot = 1,
-                        .partition_name = name,
-                        .timeout_ms = 10s,
-                        .partition_opener = opener_.get(),
-                },
-                &path))
-                << name;
     }
 }
 
@@ -921,21 +936,7 @@ TEST_F(SnapshotUpdateTest, TestRollback) {
 
     // Write some data to target partitions.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        std::string path;
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
-        ASSERT_TRUE(sm->MapUpdateSnapshot(
-                CreateLogicalPartitionParams{
-                        .block_device = fake_super,
-                        .metadata_slot = 1,
-                        .partition_name = name,
-                        .timeout_ms = 10s,
-                        .partition_opener = opener_.get(),
-                },
-                &path))
-                << name;
-        ASSERT_TRUE(WriteRandomData(path));
-        auto hash = GetHash(path);
-        ASSERT_TRUE(hash.has_value());
-        hashes_[name] = *hash;
     }
 
     // Assert that source partitions aren't affected.
@@ -1089,21 +1090,7 @@ TEST_F(SnapshotUpdateTest, RetrofitAfterRegularAb) {
 
     // Write some data to target partitions.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        std::string path;
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
-        ASSERT_TRUE(sm->MapUpdateSnapshot(
-                CreateLogicalPartitionParams{
-                        .block_device = fake_super,
-                        .metadata_slot = 1,
-                        .partition_name = name,
-                        .timeout_ms = 10s,
-                        .partition_opener = opener_.get(),
-                },
-                &path))
-                << name;
-        ASSERT_TRUE(WriteRandomData(path));
-        auto hash = GetHash(path);
-        ASSERT_TRUE(hash.has_value());
-        hashes_[name] = *hash;
     }
 
     // Assert that source partitions aren't affected.
@@ -1313,6 +1300,56 @@ TEST_F(SnapshotUpdateTest, DataWipeAfterRollback) {
     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:

fs_mgr/libsnapshot/test_helpers.cpp

@@ -62,24 +62,6 @@ std::string TestPartitionOpener::GetDeviceString(const std::string& partition_na
     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 @@ std::string ToHexString(const uint8_t* buf, size_t len) {
     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)) {

fs_mgr/libsnapshot/test_helpers.h

@@ -137,8 +137,11 @@ class SnapshotTestPropertyFetcher : public android::fs_mgr::testing::MockPropert
 // 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.
+// Write some random data to the given device.
-bool WriteRandomData(const std::string& 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);