1 // Copyright (C) 2018 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include <libsnapshot/cow_format.h>
16 #include <libsnapshot/snapshot.h>
17 
18 #include <fcntl.h>
19 #include <signal.h>
20 #include <sys/file.h>
21 #include <sys/stat.h>
22 #include <sys/types.h>
23 
24 #include <chrono>
25 #include <deque>
26 #include <future>
27 #include <iostream>
28 
29 #include <android-base/file.h>
30 #include <android-base/logging.h>
31 #include <android-base/properties.h>
32 #include <android-base/strings.h>
33 #include <android-base/unique_fd.h>
34 #include <fs_mgr/file_wait.h>
35 #include <fs_mgr/roots.h>
36 #include <fs_mgr_dm_linear.h>
37 #include <gtest/gtest.h>
38 #include <libdm/dm.h>
39 #include <libfiemap/image_manager.h>
40 #include <liblp/builder.h>
41 #include <storage_literals/storage_literals.h>
42 
43 #include <android/snapshot/snapshot.pb.h>
44 #include <libsnapshot/test_helpers.h>
45 #include "partition_cow_creator.h"
46 #include "utility.h"
47 
48 // Mock classes are not used. Header included to ensure mocked class definition aligns with the
49 // class itself.
50 #include <libsnapshot/mock_device_info.h>
51 #include <libsnapshot/mock_snapshot.h>
52 
53 namespace android {
54 namespace snapshot {
55 
56 using android::base::unique_fd;
57 using android::dm::DeviceMapper;
58 using android::dm::DmDeviceState;
59 using android::fiemap::FiemapStatus;
60 using android::fiemap::IImageManager;
61 using android::fs_mgr::BlockDeviceInfo;
62 using android::fs_mgr::CreateLogicalPartitionParams;
63 using android::fs_mgr::DestroyLogicalPartition;
64 using android::fs_mgr::EnsurePathMounted;
65 using android::fs_mgr::EnsurePathUnmounted;
66 using android::fs_mgr::Extent;
67 using android::fs_mgr::Fstab;
68 using android::fs_mgr::GetPartitionGroupName;
69 using android::fs_mgr::GetPartitionName;
70 using android::fs_mgr::Interval;
71 using android::fs_mgr::MetadataBuilder;
72 using android::fs_mgr::SlotSuffixForSlotNumber;
73 using chromeos_update_engine::DeltaArchiveManifest;
74 using chromeos_update_engine::DynamicPartitionGroup;
75 using chromeos_update_engine::PartitionUpdate;
76 using namespace ::testing;
77 using namespace android::storage_literals;
78 using namespace std::chrono_literals;
79 using namespace std::string_literals;
80 
81 // Global states. See test_helpers.h.
82 std::unique_ptr<SnapshotManager> sm;
83 TestDeviceInfo* test_device = nullptr;
84 std::string fake_super;
85 
86 void MountMetadata();
87 
88 class SnapshotTest : public ::testing::Test {
89   public:
SnapshotTest()90     SnapshotTest() : dm_(DeviceMapper::Instance()) {}
91 
92     // This is exposed for main.
Cleanup()93     void Cleanup() {
94         InitializeState();
95         CleanupTestArtifacts();
96     }
97 
98   protected:
SetUp()99     void SetUp() override {
100         SKIP_IF_NON_VIRTUAL_AB();
101 
102         SnapshotTestPropertyFetcher::SetUp();
103         InitializeState();
104         CleanupTestArtifacts();
105         FormatFakeSuper();
106         MountMetadata();
107         ASSERT_TRUE(sm->BeginUpdate());
108     }
109 
TearDown()110     void TearDown() override {
111         RETURN_IF_NON_VIRTUAL_AB();
112 
113         lock_ = nullptr;
114 
115         CleanupTestArtifacts();
116         SnapshotTestPropertyFetcher::TearDown();
117     }
118 
InitializeState()119     void InitializeState() {
120         ASSERT_TRUE(sm->EnsureImageManager());
121         image_manager_ = sm->image_manager();
122 
123         test_device->set_slot_suffix("_a");
124 
125         sm->set_use_first_stage_snapuserd(false);
126     }
127 
CleanupTestArtifacts()128     void CleanupTestArtifacts() {
129         // Normally cancelling inside a merge is not allowed. Since these
130         // are tests, we don't care, destroy everything that might exist.
131         // Note we hardcode this list because of an annoying quirk: when
132         // completing a merge, the snapshot stops existing, so we can't
133         // get an accurate list to remove.
134         lock_ = nullptr;
135 
136         std::vector<std::string> snapshots = {"test-snapshot", "test_partition_a",
137                                               "test_partition_b"};
138         for (const auto& snapshot : snapshots) {
139             ASSERT_TRUE(DeleteSnapshotDevice(snapshot));
140             DeleteBackingImage(image_manager_, snapshot + "-cow-img");
141 
142             auto status_file = sm->GetSnapshotStatusFilePath(snapshot);
143             android::base::RemoveFileIfExists(status_file);
144         }
145 
146         // Remove stale partitions in fake super.
147         std::vector<std::string> partitions = {
148                 "base-device",
149                 "test_partition_b",
150                 "test_partition_b-base",
151                 "test_partition_b-base",
152         };
153         for (const auto& partition : partitions) {
154             DeleteDevice(partition);
155         }
156 
157         if (sm->GetUpdateState() != UpdateState::None) {
158             auto state_file = sm->GetStateFilePath();
159             unlink(state_file.c_str());
160         }
161     }
162 
AcquireLock()163     bool AcquireLock() {
164         lock_ = sm->LockExclusive();
165         return !!lock_;
166     }
167 
168     // This is so main() can instantiate this to invoke Cleanup.
TestBody()169     virtual void TestBody() override {}
170 
FormatFakeSuper()171     void FormatFakeSuper() {
172         BlockDeviceInfo super_device("super", kSuperSize, 0, 0, 4096);
173         std::vector<BlockDeviceInfo> devices = {super_device};
174 
175         auto builder = MetadataBuilder::New(devices, "super", 65536, 2);
176         ASSERT_NE(builder, nullptr);
177 
178         auto metadata = builder->Export();
179         ASSERT_NE(metadata, nullptr);
180 
181         TestPartitionOpener opener(fake_super);
182         ASSERT_TRUE(FlashPartitionTable(opener, fake_super, *metadata.get()));
183     }
184 
185     // If |path| is non-null, the partition will be mapped after creation.
CreatePartition(const std::string & name,uint64_t size,std::string * path=nullptr,const std::optional<std::string> group={})186     bool CreatePartition(const std::string& name, uint64_t size, std::string* path = nullptr,
187                          const std::optional<std::string> group = {}) {
188         TestPartitionOpener opener(fake_super);
189         auto builder = MetadataBuilder::New(opener, "super", 0);
190         if (!builder) return false;
191 
192         std::string partition_group = std::string(android::fs_mgr::kDefaultGroup);
193         if (group) {
194             partition_group = *group;
195         }
196         return CreatePartition(builder.get(), name, size, path, partition_group);
197     }
198 
CreatePartition(MetadataBuilder * builder,const std::string & name,uint64_t size,std::string * path,const std::string & group)199     bool CreatePartition(MetadataBuilder* builder, const std::string& name, uint64_t size,
200                          std::string* path, const std::string& group) {
201         auto partition = builder->AddPartition(name, group, 0);
202         if (!partition) return false;
203         if (!builder->ResizePartition(partition, size)) {
204             return false;
205         }
206 
207         // Update the source slot.
208         auto metadata = builder->Export();
209         if (!metadata) return false;
210 
211         TestPartitionOpener opener(fake_super);
212         if (!UpdatePartitionTable(opener, "super", *metadata.get(), 0)) {
213             return false;
214         }
215 
216         if (!path) return true;
217 
218         CreateLogicalPartitionParams params = {
219                 .block_device = fake_super,
220                 .metadata = metadata.get(),
221                 .partition_name = name,
222                 .force_writable = true,
223                 .timeout_ms = 10s,
224         };
225         return CreateLogicalPartition(params, path);
226     }
227 
MapUpdateSnapshot(const std::string & name,std::unique_ptr<ISnapshotWriter> * writer)228     AssertionResult MapUpdateSnapshot(const std::string& name,
229                                       std::unique_ptr<ISnapshotWriter>* writer) {
230         TestPartitionOpener opener(fake_super);
231         CreateLogicalPartitionParams params{
232                 .block_device = fake_super,
233                 .metadata_slot = 1,
234                 .partition_name = name,
235                 .timeout_ms = 10s,
236                 .partition_opener = &opener,
237         };
238 
239         auto old_partition = "/dev/block/mapper/" + GetOtherPartitionName(name);
240         auto result = sm->OpenSnapshotWriter(params, {old_partition});
241         if (!result) {
242             return AssertionFailure() << "Cannot open snapshot for writing: " << name;
243         }
244         if (!result->Initialize()) {
245             return AssertionFailure() << "Cannot initialize snapshot for writing: " << name;
246         }
247 
248         if (writer) {
249             *writer = std::move(result);
250         }
251         return AssertionSuccess();
252     }
253 
MapUpdateSnapshot(const std::string & name,std::string * path)254     AssertionResult MapUpdateSnapshot(const std::string& name, std::string* path) {
255         TestPartitionOpener opener(fake_super);
256         CreateLogicalPartitionParams params{
257                 .block_device = fake_super,
258                 .metadata_slot = 1,
259                 .partition_name = name,
260                 .timeout_ms = 10s,
261                 .partition_opener = &opener,
262         };
263 
264         auto result = sm->MapUpdateSnapshot(params, path);
265         if (!result) {
266             return AssertionFailure() << "Cannot open snapshot for writing: " << name;
267         }
268         return AssertionSuccess();
269     }
270 
DeleteSnapshotDevice(const std::string & snapshot)271     AssertionResult DeleteSnapshotDevice(const std::string& snapshot) {
272         AssertionResult res = AssertionSuccess();
273         if (!(res = DeleteDevice(snapshot))) return res;
274         if (!sm->UnmapDmUserDevice(snapshot)) {
275             return AssertionFailure() << "Cannot delete dm-user device for " << snapshot;
276         }
277         if (!(res = DeleteDevice(snapshot + "-inner"))) return res;
278         if (!(res = DeleteDevice(snapshot + "-cow"))) return res;
279         if (!image_manager_->UnmapImageIfExists(snapshot + "-cow-img")) {
280             return AssertionFailure() << "Cannot unmap image " << snapshot << "-cow-img";
281         }
282         if (!(res = DeleteDevice(snapshot + "-base"))) return res;
283         if (!(res = DeleteDevice(snapshot + "-src"))) return res;
284         return AssertionSuccess();
285     }
286 
DeleteDevice(const std::string & device)287     AssertionResult DeleteDevice(const std::string& device) {
288         if (!dm_.DeleteDeviceIfExists(device)) {
289             return AssertionFailure() << "Can't delete " << device;
290         }
291         return AssertionSuccess();
292     }
293 
CreateCowImage(const std::string & name)294     AssertionResult CreateCowImage(const std::string& name) {
295         if (!sm->CreateCowImage(lock_.get(), name)) {
296             return AssertionFailure() << "Cannot create COW image " << name;
297         }
298         std::string cow_device;
299         auto map_res = MapCowImage(name, 10s, &cow_device);
300         if (!map_res) {
301             return map_res;
302         }
303         if (!InitializeKernelCow(cow_device)) {
304             return AssertionFailure() << "Cannot zero fill " << cow_device;
305         }
306         if (!sm->UnmapCowImage(name)) {
307             return AssertionFailure() << "Cannot unmap " << name << " after zero filling it";
308         }
309         return AssertionSuccess();
310     }
311 
MapCowImage(const std::string & name,const std::chrono::milliseconds & timeout_ms,std::string * path)312     AssertionResult MapCowImage(const std::string& name,
313                                 const std::chrono::milliseconds& timeout_ms, std::string* path) {
314         auto cow_image_path = sm->MapCowImage(name, timeout_ms);
315         if (!cow_image_path.has_value()) {
316             return AssertionFailure() << "Cannot map cow image " << name;
317         }
318         *path = *cow_image_path;
319         return AssertionSuccess();
320     }
321 
322     // Prepare A/B slot for a partition named "test_partition".
PrepareOneSnapshot(uint64_t device_size,std::unique_ptr<ISnapshotWriter> * writer=nullptr)323     AssertionResult PrepareOneSnapshot(uint64_t device_size,
324                                        std::unique_ptr<ISnapshotWriter>* writer = nullptr) {
325         lock_ = nullptr;
326 
327         DeltaArchiveManifest manifest;
328 
329         auto dynamic_partition_metadata = manifest.mutable_dynamic_partition_metadata();
330         dynamic_partition_metadata->set_vabc_enabled(IsCompressionEnabled());
331         dynamic_partition_metadata->set_cow_version(android::snapshot::kCowVersionMajor);
332 
333         auto group = dynamic_partition_metadata->add_groups();
334         group->set_name("group");
335         group->set_size(device_size * 2);
336         group->add_partition_names("test_partition");
337 
338         auto pu = manifest.add_partitions();
339         pu->set_partition_name("test_partition");
340         pu->set_estimate_cow_size(device_size);
341         SetSize(pu, device_size);
342 
343         auto extent = pu->add_operations()->add_dst_extents();
344         extent->set_start_block(0);
345         if (device_size) {
346             extent->set_num_blocks(device_size / manifest.block_size());
347         }
348 
349         TestPartitionOpener opener(fake_super);
350         auto builder = MetadataBuilder::New(opener, "super", 0);
351         if (!builder) {
352             return AssertionFailure() << "Failed to open MetadataBuilder";
353         }
354         builder->AddGroup("group_a", 16_GiB);
355         builder->AddGroup("group_b", 16_GiB);
356         if (!CreatePartition(builder.get(), "test_partition_a", device_size, nullptr, "group_a")) {
357             return AssertionFailure() << "Failed create test_partition_a";
358         }
359 
360         if (!sm->CreateUpdateSnapshots(manifest)) {
361             return AssertionFailure() << "Failed to create update snapshots";
362         }
363 
364         if (writer) {
365             auto res = MapUpdateSnapshot("test_partition_b", writer);
366             if (!res) {
367                 return res;
368             }
369         } else if (!IsCompressionEnabled()) {
370             std::string ignore;
371             if (!MapUpdateSnapshot("test_partition_b", &ignore)) {
372                 return AssertionFailure() << "Failed to map test_partition_b";
373             }
374         }
375         if (!AcquireLock()) {
376             return AssertionFailure() << "Failed to acquire lock";
377         }
378         return AssertionSuccess();
379     }
380 
381     // Simulate a reboot into the new slot.
SimulateReboot()382     AssertionResult SimulateReboot() {
383         lock_ = nullptr;
384         if (!sm->FinishedSnapshotWrites(false)) {
385             return AssertionFailure() << "Failed to finish snapshot writes";
386         }
387         if (!sm->UnmapUpdateSnapshot("test_partition_b")) {
388             return AssertionFailure() << "Failed to unmap COW for test_partition_b";
389         }
390         if (!dm_.DeleteDeviceIfExists("test_partition_b")) {
391             return AssertionFailure() << "Failed to delete test_partition_b";
392         }
393         if (!dm_.DeleteDeviceIfExists("test_partition_b-base")) {
394             return AssertionFailure() << "Failed to destroy test_partition_b-base";
395         }
396         return AssertionSuccess();
397     }
398 
NewManagerForFirstStageMount(const std::string & slot_suffix="")399     std::unique_ptr<SnapshotManager> NewManagerForFirstStageMount(
400             const std::string& slot_suffix = "_a") {
401         auto info = new TestDeviceInfo(fake_super, slot_suffix);
402         return NewManagerForFirstStageMount(info);
403     }
404 
NewManagerForFirstStageMount(TestDeviceInfo * info)405     std::unique_ptr<SnapshotManager> NewManagerForFirstStageMount(TestDeviceInfo* info) {
406         info->set_first_stage_init(true);
407         auto init = SnapshotManager::NewForFirstStageMount(info);
408         if (!init) {
409             return nullptr;
410         }
411         init->SetUeventRegenCallback([](const std::string& device) -> bool {
412             return android::fs_mgr::WaitForFile(device, snapshot_timeout_);
413         });
414         return init;
415     }
416 
417     static constexpr std::chrono::milliseconds snapshot_timeout_ = 5s;
418     DeviceMapper& dm_;
419     std::unique_ptr<SnapshotManager::LockedFile> lock_;
420     android::fiemap::IImageManager* image_manager_ = nullptr;
421     std::string fake_super_;
422 };
423 
TEST_F(SnapshotTest,CreateSnapshot)424 TEST_F(SnapshotTest, CreateSnapshot) {
425     ASSERT_TRUE(AcquireLock());
426 
427     PartitionCowCreator cow_creator;
428     cow_creator.compression_enabled = IsCompressionEnabled();
429     if (cow_creator.compression_enabled) {
430         cow_creator.compression_algorithm = "gz";
431     } else {
432         cow_creator.compression_algorithm = "none";
433     }
434 
435     static const uint64_t kDeviceSize = 1024 * 1024;
436     SnapshotStatus status;
437     status.set_name("test-snapshot");
438     status.set_device_size(kDeviceSize);
439     status.set_snapshot_size(kDeviceSize);
440     status.set_cow_file_size(kDeviceSize);
441     ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
442     ASSERT_TRUE(CreateCowImage("test-snapshot"));
443 
444     std::vector<std::string> snapshots;
445     ASSERT_TRUE(sm->ListSnapshots(lock_.get(), &snapshots));
446     ASSERT_EQ(snapshots.size(), 1);
447     ASSERT_EQ(snapshots[0], "test-snapshot");
448 
449     // Scope so delete can re-acquire the snapshot file lock.
450     {
451         SnapshotStatus status;
452         ASSERT_TRUE(sm->ReadSnapshotStatus(lock_.get(), "test-snapshot", &status));
453         ASSERT_EQ(status.state(), SnapshotState::CREATED);
454         ASSERT_EQ(status.device_size(), kDeviceSize);
455         ASSERT_EQ(status.snapshot_size(), kDeviceSize);
456         ASSERT_EQ(status.compression_enabled(), cow_creator.compression_enabled);
457         ASSERT_EQ(status.compression_algorithm(), cow_creator.compression_algorithm);
458     }
459 
460     ASSERT_TRUE(sm->UnmapSnapshot(lock_.get(), "test-snapshot"));
461     ASSERT_TRUE(sm->UnmapCowImage("test-snapshot"));
462     ASSERT_TRUE(sm->DeleteSnapshot(lock_.get(), "test-snapshot"));
463 }
464 
TEST_F(SnapshotTest,MapSnapshot)465 TEST_F(SnapshotTest, MapSnapshot) {
466     ASSERT_TRUE(AcquireLock());
467 
468     PartitionCowCreator cow_creator;
469     cow_creator.compression_enabled = IsCompressionEnabled();
470 
471     static const uint64_t kDeviceSize = 1024 * 1024;
472     SnapshotStatus status;
473     status.set_name("test-snapshot");
474     status.set_device_size(kDeviceSize);
475     status.set_snapshot_size(kDeviceSize);
476     status.set_cow_file_size(kDeviceSize);
477     ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
478     ASSERT_TRUE(CreateCowImage("test-snapshot"));
479 
480     std::string base_device;
481     ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
482 
483     std::string cow_device;
484     ASSERT_TRUE(MapCowImage("test-snapshot", 10s, &cow_device));
485 
486     std::string snap_device;
487     ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, cow_device, 10s,
488                                 &snap_device));
489     ASSERT_TRUE(android::base::StartsWith(snap_device, "/dev/block/dm-"));
490 }
491 
TEST_F(SnapshotTest,NoMergeBeforeReboot)492 TEST_F(SnapshotTest, NoMergeBeforeReboot) {
493     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
494 
495     // Merge should fail, since the slot hasn't changed.
496     ASSERT_FALSE(sm->InitiateMerge());
497 }
498 
TEST_F(SnapshotTest,CleanFirstStageMount)499 TEST_F(SnapshotTest, CleanFirstStageMount) {
500     // If there's no update in progress, there should be no first-stage mount
501     // needed.
502     auto sm = NewManagerForFirstStageMount();
503     ASSERT_NE(sm, nullptr);
504     ASSERT_FALSE(sm->NeedSnapshotsInFirstStageMount());
505 }
506 
TEST_F(SnapshotTest,FirstStageMountAfterRollback)507 TEST_F(SnapshotTest, FirstStageMountAfterRollback) {
508     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
509 
510     // We didn't change the slot, so we shouldn't need snapshots.
511     auto sm = NewManagerForFirstStageMount();
512     ASSERT_NE(sm, nullptr);
513     ASSERT_FALSE(sm->NeedSnapshotsInFirstStageMount());
514 
515     auto indicator = sm->GetRollbackIndicatorPath();
516     ASSERT_EQ(access(indicator.c_str(), R_OK), 0);
517 }
518 
TEST_F(SnapshotTest,Merge)519 TEST_F(SnapshotTest, Merge) {
520     ASSERT_TRUE(AcquireLock());
521 
522     static const uint64_t kDeviceSize = 1024 * 1024;
523 
524     std::unique_ptr<ISnapshotWriter> writer;
525     ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize, &writer));
526 
527     // Release the lock.
528     lock_ = nullptr;
529 
530     std::string test_string = "This is a test string.";
531     test_string.resize(writer->options().block_size);
532     ASSERT_TRUE(writer->AddRawBlocks(0, test_string.data(), test_string.size()));
533     ASSERT_TRUE(writer->Finalize());
534     writer = nullptr;
535 
536     // Done updating.
537     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
538 
539     ASSERT_TRUE(sm->UnmapUpdateSnapshot("test_partition_b"));
540 
541     test_device->set_slot_suffix("_b");
542     ASSERT_TRUE(sm->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
543     ASSERT_TRUE(sm->InitiateMerge());
544 
545     // The device should have been switched to a snapshot-merge target.
546     DeviceMapper::TargetInfo target;
547     ASSERT_TRUE(sm->IsSnapshotDevice("test_partition_b", &target));
548     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot-merge");
549 
550     // We should not be able to cancel an update now.
551     ASSERT_FALSE(sm->CancelUpdate());
552 
553     ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::MergeCompleted);
554     ASSERT_EQ(sm->GetUpdateState(), UpdateState::None);
555 
556     // The device should no longer be a snapshot or snapshot-merge.
557     ASSERT_FALSE(sm->IsSnapshotDevice("test_partition_b"));
558 
559     // Test that we can read back the string we wrote to the snapshot. Note
560     // that the base device is gone now. |snap_device| contains the correct
561     // partition.
562     unique_fd fd(open("/dev/block/mapper/test_partition_b", O_RDONLY | O_CLOEXEC));
563     ASSERT_GE(fd, 0);
564 
565     std::string buffer(test_string.size(), '\0');
566     ASSERT_TRUE(android::base::ReadFully(fd, buffer.data(), buffer.size()));
567     ASSERT_EQ(test_string, buffer);
568 }
569 
TEST_F(SnapshotTest,FirstStageMountAndMerge)570 TEST_F(SnapshotTest, FirstStageMountAndMerge) {
571     ASSERT_TRUE(AcquireLock());
572 
573     static const uint64_t kDeviceSize = 1024 * 1024;
574     ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
575     ASSERT_TRUE(SimulateReboot());
576 
577     auto init = NewManagerForFirstStageMount("_b");
578     ASSERT_NE(init, nullptr);
579     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
580     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
581 
582     ASSERT_TRUE(AcquireLock());
583 
584     // Validate that we have a snapshot device.
585     SnapshotStatus status;
586     ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
587     ASSERT_EQ(status.state(), SnapshotState::CREATED);
588     if (IsCompressionEnabled()) {
589         ASSERT_EQ(status.compression_algorithm(), "gz");
590     } else {
591         ASSERT_EQ(status.compression_algorithm(), "none");
592     }
593 
594     DeviceMapper::TargetInfo target;
595     ASSERT_TRUE(init->IsSnapshotDevice("test_partition_b", &target));
596     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot");
597 }
598 
TEST_F(SnapshotTest,FlashSuperDuringUpdate)599 TEST_F(SnapshotTest, FlashSuperDuringUpdate) {
600     ASSERT_TRUE(AcquireLock());
601 
602     static const uint64_t kDeviceSize = 1024 * 1024;
603     ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
604     ASSERT_TRUE(SimulateReboot());
605 
606     // Reflash the super partition.
607     FormatFakeSuper();
608     ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
609 
610     auto init = NewManagerForFirstStageMount("_b");
611     ASSERT_NE(init, nullptr);
612     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
613     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
614 
615     ASSERT_TRUE(AcquireLock());
616 
617     SnapshotStatus status;
618     ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
619 
620     // We should not get a snapshot device now.
621     DeviceMapper::TargetInfo target;
622     ASSERT_FALSE(init->IsSnapshotDevice("test_partition_b", &target));
623 
624     // We should see a cancelled update as well.
625     lock_ = nullptr;
626     ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::Cancelled);
627 }
628 
TEST_F(SnapshotTest,FlashSuperDuringMerge)629 TEST_F(SnapshotTest, FlashSuperDuringMerge) {
630     ASSERT_TRUE(AcquireLock());
631 
632     static const uint64_t kDeviceSize = 1024 * 1024;
633     ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
634     ASSERT_TRUE(SimulateReboot());
635 
636     auto init = NewManagerForFirstStageMount("_b");
637     ASSERT_NE(init, nullptr);
638     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
639     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
640     ASSERT_TRUE(init->InitiateMerge());
641 
642     // Now, reflash super. Note that we haven't called ProcessUpdateState, so the
643     // status is still Merging.
644     ASSERT_TRUE(DeleteSnapshotDevice("test_partition_b"));
645     ASSERT_TRUE(init->image_manager()->UnmapImageIfExists("test_partition_b-cow-img"));
646     FormatFakeSuper();
647     ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
648     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
649     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
650 
651     // Because the status is Merging, we must call ProcessUpdateState, which should
652     // detect a cancelled update.
653     ASSERT_EQ(init->ProcessUpdateState(), UpdateState::Cancelled);
654     ASSERT_EQ(init->GetUpdateState(), UpdateState::None);
655 }
656 
TEST_F(SnapshotTest,UpdateBootControlHal)657 TEST_F(SnapshotTest, UpdateBootControlHal) {
658     ASSERT_TRUE(AcquireLock());
659 
660     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::None));
661     ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
662 
663     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Initiated));
664     ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
665 
666     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Unverified));
667     ASSERT_EQ(test_device->merge_status(), MergeStatus::SNAPSHOTTED);
668 
669     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Merging));
670     ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
671 
672     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeNeedsReboot));
673     ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
674 
675     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeCompleted));
676     ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
677 
678     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeFailed));
679     ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
680 }
681 
TEST_F(SnapshotTest,MergeFailureCode)682 TEST_F(SnapshotTest, MergeFailureCode) {
683     ASSERT_TRUE(AcquireLock());
684 
685     ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeFailed,
686                                      MergeFailureCode::ListSnapshots));
687     ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
688 
689     SnapshotUpdateStatus status = sm->ReadSnapshotUpdateStatus(lock_.get());
690     ASSERT_EQ(status.state(), UpdateState::MergeFailed);
691     ASSERT_EQ(status.merge_failure_code(), MergeFailureCode::ListSnapshots);
692 }
693 
694 enum class Request { UNKNOWN, LOCK_SHARED, LOCK_EXCLUSIVE, UNLOCK, EXIT };
operator <<(std::ostream & os,Request request)695 std::ostream& operator<<(std::ostream& os, Request request) {
696     switch (request) {
697         case Request::LOCK_SHARED:
698             return os << "Shared";
699         case Request::LOCK_EXCLUSIVE:
700             return os << "Exclusive";
701         case Request::UNLOCK:
702             return os << "Unlock";
703         case Request::EXIT:
704             return os << "Exit";
705         case Request::UNKNOWN:
706             [[fallthrough]];
707         default:
708             return os << "Unknown";
709     }
710 }
711 
712 class LockTestConsumer {
713   public:
MakeRequest(Request new_request)714     AssertionResult MakeRequest(Request new_request) {
715         {
716             std::unique_lock<std::mutex> ulock(mutex_);
717             requests_.push_back(new_request);
718         }
719         cv_.notify_all();
720         return AssertionSuccess() << "Request " << new_request << " successful";
721     }
722 
723     template <typename R, typename P>
WaitFulfill(std::chrono::duration<R,P> timeout)724     AssertionResult WaitFulfill(std::chrono::duration<R, P> timeout) {
725         std::unique_lock<std::mutex> ulock(mutex_);
726         if (cv_.wait_for(ulock, timeout, [this] { return requests_.empty(); })) {
727             return AssertionSuccess() << "All requests_ fulfilled.";
728         }
729         return AssertionFailure() << "Timeout waiting for fulfilling " << requests_.size()
730                                   << " request(s), first one is "
731                                   << (requests_.empty() ? Request::UNKNOWN : requests_.front());
732     }
733 
StartHandleRequestsInBackground()734     void StartHandleRequestsInBackground() {
735         future_ = std::async(std::launch::async, &LockTestConsumer::HandleRequests, this);
736     }
737 
738   private:
HandleRequests()739     void HandleRequests() {
740         static constexpr auto consumer_timeout = 3s;
741 
742         auto next_request = Request::UNKNOWN;
743         do {
744             // Peek next request.
745             {
746                 std::unique_lock<std::mutex> ulock(mutex_);
747                 if (cv_.wait_for(ulock, consumer_timeout, [this] { return !requests_.empty(); })) {
748                     next_request = requests_.front();
749                 } else {
750                     next_request = Request::EXIT;
751                 }
752             }
753 
754             // Handle next request.
755             switch (next_request) {
756                 case Request::LOCK_SHARED: {
757                     lock_ = sm->LockShared();
758                 } break;
759                 case Request::LOCK_EXCLUSIVE: {
760                     lock_ = sm->LockExclusive();
761                 } break;
762                 case Request::EXIT:
763                     [[fallthrough]];
764                 case Request::UNLOCK: {
765                     lock_.reset();
766                 } break;
767                 case Request::UNKNOWN:
768                     [[fallthrough]];
769                 default:
770                     break;
771             }
772 
773             // Pop next request. This thread is the only thread that
774             // pops from the front of the requests_ deque.
775             {
776                 std::unique_lock<std::mutex> ulock(mutex_);
777                 if (next_request == Request::EXIT) {
778                     requests_.clear();
779                 } else {
780                     requests_.pop_front();
781                 }
782             }
783             cv_.notify_all();
784         } while (next_request != Request::EXIT);
785     }
786 
787     std::mutex mutex_;
788     std::condition_variable cv_;
789     std::deque<Request> requests_;
790     std::unique_ptr<SnapshotManager::LockedFile> lock_;
791     std::future<void> future_;
792 };
793 
794 class LockTest : public ::testing::Test {
795   public:
SetUp()796     void SetUp() {
797         SKIP_IF_NON_VIRTUAL_AB();
798         first_consumer.StartHandleRequestsInBackground();
799         second_consumer.StartHandleRequestsInBackground();
800     }
801 
TearDown()802     void TearDown() {
803         RETURN_IF_NON_VIRTUAL_AB();
804         EXPECT_TRUE(first_consumer.MakeRequest(Request::EXIT));
805         EXPECT_TRUE(second_consumer.MakeRequest(Request::EXIT));
806     }
807 
808     static constexpr auto request_timeout = 500ms;
809     LockTestConsumer first_consumer;
810     LockTestConsumer second_consumer;
811 };
812 
TEST_F(LockTest,SharedShared)813 TEST_F(LockTest, SharedShared) {
814     ASSERT_TRUE(first_consumer.MakeRequest(Request::LOCK_SHARED));
815     ASSERT_TRUE(first_consumer.WaitFulfill(request_timeout));
816     ASSERT_TRUE(second_consumer.MakeRequest(Request::LOCK_SHARED));
817     ASSERT_TRUE(second_consumer.WaitFulfill(request_timeout));
818 }
819 
820 using LockTestParam = std::pair<Request, Request>;
821 class LockTestP : public LockTest, public ::testing::WithParamInterface<LockTestParam> {};
TEST_P(LockTestP,Test)822 TEST_P(LockTestP, Test) {
823     ASSERT_TRUE(first_consumer.MakeRequest(GetParam().first));
824     ASSERT_TRUE(first_consumer.WaitFulfill(request_timeout));
825     ASSERT_TRUE(second_consumer.MakeRequest(GetParam().second));
826     ASSERT_FALSE(second_consumer.WaitFulfill(request_timeout))
827             << "Should not be able to " << GetParam().second << " while separate thread "
828             << GetParam().first;
829     ASSERT_TRUE(first_consumer.MakeRequest(Request::UNLOCK));
830     ASSERT_TRUE(second_consumer.WaitFulfill(request_timeout))
831             << "Should be able to hold lock that is released by separate thread";
832 }
833 INSTANTIATE_TEST_SUITE_P(
834         LockTest, LockTestP,
835         testing::Values(LockTestParam{Request::LOCK_EXCLUSIVE, Request::LOCK_EXCLUSIVE},
836                         LockTestParam{Request::LOCK_EXCLUSIVE, Request::LOCK_SHARED},
837                         LockTestParam{Request::LOCK_SHARED, Request::LOCK_EXCLUSIVE}),
__anon6022f21a0402(const testing::TestParamInfo<LockTestP::ParamType>& info) 838         [](const testing::TestParamInfo<LockTestP::ParamType>& info) {
839             std::stringstream ss;
840             ss << info.param.first << info.param.second;
841             return ss.str();
842         });
843 
844 class SnapshotUpdateTest : public SnapshotTest {
845   public:
SetUp()846     void SetUp() override {
847         SKIP_IF_NON_VIRTUAL_AB();
848 
849         SnapshotTest::SetUp();
850         Cleanup();
851 
852         // Cleanup() changes slot suffix, so initialize it again.
853         test_device->set_slot_suffix("_a");
854 
855         opener_ = std::make_unique<TestPartitionOpener>(fake_super);
856 
857         auto dynamic_partition_metadata = manifest_.mutable_dynamic_partition_metadata();
858         dynamic_partition_metadata->set_vabc_enabled(IsCompressionEnabled());
859         dynamic_partition_metadata->set_cow_version(android::snapshot::kCowVersionMajor);
860 
861         // Create a fake update package metadata.
862         // Not using full name "system", "vendor", "product" because these names collide with the
863         // mapped partitions on the running device.
864         // Each test modifies manifest_ slightly to indicate changes to the partition layout.
865         group_ = dynamic_partition_metadata->add_groups();
866         group_->set_name("group");
867         group_->set_size(kGroupSize);
868         group_->add_partition_names("sys");
869         group_->add_partition_names("vnd");
870         group_->add_partition_names("prd");
871         sys_ = manifest_.add_partitions();
872         sys_->set_partition_name("sys");
873         sys_->set_estimate_cow_size(2_MiB);
874         SetSize(sys_, 3_MiB);
875         vnd_ = manifest_.add_partitions();
876         vnd_->set_partition_name("vnd");
877         vnd_->set_estimate_cow_size(2_MiB);
878         SetSize(vnd_, 3_MiB);
879         prd_ = manifest_.add_partitions();
880         prd_->set_partition_name("prd");
881         prd_->set_estimate_cow_size(2_MiB);
882         SetSize(prd_, 3_MiB);
883 
884         // Initialize source partition metadata using |manifest_|.
885         src_ = MetadataBuilder::New(*opener_, "super", 0);
886         ASSERT_NE(src_, nullptr);
887         ASSERT_TRUE(FillFakeMetadata(src_.get(), manifest_, "_a"));
888         // Add sys_b which is like system_other.
889         ASSERT_TRUE(src_->AddGroup("group_b", kGroupSize));
890         auto partition = src_->AddPartition("sys_b", "group_b", 0);
891         ASSERT_NE(nullptr, partition);
892         ASSERT_TRUE(src_->ResizePartition(partition, 1_MiB));
893         auto metadata = src_->Export();
894         ASSERT_NE(nullptr, metadata);
895         ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
896 
897         // Map source partitions. Additionally, map sys_b to simulate system_other after flashing.
898         std::string path;
899         for (const auto& name : {"sys_a", "vnd_a", "prd_a", "sys_b"}) {
900             ASSERT_TRUE(CreateLogicalPartition(
901                     CreateLogicalPartitionParams{
902                             .block_device = fake_super,
903                             .metadata_slot = 0,
904                             .partition_name = name,
905                             .timeout_ms = 1s,
906                             .partition_opener = opener_.get(),
907                     },
908                     &path));
909             ASSERT_TRUE(WriteRandomData(path));
910             auto hash = GetHash(path);
911             ASSERT_TRUE(hash.has_value());
912             hashes_[name] = *hash;
913         }
914     }
TearDown()915     void TearDown() override {
916         RETURN_IF_NON_VIRTUAL_AB();
917 
918         Cleanup();
919         SnapshotTest::TearDown();
920     }
Cleanup()921     void Cleanup() {
922         if (!image_manager_) {
923             InitializeState();
924         }
925         MountMetadata();
926         for (const auto& suffix : {"_a", "_b"}) {
927             test_device->set_slot_suffix(suffix);
928             EXPECT_TRUE(sm->CancelUpdate()) << suffix;
929         }
930         EXPECT_TRUE(UnmapAll());
931     }
932 
IsPartitionUnchanged(const std::string & name)933     AssertionResult IsPartitionUnchanged(const std::string& name) {
934         std::string path;
935         if (!dm_.GetDmDevicePathByName(name, &path)) {
936             return AssertionFailure() << "Path of " << name << " cannot be determined";
937         }
938         auto hash = GetHash(path);
939         if (!hash.has_value()) {
940             return AssertionFailure() << "Cannot read partition " << name << ": " << path;
941         }
942         auto it = hashes_.find(name);
943         if (it == hashes_.end()) {
944             return AssertionFailure() << "No existing hash for " << name << ". Bad test code?";
945         }
946         if (it->second != *hash) {
947             return AssertionFailure() << "Content of " << name << " has changed";
948         }
949         return AssertionSuccess();
950     }
951 
GetSnapshotSize(const std::string & name)952     std::optional<uint64_t> GetSnapshotSize(const std::string& name) {
953         if (!AcquireLock()) {
954             return std::nullopt;
955         }
956         auto local_lock = std::move(lock_);
957 
958         SnapshotStatus status;
959         if (!sm->ReadSnapshotStatus(local_lock.get(), name, &status)) {
960             return std::nullopt;
961         }
962         return status.snapshot_size();
963     }
964 
UnmapAll()965     AssertionResult UnmapAll() {
966         for (const auto& name : {"sys", "vnd", "prd", "dlkm"}) {
967             if (!dm_.DeleteDeviceIfExists(name + "_a"s)) {
968                 return AssertionFailure() << "Cannot unmap " << name << "_a";
969             }
970             if (!DeleteSnapshotDevice(name + "_b"s)) {
971                 return AssertionFailure() << "Cannot delete snapshot " << name << "_b";
972             }
973         }
974         return AssertionSuccess();
975     }
976 
MapOneUpdateSnapshot(const std::string & name)977     AssertionResult MapOneUpdateSnapshot(const std::string& name) {
978         if (IsCompressionEnabled()) {
979             std::unique_ptr<ISnapshotWriter> writer;
980             return MapUpdateSnapshot(name, &writer);
981         } else {
982             std::string path;
983             return MapUpdateSnapshot(name, &path);
984         }
985     }
986 
WriteSnapshotAndHash(const std::string & name)987     AssertionResult WriteSnapshotAndHash(const std::string& name) {
988         if (IsCompressionEnabled()) {
989             std::unique_ptr<ISnapshotWriter> writer;
990             auto res = MapUpdateSnapshot(name, &writer);
991             if (!res) {
992                 return res;
993             }
994             if (!WriteRandomData(writer.get(), &hashes_[name])) {
995                 return AssertionFailure() << "Unable to write random data to snapshot " << name;
996             }
997             if (!writer->Finalize()) {
998                 return AssertionFailure() << "Unable to finalize COW for " << name;
999             }
1000         } else {
1001             std::string path;
1002             auto res = MapUpdateSnapshot(name, &path);
1003             if (!res) {
1004                 return res;
1005             }
1006             if (!WriteRandomData(path, std::nullopt, &hashes_[name])) {
1007                 return AssertionFailure() << "Unable to write random data to snapshot " << name;
1008             }
1009         }
1010 
1011         // Make sure updates to one device are seen by all devices.
1012         sync();
1013 
1014         return AssertionSuccess() << "Written random data to snapshot " << name
1015                                   << ", hash: " << hashes_[name];
1016     }
1017 
1018     // Generate a snapshot that moves all the upper blocks down to the start.
1019     // It doesn't really matter the order, we just want copies that reference
1020     // blocks that won't exist if the partition shrinks.
ShiftAllSnapshotBlocks(const std::string & name,uint64_t old_size)1021     AssertionResult ShiftAllSnapshotBlocks(const std::string& name, uint64_t old_size) {
1022         std::unique_ptr<ISnapshotWriter> writer;
1023         if (auto res = MapUpdateSnapshot(name, &writer); !res) {
1024             return res;
1025         }
1026         if (!writer->options().max_blocks || !*writer->options().max_blocks) {
1027             return AssertionFailure() << "No max blocks set for " << name << " writer";
1028         }
1029 
1030         uint64_t src_block = (old_size / writer->options().block_size) - 1;
1031         uint64_t dst_block = 0;
1032         uint64_t max_blocks = *writer->options().max_blocks;
1033         while (dst_block < max_blocks && dst_block < src_block) {
1034             if (!writer->AddCopy(dst_block, src_block)) {
1035                 return AssertionFailure() << "Unable to add copy for " << name << " for blocks "
1036                                           << src_block << ", " << dst_block;
1037             }
1038             dst_block++;
1039             src_block--;
1040         }
1041         if (!writer->Finalize()) {
1042             return AssertionFailure() << "Unable to finalize writer for " << name;
1043         }
1044 
1045         auto hash = HashSnapshot(writer.get());
1046         if (hash.empty()) {
1047             return AssertionFailure() << "Unable to hash snapshot writer for " << name;
1048         }
1049         hashes_[name] = hash;
1050 
1051         return AssertionSuccess();
1052     }
1053 
MapUpdateSnapshots(const std::vector<std::string> & names={"", "", ""})1054     AssertionResult MapUpdateSnapshots(const std::vector<std::string>& names = {"sys_b", "vnd_b",
1055                                                                                 "prd_b"}) {
1056         for (const auto& name : names) {
1057             auto res = MapOneUpdateSnapshot(name);
1058             if (!res) {
1059                 return res;
1060             }
1061         }
1062         return AssertionSuccess();
1063     }
1064 
1065     // Create fake install operations to grow the COW device size.
AddOperation(PartitionUpdate * partition_update,uint64_t size_bytes=0)1066     void AddOperation(PartitionUpdate* partition_update, uint64_t size_bytes = 0) {
1067         auto e = partition_update->add_operations()->add_dst_extents();
1068         e->set_start_block(0);
1069         if (size_bytes == 0) {
1070             size_bytes = GetSize(partition_update);
1071         }
1072         e->set_num_blocks(size_bytes / manifest_.block_size());
1073     }
1074 
AddOperationForPartitions(std::vector<PartitionUpdate * > partitions={})1075     void AddOperationForPartitions(std::vector<PartitionUpdate*> partitions = {}) {
1076         if (partitions.empty()) {
1077             partitions = {sys_, vnd_, prd_};
1078         }
1079         for (auto* partition : partitions) {
1080             AddOperation(partition);
1081         }
1082     }
1083 
1084     std::unique_ptr<TestPartitionOpener> opener_;
1085     DeltaArchiveManifest manifest_;
1086     std::unique_ptr<MetadataBuilder> src_;
1087     std::map<std::string, std::string> hashes_;
1088 
1089     PartitionUpdate* sys_ = nullptr;
1090     PartitionUpdate* vnd_ = nullptr;
1091     PartitionUpdate* prd_ = nullptr;
1092     DynamicPartitionGroup* group_ = nullptr;
1093 };
1094 
1095 // Test full update flow executed by update_engine. Some partitions uses super empty space,
1096 // some uses images, and some uses both.
1097 // Also test UnmapUpdateSnapshot unmaps everything.
1098 // Also test first stage mount and merge after this.
TEST_F(SnapshotUpdateTest,FullUpdateFlow)1099 TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
1100     // OTA client blindly unmaps all partitions that are possibly mapped.
1101     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1102         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
1103     }
1104 
1105     // Grow all partitions. Set |prd| large enough that |sys| and |vnd|'s COWs
1106     // fit in super, but not |prd|.
1107     constexpr uint64_t partition_size = 3788_KiB;
1108     SetSize(sys_, partition_size);
1109     SetSize(vnd_, partition_size);
1110     SetSize(prd_, 18_MiB);
1111 
1112     // Make sure |prd| does not fit in super at all. On VABC, this means we
1113     // fake an extra large COW for |vnd| to fill up super.
1114     vnd_->set_estimate_cow_size(30_MiB);
1115     prd_->set_estimate_cow_size(30_MiB);
1116 
1117     AddOperationForPartitions();
1118 
1119     // Execute the update.
1120     ASSERT_TRUE(sm->BeginUpdate());
1121     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1122 
1123     // Test that partitions prioritize using space in super.
1124     auto tgt = MetadataBuilder::New(*opener_, "super", 1);
1125     ASSERT_NE(tgt, nullptr);
1126     ASSERT_NE(nullptr, tgt->FindPartition("sys_b-cow"));
1127     ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
1128     ASSERT_EQ(nullptr, tgt->FindPartition("prd_b-cow"));
1129 
1130     // Write some data to target partitions.
1131     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1132         ASSERT_TRUE(WriteSnapshotAndHash(name));
1133     }
1134 
1135     // Assert that source partitions aren't affected.
1136     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1137         ASSERT_TRUE(IsPartitionUnchanged(name));
1138     }
1139 
1140     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1141 
1142     // Simulate shutting down the device.
1143     ASSERT_TRUE(UnmapAll());
1144 
1145     // After reboot, init does first stage mount.
1146     auto init = NewManagerForFirstStageMount("_b");
1147     ASSERT_NE(init, nullptr);
1148     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1149     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1150 
1151     auto indicator = sm->GetRollbackIndicatorPath();
1152     ASSERT_NE(access(indicator.c_str(), R_OK), 0);
1153 
1154     // Check that the target partitions have the same content.
1155     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1156         ASSERT_TRUE(IsPartitionUnchanged(name));
1157     }
1158 
1159     // Initiate the merge and wait for it to be completed.
1160     ASSERT_TRUE(init->InitiateMerge());
1161     ASSERT_EQ(init->IsSnapuserdRequired(), IsCompressionEnabled());
1162     {
1163         // We should have started in SECOND_PHASE since nothing shrinks.
1164         ASSERT_TRUE(AcquireLock());
1165         auto local_lock = std::move(lock_);
1166         auto status = init->ReadSnapshotUpdateStatus(local_lock.get());
1167         ASSERT_EQ(status.merge_phase(), MergePhase::SECOND_PHASE);
1168     }
1169     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
1170 
1171     // Make sure the second phase ran and deleted snapshots.
1172     {
1173         ASSERT_TRUE(AcquireLock());
1174         auto local_lock = std::move(lock_);
1175         std::vector<std::string> snapshots;
1176         ASSERT_TRUE(init->ListSnapshots(local_lock.get(), &snapshots));
1177         ASSERT_TRUE(snapshots.empty());
1178     }
1179 
1180     // Check that the target partitions have the same content after the merge.
1181     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1182         ASSERT_TRUE(IsPartitionUnchanged(name))
1183                 << "Content of " << name << " changes after the merge";
1184     }
1185 }
1186 
1187 // Test that shrinking and growing partitions at the same time is handled
1188 // correctly in VABC.
TEST_F(SnapshotUpdateTest,SpaceSwapUpdate)1189 TEST_F(SnapshotUpdateTest, SpaceSwapUpdate) {
1190     if (!IsCompressionEnabled()) {
1191         // b/179111359
1192         GTEST_SKIP() << "Skipping Virtual A/B Compression test";
1193     }
1194 
1195     // OTA client blindly unmaps all partitions that are possibly mapped.
1196     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1197         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
1198     }
1199 
1200     auto old_sys_size = GetSize(sys_);
1201     auto old_prd_size = GetSize(prd_);
1202 
1203     // Grow |sys| but shrink |prd|.
1204     SetSize(sys_, old_sys_size * 2);
1205     sys_->set_estimate_cow_size(8_MiB);
1206     SetSize(prd_, old_prd_size / 2);
1207     prd_->set_estimate_cow_size(1_MiB);
1208 
1209     AddOperationForPartitions();
1210 
1211     ASSERT_TRUE(sm->BeginUpdate());
1212     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1213 
1214     // Check that the old partition sizes were saved correctly.
1215     {
1216         ASSERT_TRUE(AcquireLock());
1217         auto local_lock = std::move(lock_);
1218 
1219         SnapshotStatus status;
1220         ASSERT_TRUE(sm->ReadSnapshotStatus(local_lock.get(), "prd_b", &status));
1221         ASSERT_EQ(status.old_partition_size(), 3145728);
1222         ASSERT_TRUE(sm->ReadSnapshotStatus(local_lock.get(), "sys_b", &status));
1223         ASSERT_EQ(status.old_partition_size(), 3145728);
1224     }
1225 
1226     ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
1227     ASSERT_TRUE(WriteSnapshotAndHash("vnd_b"));
1228     ASSERT_TRUE(ShiftAllSnapshotBlocks("prd_b", old_prd_size));
1229 
1230     sync();
1231 
1232     // Assert that source partitions aren't affected.
1233     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1234         ASSERT_TRUE(IsPartitionUnchanged(name));
1235     }
1236 
1237     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1238 
1239     // Simulate shutting down the device.
1240     ASSERT_TRUE(UnmapAll());
1241 
1242     // After reboot, init does first stage mount.
1243     auto init = NewManagerForFirstStageMount("_b");
1244     ASSERT_NE(init, nullptr);
1245     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1246     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1247 
1248     auto indicator = sm->GetRollbackIndicatorPath();
1249     ASSERT_NE(access(indicator.c_str(), R_OK), 0);
1250 
1251     // Check that the target partitions have the same content.
1252     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1253         ASSERT_TRUE(IsPartitionUnchanged(name));
1254     }
1255 
1256     // Initiate the merge and wait for it to be completed.
1257     ASSERT_TRUE(init->InitiateMerge());
1258     ASSERT_EQ(init->IsSnapuserdRequired(), IsCompressionEnabled());
1259     {
1260         // Check that the merge phase is FIRST_PHASE until at least one call
1261         // to ProcessUpdateState() occurs.
1262         ASSERT_TRUE(AcquireLock());
1263         auto local_lock = std::move(lock_);
1264         auto status = init->ReadSnapshotUpdateStatus(local_lock.get());
1265         ASSERT_EQ(status.merge_phase(), MergePhase::FIRST_PHASE);
1266     }
1267 
1268     // Simulate shutting down the device and creating partitions again.
1269     ASSERT_TRUE(UnmapAll());
1270     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1271 
1272     // Check that we used the correct types after rebooting mid-merge.
1273     DeviceMapper::TargetInfo target;
1274     ASSERT_TRUE(init->IsSnapshotDevice("prd_b", &target));
1275     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot-merge");
1276     ASSERT_TRUE(init->IsSnapshotDevice("sys_b", &target));
1277     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot");
1278     ASSERT_TRUE(init->IsSnapshotDevice("vnd_b", &target));
1279     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot");
1280 
1281     // Complete the merge.
1282     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
1283 
1284     // Make sure the second phase ran and deleted snapshots.
1285     {
1286         ASSERT_TRUE(AcquireLock());
1287         auto local_lock = std::move(lock_);
1288         std::vector<std::string> snapshots;
1289         ASSERT_TRUE(init->ListSnapshots(local_lock.get(), &snapshots));
1290         ASSERT_TRUE(snapshots.empty());
1291     }
1292 
1293     // Check that the target partitions have the same content after the merge.
1294     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1295         ASSERT_TRUE(IsPartitionUnchanged(name))
1296                 << "Content of " << name << " changes after the merge";
1297     }
1298 }
1299 
1300 // Test that a transient merge consistency check failure can resume properly.
TEST_F(SnapshotUpdateTest,ConsistencyCheckResume)1301 TEST_F(SnapshotUpdateTest, ConsistencyCheckResume) {
1302     if (!IsCompressionEnabled()) {
1303         // b/179111359
1304         GTEST_SKIP() << "Skipping Virtual A/B Compression test";
1305     }
1306 
1307     auto old_sys_size = GetSize(sys_);
1308     auto old_prd_size = GetSize(prd_);
1309 
1310     // Grow |sys| but shrink |prd|.
1311     SetSize(sys_, old_sys_size * 2);
1312     sys_->set_estimate_cow_size(8_MiB);
1313     SetSize(prd_, old_prd_size / 2);
1314     prd_->set_estimate_cow_size(1_MiB);
1315 
1316     AddOperationForPartitions();
1317 
1318     ASSERT_TRUE(sm->BeginUpdate());
1319     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1320     ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
1321     ASSERT_TRUE(WriteSnapshotAndHash("vnd_b"));
1322     ASSERT_TRUE(ShiftAllSnapshotBlocks("prd_b", old_prd_size));
1323 
1324     sync();
1325 
1326     // Assert that source partitions aren't affected.
1327     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1328         ASSERT_TRUE(IsPartitionUnchanged(name));
1329     }
1330 
1331     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1332 
1333     // Simulate shutting down the device.
1334     ASSERT_TRUE(UnmapAll());
1335 
1336     // After reboot, init does first stage mount.
1337     auto init = NewManagerForFirstStageMount("_b");
1338     ASSERT_NE(init, nullptr);
1339     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1340     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1341 
1342     // Check that the target partitions have the same content.
1343     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1344         ASSERT_TRUE(IsPartitionUnchanged(name));
1345     }
1346 
1347     auto old_checker = init->merge_consistency_checker();
1348 
1349     init->set_merge_consistency_checker(
1350             [](const std::string&, const SnapshotStatus&) -> MergeFailureCode {
1351                 return MergeFailureCode::WrongMergeCountConsistencyCheck;
1352             });
1353 
1354     // Initiate the merge and wait for it to be completed.
1355     ASSERT_TRUE(init->InitiateMerge());
1356     {
1357         // Check that the merge phase is FIRST_PHASE until at least one call
1358         // to ProcessUpdateState() occurs.
1359         ASSERT_TRUE(AcquireLock());
1360         auto local_lock = std::move(lock_);
1361         auto status = init->ReadSnapshotUpdateStatus(local_lock.get());
1362         ASSERT_EQ(status.merge_phase(), MergePhase::FIRST_PHASE);
1363     }
1364 
1365     // Merge should have failed.
1366     ASSERT_EQ(UpdateState::MergeFailed, init->ProcessUpdateState());
1367 
1368     // Simulate shutting down the device and creating partitions again.
1369     ASSERT_TRUE(UnmapAll());
1370 
1371     // Restore the checker.
1372     init->set_merge_consistency_checker(std::move(old_checker));
1373 
1374     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1375 
1376     // Complete the merge.
1377     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
1378 
1379     // Check that the target partitions have the same content after the merge.
1380     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1381         ASSERT_TRUE(IsPartitionUnchanged(name))
1382                 << "Content of " << name << " changes after the merge";
1383     }
1384 }
1385 
1386 // Test that if new system partitions uses empty space in super, that region is not snapshotted.
TEST_F(SnapshotUpdateTest,DirectWriteEmptySpace)1387 TEST_F(SnapshotUpdateTest, DirectWriteEmptySpace) {
1388     GTEST_SKIP() << "b/141889746";
1389     SetSize(sys_, 4_MiB);
1390     // vnd_b and prd_b are unchanged.
1391     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1392     ASSERT_EQ(3_MiB, GetSnapshotSize("sys_b").value_or(0));
1393 }
1394 
1395 // Test that if new system partitions uses space of old vendor partition, that region is
1396 // snapshotted.
TEST_F(SnapshotUpdateTest,SnapshotOldPartitions)1397 TEST_F(SnapshotUpdateTest, SnapshotOldPartitions) {
1398     SetSize(sys_, 4_MiB);  // grows
1399     SetSize(vnd_, 2_MiB);  // shrinks
1400     // prd_b is unchanged
1401     ASSERT_TRUE(sm->BeginUpdate());
1402     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1403     ASSERT_EQ(4_MiB, GetSnapshotSize("sys_b").value_or(0));
1404 }
1405 
1406 // Test that even if there seem to be empty space in target metadata, COW partition won't take
1407 // it because they are used by old partitions.
TEST_F(SnapshotUpdateTest,CowPartitionDoNotTakeOldPartitions)1408 TEST_F(SnapshotUpdateTest, CowPartitionDoNotTakeOldPartitions) {
1409     SetSize(sys_, 2_MiB);  // shrinks
1410     // vnd_b and prd_b are unchanged.
1411     ASSERT_TRUE(sm->BeginUpdate());
1412     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1413 
1414     auto tgt = MetadataBuilder::New(*opener_, "super", 1);
1415     ASSERT_NE(nullptr, tgt);
1416     auto metadata = tgt->Export();
1417     ASSERT_NE(nullptr, metadata);
1418     std::vector<std::string> written;
1419     // Write random data to all COW partitions in super
1420     for (auto p : metadata->partitions) {
1421         if (GetPartitionGroupName(metadata->groups[p.group_index]) != kCowGroupName) {
1422             continue;
1423         }
1424         std::string path;
1425         ASSERT_TRUE(CreateLogicalPartition(
1426                 CreateLogicalPartitionParams{
1427                         .block_device = fake_super,
1428                         .metadata = metadata.get(),
1429                         .partition = &p,
1430                         .timeout_ms = 1s,
1431                         .partition_opener = opener_.get(),
1432                 },
1433                 &path));
1434         ASSERT_TRUE(WriteRandomData(path));
1435         written.push_back(GetPartitionName(p));
1436     }
1437     ASSERT_FALSE(written.empty())
1438             << "No COW partitions are created even if there are empty space in super partition";
1439 
1440     // Make sure source partitions aren't affected.
1441     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1442         ASSERT_TRUE(IsPartitionUnchanged(name));
1443     }
1444 }
1445 
1446 // Test that it crashes after creating snapshot status file but before creating COW image, then
1447 // calling CreateUpdateSnapshots again works.
TEST_F(SnapshotUpdateTest,SnapshotStatusFileWithoutCow)1448 TEST_F(SnapshotUpdateTest, SnapshotStatusFileWithoutCow) {
1449     // Write some trash snapshot files to simulate leftovers from previous runs.
1450     {
1451         ASSERT_TRUE(AcquireLock());
1452         auto local_lock = std::move(lock_);
1453         SnapshotStatus status;
1454         status.set_name("sys_b");
1455         ASSERT_TRUE(sm->WriteSnapshotStatus(local_lock.get(), status));
1456         ASSERT_TRUE(image_manager_->CreateBackingImage("sys_b-cow-img", 1_MiB,
1457                                                        IImageManager::CREATE_IMAGE_DEFAULT));
1458     }
1459 
1460     // Redo the update.
1461     ASSERT_TRUE(sm->BeginUpdate());
1462     ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
1463 
1464     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1465 
1466     // Check that target partitions can be mapped.
1467     EXPECT_TRUE(MapUpdateSnapshots());
1468 }
1469 
1470 // Test that the old partitions are not modified.
TEST_F(SnapshotUpdateTest,TestRollback)1471 TEST_F(SnapshotUpdateTest, TestRollback) {
1472     // Execute the update.
1473     ASSERT_TRUE(sm->BeginUpdate());
1474     ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
1475 
1476     AddOperationForPartitions();
1477 
1478     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1479 
1480     // Write some data to target partitions.
1481     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1482         ASSERT_TRUE(WriteSnapshotAndHash(name));
1483     }
1484 
1485     // Assert that source partitions aren't affected.
1486     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1487         ASSERT_TRUE(IsPartitionUnchanged(name));
1488     }
1489 
1490     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1491 
1492     // Simulate shutting down the device.
1493     ASSERT_TRUE(UnmapAll());
1494 
1495     // After reboot, init does first stage mount.
1496     auto init = NewManagerForFirstStageMount("_b");
1497     ASSERT_NE(init, nullptr);
1498     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1499     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1500 
1501     // Check that the target partitions have the same content.
1502     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1503         ASSERT_TRUE(IsPartitionUnchanged(name));
1504     }
1505 
1506     // Simulate shutting down the device again.
1507     ASSERT_TRUE(UnmapAll());
1508     init = NewManagerForFirstStageMount("_a");
1509     ASSERT_NE(init, nullptr);
1510     ASSERT_FALSE(init->NeedSnapshotsInFirstStageMount());
1511     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1512 
1513     // Assert that the source partitions aren't affected.
1514     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1515         ASSERT_TRUE(IsPartitionUnchanged(name));
1516     }
1517 }
1518 
1519 // Test that if an update is applied but not booted into, it can be canceled.
TEST_F(SnapshotUpdateTest,CancelAfterApply)1520 TEST_F(SnapshotUpdateTest, CancelAfterApply) {
1521     ASSERT_TRUE(sm->BeginUpdate());
1522     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1523     ASSERT_TRUE(sm->CancelUpdate());
1524 }
1525 
ToIntervals(const std::vector<std::unique_ptr<Extent>> & extents)1526 static std::vector<Interval> ToIntervals(const std::vector<std::unique_ptr<Extent>>& extents) {
1527     std::vector<Interval> ret;
1528     std::transform(extents.begin(), extents.end(), std::back_inserter(ret),
1529                    [](const auto& extent) { return extent->AsLinearExtent()->AsInterval(); });
1530     return ret;
1531 }
1532 
1533 // Test that at the second update, old COW partition spaces are reclaimed.
TEST_F(SnapshotUpdateTest,ReclaimCow)1534 TEST_F(SnapshotUpdateTest, ReclaimCow) {
1535     // Make sure VABC cows are small enough that they fit in fake_super.
1536     sys_->set_estimate_cow_size(64_KiB);
1537     vnd_->set_estimate_cow_size(64_KiB);
1538     prd_->set_estimate_cow_size(64_KiB);
1539 
1540     // Execute the first update.
1541     ASSERT_TRUE(sm->BeginUpdate());
1542     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1543     ASSERT_TRUE(MapUpdateSnapshots());
1544     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1545 
1546     // Simulate shutting down the device.
1547     ASSERT_TRUE(UnmapAll());
1548 
1549     // After reboot, init does first stage mount.
1550     auto init = NewManagerForFirstStageMount("_b");
1551     ASSERT_NE(init, nullptr);
1552     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1553     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1554     init = nullptr;
1555 
1556     // Initiate the merge and wait for it to be completed.
1557     auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
1558     ASSERT_TRUE(new_sm->InitiateMerge());
1559     ASSERT_EQ(UpdateState::MergeCompleted, new_sm->ProcessUpdateState());
1560 
1561     // Execute the second update.
1562     ASSERT_TRUE(new_sm->BeginUpdate());
1563     ASSERT_TRUE(new_sm->CreateUpdateSnapshots(manifest_));
1564 
1565     // Check that the old COW space is reclaimed and does not occupy space of mapped partitions.
1566     auto src = MetadataBuilder::New(*opener_, "super", 1);
1567     ASSERT_NE(src, nullptr);
1568     auto tgt = MetadataBuilder::New(*opener_, "super", 0);
1569     ASSERT_NE(tgt, nullptr);
1570     for (const auto& cow_part_name : {"sys_a-cow", "vnd_a-cow", "prd_a-cow"}) {
1571         auto* cow_part = tgt->FindPartition(cow_part_name);
1572         ASSERT_NE(nullptr, cow_part) << cow_part_name << " does not exist in target metadata";
1573         auto cow_intervals = ToIntervals(cow_part->extents());
1574         for (const auto& old_part_name : {"sys_b", "vnd_b", "prd_b"}) {
1575             auto* old_part = src->FindPartition(old_part_name);
1576             ASSERT_NE(nullptr, old_part) << old_part_name << " does not exist in source metadata";
1577             auto old_intervals = ToIntervals(old_part->extents());
1578 
1579             auto intersect = Interval::Intersect(cow_intervals, old_intervals);
1580             ASSERT_TRUE(intersect.empty()) << "COW uses space of source partitions";
1581         }
1582     }
1583 }
1584 
TEST_F(SnapshotUpdateTest,RetrofitAfterRegularAb)1585 TEST_F(SnapshotUpdateTest, RetrofitAfterRegularAb) {
1586     constexpr auto kRetrofitGroupSize = kGroupSize / 2;
1587 
1588     // Initialize device-mapper / disk
1589     ASSERT_TRUE(UnmapAll());
1590     FormatFakeSuper();
1591 
1592     // Setup source partition metadata to have both _a and _b partitions.
1593     src_ = MetadataBuilder::New(*opener_, "super", 0);
1594     ASSERT_NE(nullptr, src_);
1595     for (const auto& suffix : {"_a"s, "_b"s}) {
1596         ASSERT_TRUE(src_->AddGroup(group_->name() + suffix, kRetrofitGroupSize));
1597         for (const auto& name : {"sys"s, "vnd"s, "prd"s}) {
1598             auto partition = src_->AddPartition(name + suffix, group_->name() + suffix, 0);
1599             ASSERT_NE(nullptr, partition);
1600             ASSERT_TRUE(src_->ResizePartition(partition, 2_MiB));
1601         }
1602     }
1603     auto metadata = src_->Export();
1604     ASSERT_NE(nullptr, metadata);
1605     ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
1606 
1607     // Flash source partitions
1608     std::string path;
1609     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1610         ASSERT_TRUE(CreateLogicalPartition(
1611                 CreateLogicalPartitionParams{
1612                         .block_device = fake_super,
1613                         .metadata_slot = 0,
1614                         .partition_name = name,
1615                         .timeout_ms = 1s,
1616                         .partition_opener = opener_.get(),
1617                 },
1618                 &path));
1619         ASSERT_TRUE(WriteRandomData(path));
1620         auto hash = GetHash(path);
1621         ASSERT_TRUE(hash.has_value());
1622         hashes_[name] = *hash;
1623     }
1624 
1625     // Setup manifest.
1626     group_->set_size(kRetrofitGroupSize);
1627     for (auto* partition : {sys_, vnd_, prd_}) {
1628         SetSize(partition, 2_MiB);
1629     }
1630     AddOperationForPartitions();
1631 
1632     ASSERT_TRUE(sm->BeginUpdate());
1633     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1634 
1635     // Test that COW image should not be created for retrofit devices; super
1636     // should be big enough.
1637     ASSERT_FALSE(image_manager_->BackingImageExists("sys_b-cow-img"));
1638     ASSERT_FALSE(image_manager_->BackingImageExists("vnd_b-cow-img"));
1639     ASSERT_FALSE(image_manager_->BackingImageExists("prd_b-cow-img"));
1640 
1641     // Write some data to target partitions.
1642     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1643         ASSERT_TRUE(WriteSnapshotAndHash(name));
1644     }
1645 
1646     // Assert that source partitions aren't affected.
1647     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
1648         ASSERT_TRUE(IsPartitionUnchanged(name));
1649     }
1650 
1651     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1652 }
1653 
TEST_F(SnapshotUpdateTest,MergeCannotRemoveCow)1654 TEST_F(SnapshotUpdateTest, MergeCannotRemoveCow) {
1655     // Make source partitions as big as possible to force COW image to be created.
1656     SetSize(sys_, 10_MiB);
1657     SetSize(vnd_, 10_MiB);
1658     SetSize(prd_, 10_MiB);
1659     sys_->set_estimate_cow_size(12_MiB);
1660     vnd_->set_estimate_cow_size(12_MiB);
1661     prd_->set_estimate_cow_size(12_MiB);
1662 
1663     src_ = MetadataBuilder::New(*opener_, "super", 0);
1664     ASSERT_NE(src_, nullptr);
1665     src_->RemoveGroupAndPartitions(group_->name() + "_a");
1666     src_->RemoveGroupAndPartitions(group_->name() + "_b");
1667     ASSERT_TRUE(FillFakeMetadata(src_.get(), manifest_, "_a"));
1668     auto metadata = src_->Export();
1669     ASSERT_NE(nullptr, metadata);
1670     ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
1671 
1672     // OTA client blindly unmaps all partitions that are possibly mapped.
1673     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1674         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
1675     }
1676 
1677     // Add operations for sys. The whole device is written.
1678     AddOperation(sys_);
1679 
1680     // Execute the update.
1681     ASSERT_TRUE(sm->BeginUpdate());
1682     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1683     ASSERT_TRUE(MapUpdateSnapshots());
1684     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1685 
1686     // Simulate shutting down the device.
1687     ASSERT_TRUE(UnmapAll());
1688 
1689     // After reboot, init does first stage mount.
1690     // Normally we should use NewManagerForFirstStageMount, but if so,
1691     // "gsid.mapped_image.sys_b-cow-img" won't be set.
1692     auto init = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
1693     ASSERT_NE(init, nullptr);
1694     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1695 
1696     // Keep an open handle to the cow device. This should cause the merge to
1697     // be incomplete.
1698     auto cow_path = android::base::GetProperty("gsid.mapped_image.sys_b-cow-img", "");
1699     unique_fd fd(open(cow_path.c_str(), O_RDONLY | O_CLOEXEC));
1700     ASSERT_GE(fd, 0);
1701 
1702     // COW cannot be removed due to open fd, so expect a soft failure.
1703     ASSERT_TRUE(init->InitiateMerge());
1704     ASSERT_EQ(UpdateState::MergeNeedsReboot, init->ProcessUpdateState());
1705 
1706     // Simulate shutting down the device.
1707     fd.reset();
1708     ASSERT_TRUE(UnmapAll());
1709 
1710     // init does first stage mount again.
1711     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1712 
1713     // sys_b should be mapped as a dm-linear device directly.
1714     ASSERT_FALSE(sm->IsSnapshotDevice("sys_b", nullptr));
1715 
1716     // Merge should be able to complete now.
1717     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
1718 }
1719 
1720 class MetadataMountedTest : public ::testing::Test {
1721   public:
1722     // This is so main() can instantiate this to invoke Cleanup.
TestBody()1723     virtual void TestBody() override {}
SetUp()1724     void SetUp() override {
1725         SKIP_IF_NON_VIRTUAL_AB();
1726         metadata_dir_ = test_device->GetMetadataDir();
1727         ASSERT_TRUE(ReadDefaultFstab(&fstab_));
1728     }
TearDown()1729     void TearDown() override {
1730         RETURN_IF_NON_VIRTUAL_AB();
1731         SetUp();
1732         // Remount /metadata
1733         test_device->set_recovery(false);
1734         EXPECT_TRUE(android::fs_mgr::EnsurePathMounted(&fstab_, metadata_dir_));
1735     }
IsMetadataMounted()1736     AssertionResult IsMetadataMounted() {
1737         Fstab mounted_fstab;
1738         if (!ReadFstabFromFile("/proc/mounts", &mounted_fstab)) {
1739             ADD_FAILURE() << "Failed to scan mounted volumes";
1740             return AssertionFailure() << "Failed to scan mounted volumes";
1741         }
1742 
1743         auto entry = GetEntryForPath(&fstab_, metadata_dir_);
1744         if (entry == nullptr) {
1745             return AssertionFailure() << "No mount point found in fstab for path " << metadata_dir_;
1746         }
1747 
1748         auto mv = GetEntryForMountPoint(&mounted_fstab, entry->mount_point);
1749         if (mv == nullptr) {
1750             return AssertionFailure() << metadata_dir_ << " is not mounted";
1751         }
1752         return AssertionSuccess() << metadata_dir_ << " is mounted";
1753     }
1754     std::string metadata_dir_;
1755     Fstab fstab_;
1756 };
1757 
MountMetadata()1758 void MountMetadata() {
1759     MetadataMountedTest().TearDown();
1760 }
1761 
TEST_F(MetadataMountedTest,Android)1762 TEST_F(MetadataMountedTest, Android) {
1763     auto device = sm->EnsureMetadataMounted();
1764     EXPECT_NE(nullptr, device);
1765     device.reset();
1766 
1767     EXPECT_TRUE(IsMetadataMounted());
1768     EXPECT_TRUE(sm->CancelUpdate()) << "Metadata dir should never be unmounted in Android mode";
1769 }
1770 
TEST_F(MetadataMountedTest,Recovery)1771 TEST_F(MetadataMountedTest, Recovery) {
1772     test_device->set_recovery(true);
1773     metadata_dir_ = test_device->GetMetadataDir();
1774 
1775     EXPECT_TRUE(android::fs_mgr::EnsurePathUnmounted(&fstab_, metadata_dir_));
1776     EXPECT_FALSE(IsMetadataMounted());
1777 
1778     auto device = sm->EnsureMetadataMounted();
1779     EXPECT_NE(nullptr, device);
1780     EXPECT_TRUE(IsMetadataMounted());
1781 
1782     device.reset();
1783     EXPECT_FALSE(IsMetadataMounted());
1784 }
1785 
1786 // Test that during a merge, we can wipe data in recovery.
TEST_F(SnapshotUpdateTest,MergeInRecovery)1787 TEST_F(SnapshotUpdateTest, MergeInRecovery) {
1788     // Execute the first update.
1789     ASSERT_TRUE(sm->BeginUpdate());
1790     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1791     ASSERT_TRUE(MapUpdateSnapshots());
1792     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1793 
1794     // Simulate shutting down the device.
1795     ASSERT_TRUE(UnmapAll());
1796 
1797     // After reboot, init does first stage mount.
1798     auto init = NewManagerForFirstStageMount("_b");
1799     ASSERT_NE(init, nullptr);
1800     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1801     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1802     init = nullptr;
1803 
1804     // Initiate the merge and then immediately stop it to simulate a reboot.
1805     auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
1806     ASSERT_TRUE(new_sm->InitiateMerge());
1807     ASSERT_TRUE(UnmapAll());
1808 
1809     // Simulate a reboot into recovery.
1810     auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
1811     test_device->set_recovery(true);
1812     new_sm = NewManagerForFirstStageMount(test_device.release());
1813 
1814     ASSERT_TRUE(new_sm->HandleImminentDataWipe());
1815     ASSERT_EQ(new_sm->GetUpdateState(), UpdateState::None);
1816 }
1817 
1818 // Test that a merge does not clear the snapshot state in fastboot.
TEST_F(SnapshotUpdateTest,MergeInFastboot)1819 TEST_F(SnapshotUpdateTest, MergeInFastboot) {
1820     // Execute the first update.
1821     ASSERT_TRUE(sm->BeginUpdate());
1822     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1823     ASSERT_TRUE(MapUpdateSnapshots());
1824     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1825 
1826     // Simulate shutting down the device.
1827     ASSERT_TRUE(UnmapAll());
1828 
1829     // After reboot, init does first stage mount.
1830     auto init = NewManagerForFirstStageMount("_b");
1831     ASSERT_NE(init, nullptr);
1832     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
1833     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1834     init = nullptr;
1835 
1836     // Initiate the merge and then immediately stop it to simulate a reboot.
1837     auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
1838     ASSERT_TRUE(new_sm->InitiateMerge());
1839     ASSERT_TRUE(UnmapAll());
1840 
1841     // Simulate a reboot into recovery.
1842     auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
1843     test_device->set_recovery(true);
1844     new_sm = NewManagerForFirstStageMount(test_device.release());
1845 
1846     ASSERT_TRUE(new_sm->FinishMergeInRecovery());
1847 
1848     auto mount = new_sm->EnsureMetadataMounted();
1849     ASSERT_TRUE(mount && mount->HasDevice());
1850     ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
1851 
1852     // Finish the merge in a normal boot.
1853     test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
1854     init = NewManagerForFirstStageMount(test_device.release());
1855     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1856     init = nullptr;
1857 
1858     test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
1859     new_sm = NewManagerForFirstStageMount(test_device.release());
1860     ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
1861     ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::None);
1862 }
1863 
1864 // Test that after an OTA, before a merge, we can wipe data in recovery.
TEST_F(SnapshotUpdateTest,DataWipeRollbackInRecovery)1865 TEST_F(SnapshotUpdateTest, DataWipeRollbackInRecovery) {
1866     // Execute the first update.
1867     ASSERT_TRUE(sm->BeginUpdate());
1868     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1869     ASSERT_TRUE(MapUpdateSnapshots());
1870     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1871 
1872     // Simulate shutting down the device.
1873     ASSERT_TRUE(UnmapAll());
1874 
1875     // Simulate a reboot into recovery.
1876     auto test_device = new TestDeviceInfo(fake_super, "_b");
1877     test_device->set_recovery(true);
1878     auto new_sm = NewManagerForFirstStageMount(test_device);
1879 
1880     ASSERT_TRUE(new_sm->HandleImminentDataWipe());
1881     // Manually mount metadata so that we can call GetUpdateState() below.
1882     MountMetadata();
1883     EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
1884     EXPECT_TRUE(test_device->IsSlotUnbootable(1));
1885     EXPECT_FALSE(test_device->IsSlotUnbootable(0));
1886 }
1887 
1888 // Test that after an OTA and a bootloader rollback with no merge, we can wipe
1889 // data in recovery.
TEST_F(SnapshotUpdateTest,DataWipeAfterRollback)1890 TEST_F(SnapshotUpdateTest, DataWipeAfterRollback) {
1891     // Execute the first update.
1892     ASSERT_TRUE(sm->BeginUpdate());
1893     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1894     ASSERT_TRUE(MapUpdateSnapshots());
1895     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
1896 
1897     // Simulate shutting down the device.
1898     ASSERT_TRUE(UnmapAll());
1899 
1900     // Simulate a rollback, with reboot into recovery.
1901     auto test_device = new TestDeviceInfo(fake_super, "_a");
1902     test_device->set_recovery(true);
1903     auto new_sm = NewManagerForFirstStageMount(test_device);
1904 
1905     ASSERT_TRUE(new_sm->HandleImminentDataWipe());
1906     EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
1907     EXPECT_FALSE(test_device->IsSlotUnbootable(0));
1908     EXPECT_FALSE(test_device->IsSlotUnbootable(1));
1909 }
1910 
1911 // Test update package that requests data wipe.
TEST_F(SnapshotUpdateTest,DataWipeRequiredInPackage)1912 TEST_F(SnapshotUpdateTest, DataWipeRequiredInPackage) {
1913     AddOperationForPartitions();
1914     // Execute the update.
1915     ASSERT_TRUE(sm->BeginUpdate());
1916     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1917 
1918     // Write some data to target partitions.
1919     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1920         ASSERT_TRUE(WriteSnapshotAndHash(name)) << name;
1921     }
1922 
1923     ASSERT_TRUE(sm->FinishedSnapshotWrites(true /* wipe */));
1924 
1925     // Simulate shutting down the device.
1926     ASSERT_TRUE(UnmapAll());
1927 
1928     // Simulate a reboot into recovery.
1929     auto test_device = new TestDeviceInfo(fake_super, "_b");
1930     test_device->set_recovery(true);
1931     auto new_sm = NewManagerForFirstStageMount(test_device);
1932 
1933     ASSERT_TRUE(new_sm->HandleImminentDataWipe());
1934     // Manually mount metadata so that we can call GetUpdateState() below.
1935     MountMetadata();
1936     EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
1937     ASSERT_FALSE(test_device->IsSlotUnbootable(1));
1938     ASSERT_FALSE(test_device->IsSlotUnbootable(0));
1939 
1940     // Now reboot into new slot.
1941     test_device = new TestDeviceInfo(fake_super, "_b");
1942     auto init = NewManagerForFirstStageMount(test_device);
1943     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
1944     // Verify that we are on the downgraded build.
1945     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1946         ASSERT_TRUE(IsPartitionUnchanged(name)) << name;
1947     }
1948 }
1949 
1950 // Test update package that requests data wipe.
TEST_F(SnapshotUpdateTest,DataWipeWithStaleSnapshots)1951 TEST_F(SnapshotUpdateTest, DataWipeWithStaleSnapshots) {
1952     AddOperationForPartitions();
1953 
1954     // Execute the update.
1955     ASSERT_TRUE(sm->BeginUpdate());
1956     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
1957 
1958     // Write some data to target partitions.
1959     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
1960         ASSERT_TRUE(WriteSnapshotAndHash(name)) << name;
1961     }
1962 
1963     // Create a stale snapshot that should not exist.
1964     {
1965         ASSERT_TRUE(AcquireLock());
1966 
1967         PartitionCowCreator cow_creator = {
1968                 .compression_enabled = IsCompressionEnabled(),
1969                 .compression_algorithm = IsCompressionEnabled() ? "gz" : "none",
1970         };
1971         SnapshotStatus status;
1972         status.set_name("sys_a");
1973         status.set_device_size(1_MiB);
1974         status.set_snapshot_size(2_MiB);
1975         status.set_cow_partition_size(2_MiB);
1976 
1977         ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
1978         lock_ = nullptr;
1979 
1980         ASSERT_TRUE(sm->EnsureImageManager());
1981         ASSERT_TRUE(sm->image_manager()->CreateBackingImage("sys_a", 1_MiB, 0));
1982     }
1983 
1984     ASSERT_TRUE(sm->FinishedSnapshotWrites(true /* wipe */));
1985 
1986     // Simulate shutting down the device.
1987     ASSERT_TRUE(UnmapAll());
1988 
1989     // Simulate a reboot into recovery.
1990     auto test_device = new TestDeviceInfo(fake_super, "_b");
1991     test_device->set_recovery(true);
1992     auto new_sm = NewManagerForFirstStageMount(test_device);
1993 
1994     ASSERT_TRUE(new_sm->HandleImminentDataWipe());
1995     // Manually mount metadata so that we can call GetUpdateState() below.
1996     MountMetadata();
1997     EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
1998     ASSERT_FALSE(test_device->IsSlotUnbootable(1));
1999     ASSERT_FALSE(test_device->IsSlotUnbootable(0));
2000 
2001     // Now reboot into new slot.
2002     test_device = new TestDeviceInfo(fake_super, "_b");
2003     auto init = NewManagerForFirstStageMount(test_device);
2004     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
2005     // Verify that we are on the downgraded build.
2006     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
2007         ASSERT_TRUE(IsPartitionUnchanged(name)) << name;
2008     }
2009 }
2010 
TEST_F(SnapshotUpdateTest,Hashtree)2011 TEST_F(SnapshotUpdateTest, Hashtree) {
2012     constexpr auto partition_size = 4_MiB;
2013     constexpr auto data_size = 3_MiB;
2014     constexpr auto hashtree_size = 512_KiB;
2015     constexpr auto fec_size = partition_size - data_size - hashtree_size;
2016 
2017     const auto block_size = manifest_.block_size();
2018     SetSize(sys_, partition_size);
2019     AddOperation(sys_, data_size);
2020 
2021     sys_->set_estimate_cow_size(partition_size + data_size);
2022 
2023     // Set hastree extents.
2024     sys_->mutable_hash_tree_data_extent()->set_start_block(0);
2025     sys_->mutable_hash_tree_data_extent()->set_num_blocks(data_size / block_size);
2026 
2027     sys_->mutable_hash_tree_extent()->set_start_block(data_size / block_size);
2028     sys_->mutable_hash_tree_extent()->set_num_blocks(hashtree_size / block_size);
2029 
2030     // Set FEC extents.
2031     sys_->mutable_fec_data_extent()->set_start_block(0);
2032     sys_->mutable_fec_data_extent()->set_num_blocks((data_size + hashtree_size) / block_size);
2033 
2034     sys_->mutable_fec_extent()->set_start_block((data_size + hashtree_size) / block_size);
2035     sys_->mutable_fec_extent()->set_num_blocks(fec_size / block_size);
2036 
2037     ASSERT_TRUE(sm->BeginUpdate());
2038     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2039 
2040     // Map and write some data to target partition.
2041     ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
2042     ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
2043 
2044     // Finish update.
2045     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
2046 
2047     // Simulate shutting down the device.
2048     ASSERT_TRUE(UnmapAll());
2049 
2050     // After reboot, init does first stage mount.
2051     auto init = NewManagerForFirstStageMount("_b");
2052     ASSERT_NE(init, nullptr);
2053     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
2054     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
2055 
2056     // Check that the target partition have the same content. Hashtree and FEC extents
2057     // should be accounted for.
2058     ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
2059 }
2060 
2061 // Test for overflow bit after update
TEST_F(SnapshotUpdateTest,Overflow)2062 TEST_F(SnapshotUpdateTest, Overflow) {
2063     if (IsCompressionEnabled()) {
2064         GTEST_SKIP() << "No overflow bit set for userspace COWs";
2065     }
2066 
2067     const auto actual_write_size = GetSize(sys_);
2068     const auto declared_write_size = actual_write_size - 1_MiB;
2069 
2070     AddOperation(sys_, declared_write_size);
2071 
2072     // Execute the update.
2073     ASSERT_TRUE(sm->BeginUpdate());
2074     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2075 
2076     // Map and write some data to target partitions.
2077     ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
2078     ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
2079 
2080     std::vector<android::dm::DeviceMapper::TargetInfo> table;
2081     ASSERT_TRUE(DeviceMapper::Instance().GetTableStatus("sys_b", &table));
2082     ASSERT_EQ(1u, table.size());
2083     EXPECT_TRUE(table[0].IsOverflowSnapshot());
2084 
2085     ASSERT_FALSE(sm->FinishedSnapshotWrites(false))
2086             << "FinishedSnapshotWrites should detect overflow of CoW device.";
2087 }
2088 
TEST_F(SnapshotUpdateTest,LowSpace)2089 TEST_F(SnapshotUpdateTest, LowSpace) {
2090     static constexpr auto kMaxFree = 10_MiB;
2091     auto userdata = std::make_unique<LowSpaceUserdata>();
2092     ASSERT_TRUE(userdata->Init(kMaxFree));
2093 
2094     // Grow all partitions to 10_MiB, total 30_MiB. This requires 30 MiB of CoW space. After
2095     // using the empty space in super (< 1 MiB), it uses 30 MiB of /userdata space.
2096     constexpr uint64_t partition_size = 10_MiB;
2097     SetSize(sys_, partition_size);
2098     SetSize(vnd_, partition_size);
2099     SetSize(prd_, partition_size);
2100     sys_->set_estimate_cow_size(partition_size);
2101     vnd_->set_estimate_cow_size(partition_size);
2102     prd_->set_estimate_cow_size(partition_size);
2103 
2104     AddOperationForPartitions();
2105 
2106     // Execute the update.
2107     ASSERT_TRUE(sm->BeginUpdate());
2108     auto res = sm->CreateUpdateSnapshots(manifest_);
2109     ASSERT_FALSE(res);
2110     ASSERT_EQ(Return::ErrorCode::NO_SPACE, res.error_code());
2111     ASSERT_GE(res.required_size(), 14_MiB);
2112     ASSERT_LT(res.required_size(), 40_MiB);
2113 }
2114 
TEST_F(SnapshotUpdateTest,AddPartition)2115 TEST_F(SnapshotUpdateTest, AddPartition) {
2116     // OTA client blindly unmaps all partitions that are possibly mapped.
2117     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
2118         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
2119     }
2120 
2121     group_->add_partition_names("dlkm");
2122 
2123     auto dlkm = manifest_.add_partitions();
2124     dlkm->set_partition_name("dlkm");
2125     dlkm->set_estimate_cow_size(2_MiB);
2126     SetSize(dlkm, 3_MiB);
2127 
2128     // Grow all partitions. Set |prd| large enough that |sys| and |vnd|'s COWs
2129     // fit in super, but not |prd|.
2130     constexpr uint64_t partition_size = 3788_KiB;
2131     SetSize(sys_, partition_size);
2132     SetSize(vnd_, partition_size);
2133     SetSize(prd_, partition_size);
2134     SetSize(dlkm, partition_size);
2135 
2136     AddOperationForPartitions({sys_, vnd_, prd_, dlkm});
2137 
2138     // Execute the update.
2139     ASSERT_TRUE(sm->BeginUpdate());
2140     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2141 
2142     // Write some data to target partitions.
2143     for (const auto& name : {"sys_b", "vnd_b", "prd_b", "dlkm_b"}) {
2144         ASSERT_TRUE(WriteSnapshotAndHash(name));
2145     }
2146 
2147     // Assert that source partitions aren't affected.
2148     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
2149         ASSERT_TRUE(IsPartitionUnchanged(name));
2150     }
2151 
2152     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
2153 
2154     // Simulate shutting down the device.
2155     ASSERT_TRUE(UnmapAll());
2156 
2157     // After reboot, init does first stage mount.
2158     auto init = NewManagerForFirstStageMount("_b");
2159     ASSERT_NE(init, nullptr);
2160 
2161     ASSERT_TRUE(init->EnsureSnapuserdConnected());
2162     init->set_use_first_stage_snapuserd(true);
2163 
2164     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
2165     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
2166 
2167     // Check that the target partitions have the same content.
2168     std::vector<std::string> partitions = {"sys_b", "vnd_b", "prd_b", "dlkm_b"};
2169     for (const auto& name : partitions) {
2170         ASSERT_TRUE(IsPartitionUnchanged(name));
2171     }
2172 
2173     ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
2174     for (const auto& name : partitions) {
2175         ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete(name + "-user-cow-init"));
2176     }
2177 
2178     // Initiate the merge and wait for it to be completed.
2179     ASSERT_TRUE(init->InitiateMerge());
2180     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
2181 
2182     // Check that the target partitions have the same content after the merge.
2183     for (const auto& name : {"sys_b", "vnd_b", "prd_b", "dlkm_b"}) {
2184         ASSERT_TRUE(IsPartitionUnchanged(name))
2185                 << "Content of " << name << " changes after the merge";
2186     }
2187 }
2188 
2189 class AutoKill final {
2190   public:
AutoKill(pid_t pid)2191     explicit AutoKill(pid_t pid) : pid_(pid) {}
~AutoKill()2192     ~AutoKill() {
2193         if (pid_ > 0) kill(pid_, SIGKILL);
2194     }
2195 
valid() const2196     bool valid() const { return pid_ > 0; }
2197 
2198   private:
2199     pid_t pid_;
2200 };
2201 
TEST_F(SnapshotUpdateTest,DaemonTransition)2202 TEST_F(SnapshotUpdateTest, DaemonTransition) {
2203     if (!IsCompressionEnabled()) {
2204         GTEST_SKIP() << "Skipping Virtual A/B Compression test";
2205     }
2206 
2207     // Ensure a connection to the second-stage daemon, but use the first-stage
2208     // code paths thereafter.
2209     ASSERT_TRUE(sm->EnsureSnapuserdConnected());
2210     sm->set_use_first_stage_snapuserd(true);
2211 
2212     AddOperationForPartitions();
2213     // Execute the update.
2214     ASSERT_TRUE(sm->BeginUpdate());
2215     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2216     ASSERT_TRUE(MapUpdateSnapshots());
2217     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
2218     ASSERT_TRUE(UnmapAll());
2219 
2220     auto init = NewManagerForFirstStageMount("_b");
2221     ASSERT_NE(init, nullptr);
2222 
2223     ASSERT_TRUE(init->EnsureSnapuserdConnected());
2224     init->set_use_first_stage_snapuserd(true);
2225 
2226     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
2227     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
2228 
2229     ASSERT_EQ(access("/dev/dm-user/sys_b-user-cow-init", F_OK), 0);
2230     ASSERT_EQ(access("/dev/dm-user/sys_b-user-cow", F_OK), -1);
2231 
2232     ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
2233 
2234     // :TODO: this is a workaround to ensure the handler list stays empty. We
2235     // should make this test more like actual init, and spawn two copies of
2236     // snapuserd, given how many other tests we now have for normal snapuserd.
2237     ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("sys_b-user-cow-init"));
2238     ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("vnd_b-user-cow-init"));
2239     ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("prd_b-user-cow-init"));
2240 
2241     // The control device should have been renamed.
2242     ASSERT_TRUE(android::fs_mgr::WaitForFileDeleted("/dev/dm-user/sys_b-user-cow-init", 10s));
2243     ASSERT_EQ(access("/dev/dm-user/sys_b-user-cow", F_OK), 0);
2244 }
2245 
TEST_F(SnapshotUpdateTest,MapAllSnapshots)2246 TEST_F(SnapshotUpdateTest, MapAllSnapshots) {
2247     AddOperationForPartitions();
2248     // Execute the update.
2249     ASSERT_TRUE(sm->BeginUpdate());
2250     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2251     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
2252         ASSERT_TRUE(WriteSnapshotAndHash(name));
2253     }
2254     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
2255     ASSERT_TRUE(sm->MapAllSnapshots(10s));
2256 
2257     // Read bytes back and verify they match the cache.
2258     ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
2259 
2260     ASSERT_TRUE(sm->UnmapAllSnapshots());
2261 }
2262 
TEST_F(SnapshotUpdateTest,CancelOnTargetSlot)2263 TEST_F(SnapshotUpdateTest, CancelOnTargetSlot) {
2264     AddOperationForPartitions();
2265 
2266     // Execute the update from B->A.
2267     test_device->set_slot_suffix("_b");
2268     ASSERT_TRUE(sm->BeginUpdate());
2269     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2270 
2271     std::string path;
2272     ASSERT_TRUE(CreateLogicalPartition(
2273             CreateLogicalPartitionParams{
2274                     .block_device = fake_super,
2275                     .metadata_slot = 0,
2276                     .partition_name = "sys_a",
2277                     .timeout_ms = 1s,
2278                     .partition_opener = opener_.get(),
2279             },
2280             &path));
2281 
2282     // Hold sys_a open so it can't be unmapped.
2283     unique_fd fd(open(path.c_str(), O_RDONLY));
2284 
2285     // Switch back to "A", make sure we can cancel. Instead of unmapping sys_a
2286     // we should simply delete the old snapshots.
2287     test_device->set_slot_suffix("_a");
2288     ASSERT_TRUE(sm->BeginUpdate());
2289 }
2290 
2291 class FlashAfterUpdateTest : public SnapshotUpdateTest,
2292                              public WithParamInterface<std::tuple<uint32_t, bool>> {
2293   public:
InitiateMerge(const std::string & slot_suffix)2294     AssertionResult InitiateMerge(const std::string& slot_suffix) {
2295         auto sm = SnapshotManager::New(new TestDeviceInfo(fake_super, slot_suffix));
2296         if (!sm->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_)) {
2297             return AssertionFailure() << "Cannot CreateLogicalAndSnapshotPartitions";
2298         }
2299         if (!sm->InitiateMerge()) {
2300             return AssertionFailure() << "Cannot initiate merge";
2301         }
2302         return AssertionSuccess();
2303     }
2304 };
2305 
TEST_P(FlashAfterUpdateTest,FlashSlotAfterUpdate)2306 TEST_P(FlashAfterUpdateTest, FlashSlotAfterUpdate) {
2307     // OTA client blindly unmaps all partitions that are possibly mapped.
2308     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
2309         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
2310     }
2311 
2312     // Execute the update.
2313     ASSERT_TRUE(sm->BeginUpdate());
2314     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
2315     ASSERT_TRUE(MapUpdateSnapshots());
2316     ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
2317 
2318     // Simulate shutting down the device.
2319     ASSERT_TRUE(UnmapAll());
2320 
2321     bool after_merge = std::get<1>(GetParam());
2322     if (after_merge) {
2323         ASSERT_TRUE(InitiateMerge("_b"));
2324         // Simulate shutting down the device after merge has initiated.
2325         ASSERT_TRUE(UnmapAll());
2326     }
2327 
2328     auto flashed_slot = std::get<0>(GetParam());
2329     auto flashed_slot_suffix = SlotSuffixForSlotNumber(flashed_slot);
2330 
2331     // Simulate flashing |flashed_slot|. This clears the UPDATED flag.
2332     auto flashed_builder = MetadataBuilder::New(*opener_, "super", flashed_slot);
2333     ASSERT_NE(flashed_builder, nullptr);
2334     flashed_builder->RemoveGroupAndPartitions(group_->name() + flashed_slot_suffix);
2335     flashed_builder->RemoveGroupAndPartitions(kCowGroupName);
2336     ASSERT_TRUE(FillFakeMetadata(flashed_builder.get(), manifest_, flashed_slot_suffix));
2337 
2338     // Deliberately remove a partition from this build so that
2339     // InitiateMerge do not switch state to "merging". This is possible in
2340     // practice because the list of dynamic partitions may change.
2341     ASSERT_NE(nullptr, flashed_builder->FindPartition("prd" + flashed_slot_suffix));
2342     flashed_builder->RemovePartition("prd" + flashed_slot_suffix);
2343 
2344     // Note that fastbootd always updates the partition table of both slots.
2345     auto flashed_metadata = flashed_builder->Export();
2346     ASSERT_NE(nullptr, flashed_metadata);
2347     ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *flashed_metadata, 0));
2348     ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *flashed_metadata, 1));
2349 
2350     std::string path;
2351     for (const auto& name : {"sys", "vnd"}) {
2352         ASSERT_TRUE(CreateLogicalPartition(
2353                 CreateLogicalPartitionParams{
2354                         .block_device = fake_super,
2355                         .metadata_slot = flashed_slot,
2356                         .partition_name = name + flashed_slot_suffix,
2357                         .timeout_ms = 1s,
2358                         .partition_opener = opener_.get(),
2359                 },
2360                 &path));
2361         ASSERT_TRUE(WriteRandomData(path));
2362         auto hash = GetHash(path);
2363         ASSERT_TRUE(hash.has_value());
2364         hashes_[name + flashed_slot_suffix] = *hash;
2365     }
2366 
2367     // Simulate shutting down the device after flash.
2368     ASSERT_TRUE(UnmapAll());
2369 
2370     // Simulate reboot. After reboot, init does first stage mount.
2371     auto init = NewManagerForFirstStageMount(flashed_slot_suffix);
2372     ASSERT_NE(init, nullptr);
2373 
2374     if (flashed_slot && after_merge) {
2375         ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
2376     }
2377     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
2378 
2379     // Check that the target partitions have the same content.
2380     for (const auto& name : {"sys", "vnd"}) {
2381         ASSERT_TRUE(IsPartitionUnchanged(name + flashed_slot_suffix));
2382     }
2383 
2384     // There should be no snapshot to merge.
2385     auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, flashed_slot_suffix));
2386     if (flashed_slot == 0 && after_merge) {
2387         ASSERT_EQ(UpdateState::MergeCompleted, new_sm->ProcessUpdateState());
2388     } else {
2389         // update_engine calls ProcessUpdateState first -- should see Cancelled.
2390         ASSERT_EQ(UpdateState::Cancelled, new_sm->ProcessUpdateState());
2391     }
2392 
2393     // Next OTA calls CancelUpdate no matter what.
2394     ASSERT_TRUE(new_sm->CancelUpdate());
2395 }
2396 
2397 INSTANTIATE_TEST_SUITE_P(Snapshot, FlashAfterUpdateTest, Combine(Values(0, 1), Bool()),
__anon6022f21a0702(const TestParamInfo<FlashAfterUpdateTest::ParamType>& info) 2398                          [](const TestParamInfo<FlashAfterUpdateTest::ParamType>& info) {
2399                              return "Flash"s + (std::get<0>(info.param) ? "New"s : "Old"s) +
2400                                     "Slot"s + (std::get<1>(info.param) ? "After"s : "Before"s) +
2401                                     "Merge"s;
2402                          });
2403 
2404 // Test behavior of ImageManager::Create on low space scenario. These tests assumes image manager
2405 // uses /data as backup device.
2406 class ImageManagerTest : public SnapshotTest, public WithParamInterface<uint64_t> {
2407   protected:
SetUp()2408     void SetUp() override {
2409         SKIP_IF_NON_VIRTUAL_AB();
2410         SnapshotTest::SetUp();
2411         userdata_ = std::make_unique<LowSpaceUserdata>();
2412         ASSERT_TRUE(userdata_->Init(GetParam()));
2413     }
TearDown()2414     void TearDown() override {
2415         RETURN_IF_NON_VIRTUAL_AB();
2416         return;  // BUG(149738928)
2417 
2418         EXPECT_TRUE(!image_manager_->BackingImageExists(kImageName) ||
2419                     image_manager_->DeleteBackingImage(kImageName));
2420     }
2421     static constexpr const char* kImageName = "my_image";
2422     std::unique_ptr<LowSpaceUserdata> userdata_;
2423 };
2424 
TEST_P(ImageManagerTest,CreateImageEnoughAvailSpace)2425 TEST_P(ImageManagerTest, CreateImageEnoughAvailSpace) {
2426     if (userdata_->available_space() == 0) {
2427         GTEST_SKIP() << "/data is full (" << userdata_->available_space()
2428                      << " bytes available), skipping";
2429     }
2430     ASSERT_TRUE(image_manager_->CreateBackingImage(kImageName, userdata_->available_space(),
2431                                                    IImageManager::CREATE_IMAGE_DEFAULT))
2432             << "Should be able to create image with size = " << userdata_->available_space()
2433             << " bytes";
2434     ASSERT_TRUE(image_manager_->DeleteBackingImage(kImageName))
2435             << "Should be able to delete created image";
2436 }
2437 
TEST_P(ImageManagerTest,CreateImageNoSpace)2438 TEST_P(ImageManagerTest, CreateImageNoSpace) {
2439     uint64_t to_allocate = userdata_->free_space() + userdata_->bsize();
2440     auto res = image_manager_->CreateBackingImage(kImageName, to_allocate,
2441                                                   IImageManager::CREATE_IMAGE_DEFAULT);
2442     ASSERT_FALSE(res) << "Should not be able to create image with size = " << to_allocate
2443                       << " bytes because only " << userdata_->free_space() << " bytes are free";
2444     ASSERT_EQ(FiemapStatus::ErrorCode::NO_SPACE, res.error_code()) << res.string();
2445 }
2446 
ImageManagerTestParams()2447 std::vector<uint64_t> ImageManagerTestParams() {
2448     std::vector<uint64_t> ret;
2449     for (uint64_t size = 1_MiB; size <= 512_MiB; size *= 2) {
2450         ret.push_back(size);
2451     }
2452     return ret;
2453 }
2454 
2455 INSTANTIATE_TEST_SUITE_P(ImageManagerTest, ImageManagerTest, ValuesIn(ImageManagerTestParams()));
2456 
Mkdir(const std::string & path)2457 bool Mkdir(const std::string& path) {
2458     if (mkdir(path.c_str(), 0700) && errno != EEXIST) {
2459         std::cerr << "Could not mkdir " << path << ": " << strerror(errno) << std::endl;
2460         return false;
2461     }
2462     return true;
2463 }
2464 
2465 class SnapshotTestEnvironment : public ::testing::Environment {
2466   public:
~SnapshotTestEnvironment()2467     ~SnapshotTestEnvironment() override {}
2468     void SetUp() override;
2469     void TearDown() override;
2470 
2471   private:
2472     bool CreateFakeSuper();
2473 
2474     std::unique_ptr<IImageManager> super_images_;
2475 };
2476 
CreateFakeSuper()2477 bool SnapshotTestEnvironment::CreateFakeSuper() {
2478     // Create and map the fake super partition.
2479     static constexpr int kImageFlags =
2480             IImageManager::CREATE_IMAGE_DEFAULT | IImageManager::CREATE_IMAGE_ZERO_FILL;
2481     if (!super_images_->CreateBackingImage("fake-super", kSuperSize, kImageFlags)) {
2482         LOG(ERROR) << "Could not create fake super partition";
2483         return false;
2484     }
2485     if (!super_images_->MapImageDevice("fake-super", 10s, &fake_super)) {
2486         LOG(ERROR) << "Could not map fake super partition";
2487         return false;
2488     }
2489     test_device->set_fake_super(fake_super);
2490     return true;
2491 }
2492 
SetUp()2493 void SnapshotTestEnvironment::SetUp() {
2494     // b/163082876: GTEST_SKIP in Environment will make atest report incorrect results. Until
2495     // that is fixed, don't call GTEST_SKIP here, but instead call GTEST_SKIP in individual test
2496     // suites.
2497     RETURN_IF_NON_VIRTUAL_AB_MSG("Virtual A/B is not enabled, skipping global setup.\n");
2498 
2499     std::vector<std::string> paths = {
2500             // clang-format off
2501             "/data/gsi/ota/test",
2502             "/data/gsi/ota/test/super",
2503             "/metadata/gsi/ota/test",
2504             "/metadata/gsi/ota/test/super",
2505             "/metadata/ota/test",
2506             "/metadata/ota/test/snapshots",
2507             // clang-format on
2508     };
2509     for (const auto& path : paths) {
2510         ASSERT_TRUE(Mkdir(path));
2511     }
2512 
2513     // Create this once, otherwise, gsid will start/stop between each test.
2514     test_device = new TestDeviceInfo();
2515     sm = SnapshotManager::New(test_device);
2516     ASSERT_NE(nullptr, sm) << "Could not create snapshot manager";
2517 
2518     // Use a separate image manager for our fake super partition.
2519     super_images_ = IImageManager::Open("ota/test/super", 10s);
2520     ASSERT_NE(nullptr, super_images_) << "Could not create image manager";
2521 
2522     // Map the old image if one exists so we can safely unmap everything that
2523     // depends on it.
2524     bool recreate_fake_super;
2525     if (super_images_->BackingImageExists("fake-super")) {
2526         if (super_images_->IsImageMapped("fake-super")) {
2527             ASSERT_TRUE(super_images_->GetMappedImageDevice("fake-super", &fake_super));
2528         } else {
2529             ASSERT_TRUE(super_images_->MapImageDevice("fake-super", 10s, &fake_super));
2530         }
2531         test_device->set_fake_super(fake_super);
2532         recreate_fake_super = true;
2533     } else {
2534         ASSERT_TRUE(CreateFakeSuper());
2535         recreate_fake_super = false;
2536     }
2537 
2538     // Clean up previous run.
2539     MetadataMountedTest().TearDown();
2540     SnapshotUpdateTest().Cleanup();
2541     SnapshotTest().Cleanup();
2542 
2543     if (recreate_fake_super) {
2544         // Clean up any old copy.
2545         DeleteBackingImage(super_images_.get(), "fake-super");
2546         ASSERT_TRUE(CreateFakeSuper());
2547     }
2548 }
2549 
TearDown()2550 void SnapshotTestEnvironment::TearDown() {
2551     RETURN_IF_NON_VIRTUAL_AB();
2552     if (super_images_ != nullptr) {
2553         DeleteBackingImage(super_images_.get(), "fake-super");
2554     }
2555 }
2556 
2557 }  // namespace snapshot
2558 }  // namespace android
2559 
main(int argc,char ** argv)2560 int main(int argc, char** argv) {
2561     ::testing::InitGoogleTest(&argc, argv);
2562     ::testing::AddGlobalTestEnvironment(new ::android::snapshot::SnapshotTestEnvironment());
2563     return RUN_ALL_TESTS();
2564 }
2565