1 //
2 // Copyright (C) 2021 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 <unistd.h>
18
19 #include <android-base/file.h>
20 #include <android-base/mapped_file.h>
21 #include <android-base/properties.h>
22 #include <bsdiff/bsdiff.h>
23 #include <gtest/gtest.h>
24 #include <libsnapshot/cow_writer.h>
25 #include <libsnapshot/mock_cow_writer.h>
26
27 #include "update_engine/common/error_code.h"
28 #include "update_engine/common/hash_calculator.h"
29 #include "update_engine/common/mock_dynamic_partition_control.h"
30 #include "update_engine/common/utils.h"
31 #include "update_engine/payload_consumer/vabc_partition_writer.h"
32 #include "update_engine/payload_generator/delta_diff_generator.h"
33 #include "update_engine/payload_generator/extent_ranges.h"
34 #include "update_engine/update_metadata.pb.h"
35
36 namespace chromeos_update_engine {
37
38 using android::snapshot::CowOptions;
39 using testing::_;
40 using testing::Args;
41 using testing::ElementsAreArray;
42 using testing::Invoke;
43 using testing::Return;
44 using testing::Sequence;
45 using utils::GetReadonlyZeroBlock;
46
47 namespace {
48
49 static constexpr auto& fake_part_name = "fake_part";
50 static constexpr size_t FAKE_PART_SIZE = 4096 * 50;
51 class VABCPartitionWriterTest : public ::testing::Test {
52 public:
SetUp()53 void SetUp() override {
54 ftruncate(source_part_.fd, FAKE_PART_SIZE);
55 ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
56 .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::NONE)));
57 }
58
59 protected:
60 void AddBlockTest(bool xor_enabled);
AddMergeOp(PartitionUpdate * partition,std::array<size_t,2> src_extent,std::array<size_t,2> dst_extent,CowMergeOperation_Type type)61 CowMergeOperation* AddMergeOp(PartitionUpdate* partition,
62 std::array<size_t, 2> src_extent,
63 std::array<size_t, 2> dst_extent,
64 CowMergeOperation_Type type) {
65 auto merge_op = partition->add_merge_operations();
66 auto src = merge_op->mutable_src_extent();
67 src->set_start_block(src_extent[0]);
68 src->set_num_blocks(src_extent[1]);
69 auto dst = merge_op->mutable_dst_extent();
70 dst->set_start_block(dst_extent[0]);
71 dst->set_num_blocks(dst_extent[1]);
72 merge_op->set_type(type);
73 return merge_op;
74 }
75
76 android::snapshot::CowOptions options_ = {
77 .block_size = static_cast<uint32_t>(kBlockSize)};
78 android::snapshot::MockCowWriter cow_writer_;
79 MockDynamicPartitionControl dynamic_control_;
80 PartitionUpdate partition_update_;
81 InstallPlan install_plan_;
82 TemporaryFile source_part_;
83 InstallPlan::Partition install_part_{.name = fake_part_name,
84 .source_path = source_part_.path,
85 .source_size = FAKE_PART_SIZE};
86 };
87
TEST_F(VABCPartitionWriterTest,MergeSequenceWriteTest)88 TEST_F(VABCPartitionWriterTest, MergeSequenceWriteTest) {
89 AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
90 AddMergeOp(&partition_update_, {12, 2}, {13, 2}, CowMergeOperation::COW_XOR);
91 AddMergeOp(&partition_update_, {15, 1}, {20, 1}, CowMergeOperation::COW_COPY);
92 AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
93 AddMergeOp(&partition_update_, {42, 5}, {40, 5}, CowMergeOperation::COW_XOR);
94 VABCPartitionWriter writer_{
95 partition_update_, install_part_, &dynamic_control_, kBlockSize};
96 EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
97 .WillOnce(Invoke([](const std::string&,
98 const std::optional<std::string>&,
99 std::optional<uint64_t>) {
100 auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
101 auto expected_merge_sequence = {10, 14, 13, 20, 25, 40, 41, 42, 43, 44};
102 EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
103 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
104 .WillOnce(Return(true));
105 ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
106 ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
107 return cow_writer;
108 }));
109 EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
110 .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
111 ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
112 }
113
TEST_F(VABCPartitionWriterTest,MergeSequenceXorSameBlock)114 TEST_F(VABCPartitionWriterTest, MergeSequenceXorSameBlock) {
115 AddMergeOp(&partition_update_, {19, 4}, {19, 3}, CowMergeOperation::COW_XOR)
116 ->set_src_offset(1);
117 VABCPartitionWriter writer_{
118 partition_update_, install_part_, &dynamic_control_, kBlockSize};
119 EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
120 .WillOnce(
121 Invoke([](const std::string&,
122 const std::optional<std::string>&,
123 std::optional<uint64_t>) {
124 auto cow_writer =
125 std::make_unique<android::snapshot::MockCowWriter>();
126 auto expected_merge_sequence = {19, 20, 21};
127 EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
128 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
129 .WillOnce(Return(true));
130 ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
131 ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
132 return cow_writer;
133 }));
134 EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
135 .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
136 ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
137 }
138
TEST_F(VABCPartitionWriterTest,AddBlockTestXor)139 TEST_F(VABCPartitionWriterTest, AddBlockTestXor) {
140 return AddBlockTest(true);
141 }
142
TEST_F(VABCPartitionWriterTest,AddBlockTestNoXor)143 TEST_F(VABCPartitionWriterTest, AddBlockTestNoXor) {
144 return AddBlockTest(false);
145 }
146
AddBlockTest(bool xor_enabled)147 void VABCPartitionWriterTest::AddBlockTest(bool xor_enabled) {
148 if (xor_enabled) {
149 ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
150 .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
151 } else {
152 ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
153 .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::NONE)));
154 }
155 InstallOperation& install_op = *partition_update_.add_operations();
156 install_op.set_type(InstallOperation::SOURCE_COPY);
157 *install_op.add_src_extents() = ExtentForRange(5, 1);
158 *install_op.add_src_extents() = ExtentForRange(10, 1);
159 *install_op.add_src_extents() = ExtentForRange(15, 2);
160 *install_op.add_src_extents() = ExtentForRange(20, 2);
161
162 *install_op.add_dst_extents() = ExtentForRange(10, 1);
163 *install_op.add_dst_extents() = ExtentForRange(15, 1);
164 *install_op.add_dst_extents() = ExtentForRange(20, 2);
165 *install_op.add_dst_extents() = ExtentForRange(25, 2);
166 AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
167 AddMergeOp(&partition_update_, {10, 1}, {15, 1}, CowMergeOperation::COW_COPY);
168 AddMergeOp(&partition_update_, {15, 2}, {20, 2}, CowMergeOperation::COW_COPY);
169 AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
170 VABCPartitionWriter writer_{
171 partition_update_, install_part_, &dynamic_control_, kBlockSize};
172 EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
173 .WillOnce(Invoke([xor_enabled](const std::string&,
174 const std::optional<std::string>&,
175 std::optional<uint64_t>) {
176 auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
177 ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
178 ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
179 if (xor_enabled) {
180 EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
181 .WillOnce(Return(true));
182 EXPECT_CALL(*cow_writer, AddCopy(10, 5, 1));
183 EXPECT_CALL(*cow_writer, AddCopy(15, 10, 1));
184 // libsnapshot want blocks in reverser order, so 21 goes before 20
185 EXPECT_CALL(*cow_writer, AddCopy(20, 15, 2));
186
187 EXPECT_CALL(*cow_writer, AddCopy(25, 20, 1));
188 EXPECT_CALL(*cow_writer, AddRawBlocks(26, _, 4096))
189 .WillOnce(Return(true));
190 EXPECT_CALL(*cow_writer, Finalize());
191 } else {
192 Sequence s;
193 EXPECT_CALL(*cow_writer, AddCopy(10, 5, 1)).InSequence(s);
194 EXPECT_CALL(*cow_writer, AddCopy(15, 10, 1)).InSequence(s);
195 // libsnapshot want blocks in reverser order, so 21 goes before 20
196 EXPECT_CALL(*cow_writer, AddCopy(20, 15, 2)).InSequence(s);
197
198 EXPECT_CALL(*cow_writer, AddCopy(25, 20, 1)).InSequence(s);
199 EXPECT_CALL(*cow_writer, AddRawBlocks(26, _, 4096))
200 .InSequence(s)
201 .WillOnce(Return(true));
202 }
203 return cow_writer;
204 }));
205 ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
206 ErrorCode error{};
207 ASSERT_TRUE(writer_.PerformSourceCopyOperation(install_op, &error));
208 }
209
GetNoopBSDIFF(size_t data_size)210 std::string GetNoopBSDIFF(size_t data_size) {
211 auto zeros = GetReadonlyZeroBlock(data_size);
212 TemporaryFile patch_file;
213 int error = bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(zeros->data()),
214 zeros->size(),
215 reinterpret_cast<const uint8_t*>(zeros->data()),
216 zeros->size(),
217 patch_file.path,
218 nullptr);
219 if (error) {
220 LOG(ERROR) << "Failed to generate BSDIFF patch " << error;
221 return {};
222 }
223 std::string patch_data;
224 if (!utils::ReadFile(patch_file.path, &patch_data)) {
225 return {};
226 }
227 return patch_data;
228 }
229
TEST_F(VABCPartitionWriterTest,StreamXORBlockTest)230 TEST_F(VABCPartitionWriterTest, StreamXORBlockTest) {
231 AddMergeOp(&partition_update_, {5, 2}, {10, 2}, CowMergeOperation::COW_XOR);
232 AddMergeOp(&partition_update_, {8, 2}, {13, 2}, CowMergeOperation::COW_XOR);
233 auto install_op = partition_update_.add_operations();
234 *install_op->add_src_extents() = ExtentForRange(5, 5);
235 *install_op->add_dst_extents() = ExtentForRange(10, 5);
236 install_op->set_type(InstallOperation::SOURCE_BSDIFF);
237 auto data_hash = install_op->mutable_src_sha256_hash();
238 auto zeros = GetReadonlyZeroBlock(kBlockSize * 5);
239 brillo::Blob expected_hash;
240 truncate64(source_part_.path, kBlockSize * 20);
241 HashCalculator::RawHashOfBytes(zeros->data(), zeros->size(), &expected_hash);
242 data_hash->assign(reinterpret_cast<const char*>(expected_hash.data()),
243 expected_hash.size());
244
245 EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
246 .WillOnce(Invoke([](const std::string&,
247 const std::optional<std::string>&,
248 std::optional<uint64_t>) {
249 auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
250 ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
251 auto expected_merge_sequence = {10, 11, 13, 14};
252 auto expected_merge_sequence_rev = {11, 10, 14, 13};
253 const bool is_ascending = android::base::GetBoolProperty(
254 "ro.virtual_ab.userspace.snapshots.enabled", false);
255 if (!is_ascending) {
256 EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
257 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence_rev)))
258 .WillOnce(Return(true));
259 } else {
260 EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
261 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
262 .WillOnce(Return(true));
263 }
264 EXPECT_CALL(*cow_writer, AddCopy(_, _, _)).Times(0);
265 EXPECT_CALL(*cow_writer, AddRawBlocks(_, _, _)).WillOnce(Return(true));
266 EXPECT_CALL(*cow_writer, AddXorBlocks(10, _, kBlockSize * 2, 5, 0))
267 .WillOnce(Return(true));
268 EXPECT_CALL(*cow_writer, AddXorBlocks(13, _, kBlockSize * 2, 8, 0))
269 .WillOnce(Return(true));
270 return cow_writer;
271 }));
272 EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
273 .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
274 VABCPartitionWriter writer_{
275 partition_update_, install_part_, &dynamic_control_, kBlockSize};
276 ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
277 const auto patch_data = GetNoopBSDIFF(kBlockSize * 5);
278 ASSERT_GT(patch_data.size(), 0UL);
279 ASSERT_TRUE(writer_.PerformDiffOperation(
280 *install_op, nullptr, patch_data.data(), patch_data.size()));
281 }
282
283 } // namespace
284
285 } // namespace chromeos_update_engine
286