libsnapshot: Allow forward merge on FDR

If forward merge indicator exists at unverified state during FDR,
we are allowed to initiate merge and finish it before wiping userdata.

Test: downgrades
Test: update with data wipe
Test: vts_libsnapshot_test
Bug: 152094219
Change-Id: I151f1af4aca446cc52e669c5c3ea2b2e55acb468
(cherry picked from commit 973b8c10d1a74494e87bdda39bcfabf646abf0dd)
Merged-In: I151f1af4aca446cc52e669c5c3ea2b2e55acb468
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 08d17d1..1daa83b 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -529,6 +529,13 @@
     // allow forward merge on FDR.
     bool UpdateForwardMergeIndicator(bool wipe);
 
+    // Helper for HandleImminentDataWipe.
+    // Call ProcessUpdateState and handle states with special rules before data wipe. Specifically,
+    // if |allow_forward_merge| and allow-forward-merge indicator exists, initiate merge if
+    // necessary.
+    bool ProcessUpdateStateOnDataWipe(bool allow_forward_merge,
+                                      const std::function<bool()>& callback);
+
     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 f7df181..c9fa28e 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -2522,19 +2522,39 @@
         return false;
     }
 
-    UpdateState state = ProcessUpdateState([&]() -> bool {
+    auto process_callback = [&]() -> bool {
         if (callback) {
             callback();
         }
         return true;
-    });
+    };
+    if (!ProcessUpdateStateOnDataWipe(true /* allow_forward_merge */, process_callback)) {
+        return false;
+    }
+
+    // Nothing should be depending on partitions now, so unmap them all.
+    if (!UnmapAllPartitions()) {
+        LOG(ERROR) << "Unable to unmap all partitions; fastboot may fail to flash.";
+    }
+    return true;
+}
+
+bool SnapshotManager::ProcessUpdateStateOnDataWipe(bool allow_forward_merge,
+                                                   const std::function<bool()>& callback) {
+    auto slot_number = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    UpdateState state = ProcessUpdateState(callback);
     LOG(INFO) << "Update state in recovery: " << state;
     switch (state) {
         case UpdateState::MergeFailed:
             LOG(ERROR) << "Unrecoverable merge failure detected.";
             return false;
         case UpdateState::Unverified: {
-            // If an OTA was just applied but has not yet started merging, we
+            // If an OTA was just applied but has not yet started merging:
+            //
+            // - if forward merge is allowed, initiate merge and call
+            // ProcessUpdateState again.
+            //
+            // - if forward merge is not allowed, we
             // have no choice but to revert slots, because the current slot will
             // immediately become unbootable. Rather than wait for the device
             // to reboot N times until a rollback, we proactively disable the
@@ -2544,8 +2564,17 @@
             // as an error here.
             auto slot = GetCurrentSlot();
             if (slot == Slot::Target) {
+                if (allow_forward_merge &&
+                    access(GetForwardMergeIndicatorPath().c_str(), F_OK) == 0) {
+                    LOG(INFO) << "Forward merge allowed, initiating merge now.";
+                    return InitiateMerge() &&
+                           ProcessUpdateStateOnDataWipe(false /* allow_forward_merge */, callback);
+                }
+
                 LOG(ERROR) << "Reverting to old slot since update will be deleted.";
                 device_->SetSlotAsUnbootable(slot_number);
+            } else {
+                LOG(INFO) << "Booting from " << slot << " slot, no action is taken.";
             }
             break;
         }
@@ -2557,11 +2586,6 @@
         default:
             break;
     }
-
-    // Nothing should be depending on partitions now, so unmap them all.
-    if (!UnmapAllPartitions()) {
-        LOG(ERROR) << "Unable to unmap all partitions; fastboot may fail to flash.";
-    }
     return true;
 }
 
@@ -2643,10 +2667,15 @@
 }
 
 bool SnapshotManager::UpdateForwardMergeIndicator(bool wipe) {
+    auto path = GetForwardMergeIndicatorPath();
+
     if (!wipe) {
+        LOG(INFO) << "Wipe is not scheduled. Deleting forward merge indicator.";
         return RemoveFileIfExists(path);
     }
 
+    // TODO(b/152094219): Don't forward merge if no CoW file is allocated.
+
     LOG(INFO) << "Wipe will be scheduled. Allowing forward merge of snapshots.";
     if (!android::base::WriteStringToFile("1", path)) {
         PLOG(ERROR) << "Unable to write forward merge indicator: " << path;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 862add1..f82c082 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -1501,6 +1501,45 @@
     EXPECT_FALSE(test_device->IsSlotUnbootable(1));
 }
 
+// Test update package that requests data wipe.
+TEST_F(SnapshotUpdateTest, DataWipeRequiredInPackage) {
+    AddOperationForPartitions();
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name)) << name;
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(true /* wipe */));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_b");
+    test_device->set_recovery(true);
+    auto new_sm = SnapshotManager::NewForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    // Manually mount metadata so that we can call GetUpdateState() below.
+    MountMetadata();
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    ASSERT_FALSE(test_device->IsSlotUnbootable(1));
+    ASSERT_FALSE(test_device->IsSlotUnbootable(0));
+
+    // Now reboot into new slot.
+    test_device = new TestDeviceInfo(fake_super, "_b");
+    auto init = SnapshotManager::NewForFirstStageMount(test_device);
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    // Verify that we are on the downgraded build.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name)) << name;
+    }
+}
+
 TEST_F(SnapshotUpdateTest, Hashtree) {
     constexpr auto partition_size = 4_MiB;
     constexpr auto data_size = 3_MiB;