Merge "libsnapshot: snaity check overflow bit at the end of update"
diff --git a/fs_mgr/libdm/dm.cpp b/fs_mgr/libdm/dm.cpp
index 809318c..254fbed 100644
--- a/fs_mgr/libdm/dm.cpp
+++ b/fs_mgr/libdm/dm.cpp
@@ -575,5 +575,9 @@
     return "/dev/block/" + sub_device_name;
 }
 
+bool DeviceMapper::TargetInfo::IsOverflowSnapshot() const {
+    return spec.target_type == "snapshot"s && data == "Overflow"s;
+}
+
 }  // namespace dm
 }  // namespace android
diff --git a/fs_mgr/libdm/include/libdm/dm.h b/fs_mgr/libdm/include/libdm/dm.h
index 418210c..abe9c4c 100644
--- a/fs_mgr/libdm/include/libdm/dm.h
+++ b/fs_mgr/libdm/include/libdm/dm.h
@@ -205,6 +205,8 @@
         TargetInfo() {}
         TargetInfo(const struct dm_target_spec& spec, const std::string& data)
             : spec(spec), data(data) {}
+
+        bool IsOverflowSnapshot() const;
     };
     bool GetTableStatus(const std::string& name, std::vector<TargetInfo>* table);
 
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 7450d19..5738b96 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -155,6 +155,7 @@
     // Mark snapshot writes as having completed. After this, new snapshots cannot
     // be created, and the device must either cancel the OTA (either before
     // rebooting or after rolling back), or merge the OTA.
+    // Before calling this function, all snapshots must be mapped.
     bool FinishedSnapshotWrites();
 
   private:
@@ -490,6 +491,11 @@
     // This should only be called in recovery.
     bool UnmapAllPartitions();
 
+    // Sanity check no snapshot overflows. Note that this returns false negatives if the snapshot
+    // overflows, then is remapped and not written afterwards. Hence, the function may only serve
+    // as a sanity check.
+    bool EnsureNoOverflowSnapshot(LockedFile* lock);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 830495c..f38db43 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -214,6 +214,11 @@
         return false;
     }
 
+    if (!EnsureNoOverflowSnapshot(lock.get())) {
+        LOG(ERROR) << "Cannot ensure there are no overflow snapshots.";
+        return false;
+    }
+
     // This file acts as both a quick indicator for init (it can use access(2)
     // to decide how to do first-stage mounts), and it stores the old slot, so
     // we can tell whether or not we performed a rollback.
@@ -2303,5 +2308,36 @@
     return true;
 }
 
+bool SnapshotManager::EnsureNoOverflowSnapshot(LockedFile* lock) {
+    CHECK(lock);
+
+    std::vector<std::string> snapshots;
+    if (!ListSnapshots(lock, &snapshots)) {
+        LOG(ERROR) << "Could not list snapshots.";
+        return false;
+    }
+
+    auto& dm = DeviceMapper::Instance();
+    for (const auto& snapshot : snapshots) {
+        std::vector<DeviceMapper::TargetInfo> targets;
+        if (!dm.GetTableStatus(snapshot, &targets)) {
+            LOG(ERROR) << "Could not read snapshot device table: " << snapshot;
+            return false;
+        }
+        if (targets.size() != 1) {
+            LOG(ERROR) << "Unexpected device-mapper table for snapshot: " << snapshot
+                       << ", size = " << targets.size();
+            return false;
+        }
+        if (targets[0].IsOverflowSnapshot()) {
+            LOG(ERROR) << "Detected overflow in snapshot " << snapshot
+                       << ", CoW device size computation is wrong!";
+            return false;
+        }
+    }
+
+    return true;
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index ff943f2..964b21a 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -273,6 +273,61 @@
         return AssertionSuccess();
     }
 
+    // Prepare A/B slot for a partition named "test_partition".
+    AssertionResult PrepareOneSnapshot(uint64_t device_size,
+                                       std::string* out_snap_device = nullptr) {
+        std::string base_device, cow_device, snap_device;
+        if (!CreatePartition("test_partition_a", device_size)) {
+            return AssertionFailure();
+        }
+        if (!MapUpdatePartitions()) {
+            return AssertionFailure();
+        }
+        if (!dm_.GetDmDevicePathByName("test_partition_b-base", &base_device)) {
+            return AssertionFailure();
+        }
+        SnapshotStatus status;
+        status.set_name("test_partition_b");
+        status.set_device_size(device_size);
+        status.set_snapshot_size(device_size);
+        status.set_cow_file_size(device_size);
+        if (!sm->CreateSnapshot(lock_.get(), &status)) {
+            return AssertionFailure();
+        }
+        if (!CreateCowImage("test_partition_b")) {
+            return AssertionFailure();
+        }
+        if (!MapCowImage("test_partition_b", 10s, &cow_device)) {
+            return AssertionFailure();
+        }
+        if (!sm->MapSnapshot(lock_.get(), "test_partition_b", base_device, cow_device, 10s,
+                             &snap_device)) {
+            return AssertionFailure();
+        }
+        if (out_snap_device) {
+            *out_snap_device = std::move(snap_device);
+        }
+        return AssertionSuccess();
+    }
+
+    // Simulate a reboot into the new slot.
+    AssertionResult SimulateReboot() {
+        lock_ = nullptr;
+        if (!sm->FinishedSnapshotWrites()) {
+            return AssertionFailure();
+        }
+        if (!dm_.DeleteDevice("test_partition_b")) {
+            return AssertionFailure();
+        }
+        if (!DestroyLogicalPartition("test_partition_b-base")) {
+            return AssertionFailure();
+        }
+        if (!sm->UnmapCowImage("test_partition_b")) {
+            return AssertionFailure();
+        }
+        return AssertionSuccess();
+    }
+
     DeviceMapper& dm_;
     std::unique_ptr<SnapshotManager::LockedFile> lock_;
     android::fiemap::IImageManager* image_manager_ = nullptr;
@@ -389,21 +444,8 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-
-    std::string base_device, cow_device, snap_device;
-    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
-    ASSERT_TRUE(MapUpdatePartitions());
-    ASSERT_TRUE(dm_.GetDmDevicePathByName("test_partition_b-base", &base_device));
-    SnapshotStatus status;
-    status.set_name("test_partition_b");
-    status.set_device_size(kDeviceSize);
-    status.set_snapshot_size(kDeviceSize);
-    status.set_cow_file_size(kDeviceSize);
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &status));
-    ASSERT_TRUE(CreateCowImage("test_partition_b"));
-    ASSERT_TRUE(MapCowImage("test_partition_b", 10s, &cow_device));
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test_partition_b", base_device, cow_device, 10s,
-                                &snap_device));
+    std::string snap_device;
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize, &snap_device));
 
     std::string test_string = "This is a test string.";
     {
@@ -455,21 +497,8 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-
-    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
-    ASSERT_TRUE(MapUpdatePartitions());
-    SnapshotStatus status;
-    status.set_name("test_partition_b");
-    status.set_device_size(kDeviceSize);
-    status.set_snapshot_size(kDeviceSize);
-    status.set_cow_file_size(kDeviceSize);
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &status));
-    ASSERT_TRUE(CreateCowImage("test_partition_b"));
-
-    // Simulate a reboot into the new slot.
-    lock_ = nullptr;
-    ASSERT_TRUE(sm->FinishedSnapshotWrites());
-    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
 
     auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
@@ -479,6 +508,7 @@
     ASSERT_TRUE(AcquireLock());
 
     // Validate that we have a snapshot device.
+    SnapshotStatus status;
     ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
     ASSERT_EQ(status.state(), SnapshotState::CREATED);
 
@@ -492,21 +522,8 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-
-    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
-    ASSERT_TRUE(MapUpdatePartitions());
-    SnapshotStatus status;
-    status.set_name("test_partition_b");
-    status.set_device_size(kDeviceSize);
-    status.set_snapshot_size(kDeviceSize);
-    status.set_cow_file_size(kDeviceSize);
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &status));
-    ASSERT_TRUE(CreateCowImage("test_partition_b"));
-
-    // Simulate a reboot into the new slot.
-    lock_ = nullptr;
-    ASSERT_TRUE(sm->FinishedSnapshotWrites());
-    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
 
     // Reflash the super partition.
     FormatFakeSuper();
@@ -519,6 +536,7 @@
 
     ASSERT_TRUE(AcquireLock());
 
+    SnapshotStatus status;
     ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
 
     // We should not get a snapshot device now.
@@ -535,21 +553,8 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-
-    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
-    ASSERT_TRUE(MapUpdatePartitions());
-    SnapshotStatus status;
-    status.set_name("test_partition_b");
-    status.set_device_size(kDeviceSize);
-    status.set_snapshot_size(kDeviceSize);
-    status.set_cow_file_size(kDeviceSize);
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &status));
-    ASSERT_TRUE(CreateCowImage("test_partition_b"));
-
-    // Simulate a reboot into the new slot.
-    lock_ = nullptr;
-    ASSERT_TRUE(sm->FinishedSnapshotWrites());
-    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
 
     auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
@@ -905,6 +910,17 @@
                                   << ", hash: " << hashes_[name];
     }
 
+    AssertionResult MapUpdateSnapshots(const std::vector<std::string>& names = {"sys_b", "vnd_b",
+                                                                                "prd_b"}) {
+        for (const auto& name : names) {
+            auto res = MapUpdateSnapshot(name);
+            if (!res) {
+                return res;
+            }
+        }
+        return AssertionSuccess();
+    }
+
     std::unique_ptr<TestPartitionOpener> opener_;
     DeltaArchiveManifest manifest_;
     std::unique_ptr<MetadataBuilder> src_;
@@ -1064,9 +1080,7 @@
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
 
     // Check that target partitions can be mapped.
-    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        EXPECT_TRUE(MapUpdateSnapshot(name));
-    }
+    EXPECT_TRUE(MapUpdateSnapshots());
 }
 
 // Test that the old partitions are not modified.
@@ -1142,6 +1156,7 @@
     // Execute the first update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.
@@ -1277,6 +1292,7 @@
     // Execute the update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.
@@ -1379,6 +1395,7 @@
     // Execute the first update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.
@@ -1410,6 +1427,7 @@
     // Execute the first update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.
@@ -1434,6 +1452,7 @@
     // Execute the first update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.
@@ -1480,7 +1499,8 @@
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
 
-    // Write some data to target partition.
+    // Map and write some data to target partition.
+    ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
     ASSERT_TRUE(WriteSnapshotAndHash("sys_b", partition_size));
 
     // Finish update.
@@ -1500,6 +1520,32 @@
     ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
 }
 
+// Test for overflow bit after update
+TEST_F(SnapshotUpdateTest, Overflow) {
+    const auto actual_write_size = GetSize(sys_);
+    const auto declared_write_size = actual_write_size - 1_MiB;
+
+    auto e = sys_->add_operations()->add_dst_extents();
+    e->set_start_block(0);
+    e->set_num_blocks(declared_write_size / manifest_.block_size());
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Map and write some data to target partitions.
+    ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
+    ASSERT_TRUE(WriteSnapshotAndHash("sys_b", actual_write_size));
+
+    std::vector<android::dm::DeviceMapper::TargetInfo> table;
+    ASSERT_TRUE(DeviceMapper::Instance().GetTableStatus("sys_b", &table));
+    ASSERT_EQ(1u, table.size());
+    EXPECT_TRUE(table[0].IsOverflowSnapshot());
+
+    ASSERT_FALSE(sm->FinishedSnapshotWrites())
+            << "FinishedSnapshotWrites should detect overflow of CoW device.";
+}
+
 class FlashAfterUpdateTest : public SnapshotUpdateTest,
                              public WithParamInterface<std::tuple<uint32_t, bool>> {
   public:
@@ -1524,7 +1570,7 @@
     // Execute the update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
-
+    ASSERT_TRUE(MapUpdateSnapshots());
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
 
     // Simulate shutting down the device.