1 //
2 // Copyright (C) 2020 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 <array>
18 #include <cstring>
19 #include <map>
20 #include <memory>
21 #include <numeric>
22 #include <string>
23 #include <vector>
24
25 #include <gtest/gtest.h>
26 #include <google/protobuf/message_lite.h>
27 #include <libsnapshot/cow_writer.h>
28
29 #include "update_engine/payload_consumer/snapshot_extent_writer.h"
30 #include "update_engine/payload_generator/delta_diff_generator.h"
31 #include "update_engine/update_metadata.pb.h"
32
33 namespace chromeos_update_engine {
34
35 class FakeCowWriter : public android::snapshot::ICowWriter {
36 public:
37 struct CowOp {
38 enum { COW_COPY, COW_REPLACE, COW_ZERO } type;
39 std::vector<unsigned char> data;
40 union {
41 size_t source_block;
42 size_t num_blocks;
43 };
44 };
45 using ICowWriter::ICowWriter;
46 ~FakeCowWriter() = default;
47
AddCopy(uint64_t new_block,uint64_t old_block,uint64_t num_blocks)48 bool AddCopy(uint64_t new_block,
49 uint64_t old_block,
50 uint64_t num_blocks) override {
51 for (size_t i = 0; i < num_blocks; i++) {
52 operations_[new_block + i] = {
53 .type = CowOp::COW_COPY,
54 .source_block = static_cast<size_t>(old_block + i)};
55 }
56 return true;
57 }
AddRawBlocks(uint64_t new_block_start,const void * data,size_t size)58 bool AddRawBlocks(uint64_t new_block_start,
59 const void* data,
60 size_t size) override {
61 auto&& op = operations_[new_block_start];
62 const auto uint8_ptr = static_cast<const unsigned char*>(data);
63 op.data.insert(op.data.end(), uint8_ptr, uint8_ptr + size);
64 return true;
65 }
AddZeroBlocks(uint64_t new_block_start,uint64_t num_blocks)66 bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override {
67 operations_[new_block_start] = {.type = CowOp::COW_ZERO};
68 return true;
69 }
AddXorBlocks(uint32_t new_block_start,const void * data,size_t size,uint32_t old_block,uint16_t offset)70 bool AddXorBlocks(uint32_t new_block_start,
71 const void* data,
72 size_t size,
73 uint32_t old_block,
74 uint16_t offset) override {
75 return false;
76 }
Finalize()77 bool Finalize() override {
78 finalize_called_ = true;
79 return true;
80 }
81
AddLabel(uint64_t label)82 bool AddLabel(uint64_t label) {
83 label_count_++;
84 return true;
85 }
86
AddSequenceData(size_t num_ops,const uint32_t * data)87 bool AddSequenceData(size_t num_ops, const uint32_t* data) override {
88 return false;
89 }
90
GetBlockSize() const91 uint32_t GetBlockSize() const override { return 4096; }
GetMaxBlocks() const92 std::optional<uint32_t> GetMaxBlocks() const override { return {}; }
93
OpenReader()94 std::unique_ptr<android::snapshot::ICowReader> OpenReader() override {
95 return nullptr;
96 }
OpenFileDescriptor(const std::optional<std::string> &)97 std::unique_ptr<FileDescriptor> OpenFileDescriptor(
98 const std::optional<std::string>&) override {
99 return nullptr;
100 }
101
102 // Return number of bytes the cow image occupies on disk.
GetCowSize() const103 uint64_t GetCowSize() const {
104 return std::accumulate(
105 operations_.begin(), operations_.end(), 0, [](auto&& acc, auto&& op) {
106 return acc + op.second.data.size();
107 });
108 }
GetCowSizeInfo() const109 android::snapshot::CowSizeInfo GetCowSizeInfo() const override {
110 return android::snapshot::CowSizeInfo{GetCowSize(), 0};
111 }
Contains(size_t block)112 bool Contains(size_t block) {
113 return operations_.find(block) != operations_.end();
114 }
115 bool finalize_called_ = true;
116 size_t label_count_ = 0;
117 std::map<size_t, CowOp> operations_;
118 };
119
120 class SnapshotExtentWriterTest : public ::testing::Test {
121 public:
SetUp()122 void SetUp() override {}
123
124 protected:
125 android::snapshot::CowOptions options_ = {
126 .block_size = static_cast<uint32_t>(kBlockSize)};
127 FakeCowWriter cow_writer_;
128 SnapshotExtentWriter writer_{&cow_writer_};
129 };
130
AddExtent(google::protobuf::RepeatedPtrField<Extent> * extents,size_t start_block,size_t num_blocks)131 void AddExtent(google::protobuf::RepeatedPtrField<Extent>* extents,
132 size_t start_block,
133 size_t num_blocks) {
134 auto&& extent = extents->Add();
135 extent->set_start_block(start_block);
136 extent->set_num_blocks(num_blocks);
137 }
138
TEST_F(SnapshotExtentWriterTest,BufferWrites)139 TEST_F(SnapshotExtentWriterTest, BufferWrites) {
140 google::protobuf::RepeatedPtrField<Extent> extents;
141 AddExtent(&extents, 123, 1);
142 writer_.Init(extents, kBlockSize);
143
144 std::vector<uint8_t> buf(kBlockSize, 0);
145 buf[123] = 231;
146 buf[231] = 123;
147 buf[buf.size() - 1] = 255;
148
149 writer_.Write(buf.data(), kBlockSize - 1);
150 ASSERT_TRUE(cow_writer_.operations_.empty())
151 << "Haven't send data of a complete block yet, CowWriter should not be "
152 "invoked.";
153 writer_.Write(buf.data() + kBlockSize - 1, 1);
154 ASSERT_TRUE(cow_writer_.Contains(123))
155 << "Once a block of data is sent to SnapshotExtentWriter, it should "
156 "forward data to cow_writer.";
157 ASSERT_EQ(cow_writer_.operations_.size(), 1U);
158 ASSERT_EQ(buf, cow_writer_.operations_[123].data);
159 }
160
TEST_F(SnapshotExtentWriterTest,NonBufferedWrites)161 TEST_F(SnapshotExtentWriterTest, NonBufferedWrites) {
162 google::protobuf::RepeatedPtrField<Extent> extents;
163 AddExtent(&extents, 123, 1);
164 AddExtent(&extents, 125, 1);
165 writer_.Init(extents, kBlockSize);
166
167 std::vector<uint8_t> buf(kBlockSize * 2, 0);
168 buf[123] = 231;
169 buf[231] = 123;
170 buf[buf.size() - 1] = 255;
171
172 writer_.Write(buf.data(), buf.size());
173 ASSERT_TRUE(cow_writer_.Contains(123));
174 ASSERT_TRUE(cow_writer_.Contains(125));
175
176 ASSERT_EQ(cow_writer_.operations_.size(), 2U);
177 auto actual_data = cow_writer_.operations_[123].data;
178 actual_data.insert(actual_data.end(),
179 cow_writer_.operations_[125].data.begin(),
180 cow_writer_.operations_[125].data.end());
181 ASSERT_EQ(buf, actual_data);
182 }
183
TEST_F(SnapshotExtentWriterTest,WriteAcrossBlockBoundary)184 TEST_F(SnapshotExtentWriterTest, WriteAcrossBlockBoundary) {
185 google::protobuf::RepeatedPtrField<Extent> extents;
186 AddExtent(&extents, 123, 1);
187 AddExtent(&extents, 125, 2);
188 writer_.Init(extents, kBlockSize);
189
190 std::vector<uint8_t> buf(kBlockSize * 3);
191 std::memset(buf.data(), 0, buf.size());
192 buf[123] = 231;
193 buf[231] = 123;
194 buf[buf.size() - 1] = 255;
195 buf[kBlockSize - 1] = 254;
196
197 writer_.Write(buf.data(), kBlockSize - 1);
198 ASSERT_TRUE(cow_writer_.operations_.empty())
199 << "Haven't send data of a complete block yet, CowWriter should not be "
200 "invoked.";
201 writer_.Write(buf.data() + kBlockSize - 1, 1 + kBlockSize * 2);
202 ASSERT_TRUE(cow_writer_.Contains(123));
203 ASSERT_TRUE(cow_writer_.Contains(125));
204
205 ASSERT_EQ(cow_writer_.operations_.size(), 2U);
206 auto actual_data = cow_writer_.operations_[123].data;
207 actual_data.insert(actual_data.end(),
208 cow_writer_.operations_[125].data.begin(),
209 cow_writer_.operations_[125].data.end());
210 ASSERT_EQ(buf, actual_data);
211 }
212 } // namespace chromeos_update_engine
213