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 <optional>
16 #include <tuple>
17 
18 #include <gmock/gmock.h>
19 #include <gtest/gtest.h>
20 #include <libdm/dm.h>
21 #include <liblp/builder.h>
22 #include <liblp/property_fetcher.h>
23 
24 #include <libsnapshot/test_helpers.h>
25 
26 #include "dm_snapshot_internals.h"
27 #include "partition_cow_creator.h"
28 #include "utility.h"
29 
30 using namespace android::fs_mgr;
31 
32 using chromeos_update_engine::InstallOperation;
33 using UeExtent = chromeos_update_engine::Extent;
34 using google::protobuf::RepeatedPtrField;
35 using ::testing::Matches;
36 using ::testing::Pointwise;
37 using ::testing::Truly;
38 
39 namespace android {
40 namespace snapshot {
41 
42 class PartitionCowCreatorTest : public ::testing::Test {
43   public:
SetUp()44     void SetUp() override {
45         SKIP_IF_NON_VIRTUAL_AB();
46         SnapshotTestPropertyFetcher::SetUp();
47     }
TearDown()48     void TearDown() override {
49         RETURN_IF_NON_VIRTUAL_AB();
50         SnapshotTestPropertyFetcher::TearDown();
51     }
52 };
53 
TEST_F(PartitionCowCreatorTest,IntersectSelf)54 TEST_F(PartitionCowCreatorTest, IntersectSelf) {
55     constexpr uint64_t super_size = 1_MiB;
56     constexpr uint64_t partition_size = 40_KiB;
57 
58     auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
59     ASSERT_NE(builder_a, nullptr);
60     auto system_a = builder_a->AddPartition("system_a", LP_PARTITION_ATTR_READONLY);
61     ASSERT_NE(system_a, nullptr);
62     ASSERT_TRUE(builder_a->ResizePartition(system_a, partition_size));
63 
64     auto builder_b = MetadataBuilder::New(super_size, 1_KiB, 2);
65     ASSERT_NE(builder_b, nullptr);
66     auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
67     ASSERT_NE(system_b, nullptr);
68     ASSERT_TRUE(builder_b->ResizePartition(system_b, partition_size));
69 
70     PartitionCowCreator creator{.target_metadata = builder_b.get(),
71                                 .target_suffix = "_b",
72                                 .target_partition = system_b,
73                                 .current_metadata = builder_a.get(),
74                                 .current_suffix = "_a"};
75     auto ret = creator.Run();
76     ASSERT_TRUE(ret.has_value());
77     ASSERT_EQ(partition_size, ret->snapshot_status.device_size());
78     ASSERT_EQ(partition_size, ret->snapshot_status.snapshot_size());
79 }
80 
TEST_F(PartitionCowCreatorTest,Holes)81 TEST_F(PartitionCowCreatorTest, Holes) {
82     const auto& opener = test_device->GetPartitionOpener();
83 
84     constexpr auto slack_space = 1_MiB;
85     constexpr auto big_size = (kSuperSize - slack_space) / 2;
86     constexpr auto small_size = big_size / 2;
87 
88     BlockDeviceInfo super_device("super", kSuperSize, 0, 0, 4_KiB);
89     std::vector<BlockDeviceInfo> devices = {super_device};
90     auto source = MetadataBuilder::New(devices, "super", 1_KiB, 2);
91     auto system = source->AddPartition("system_a", 0);
92     ASSERT_NE(nullptr, system);
93     ASSERT_TRUE(source->ResizePartition(system, big_size));
94     auto vendor = source->AddPartition("vendor_a", 0);
95     ASSERT_NE(nullptr, vendor);
96     ASSERT_TRUE(source->ResizePartition(vendor, big_size));
97     // Create a hole between system and vendor
98     ASSERT_TRUE(source->ResizePartition(system, small_size));
99     auto source_metadata = source->Export();
100     ASSERT_NE(nullptr, source_metadata);
101     ASSERT_TRUE(FlashPartitionTable(opener, fake_super, *source_metadata.get()));
102 
103     auto target = MetadataBuilder::NewForUpdate(opener, "super", 0, 1);
104     // Shrink vendor
105     vendor = target->FindPartition("vendor_b");
106     ASSERT_NE(nullptr, vendor);
107     ASSERT_TRUE(target->ResizePartition(vendor, small_size));
108     // Grow system to take hole & saved space from vendor
109     system = target->FindPartition("system_b");
110     ASSERT_NE(nullptr, system);
111     ASSERT_TRUE(target->ResizePartition(system, big_size * 2 - small_size));
112 
113     PartitionCowCreator creator{.target_metadata = target.get(),
114                                 .target_suffix = "_b",
115                                 .target_partition = system,
116                                 .current_metadata = source.get(),
117                                 .current_suffix = "_a"};
118     auto ret = creator.Run();
119     ASSERT_TRUE(ret.has_value());
120 }
121 
TEST_F(PartitionCowCreatorTest,CowSize)122 TEST_F(PartitionCowCreatorTest, CowSize) {
123     using InstallOperation = chromeos_update_engine::InstallOperation;
124     using RepeatedInstallOperationPtr = google::protobuf::RepeatedPtrField<InstallOperation>;
125     using Extent = chromeos_update_engine::Extent;
126 
127     constexpr uint64_t super_size = 50_MiB;
128     constexpr uint64_t partition_size = 40_MiB;
129 
130     auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
131     ASSERT_NE(builder_a, nullptr);
132     auto system_a = builder_a->AddPartition("system_a", LP_PARTITION_ATTR_READONLY);
133     ASSERT_NE(system_a, nullptr);
134     ASSERT_TRUE(builder_a->ResizePartition(system_a, partition_size));
135 
136     auto builder_b = MetadataBuilder::New(super_size, 1_KiB, 2);
137     ASSERT_NE(builder_b, nullptr);
138     auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
139     ASSERT_NE(system_b, nullptr);
140     ASSERT_TRUE(builder_b->ResizePartition(system_b, partition_size));
141 
142     const uint64_t block_size = builder_b->logical_block_size();
143     const uint64_t chunk_size = kSnapshotChunkSize * dm::kSectorSize;
144     ASSERT_EQ(chunk_size, block_size);
145 
146     auto cow_device_size = [](const std::vector<InstallOperation>& iopv, MetadataBuilder* builder_a,
147                               MetadataBuilder* builder_b, Partition* system_b) {
148         PartitionUpdate update;
149         *update.mutable_operations() = RepeatedInstallOperationPtr(iopv.begin(), iopv.end());
150 
151         PartitionCowCreator creator{.target_metadata = builder_b,
152                                     .target_suffix = "_b",
153                                     .target_partition = system_b,
154                                     .current_metadata = builder_a,
155                                     .current_suffix = "_a",
156                                     .update = &update};
157 
158         auto ret = creator.Run();
159 
160         if (ret.has_value()) {
161             return ret->snapshot_status.cow_file_size() + ret->snapshot_status.cow_partition_size();
162         }
163         return std::numeric_limits<uint64_t>::max();
164     };
165 
166     std::vector<InstallOperation> iopv;
167     InstallOperation iop;
168     Extent* e;
169 
170     // No data written, no operations performed
171     ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
172 
173     // No data written
174     e = iop.add_dst_extents();
175     e->set_start_block(0);
176     e->set_num_blocks(0);
177     iopv.push_back(iop);
178     ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
179 
180     e = iop.add_dst_extents();
181     e->set_start_block(1);
182     e->set_num_blocks(0);
183     iopv.push_back(iop);
184     ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
185 
186     // Fill the first block
187     e = iop.add_dst_extents();
188     e->set_start_block(0);
189     e->set_num_blocks(1);
190     iopv.push_back(iop);
191     ASSERT_EQ(3 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
192 
193     // Fill the second block
194     e = iop.add_dst_extents();
195     e->set_start_block(1);
196     e->set_num_blocks(1);
197     iopv.push_back(iop);
198     ASSERT_EQ(4 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
199 
200     // Jump to 5th block and write 2
201     e = iop.add_dst_extents();
202     e->set_start_block(5);
203     e->set_num_blocks(2);
204     iopv.push_back(iop);
205     ASSERT_EQ(6 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
206 }
207 
TEST_F(PartitionCowCreatorTest,Zero)208 TEST_F(PartitionCowCreatorTest, Zero) {
209     constexpr uint64_t super_size = 1_MiB;
210     auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
211     ASSERT_NE(builder_a, nullptr);
212 
213     auto builder_b = MetadataBuilder::New(super_size, 1_KiB, 2);
214     ASSERT_NE(builder_b, nullptr);
215     auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
216     ASSERT_NE(system_b, nullptr);
217 
218     PartitionCowCreator creator{.target_metadata = builder_b.get(),
219                                 .target_suffix = "_b",
220                                 .target_partition = system_b,
221                                 .current_metadata = builder_a.get(),
222                                 .current_suffix = "_a",
223                                 .update = nullptr};
224 
225     auto ret = creator.Run();
226 
227     ASSERT_EQ(0u, ret->snapshot_status.device_size());
228     ASSERT_EQ(0u, ret->snapshot_status.snapshot_size());
229     ASSERT_EQ(0u, ret->snapshot_status.cow_file_size());
230     ASSERT_EQ(0u, ret->snapshot_status.cow_partition_size());
231 }
232 
TEST_F(PartitionCowCreatorTest,CompressionEnabled)233 TEST_F(PartitionCowCreatorTest, CompressionEnabled) {
234     constexpr uint64_t super_size = 1_MiB;
235     auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
236     ASSERT_NE(builder_a, nullptr);
237 
238     auto builder_b = MetadataBuilder::New(super_size, 1_KiB, 2);
239     ASSERT_NE(builder_b, nullptr);
240     auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
241     ASSERT_NE(system_b, nullptr);
242     ASSERT_TRUE(builder_b->ResizePartition(system_b, 128_KiB));
243 
244     PartitionUpdate update;
245     update.set_estimate_cow_size(256_KiB);
246 
247     PartitionCowCreator creator{.target_metadata = builder_b.get(),
248                                 .target_suffix = "_b",
249                                 .target_partition = system_b,
250                                 .current_metadata = builder_a.get(),
251                                 .current_suffix = "_a",
252                                 .using_snapuserd = true,
253                                 .update = &update};
254 
255     auto ret = creator.Run();
256     ASSERT_TRUE(ret.has_value());
257     ASSERT_EQ(ret->snapshot_status.cow_file_size(), 1458176);
258 }
259 
TEST_F(PartitionCowCreatorTest,CompressionWithNoManifest)260 TEST_F(PartitionCowCreatorTest, CompressionWithNoManifest) {
261     constexpr uint64_t super_size = 1_MiB;
262     auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
263     ASSERT_NE(builder_a, nullptr);
264 
265     auto builder_b = MetadataBuilder::New(super_size, 1_KiB, 2);
266     ASSERT_NE(builder_b, nullptr);
267     auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
268     ASSERT_NE(system_b, nullptr);
269     ASSERT_TRUE(builder_b->ResizePartition(system_b, 128_KiB));
270 
271     PartitionUpdate update;
272 
273     PartitionCowCreator creator{.target_metadata = builder_b.get(),
274                                 .target_suffix = "_b",
275                                 .target_partition = system_b,
276                                 .current_metadata = builder_a.get(),
277                                 .current_suffix = "_a",
278                                 .using_snapuserd = true,
279                                 .update = nullptr};
280 
281     auto ret = creator.Run();
282     ASSERT_FALSE(ret.has_value());
283 }
284 
TEST(DmSnapshotInternals,CowSizeCalculator)285 TEST(DmSnapshotInternals, CowSizeCalculator) {
286     SKIP_IF_NON_VIRTUAL_AB();
287 
288     DmSnapCowSizeCalculator cc(512, 8);
289     unsigned long int b;
290 
291     // Empty COW
292     ASSERT_EQ(cc.cow_size_sectors(), 16);
293 
294     // First chunk written
295     for (b = 0; b < 4_KiB; ++b) {
296         cc.WriteByte(b);
297         ASSERT_EQ(cc.cow_size_sectors(), 24);
298     }
299 
300     // Second chunk written
301     for (b = 4_KiB; b < 8_KiB; ++b) {
302         cc.WriteByte(b);
303         ASSERT_EQ(cc.cow_size_sectors(), 32);
304     }
305 
306     // Leave a hole and write 5th chunk
307     for (b = 16_KiB; b < 20_KiB; ++b) {
308         cc.WriteByte(b);
309         ASSERT_EQ(cc.cow_size_sectors(), 40);
310     }
311 
312     // Write a byte that would surely overflow the counter
313     cc.WriteChunk(std::numeric_limits<uint64_t>::max());
314     ASSERT_FALSE(cc.cow_size_sectors().has_value());
315 }
316 
BlocksToExtents(const std::vector<uint64_t> & blocks,google::protobuf::RepeatedPtrField<UeExtent> * extents)317 void BlocksToExtents(const std::vector<uint64_t>& blocks,
318                      google::protobuf::RepeatedPtrField<UeExtent>* extents) {
319     for (uint64_t block : blocks) {
320         AppendExtent(extents, block, 1);
321     }
322 }
323 
324 template <typename T>
ExtentsToBlocks(const T & extents)325 std::vector<uint64_t> ExtentsToBlocks(const T& extents) {
326     std::vector<uint64_t> blocks;
327     for (const auto& extent : extents) {
328         for (uint64_t offset = 0; offset < extent.num_blocks(); ++offset) {
329             blocks.push_back(extent.start_block() + offset);
330         }
331     }
332     return blocks;
333 }
334 
CreateCopyOp(const std::vector<uint64_t> & src_blocks,const std::vector<uint64_t> & dst_blocks)335 InstallOperation CreateCopyOp(const std::vector<uint64_t>& src_blocks,
336                               const std::vector<uint64_t>& dst_blocks) {
337     InstallOperation op;
338     op.set_type(InstallOperation::SOURCE_COPY);
339     BlocksToExtents(src_blocks, op.mutable_src_extents());
340     BlocksToExtents(dst_blocks, op.mutable_dst_extents());
341     return op;
342 }
343 
344 // ExtentEqual(tuple<UeExtent, UeExtent>)
345 MATCHER(ExtentEqual, "") {
346     auto&& [a, b] = arg;
347     return a.start_block() == b.start_block() && a.num_blocks() == b.num_blocks();
348 }
349 
350 struct OptimizeOperationTestParam {
351     InstallOperation input;
352     std::optional<InstallOperation> expected_output;
353 };
354 
355 class OptimizeOperationTest : public ::testing::TestWithParam<OptimizeOperationTestParam> {
SetUp()356     void SetUp() override { SKIP_IF_NON_VIRTUAL_AB(); }
357 };
TEST_P(OptimizeOperationTest,Test)358 TEST_P(OptimizeOperationTest, Test) {
359     InstallOperation actual_output;
360     EXPECT_EQ(GetParam().expected_output.has_value(),
361               OptimizeSourceCopyOperation(GetParam().input, &actual_output))
362             << "OptimizeSourceCopyOperation should "
363             << (GetParam().expected_output.has_value() ? "succeed" : "fail");
364     if (!GetParam().expected_output.has_value()) return;
365     EXPECT_THAT(actual_output.src_extents(),
366                 Pointwise(ExtentEqual(), GetParam().expected_output->src_extents()));
367     EXPECT_THAT(actual_output.dst_extents(),
368                 Pointwise(ExtentEqual(), GetParam().expected_output->dst_extents()));
369 }
370 
GetOptimizeOperationTestParams()371 std::vector<OptimizeOperationTestParam> GetOptimizeOperationTestParams() {
372     return {
373             {CreateCopyOp({}, {}), CreateCopyOp({}, {})},
374             {CreateCopyOp({1, 2, 4}, {1, 2, 4}), CreateCopyOp({}, {})},
375             {CreateCopyOp({1, 2, 3}, {4, 5, 6}), std::nullopt},
376             {CreateCopyOp({3, 2}, {1, 2}), CreateCopyOp({3}, {1})},
377             {CreateCopyOp({5, 6, 3, 4, 1, 2}, {1, 2, 3, 4, 5, 6}),
378              CreateCopyOp({5, 6, 1, 2}, {1, 2, 5, 6})},
379             {CreateCopyOp({1, 2, 3, 5, 5, 6}, {5, 6, 3, 4, 1, 2}),
380              CreateCopyOp({1, 2, 5, 5, 6}, {5, 6, 4, 1, 2})},
381             {CreateCopyOp({1, 2, 5, 6, 9, 10}, {1, 4, 5, 6, 7, 8}),
382              CreateCopyOp({2, 9, 10}, {4, 7, 8})},
383             {CreateCopyOp({2, 3, 3, 4, 4}, {1, 2, 3, 4, 5}), CreateCopyOp({2, 3, 4}, {1, 2, 5})},
384     };
385 }
386 
387 INSTANTIATE_TEST_CASE_P(Snapshot, OptimizeOperationTest,
388                         ::testing::ValuesIn(GetOptimizeOperationTestParams()));
389 
390 }  // namespace snapshot
391 }  // namespace android
392