Merge "fastboot: Fix snapshot-update merge behavior." into rvc-dev am: 4adbeb470f

Change-Id: I89477c995a82f6f4fbd38d3393538feb1c4e3c16

fastboot/device/commands.cpp

@@ -625,7 +625,7 @@ bool SnapshotUpdateHandler(FastbootDevice* device, const std::vector<std::string
         if (!sm) {
             return device->WriteFail("Unable to create SnapshotManager");
         }
-        if (!sm->HandleImminentDataWipe()) {
+        if (!sm->FinishMergeInRecovery()) {
             return device->WriteFail("Unable to finish snapshot merge");
         }
     } else {

fs_mgr/libsnapshot/include/libsnapshot/snapshot.h

@@ -223,6 +223,10 @@ class SnapshotManager final {
     // optional callback fires periodically to query progress via GetUpdateState.
     bool HandleImminentDataWipe(const std::function<void()>& callback = {});
 
+    // Force a merge to complete in recovery. This is similar to HandleImminentDataWipe
+    // but does not expect a data wipe after.
+    bool FinishMergeInRecovery();
+
     // This method is only allowed in recovery and is used as a helper to
     // initialize the snapshot devices as a requirement to mount a snapshotted
     // /system in recovery.
@@ -541,6 +545,7 @@ class SnapshotManager final {
     std::unique_ptr<IDeviceInfo> device_;
     std::unique_ptr<IImageManager> images_;
     bool has_local_image_manager_ = false;
+    bool in_factory_data_reset_ = false;
 };
 
 }  // namespace snapshot

fs_mgr/libsnapshot/snapshot.cpp

@@ -577,8 +577,16 @@ bool SnapshotManager::InitiateMerge() {
         return false;
     }
 
+    auto other_suffix = device_->GetOtherSlotSuffix();
+
     auto& dm = DeviceMapper::Instance();
     for (const auto& snapshot : snapshots) {
+        if (android::base::EndsWith(snapshot, other_suffix)) {
+            // Allow the merge to continue, but log this unexpected case.
+            LOG(ERROR) << "Unexpected snapshot found during merge: " << snapshot;
+            continue;
+        }
+
         // The device has to be mapped, since everything should be merged at
         // the same time. This is a fairly serious error. We could forcefully
         // map everything here, but it should have been mapped during first-
@@ -1008,6 +1016,15 @@ std::string SnapshotManager::GetForwardMergeIndicatorPath() {
 }
 
 void SnapshotManager::AcknowledgeMergeSuccess(LockedFile* lock) {
+    // It's not possible to remove update state in recovery, so write an
+    // indicator that cleanup is needed on reboot. If a factory data reset
+    // was requested, it doesn't matter, everything will get wiped anyway.
+    // To make testing easier we consider a /data wipe as cleaned up.
+    if (device_->IsRecovery() && !in_factory_data_reset_) {
+        WriteUpdateState(lock, UpdateState::MergeCompleted);
+        return;
+    }
+
     RemoveAllUpdateState(lock);
 }
 
@@ -2528,7 +2545,43 @@ bool SnapshotManager::HandleImminentDataWipe(const std::function<void()>& callba
         }
         return true;
     };
-    if (!ProcessUpdateStateOnDataWipe(true /* allow_forward_merge */, process_callback)) {
+
+    in_factory_data_reset_ = true;
+    bool ok = ProcessUpdateStateOnDataWipe(true /* allow_forward_merge */, process_callback);
+    in_factory_data_reset_ = false;
+
+    if (!ok) {
+        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::FinishMergeInRecovery() {
+    if (!device_->IsRecovery()) {
+        LOG(ERROR) << "Data wipes are only allowed in recovery.";
+        return false;
+    }
+
+    auto mount = EnsureMetadataMounted();
+    if (!mount || !mount->HasDevice()) {
+        return false;
+    }
+
+    auto slot_number = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    auto super_path = device_->GetSuperDevice(slot_number);
+    if (!CreateLogicalAndSnapshotPartitions(super_path)) {
+        LOG(ERROR) << "Unable to map partitions to complete merge.";
+        return false;
+    }
+
+    UpdateState state = ProcessUpdateState();
+    if (state != UpdateState::MergeCompleted) {
+        LOG(ERROR) << "Merge returned unexpected status: " << state;
         return false;
     }
 

fs_mgr/libsnapshot/snapshot_test.cpp

@@ -1454,6 +1454,52 @@ TEST_F(SnapshotUpdateTest, MergeInRecovery) {
     ASSERT_EQ(new_sm->GetUpdateState(), UpdateState::None);
 }
 
+// Test that a merge does not clear the snapshot state in fastboot.
+TEST_F(SnapshotUpdateTest, MergeInFastboot) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // 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", snapshot_timeout_));
+    init = nullptr;
+
+    // Initiate the merge and then immediately stop it to simulate a reboot.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    test_device->set_recovery(true);
+    new_sm = SnapshotManager::NewForFirstStageMount(test_device.release());
+
+    ASSERT_TRUE(new_sm->FinishMergeInRecovery());
+
+    auto mount = new_sm->EnsureMetadataMounted();
+    ASSERT_TRUE(mount && mount->HasDevice());
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
+
+    // Finish the merge in a normal boot.
+    test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    init = SnapshotManager::NewForFirstStageMount(test_device.release());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    init = nullptr;
+
+    test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    new_sm = SnapshotManager::NewForFirstStageMount(test_device.release());
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::None);
+}
+
 // Test that after an OTA, before a merge, we can wipe data in recovery.
 TEST_F(SnapshotUpdateTest, DataWipeRollbackInRecovery) {
     // Execute the first update.