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