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