1 //
2 // Copyright (C) 2019 The Android Open Source Project
3 //
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
7 //
8 //      http://www.apache.org/licenses/LICENSE-2.0
9 //
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15 //
16 
17 #include "update_engine/dynamic_partition_control_android.h"
18 
19 #include <set>
20 #include <vector>
21 
22 #include <base/logging.h>
23 #include <base/strings/string_util.h>
24 #include <gmock/gmock.h>
25 #include <gtest/gtest.h>
26 #include <libavb/libavb.h>
27 #include <libsnapshot/mock_snapshot.h>
28 
29 #include "update_engine/common/mock_prefs.h"
30 #include "update_engine/common/test_utils.h"
31 #include "update_engine/dynamic_partition_test_utils.h"
32 #include "update_engine/mock_dynamic_partition_control.h"
33 
34 using android::dm::DmDeviceState;
35 using android::snapshot::MockSnapshotManager;
36 using chromeos_update_engine::test_utils::ScopedLoopbackDeviceBinder;
37 using chromeos_update_engine::test_utils::ScopedTempFile;
38 using std::string;
39 using testing::_;
40 using testing::AnyNumber;
41 using testing::AnyOf;
42 using testing::Invoke;
43 using testing::NiceMock;
44 using testing::Not;
45 using testing::Optional;
46 using testing::Return;
47 
48 namespace chromeos_update_engine {
49 
50 class DynamicPartitionControlAndroidTest : public ::testing::Test {
51  public:
SetUp()52   void SetUp() override {
53     module_ = std::make_unique<NiceMock<MockDynamicPartitionControlAndroid>>();
54 
55     ON_CALL(dynamicControl(), GetDynamicPartitionsFeatureFlag())
56         .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
57     ON_CALL(dynamicControl(), GetVirtualAbFeatureFlag())
58         .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::NONE)));
59 
60     ON_CALL(dynamicControl(), GetDeviceDir(_))
61         .WillByDefault(Invoke([](auto path) {
62           *path = kFakeDevicePath;
63           return true;
64         }));
65 
66     ON_CALL(dynamicControl(), GetSuperPartitionName(_))
67         .WillByDefault(Return(kFakeSuper));
68 
69     ON_CALL(dynamicControl(), GetDmDevicePathByName(_, _))
70         .WillByDefault(Invoke([](auto partition_name_suffix, auto device) {
71           *device = GetDmDevice(partition_name_suffix);
72           return true;
73         }));
74 
75     ON_CALL(dynamicControl(), EraseSystemOtherAvbFooter(_, _))
76         .WillByDefault(Return(true));
77 
78     ON_CALL(dynamicControl(), IsRecovery()).WillByDefault(Return(false));
79 
80     ON_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
81         .WillByDefault(Invoke([&](uint32_t source_slot,
82                                   uint32_t target_slot,
83                                   const DeltaArchiveManifest& manifest,
84                                   bool delete_source) {
85           return dynamicControl().RealPrepareDynamicPartitionsForUpdate(
86               source_slot, target_slot, manifest, delete_source);
87         }));
88   }
89 
90   // Return the mocked DynamicPartitionControlInterface.
dynamicControl()91   NiceMock<MockDynamicPartitionControlAndroid>& dynamicControl() {
92     return static_cast<NiceMock<MockDynamicPartitionControlAndroid>&>(*module_);
93   }
94 
GetSuperDevice(uint32_t slot)95   std::string GetSuperDevice(uint32_t slot) {
96     return GetDevice(dynamicControl().GetSuperPartitionName(slot));
97   }
98 
source()99   uint32_t source() { return slots_.source; }
target()100   uint32_t target() { return slots_.target; }
101 
102   // Return partition names with suffix of source().
S(const std::string & name)103   std::string S(const std::string& name) {
104     return name + kSlotSuffixes[source()];
105   }
106 
107   // Return partition names with suffix of target().
T(const std::string & name)108   std::string T(const std::string& name) {
109     return name + kSlotSuffixes[target()];
110   }
111 
112   // Set the fake metadata to return when LoadMetadataBuilder is called on
113   // |slot|.
SetMetadata(uint32_t slot,const PartitionSuffixSizes & sizes,uint32_t partition_attr=0)114   void SetMetadata(uint32_t slot,
115                    const PartitionSuffixSizes& sizes,
116                    uint32_t partition_attr = 0) {
117     EXPECT_CALL(dynamicControl(),
118                 LoadMetadataBuilder(GetSuperDevice(slot), slot))
119         .Times(AnyNumber())
120         .WillRepeatedly(Invoke([sizes, partition_attr](auto, auto) {
121           return NewFakeMetadata(PartitionSuffixSizesToManifest(sizes),
122                                  partition_attr);
123         }));
124 
125     EXPECT_CALL(dynamicControl(),
126                 LoadMetadataBuilder(GetSuperDevice(slot), slot, _))
127         .Times(AnyNumber())
128         .WillRepeatedly(Invoke([sizes, partition_attr](auto, auto, auto) {
129           return NewFakeMetadata(PartitionSuffixSizesToManifest(sizes),
130                                  partition_attr);
131         }));
132   }
133 
ExpectStoreMetadata(const PartitionSuffixSizes & partition_sizes)134   void ExpectStoreMetadata(const PartitionSuffixSizes& partition_sizes) {
135     EXPECT_CALL(dynamicControl(),
136                 StoreMetadata(GetSuperDevice(target()),
137                               MetadataMatches(partition_sizes),
138                               target()))
139         .WillOnce(Return(true));
140   }
141 
142   // Expect that UnmapPartitionOnDeviceMapper is called on target() metadata
143   // slot with each partition in |partitions|.
ExpectUnmap(const std::set<std::string> & partitions)144   void ExpectUnmap(const std::set<std::string>& partitions) {
145     // Error when UnmapPartitionOnDeviceMapper is called on unknown arguments.
146     ON_CALL(dynamicControl(), UnmapPartitionOnDeviceMapper(_))
147         .WillByDefault(Return(false));
148 
149     for (const auto& partition : partitions) {
150       EXPECT_CALL(dynamicControl(), UnmapPartitionOnDeviceMapper(partition))
151           .WillOnce(Return(true));
152     }
153   }
PreparePartitionsForUpdate(const PartitionSizes & partition_sizes)154   bool PreparePartitionsForUpdate(const PartitionSizes& partition_sizes) {
155     return dynamicControl().PreparePartitionsForUpdate(
156         source(),
157         target(),
158         PartitionSizesToManifest(partition_sizes),
159         true,
160         nullptr);
161   }
SetSlots(const TestParam & slots)162   void SetSlots(const TestParam& slots) { slots_ = slots; }
163 
SetSnapshotEnabled(bool enabled)164   void SetSnapshotEnabled(bool enabled) {
165     dynamicControl().target_supports_snapshot_ = enabled;
166   }
167 
168   struct Listener : public ::testing::MatchResultListener {
Listenerchromeos_update_engine::DynamicPartitionControlAndroidTest::Listener169     explicit Listener(std::ostream* os) : MatchResultListener(os) {}
170   };
171 
UpdatePartitionMetadata(const PartitionSuffixSizes & source_metadata,const PartitionSizes & update_metadata,const PartitionSuffixSizes & expected)172   testing::AssertionResult UpdatePartitionMetadata(
173       const PartitionSuffixSizes& source_metadata,
174       const PartitionSizes& update_metadata,
175       const PartitionSuffixSizes& expected) {
176     return UpdatePartitionMetadata(
177         PartitionSuffixSizesToManifest(source_metadata),
178         PartitionSizesToManifest(update_metadata),
179         PartitionSuffixSizesToManifest(expected));
180   }
UpdatePartitionMetadata(const DeltaArchiveManifest & source_manifest,const DeltaArchiveManifest & update_manifest,const DeltaArchiveManifest & expected)181   testing::AssertionResult UpdatePartitionMetadata(
182       const DeltaArchiveManifest& source_manifest,
183       const DeltaArchiveManifest& update_manifest,
184       const DeltaArchiveManifest& expected) {
185     return UpdatePartitionMetadata(
186         source_manifest, update_manifest, MetadataMatches(expected));
187   }
UpdatePartitionMetadata(const DeltaArchiveManifest & source_manifest,const DeltaArchiveManifest & update_manifest,const Matcher<MetadataBuilder * > & matcher)188   testing::AssertionResult UpdatePartitionMetadata(
189       const DeltaArchiveManifest& source_manifest,
190       const DeltaArchiveManifest& update_manifest,
191       const Matcher<MetadataBuilder*>& matcher) {
192     auto super_metadata = NewFakeMetadata(source_manifest);
193     if (!module_->UpdatePartitionMetadata(
194             super_metadata.get(), target(), update_manifest)) {
195       return testing::AssertionFailure()
196              << "UpdatePartitionMetadataInternal failed";
197     }
198     std::stringstream ss;
199     Listener listener(&ss);
200     if (matcher.MatchAndExplain(super_metadata.get(), &listener)) {
201       return testing::AssertionSuccess() << ss.str();
202     } else {
203       return testing::AssertionFailure() << ss.str();
204     }
205   }
206 
207   std::unique_ptr<DynamicPartitionControlAndroid> module_;
208   TestParam slots_;
209 };
210 
211 class DynamicPartitionControlAndroidTestP
212     : public DynamicPartitionControlAndroidTest,
213       public ::testing::WithParamInterface<TestParam> {
214  public:
SetUp()215   void SetUp() override {
216     DynamicPartitionControlAndroidTest::SetUp();
217     SetSlots(GetParam());
218   }
219 };
220 
221 // Test resize case. Grow if target metadata contains a partition with a size
222 // less than expected.
TEST_P(DynamicPartitionControlAndroidTestP,NeedGrowIfSizeNotMatchWhenResizing)223 TEST_P(DynamicPartitionControlAndroidTestP,
224        NeedGrowIfSizeNotMatchWhenResizing) {
225   PartitionSuffixSizes source_metadata{{S("system"), 2_GiB},
226                                        {S("vendor"), 1_GiB},
227                                        {T("system"), 2_GiB},
228                                        {T("vendor"), 1_GiB}};
229   PartitionSuffixSizes expected{{S("system"), 2_GiB},
230                                 {S("vendor"), 1_GiB},
231                                 {T("system"), 3_GiB},
232                                 {T("vendor"), 1_GiB}};
233   PartitionSizes update_metadata{{"system", 3_GiB}, {"vendor", 1_GiB}};
234   EXPECT_TRUE(
235       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
236 }
237 
238 // Test resize case. Shrink if target metadata contains a partition with a size
239 // greater than expected.
TEST_P(DynamicPartitionControlAndroidTestP,NeedShrinkIfSizeNotMatchWhenResizing)240 TEST_P(DynamicPartitionControlAndroidTestP,
241        NeedShrinkIfSizeNotMatchWhenResizing) {
242   PartitionSuffixSizes source_metadata{{S("system"), 2_GiB},
243                                        {S("vendor"), 1_GiB},
244                                        {T("system"), 2_GiB},
245                                        {T("vendor"), 1_GiB}};
246   PartitionSuffixSizes expected{{S("system"), 2_GiB},
247                                 {S("vendor"), 1_GiB},
248                                 {T("system"), 2_GiB},
249                                 {T("vendor"), 150_MiB}};
250   PartitionSizes update_metadata{{"system", 2_GiB}, {"vendor", 150_MiB}};
251   EXPECT_TRUE(
252       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
253 }
254 
255 // Test adding partitions on the first run.
TEST_P(DynamicPartitionControlAndroidTestP,AddPartitionToEmptyMetadata)256 TEST_P(DynamicPartitionControlAndroidTestP, AddPartitionToEmptyMetadata) {
257   PartitionSuffixSizes source_metadata{};
258   PartitionSuffixSizes expected{{T("system"), 2_GiB}, {T("vendor"), 1_GiB}};
259   PartitionSizes update_metadata{{"system", 2_GiB}, {"vendor", 1_GiB}};
260   EXPECT_TRUE(
261       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
262 }
263 
264 // Test subsequent add case.
TEST_P(DynamicPartitionControlAndroidTestP,AddAdditionalPartition)265 TEST_P(DynamicPartitionControlAndroidTestP, AddAdditionalPartition) {
266   PartitionSuffixSizes source_metadata{{S("system"), 2_GiB},
267                                        {T("system"), 2_GiB}};
268   PartitionSuffixSizes expected{
269       {S("system"), 2_GiB}, {T("system"), 2_GiB}, {T("vendor"), 1_GiB}};
270   PartitionSizes update_metadata{{"system", 2_GiB}, {"vendor", 1_GiB}};
271   EXPECT_TRUE(
272       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
273 }
274 
275 // Test delete one partition.
TEST_P(DynamicPartitionControlAndroidTestP,DeletePartition)276 TEST_P(DynamicPartitionControlAndroidTestP, DeletePartition) {
277   PartitionSuffixSizes source_metadata{{S("system"), 2_GiB},
278                                        {S("vendor"), 1_GiB},
279                                        {T("system"), 2_GiB},
280                                        {T("vendor"), 1_GiB}};
281   // No T("vendor")
282   PartitionSuffixSizes expected{
283       {S("system"), 2_GiB}, {S("vendor"), 1_GiB}, {T("system"), 2_GiB}};
284   PartitionSizes update_metadata{{"system", 2_GiB}};
285   EXPECT_TRUE(
286       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
287 }
288 
289 // Test delete all partitions.
TEST_P(DynamicPartitionControlAndroidTestP,DeleteAll)290 TEST_P(DynamicPartitionControlAndroidTestP, DeleteAll) {
291   PartitionSuffixSizes source_metadata{{S("system"), 2_GiB},
292                                        {S("vendor"), 1_GiB},
293                                        {T("system"), 2_GiB},
294                                        {T("vendor"), 1_GiB}};
295   PartitionSuffixSizes expected{{S("system"), 2_GiB}, {S("vendor"), 1_GiB}};
296   PartitionSizes update_metadata{};
297   EXPECT_TRUE(
298       UpdatePartitionMetadata(source_metadata, update_metadata, expected));
299 }
300 
301 // Test corrupt source metadata case.
TEST_P(DynamicPartitionControlAndroidTestP,CorruptedSourceMetadata)302 TEST_P(DynamicPartitionControlAndroidTestP, CorruptedSourceMetadata) {
303   EXPECT_CALL(dynamicControl(),
304               LoadMetadataBuilder(GetSuperDevice(source()), source(), _))
305       .WillOnce(Invoke([](auto, auto, auto) { return nullptr; }));
306   ExpectUnmap({T("system")});
307 
308   EXPECT_FALSE(PreparePartitionsForUpdate({{"system", 1_GiB}}))
309       << "Should not be able to continue with corrupt source metadata";
310 }
311 
312 // Test that UpdatePartitionMetadata fails if there is not enough space on the
313 // device.
TEST_P(DynamicPartitionControlAndroidTestP,NotEnoughSpace)314 TEST_P(DynamicPartitionControlAndroidTestP, NotEnoughSpace) {
315   PartitionSuffixSizes source_metadata{{S("system"), 3_GiB},
316                                        {S("vendor"), 2_GiB},
317                                        {T("system"), 0},
318                                        {T("vendor"), 0}};
319   PartitionSizes update_metadata{{"system", 3_GiB}, {"vendor", 3_GiB}};
320 
321   EXPECT_FALSE(UpdatePartitionMetadata(source_metadata, update_metadata, {}))
322       << "Should not be able to fit 11GiB data into 10GiB space";
323 }
324 
TEST_P(DynamicPartitionControlAndroidTestP,NotEnoughSpaceForSlot)325 TEST_P(DynamicPartitionControlAndroidTestP, NotEnoughSpaceForSlot) {
326   PartitionSuffixSizes source_metadata{{S("system"), 1_GiB},
327                                        {S("vendor"), 1_GiB},
328                                        {T("system"), 0},
329                                        {T("vendor"), 0}};
330   PartitionSizes update_metadata{{"system", 3_GiB}, {"vendor", 3_GiB}};
331   EXPECT_FALSE(UpdatePartitionMetadata(source_metadata, update_metadata, {}))
332       << "Should not be able to grow over size of super / 2";
333 }
334 
TEST_P(DynamicPartitionControlAndroidTestP,ApplyRetrofitUpdateOnDynamicPartitionsEnabledBuild)335 TEST_P(DynamicPartitionControlAndroidTestP,
336        ApplyRetrofitUpdateOnDynamicPartitionsEnabledBuild) {
337   ON_CALL(dynamicControl(), GetDynamicPartitionsFeatureFlag())
338       .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::RETROFIT)));
339   // Static partition {system,bar}_{a,b} exists.
340   EXPECT_CALL(dynamicControl(),
341               DeviceExists(AnyOf(GetDevice(S("bar")),
342                                  GetDevice(T("bar")),
343                                  GetDevice(S("system")),
344                                  GetDevice(T("system")))))
345       .WillRepeatedly(Return(true));
346 
347   SetMetadata(source(),
348               {{S("system"), 2_GiB},
349                {S("vendor"), 1_GiB},
350                {T("system"), 2_GiB},
351                {T("vendor"), 1_GiB}});
352 
353   // Not calling through
354   // DynamicPartitionControlAndroidTest::PreparePartitionsForUpdate(), since we
355   // don't want any default group in the PartitionMetadata.
356   EXPECT_TRUE(dynamicControl().PreparePartitionsForUpdate(
357       source(), target(), {}, true, nullptr));
358 
359   // Should use dynamic source partitions.
360   EXPECT_CALL(dynamicControl(), GetState(S("system")))
361       .Times(1)
362       .WillOnce(Return(DmDeviceState::ACTIVE));
363   string system_device;
364   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
365       "system", source(), source(), &system_device));
366   EXPECT_EQ(GetDmDevice(S("system")), system_device);
367 
368   // Should use static target partitions without querying dynamic control.
369   EXPECT_CALL(dynamicControl(), GetState(T("system"))).Times(0);
370   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
371       "system", target(), source(), &system_device));
372   EXPECT_EQ(GetDevice(T("system")), system_device);
373 
374   // Static partition "bar".
375   EXPECT_CALL(dynamicControl(), GetState(S("bar"))).Times(0);
376   std::string bar_device;
377   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
378       "bar", source(), source(), &bar_device));
379   EXPECT_EQ(GetDevice(S("bar")), bar_device);
380 
381   EXPECT_CALL(dynamicControl(), GetState(T("bar"))).Times(0);
382   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
383       "bar", target(), source(), &bar_device));
384   EXPECT_EQ(GetDevice(T("bar")), bar_device);
385 }
386 
TEST_P(DynamicPartitionControlAndroidTestP,GetPartitionDeviceWhenResumingUpdate)387 TEST_P(DynamicPartitionControlAndroidTestP,
388        GetPartitionDeviceWhenResumingUpdate) {
389   // Static partition bar_{a,b} exists.
390   EXPECT_CALL(dynamicControl(),
391               DeviceExists(AnyOf(GetDevice(S("bar")), GetDevice(T("bar")))))
392       .WillRepeatedly(Return(true));
393 
394   // Both of the two slots contain valid partition metadata, since this is
395   // resuming an update.
396   SetMetadata(source(),
397               {{S("system"), 2_GiB},
398                {S("vendor"), 1_GiB},
399                {T("system"), 2_GiB},
400                {T("vendor"), 1_GiB}});
401   SetMetadata(target(),
402               {{S("system"), 2_GiB},
403                {S("vendor"), 1_GiB},
404                {T("system"), 2_GiB},
405                {T("vendor"), 1_GiB}});
406 
407   EXPECT_TRUE(dynamicControl().PreparePartitionsForUpdate(
408       source(),
409       target(),
410       PartitionSizesToManifest({{"system", 2_GiB}, {"vendor", 1_GiB}}),
411       false,
412       nullptr));
413 
414   // Dynamic partition "system".
415   EXPECT_CALL(dynamicControl(), GetState(S("system")))
416       .Times(1)
417       .WillOnce(Return(DmDeviceState::ACTIVE));
418   string system_device;
419   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
420       "system", source(), source(), &system_device));
421   EXPECT_EQ(GetDmDevice(S("system")), system_device);
422 
423   EXPECT_CALL(dynamicControl(), GetState(T("system")))
424       .Times(AnyNumber())
425       .WillOnce(Return(DmDeviceState::ACTIVE));
426   EXPECT_CALL(dynamicControl(),
427               MapPartitionOnDeviceMapper(
428                   GetSuperDevice(target()), T("system"), target(), _, _))
429       .Times(AnyNumber())
430       .WillRepeatedly(
431           Invoke([](const auto&, const auto& name, auto, auto, auto* device) {
432             *device = "/fake/remapped/" + name;
433             return true;
434           }));
435   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
436       "system", target(), source(), &system_device));
437   EXPECT_EQ("/fake/remapped/" + T("system"), system_device);
438 
439   // Static partition "bar".
440   EXPECT_CALL(dynamicControl(), GetState(S("bar"))).Times(0);
441   std::string bar_device;
442   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
443       "bar", source(), source(), &bar_device));
444   EXPECT_EQ(GetDevice(S("bar")), bar_device);
445 
446   EXPECT_CALL(dynamicControl(), GetState(T("bar"))).Times(0);
447   EXPECT_TRUE(dynamicControl().GetPartitionDevice(
448       "bar", target(), source(), &bar_device));
449   EXPECT_EQ(GetDevice(T("bar")), bar_device);
450 }
451 
452 INSTANTIATE_TEST_CASE_P(DynamicPartitionControlAndroidTest,
453                         DynamicPartitionControlAndroidTestP,
454                         testing::Values(TestParam{0, 1}, TestParam{1, 0}));
455 
456 class DynamicPartitionControlAndroidGroupTestP
457     : public DynamicPartitionControlAndroidTestP {
458  public:
459   DeltaArchiveManifest source_manifest;
SetUp()460   void SetUp() override {
461     DynamicPartitionControlAndroidTestP::SetUp();
462     AddGroupAndPartition(
463         &source_manifest, S("android"), 3_GiB, S("system"), 2_GiB);
464     AddGroupAndPartition(&source_manifest, S("oem"), 2_GiB, S("vendor"), 1_GiB);
465     AddGroupAndPartition(&source_manifest, T("android"), 3_GiB, T("system"), 0);
466     AddGroupAndPartition(&source_manifest, T("oem"), 2_GiB, T("vendor"), 0);
467   }
468 
AddGroupAndPartition(DeltaArchiveManifest * manifest,const string & group,uint64_t group_size,const string & partition,uint64_t partition_size)469   void AddGroupAndPartition(DeltaArchiveManifest* manifest,
470                             const string& group,
471                             uint64_t group_size,
472                             const string& partition,
473                             uint64_t partition_size) {
474     auto* g = AddGroup(manifest, group, group_size);
475     AddPartition(manifest, g, partition, partition_size);
476   }
477 };
478 
479 // Allow to resize within group.
TEST_P(DynamicPartitionControlAndroidGroupTestP,ResizeWithinGroup)480 TEST_P(DynamicPartitionControlAndroidGroupTestP, ResizeWithinGroup) {
481   DeltaArchiveManifest expected;
482   AddGroupAndPartition(&expected, T("android"), 3_GiB, T("system"), 3_GiB);
483   AddGroupAndPartition(&expected, T("oem"), 2_GiB, T("vendor"), 2_GiB);
484 
485   DeltaArchiveManifest update_manifest;
486   AddGroupAndPartition(&update_manifest, "android", 3_GiB, "system", 3_GiB);
487   AddGroupAndPartition(&update_manifest, "oem", 2_GiB, "vendor", 2_GiB);
488 
489   EXPECT_TRUE(
490       UpdatePartitionMetadata(source_manifest, update_manifest, expected));
491 }
492 
TEST_P(DynamicPartitionControlAndroidGroupTestP,NotEnoughSpaceForGroup)493 TEST_P(DynamicPartitionControlAndroidGroupTestP, NotEnoughSpaceForGroup) {
494   DeltaArchiveManifest update_manifest;
495   AddGroupAndPartition(&update_manifest, "android", 3_GiB, "system", 1_GiB),
496       AddGroupAndPartition(&update_manifest, "oem", 2_GiB, "vendor", 3_GiB);
497   EXPECT_FALSE(UpdatePartitionMetadata(source_manifest, update_manifest, {}))
498       << "Should not be able to grow over maximum size of group";
499 }
500 
TEST_P(DynamicPartitionControlAndroidGroupTestP,GroupTooBig)501 TEST_P(DynamicPartitionControlAndroidGroupTestP, GroupTooBig) {
502   DeltaArchiveManifest update_manifest;
503   AddGroup(&update_manifest, "android", 3_GiB);
504   AddGroup(&update_manifest, "oem", 3_GiB);
505   EXPECT_FALSE(UpdatePartitionMetadata(source_manifest, update_manifest, {}))
506       << "Should not be able to grow over size of super / 2";
507 }
508 
TEST_P(DynamicPartitionControlAndroidGroupTestP,AddPartitionToGroup)509 TEST_P(DynamicPartitionControlAndroidGroupTestP, AddPartitionToGroup) {
510   DeltaArchiveManifest expected;
511   auto* g = AddGroup(&expected, T("android"), 3_GiB);
512   AddPartition(&expected, g, T("system"), 2_GiB);
513   AddPartition(&expected, g, T("system_ext"), 1_GiB);
514 
515   DeltaArchiveManifest update_manifest;
516   g = AddGroup(&update_manifest, "android", 3_GiB);
517   AddPartition(&update_manifest, g, "system", 2_GiB);
518   AddPartition(&update_manifest, g, "system_ext", 1_GiB);
519   AddGroupAndPartition(&update_manifest, "oem", 2_GiB, "vendor", 2_GiB);
520 
521   EXPECT_TRUE(
522       UpdatePartitionMetadata(source_manifest, update_manifest, expected));
523 }
524 
TEST_P(DynamicPartitionControlAndroidGroupTestP,RemovePartitionFromGroup)525 TEST_P(DynamicPartitionControlAndroidGroupTestP, RemovePartitionFromGroup) {
526   DeltaArchiveManifest expected;
527   AddGroup(&expected, T("android"), 3_GiB);
528 
529   DeltaArchiveManifest update_manifest;
530   AddGroup(&update_manifest, "android", 3_GiB);
531   AddGroupAndPartition(&update_manifest, "oem", 2_GiB, "vendor", 2_GiB);
532 
533   EXPECT_TRUE(
534       UpdatePartitionMetadata(source_manifest, update_manifest, expected));
535 }
536 
TEST_P(DynamicPartitionControlAndroidGroupTestP,AddGroup)537 TEST_P(DynamicPartitionControlAndroidGroupTestP, AddGroup) {
538   DeltaArchiveManifest expected;
539   AddGroupAndPartition(
540       &expected, T("new_group"), 2_GiB, T("new_partition"), 2_GiB);
541 
542   DeltaArchiveManifest update_manifest;
543   AddGroupAndPartition(&update_manifest, "android", 2_GiB, "system", 2_GiB);
544   AddGroupAndPartition(&update_manifest, "oem", 1_GiB, "vendor", 1_GiB);
545   AddGroupAndPartition(
546       &update_manifest, "new_group", 2_GiB, "new_partition", 2_GiB);
547   EXPECT_TRUE(
548       UpdatePartitionMetadata(source_manifest, update_manifest, expected));
549 }
550 
TEST_P(DynamicPartitionControlAndroidGroupTestP,RemoveGroup)551 TEST_P(DynamicPartitionControlAndroidGroupTestP, RemoveGroup) {
552   DeltaArchiveManifest update_manifest;
553   AddGroupAndPartition(&update_manifest, "android", 2_GiB, "system", 2_GiB);
554 
555   EXPECT_TRUE(UpdatePartitionMetadata(
556       source_manifest, update_manifest, Not(HasGroup(T("oem")))));
557 }
558 
TEST_P(DynamicPartitionControlAndroidGroupTestP,ResizeGroup)559 TEST_P(DynamicPartitionControlAndroidGroupTestP, ResizeGroup) {
560   DeltaArchiveManifest expected;
561   AddGroupAndPartition(&expected, T("android"), 2_GiB, T("system"), 2_GiB);
562   AddGroupAndPartition(&expected, T("oem"), 3_GiB, T("vendor"), 3_GiB);
563   DeltaArchiveManifest update_manifest;
564   AddGroupAndPartition(&update_manifest, "android", 2_GiB, "system", 2_GiB),
565       AddGroupAndPartition(&update_manifest, "oem", 3_GiB, "vendor", 3_GiB);
566   EXPECT_TRUE(
567       UpdatePartitionMetadata(source_manifest, update_manifest, expected));
568 }
569 
570 INSTANTIATE_TEST_CASE_P(DynamicPartitionControlAndroidTest,
571                         DynamicPartitionControlAndroidGroupTestP,
572                         testing::Values(TestParam{0, 1}, TestParam{1, 0}));
573 
update_sizes_0()574 const PartitionSuffixSizes update_sizes_0() {
575   // Initial state is 0 for "other" slot.
576   return {
577       {"grown_a", 2_GiB},
578       {"shrunk_a", 1_GiB},
579       {"same_a", 100_MiB},
580       {"deleted_a", 150_MiB},
581       // no added_a
582       {"grown_b", 200_MiB},
583       // simulate system_other
584       {"shrunk_b", 0},
585       {"same_b", 0},
586       {"deleted_b", 0},
587       // no added_b
588   };
589 }
590 
update_sizes_1()591 const PartitionSuffixSizes update_sizes_1() {
592   return {
593       {"grown_a", 2_GiB},
594       {"shrunk_a", 1_GiB},
595       {"same_a", 100_MiB},
596       {"deleted_a", 150_MiB},
597       // no added_a
598       {"grown_b", 3_GiB},
599       {"shrunk_b", 150_MiB},
600       {"same_b", 100_MiB},
601       {"added_b", 150_MiB},
602       // no deleted_b
603   };
604 }
605 
update_sizes_2()606 const PartitionSuffixSizes update_sizes_2() {
607   return {
608       {"grown_a", 4_GiB},
609       {"shrunk_a", 100_MiB},
610       {"same_a", 100_MiB},
611       {"deleted_a", 64_MiB},
612       // no added_a
613       {"grown_b", 3_GiB},
614       {"shrunk_b", 150_MiB},
615       {"same_b", 100_MiB},
616       {"added_b", 150_MiB},
617       // no deleted_b
618   };
619 }
620 
621 // Test case for first update after the device is manufactured, in which
622 // case the "other" slot is likely of size "0" (except system, which is
623 // non-zero because of system_other partition)
TEST_F(DynamicPartitionControlAndroidTest,SimulatedFirstUpdate)624 TEST_F(DynamicPartitionControlAndroidTest, SimulatedFirstUpdate) {
625   SetSlots({0, 1});
626 
627   SetMetadata(source(), update_sizes_0());
628   SetMetadata(target(), update_sizes_0());
629   ExpectStoreMetadata(update_sizes_1());
630   ExpectUnmap({"grown_b", "shrunk_b", "same_b", "added_b"});
631 
632   EXPECT_TRUE(PreparePartitionsForUpdate({{"grown", 3_GiB},
633                                           {"shrunk", 150_MiB},
634                                           {"same", 100_MiB},
635                                           {"added", 150_MiB}}));
636 }
637 
638 // After first update, test for the second update. In the second update, the
639 // "added" partition is deleted and "deleted" partition is re-added.
TEST_F(DynamicPartitionControlAndroidTest,SimulatedSecondUpdate)640 TEST_F(DynamicPartitionControlAndroidTest, SimulatedSecondUpdate) {
641   SetSlots({1, 0});
642 
643   SetMetadata(source(), update_sizes_1());
644   SetMetadata(target(), update_sizes_0());
645 
646   ExpectStoreMetadata(update_sizes_2());
647   ExpectUnmap({"grown_a", "shrunk_a", "same_a", "deleted_a"});
648 
649   EXPECT_TRUE(PreparePartitionsForUpdate({{"grown", 4_GiB},
650                                           {"shrunk", 100_MiB},
651                                           {"same", 100_MiB},
652                                           {"deleted", 64_MiB}}));
653 }
654 
TEST_F(DynamicPartitionControlAndroidTest,ApplyingToCurrentSlot)655 TEST_F(DynamicPartitionControlAndroidTest, ApplyingToCurrentSlot) {
656   SetSlots({1, 1});
657   EXPECT_FALSE(PreparePartitionsForUpdate({}))
658       << "Should not be able to apply to current slot.";
659 }
660 
TEST_P(DynamicPartitionControlAndroidTestP,OptimizeOperationTest)661 TEST_P(DynamicPartitionControlAndroidTestP, OptimizeOperationTest) {
662   ASSERT_TRUE(dynamicControl().PreparePartitionsForUpdate(
663       source(),
664       target(),
665       PartitionSizesToManifest({{"foo", 4_MiB}}),
666       false,
667       nullptr));
668   dynamicControl().set_fake_mapped_devices({T("foo")});
669 
670   InstallOperation iop;
671   InstallOperation optimized;
672   Extent *se, *de;
673 
674   // Not a SOURCE_COPY operation, cannot skip.
675   iop.set_type(InstallOperation::REPLACE);
676   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
677 
678   iop.set_type(InstallOperation::SOURCE_COPY);
679 
680   // By default GetVirtualAbFeatureFlag is disabled. Cannot skip operation.
681   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
682 
683   // Enable GetVirtualAbFeatureFlag in the mock interface.
684   ON_CALL(dynamicControl(), GetVirtualAbFeatureFlag())
685       .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
686 
687   // By default target_supports_snapshot_ is set to false. Cannot skip
688   // operation.
689   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
690 
691   SetSnapshotEnabled(true);
692 
693   // Empty source and destination. Skip.
694   EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
695   EXPECT_TRUE(optimized.src_extents().empty());
696   EXPECT_TRUE(optimized.dst_extents().empty());
697 
698   se = iop.add_src_extents();
699   se->set_start_block(0);
700   se->set_num_blocks(1);
701 
702   // There is something in sources, but destinations are empty. Cannot skip.
703   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
704 
705   InstallOperation iop2;
706 
707   de = iop2.add_dst_extents();
708   de->set_start_block(0);
709   de->set_num_blocks(1);
710 
711   // There is something in destinations, but sources are empty. Cannot skip.
712   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop2, &optimized));
713 
714   de = iop.add_dst_extents();
715   de->set_start_block(0);
716   de->set_num_blocks(1);
717 
718   // Sources and destinations are identical. Skip.
719   EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
720   EXPECT_TRUE(optimized.src_extents().empty());
721   EXPECT_TRUE(optimized.dst_extents().empty());
722 
723   se = iop.add_src_extents();
724   se->set_start_block(1);
725   se->set_num_blocks(5);
726 
727   // There is something in source, but not in destination. Cannot skip.
728   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
729 
730   de = iop.add_dst_extents();
731   de->set_start_block(1);
732   de->set_num_blocks(5);
733 
734   // There is source and destination are equal. Skip.
735   EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
736   EXPECT_TRUE(optimized.src_extents().empty());
737   EXPECT_TRUE(optimized.dst_extents().empty());
738 
739   de = iop.add_dst_extents();
740   de->set_start_block(6);
741   de->set_num_blocks(5);
742 
743   // There is something extra in dest. Cannot skip.
744   EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
745 
746   se = iop.add_src_extents();
747   se->set_start_block(6);
748   se->set_num_blocks(5);
749 
750   // Source and dest are identical again. Skip.
751   EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
752   EXPECT_TRUE(optimized.src_extents().empty());
753   EXPECT_TRUE(optimized.dst_extents().empty());
754 
755   iop.Clear();
756   iop.set_type(InstallOperation::SOURCE_COPY);
757   se = iop.add_src_extents();
758   se->set_start_block(1);
759   se->set_num_blocks(1);
760   se = iop.add_src_extents();
761   se->set_start_block(3);
762   se->set_num_blocks(2);
763   se = iop.add_src_extents();
764   se->set_start_block(7);
765   se->set_num_blocks(2);
766   de = iop.add_dst_extents();
767   de->set_start_block(2);
768   de->set_num_blocks(5);
769 
770   // [1, 3, 4, 7, 8] -> [2, 3, 4, 5, 6] should return [1, 7, 8] -> [2, 5, 6]
771   EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
772   ASSERT_EQ(2, optimized.src_extents_size());
773   ASSERT_EQ(2, optimized.dst_extents_size());
774   EXPECT_EQ(1u, optimized.src_extents(0).start_block());
775   EXPECT_EQ(1u, optimized.src_extents(0).num_blocks());
776   EXPECT_EQ(2u, optimized.dst_extents(0).start_block());
777   EXPECT_EQ(1u, optimized.dst_extents(0).num_blocks());
778   EXPECT_EQ(7u, optimized.src_extents(1).start_block());
779   EXPECT_EQ(2u, optimized.src_extents(1).num_blocks());
780   EXPECT_EQ(5u, optimized.dst_extents(1).start_block());
781   EXPECT_EQ(2u, optimized.dst_extents(1).num_blocks());
782 
783   // Don't skip for static partitions.
784   EXPECT_FALSE(dynamicControl().OptimizeOperation("bar", iop, &optimized));
785 }
786 
TEST_F(DynamicPartitionControlAndroidTest,ResetUpdate)787 TEST_F(DynamicPartitionControlAndroidTest, ResetUpdate) {
788   MockPrefs prefs;
789   ASSERT_TRUE(dynamicControl().ResetUpdate(&prefs));
790 }
791 
TEST_F(DynamicPartitionControlAndroidTest,IsAvbNotEnabledInFstab)792 TEST_F(DynamicPartitionControlAndroidTest, IsAvbNotEnabledInFstab) {
793   // clang-format off
794   std::string fstab_content =
795       "system /postinstall ext4 ro,nosuid,nodev,noexec slotselect_other,logical\n"  // NOLINT(whitespace/line_length)
796       "/dev/block/by-name/system /postinstall ext4 ro,nosuid,nodev,noexec slotselect_other\n";  // NOLINT(whitespace/line_length)
797   // clang-format on
798   ScopedTempFile fstab;
799   ASSERT_TRUE(test_utils::WriteFileString(fstab.path(), fstab_content));
800   ASSERT_THAT(dynamicControl().RealIsAvbEnabledInFstab(fstab.path()),
801               Optional(false));
802 }
803 
TEST_F(DynamicPartitionControlAndroidTest,IsAvbEnabledInFstab)804 TEST_F(DynamicPartitionControlAndroidTest, IsAvbEnabledInFstab) {
805   // clang-format off
806   std::string fstab_content =
807       "system /postinstall ext4 ro,nosuid,nodev,noexec slotselect_other,logical,avb_keys=/foo\n";  // NOLINT(whitespace/line_length)
808   // clang-format on
809   ScopedTempFile fstab;
810   ASSERT_TRUE(test_utils::WriteFileString(fstab.path(), fstab_content));
811   ASSERT_THAT(dynamicControl().RealIsAvbEnabledInFstab(fstab.path()),
812               Optional(true));
813 }
814 
TEST_P(DynamicPartitionControlAndroidTestP,AvbNotEnabledOnSystemOther)815 TEST_P(DynamicPartitionControlAndroidTestP, AvbNotEnabledOnSystemOther) {
816   ON_CALL(dynamicControl(), GetSystemOtherPath(_, _, _, _, _))
817       .WillByDefault(Invoke([&](auto source_slot,
818                                 auto target_slot,
819                                 const auto& name,
820                                 auto path,
821                                 auto should_unmap) {
822         return dynamicControl().RealGetSystemOtherPath(
823             source_slot, target_slot, name, path, should_unmap);
824       }));
825   ON_CALL(dynamicControl(), IsAvbEnabledOnSystemOther())
826       .WillByDefault(Return(false));
827   EXPECT_TRUE(
828       dynamicControl().RealEraseSystemOtherAvbFooter(source(), target()));
829 }
830 
TEST_P(DynamicPartitionControlAndroidTestP,NoSystemOtherToErase)831 TEST_P(DynamicPartitionControlAndroidTestP, NoSystemOtherToErase) {
832   SetMetadata(source(), {{S("system"), 100_MiB}});
833   ON_CALL(dynamicControl(), IsAvbEnabledOnSystemOther())
834       .WillByDefault(Return(true));
835   std::string path;
836   bool should_unmap;
837   ASSERT_TRUE(dynamicControl().RealGetSystemOtherPath(
838       source(), target(), T("system"), &path, &should_unmap));
839   ASSERT_TRUE(path.empty()) << path;
840   ASSERT_FALSE(should_unmap);
841   ON_CALL(dynamicControl(), GetSystemOtherPath(_, _, _, _, _))
842       .WillByDefault(Invoke([&](auto source_slot,
843                                 auto target_slot,
844                                 const auto& name,
845                                 auto path,
846                                 auto should_unmap) {
847         return dynamicControl().RealGetSystemOtherPath(
848             source_slot, target_slot, name, path, should_unmap);
849       }));
850   EXPECT_TRUE(
851       dynamicControl().RealEraseSystemOtherAvbFooter(source(), target()));
852 }
853 
TEST_P(DynamicPartitionControlAndroidTestP,SkipEraseUpdatedSystemOther)854 TEST_P(DynamicPartitionControlAndroidTestP, SkipEraseUpdatedSystemOther) {
855   PartitionSuffixSizes sizes{{S("system"), 100_MiB}, {T("system"), 100_MiB}};
856   SetMetadata(source(), sizes, LP_PARTITION_ATTR_UPDATED);
857   ON_CALL(dynamicControl(), IsAvbEnabledOnSystemOther())
858       .WillByDefault(Return(true));
859   std::string path;
860   bool should_unmap;
861   ASSERT_TRUE(dynamicControl().RealGetSystemOtherPath(
862       source(), target(), T("system"), &path, &should_unmap));
863   ASSERT_TRUE(path.empty()) << path;
864   ASSERT_FALSE(should_unmap);
865   ON_CALL(dynamicControl(), GetSystemOtherPath(_, _, _, _, _))
866       .WillByDefault(Invoke([&](auto source_slot,
867                                 auto target_slot,
868                                 const auto& name,
869                                 auto path,
870                                 auto should_unmap) {
871         return dynamicControl().RealGetSystemOtherPath(
872             source_slot, target_slot, name, path, should_unmap);
873       }));
874   EXPECT_TRUE(
875       dynamicControl().RealEraseSystemOtherAvbFooter(source(), target()));
876 }
877 
TEST_P(DynamicPartitionControlAndroidTestP,EraseSystemOtherAvbFooter)878 TEST_P(DynamicPartitionControlAndroidTestP, EraseSystemOtherAvbFooter) {
879   constexpr uint64_t file_size = 1_MiB;
880   static_assert(file_size > AVB_FOOTER_SIZE);
881   ScopedTempFile system_other;
882   brillo::Blob original(file_size, 'X');
883   ASSERT_TRUE(test_utils::WriteFileVector(system_other.path(), original));
884   std::string mnt_path;
885   ScopedLoopbackDeviceBinder dev(system_other.path(), true, &mnt_path);
886   ASSERT_TRUE(dev.is_bound());
887 
888   brillo::Blob device_content;
889   ASSERT_TRUE(utils::ReadFile(mnt_path, &device_content));
890   ASSERT_EQ(original, device_content);
891 
892   PartitionSuffixSizes sizes{{S("system"), 100_MiB}, {T("system"), file_size}};
893   SetMetadata(source(), sizes);
894   ON_CALL(dynamicControl(), IsAvbEnabledOnSystemOther())
895       .WillByDefault(Return(true));
896   EXPECT_CALL(dynamicControl(),
897               GetSystemOtherPath(source(), target(), T("system"), _, _))
898       .WillRepeatedly(
899           Invoke([&](auto, auto, const auto&, auto path, auto should_unmap) {
900             *path = mnt_path;
901             *should_unmap = false;
902             return true;
903           }));
904   ASSERT_TRUE(
905       dynamicControl().RealEraseSystemOtherAvbFooter(source(), target()));
906 
907   device_content.clear();
908   ASSERT_TRUE(utils::ReadFile(mnt_path, &device_content));
909   brillo::Blob new_expected(original);
910   // Clear the last AVB_FOOTER_SIZE bytes.
911   new_expected.resize(file_size - AVB_FOOTER_SIZE);
912   new_expected.resize(file_size, '\0');
913   ASSERT_EQ(new_expected, device_content);
914 }
915 
916 class FakeAutoDevice : public android::snapshot::AutoDevice {
917  public:
FakeAutoDevice()918   FakeAutoDevice() : AutoDevice("") {}
919 };
920 
921 class SnapshotPartitionTestP : public DynamicPartitionControlAndroidTestP {
922  public:
SetUp()923   void SetUp() override {
924     DynamicPartitionControlAndroidTestP::SetUp();
925     ON_CALL(dynamicControl(), GetVirtualAbFeatureFlag())
926         .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
927 
928     snapshot_ = new NiceMock<MockSnapshotManager>();
929     dynamicControl().snapshot_.reset(snapshot_);  // takes ownership
930     EXPECT_CALL(*snapshot_, BeginUpdate()).WillOnce(Return(true));
931     EXPECT_CALL(*snapshot_, EnsureMetadataMounted())
932         .WillRepeatedly(
933             Invoke([]() { return std::make_unique<FakeAutoDevice>(); }));
934 
935     manifest_ =
936         PartitionSizesToManifest({{"system", 3_GiB}, {"vendor", 1_GiB}});
937   }
ExpectCreateUpdateSnapshots(android::snapshot::Return val)938   void ExpectCreateUpdateSnapshots(android::snapshot::Return val) {
939     manifest_.mutable_dynamic_partition_metadata()->set_snapshot_enabled(true);
940     EXPECT_CALL(*snapshot_, CreateUpdateSnapshots(_))
941         .WillRepeatedly(Invoke([&, val](const auto& manifest) {
942           // Deep comparison requires full protobuf library. Comparing the
943           // pointers are sufficient.
944           EXPECT_EQ(&manifest_, &manifest);
945           LOG(WARNING) << "CreateUpdateSnapshots returning " << val.string();
946           return val;
947         }));
948   }
PreparePartitionsForUpdate(uint64_t * required_size)949   bool PreparePartitionsForUpdate(uint64_t* required_size) {
950     return dynamicControl().PreparePartitionsForUpdate(
951         source(), target(), manifest_, true /* update */, required_size);
952   }
953   MockSnapshotManager* snapshot_ = nullptr;
954   DeltaArchiveManifest manifest_;
955 };
956 
957 // Test happy path of PreparePartitionsForUpdate on a Virtual A/B device.
TEST_P(SnapshotPartitionTestP,PreparePartitions)958 TEST_P(SnapshotPartitionTestP, PreparePartitions) {
959   ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
960   uint64_t required_size = 0;
961   EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
962   EXPECT_EQ(0u, required_size);
963 }
964 
965 // Test that if not enough space, required size returned by SnapshotManager is
966 // passed up.
TEST_P(SnapshotPartitionTestP,PreparePartitionsNoSpace)967 TEST_P(SnapshotPartitionTestP, PreparePartitionsNoSpace) {
968   ExpectCreateUpdateSnapshots(android::snapshot::Return::NoSpace(1_GiB));
969   uint64_t required_size = 0;
970   EXPECT_FALSE(PreparePartitionsForUpdate(&required_size));
971   EXPECT_EQ(1_GiB, required_size);
972 }
973 
974 // Test that in recovery, use empty space in super partition for a snapshot
975 // update first.
TEST_P(SnapshotPartitionTestP,RecoveryUseSuperEmpty)976 TEST_P(SnapshotPartitionTestP, RecoveryUseSuperEmpty) {
977   ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
978   EXPECT_CALL(dynamicControl(), IsRecovery()).WillRepeatedly(Return(true));
979   // Must not call PrepareDynamicPartitionsForUpdate if
980   // PrepareSnapshotPartitionsForUpdate succeeds.
981   EXPECT_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
982       .Times(0);
983   uint64_t required_size = 0;
984   EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
985   EXPECT_EQ(0u, required_size);
986 }
987 
988 // Test that in recovery, if CreateUpdateSnapshots throws an error, try
989 // the flashing path for full updates.
TEST_P(SnapshotPartitionTestP,RecoveryErrorShouldDeleteSource)990 TEST_P(SnapshotPartitionTestP, RecoveryErrorShouldDeleteSource) {
991   // Expectation on PreparePartitionsForUpdate
992   ExpectCreateUpdateSnapshots(android::snapshot::Return::NoSpace(1_GiB));
993   EXPECT_CALL(dynamicControl(), IsRecovery()).WillRepeatedly(Return(true));
994   EXPECT_CALL(*snapshot_, CancelUpdate()).WillOnce(Return(true));
995   EXPECT_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
996       .WillRepeatedly(Invoke([&](auto source_slot,
997                                  auto target_slot,
998                                  const auto& manifest,
999                                  auto delete_source) {
1000         EXPECT_EQ(source(), source_slot);
1001         EXPECT_EQ(target(), target_slot);
1002         // Deep comparison requires full protobuf library. Comparing the
1003         // pointers are sufficient.
1004         EXPECT_EQ(&manifest_, &manifest);
1005         EXPECT_TRUE(delete_source);
1006         return dynamicControl().RealPrepareDynamicPartitionsForUpdate(
1007             source_slot, target_slot, manifest, delete_source);
1008       }));
1009   // Expectation on PrepareDynamicPartitionsForUpdate
1010   SetMetadata(source(), {{S("system"), 2_GiB}, {S("vendor"), 1_GiB}});
1011   ExpectUnmap({T("system"), T("vendor")});
1012   // Expect that the source partitions aren't present in target super metadata.
1013   ExpectStoreMetadata({{T("system"), 3_GiB}, {T("vendor"), 1_GiB}});
1014 
1015   uint64_t required_size = 0;
1016   EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
1017   EXPECT_EQ(0u, required_size);
1018 }
1019 
1020 INSTANTIATE_TEST_CASE_P(DynamicPartitionControlAndroidTest,
1021                         SnapshotPartitionTestP,
1022                         testing::Values(TestParam{0, 1}, TestParam{1, 0}));
1023 
1024 }  // namespace chromeos_update_engine
1025