1 //
2 // Unless required by applicable law or agreed to in writing, software
3 // distributed under the License is distributed on an "AS IS" BASIS,
4 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
5 // See the License for the specific language governing permissions and
6 // limitations under the License.
7 
8 #include <sys/stat.h>
9 
10 #include <cstdio>
11 #include <limits>
12 #include <memory>
13 
14 #include <android-base/file.h>
15 #include <android-base/logging.h>
16 #include <android-base/unique_fd.h>
17 #include <gtest/gtest.h>
18 #include <libsnapshot/cow_format.h>
19 #include <libsnapshot/cow_reader.h>
20 #include <libsnapshot/cow_writer.h>
21 #include <storage_literals/storage_literals.h>
22 #include "writer_v2.h"
23 #include "writer_v3.h"
24 
25 using android::base::unique_fd;
26 using testing::AssertionFailure;
27 using testing::AssertionResult;
28 using testing::AssertionSuccess;
29 
30 namespace android {
31 namespace snapshot {
32 
33 using namespace android::storage_literals;
34 using ::testing::TestWithParam;
35 
36 class CowTestV3 : public ::testing::Test {
37   protected:
SetUp()38     virtual void SetUp() override {
39         cow_ = std::make_unique<TemporaryFile>();
40         ASSERT_GE(cow_->fd, 0) << strerror(errno);
41     }
42 
TearDown()43     virtual void TearDown() override { cow_ = nullptr; }
44 
GetCowFd()45     unique_fd GetCowFd() { return unique_fd{dup(cow_->fd)}; }
46 
47     std::unique_ptr<TemporaryFile> cow_;
48 };
49 
50 // Helper to check read sizes.
ReadData(CowReader & reader,const CowOperation * op,void * buffer,size_t size)51 static inline bool ReadData(CowReader& reader, const CowOperation* op, void* buffer, size_t size) {
52     return reader.ReadData(op, buffer, size) == size;
53 }
54 
TEST_F(CowTestV3,CowHeaderV2Test)55 TEST_F(CowTestV3, CowHeaderV2Test) {
56     CowOptions options;
57     options.cluster_ops = 5;
58     options.num_merge_ops = 1;
59     options.block_size = 4096;
60     std::string data = "This is some data, believe it";
61     data.resize(options.block_size, '\0');
62     auto writer_v2 = std::make_unique<CowWriterV2>(options, GetCowFd());
63     ASSERT_TRUE(writer_v2->Initialize());
64     ASSERT_TRUE(writer_v2->Finalize());
65 
66     CowReader reader;
67     ASSERT_TRUE(reader.Parse(cow_->fd));
68 
69     const auto& header = reader.GetHeader();
70     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
71     ASSERT_EQ(header.prefix.major_version, 2);
72     ASSERT_EQ(header.prefix.minor_version, 0);
73     ASSERT_EQ(header.block_size, options.block_size);
74     ASSERT_EQ(header.cluster_ops, options.cluster_ops);
75 }
76 
TEST_F(CowTestV3,Header)77 TEST_F(CowTestV3, Header) {
78     CowOptions options;
79     options.op_count_max = 15;
80     auto writer = CreateCowWriter(3, options, GetCowFd());
81     ASSERT_TRUE(writer->Finalize());
82 
83     CowReader reader;
84     ASSERT_TRUE(reader.Parse(cow_->fd));
85 
86     const auto& header = reader.GetHeader();
87     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
88     ASSERT_EQ(header.prefix.major_version, 3);
89     ASSERT_EQ(header.prefix.minor_version, 0);
90     ASSERT_EQ(header.block_size, options.block_size);
91     ASSERT_EQ(header.cluster_ops, 0);
92 }
93 
TEST_F(CowTestV3,MaxOp)94 TEST_F(CowTestV3, MaxOp) {
95     CowOptions options;
96     options.op_count_max = 20;
97     auto writer = CreateCowWriter(3, options, GetCowFd());
98     ASSERT_FALSE(writer->AddZeroBlocks(1, 21));
99     ASSERT_TRUE(writer->AddZeroBlocks(1, 20));
100     std::string data = "This is some data, believe it";
101     data.resize(options.block_size, '\0');
102 
103     ASSERT_FALSE(writer->AddRawBlocks(5, data.data(), data.size()));
104 
105     ASSERT_TRUE(writer->Finalize());
106 
107     CowReader reader;
108     ASSERT_TRUE(reader.Parse(cow_->fd));
109     ASSERT_EQ(reader.header_v3().op_count, 20);
110 }
111 
TEST_F(CowTestV3,MaxOpSingleThreadCompression)112 TEST_F(CowTestV3, MaxOpSingleThreadCompression) {
113     CowOptions options;
114     options.op_count_max = 20;
115     options.num_compress_threads = 1;
116     options.compression_factor = 4096 * 8;
117     options.compression = "lz4";
118 
119     auto writer = CreateCowWriter(3, options, GetCowFd());
120     ASSERT_TRUE(writer->AddZeroBlocks(1, 20));
121     std::string data = "This is some data, believe it";
122     data.resize(options.block_size, '\0');
123 
124     ASSERT_FALSE(writer->AddRawBlocks(5, data.data(), data.size()));
125 
126     ASSERT_TRUE(writer->Finalize());
127 }
128 
TEST_F(CowTestV3,MaxOpMultiThreadCompression)129 TEST_F(CowTestV3, MaxOpMultiThreadCompression) {
130     CowOptions options;
131     options.op_count_max = 20;
132     options.num_compress_threads = 2;
133     options.compression_factor = 4096 * 8;
134     options.compression = "lz4";
135 
136     auto writer = CreateCowWriter(3, options, GetCowFd());
137     ASSERT_TRUE(writer->AddZeroBlocks(1, 20));
138     std::string data = "This is some data, believe it";
139     data.resize(options.block_size, '\0');
140 
141     ASSERT_FALSE(writer->AddRawBlocks(5, data.data(), data.size()));
142 
143     ASSERT_TRUE(writer->Finalize());
144 }
145 
TEST_F(CowTestV3,ZeroOp)146 TEST_F(CowTestV3, ZeroOp) {
147     CowOptions options;
148     options.op_count_max = 20;
149     auto writer = CreateCowWriter(3, options, GetCowFd());
150     ASSERT_TRUE(writer->AddZeroBlocks(1, 2));
151     ASSERT_TRUE(writer->Finalize());
152 
153     CowReader reader;
154     ASSERT_TRUE(reader.Parse(cow_->fd));
155     ASSERT_EQ(reader.header_v3().op_count, 2);
156 
157     auto iter = reader.GetOpIter();
158     ASSERT_NE(iter, nullptr);
159     ASSERT_FALSE(iter->AtEnd());
160 
161     auto op = iter->Get();
162     ASSERT_EQ(op->type(), kCowZeroOp);
163     ASSERT_EQ(op->data_length, 0);
164     ASSERT_EQ(op->new_block, 1);
165     ASSERT_EQ(op->source(), 0);
166 
167     iter->Next();
168     ASSERT_FALSE(iter->AtEnd());
169     op = iter->Get();
170 
171     ASSERT_EQ(op->type(), kCowZeroOp);
172     ASSERT_EQ(op->data_length, 0);
173     ASSERT_EQ(op->new_block, 2);
174     ASSERT_EQ(op->source(), 0);
175 }
176 
TEST_F(CowTestV3,ReplaceOp)177 TEST_F(CowTestV3, ReplaceOp) {
178     CowOptions options;
179     options.op_count_max = 20;
180     options.scratch_space = false;
181     auto writer = CreateCowWriter(3, options, GetCowFd());
182     std::string data = "This is some data, believe it";
183     data.resize(options.block_size, '\0');
184 
185     ASSERT_TRUE(writer->AddRawBlocks(5, data.data(), data.size()));
186     ASSERT_TRUE(writer->Finalize());
187 
188     CowReader reader;
189     ASSERT_TRUE(reader.Parse(cow_->fd));
190 
191     const auto& header = reader.header_v3();
192     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
193     ASSERT_EQ(header.prefix.major_version, 3);
194     ASSERT_EQ(header.prefix.minor_version, kCowVersionMinor);
195     ASSERT_EQ(header.block_size, options.block_size);
196     ASSERT_EQ(header.op_count, 1);
197 
198     auto iter = reader.GetOpIter();
199     ASSERT_NE(iter, nullptr);
200     ASSERT_FALSE(iter->AtEnd());
201 
202     auto op = iter->Get();
203     std::string sink(data.size(), '\0');
204 
205     ASSERT_EQ(op->type(), kCowReplaceOp);
206     ASSERT_EQ(op->data_length, 4096);
207     ASSERT_EQ(op->new_block, 5);
208     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
209     ASSERT_EQ(sink, data);
210 }
211 
TEST_F(CowTestV3,BigReplaceOp)212 TEST_F(CowTestV3, BigReplaceOp) {
213     CowOptions options;
214     options.op_count_max = 10000;
215     options.batch_write = true;
216     options.cluster_ops = 2048;
217 
218     auto writer = CreateCowWriter(3, options, GetCowFd());
219     std::string data = "This is some data, believe it";
220     data.resize(options.block_size * 4096, '\0');
221     for (int i = 0; i < data.size(); i++) {
222         data[i] = static_cast<char>('A' + i / options.block_size);
223     }
224     ASSERT_TRUE(writer->AddRawBlocks(5, data.data(), data.size()));
225     ASSERT_TRUE(writer->Finalize());
226 
227     CowReader reader;
228     ASSERT_TRUE(reader.Parse(cow_->fd));
229 
230     const auto& header = reader.header_v3();
231     ASSERT_EQ(header.op_count, 4096);
232 
233     auto iter = reader.GetOpIter();
234     ASSERT_NE(iter, nullptr);
235     ASSERT_FALSE(iter->AtEnd());
236 
237     size_t i = 0;
238 
239     while (!iter->AtEnd()) {
240         auto op = iter->Get();
241         std::string sink(options.block_size, '\0');
242         ASSERT_EQ(op->type(), kCowReplaceOp);
243         ASSERT_EQ(op->data_length, options.block_size);
244         ASSERT_EQ(op->new_block, 5 + i);
245         ASSERT_TRUE(ReadData(reader, op, sink.data(), options.block_size));
246         ASSERT_EQ(std::string_view(sink),
247                   std::string_view(data).substr(i * options.block_size, options.block_size))
248                 << " readback data for " << i << "th block does not match";
249         iter->Next();
250         i++;
251     }
252 }
253 
TEST_F(CowTestV3,ConsecutiveReplaceOp)254 TEST_F(CowTestV3, ConsecutiveReplaceOp) {
255     CowOptions options;
256     options.op_count_max = 20;
257     options.scratch_space = false;
258     auto writer = CreateCowWriter(3, options, GetCowFd());
259     std::string data;
260     data.resize(options.block_size * 5);
261     for (int i = 0; i < data.size(); i++) {
262         data[i] = static_cast<char>('A' + i / options.block_size);
263     }
264 
265     ASSERT_TRUE(writer->AddRawBlocks(5, data.data(), data.size()));
266     ASSERT_TRUE(writer->Finalize());
267 
268     CowReader reader;
269     ASSERT_TRUE(reader.Parse(cow_->fd));
270 
271     const auto& header = reader.header_v3();
272     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
273     ASSERT_EQ(header.prefix.major_version, 3);
274     ASSERT_EQ(header.prefix.minor_version, kCowVersionMinor);
275     ASSERT_EQ(header.block_size, options.block_size);
276     ASSERT_EQ(header.op_count, 5);
277 
278     auto iter = reader.GetOpIter();
279     ASSERT_NE(iter, nullptr);
280     ASSERT_FALSE(iter->AtEnd());
281 
282     size_t i = 0;
283 
284     while (!iter->AtEnd()) {
285         auto op = iter->Get();
286         std::string sink(options.block_size, '\0');
287         ASSERT_EQ(op->type(), kCowReplaceOp);
288         ASSERT_EQ(op->data_length, options.block_size);
289         ASSERT_EQ(op->new_block, 5 + i);
290         ASSERT_TRUE(ReadData(reader, op, sink.data(), options.block_size));
291         ASSERT_EQ(std::string_view(sink),
292                   std::string_view(data).substr(i * options.block_size, options.block_size))
293                 << " readback data for " << i << "th block does not match";
294         iter->Next();
295         i++;
296     }
297 
298     ASSERT_EQ(i, 5);
299 }
300 
TEST_F(CowTestV3,CopyOp)301 TEST_F(CowTestV3, CopyOp) {
302     CowOptions options;
303     options.op_count_max = 100;
304     auto writer = CreateCowWriter(3, options, GetCowFd());
305 
306     ASSERT_TRUE(writer->AddCopy(10, 1000, 100));
307     ASSERT_TRUE(writer->Finalize());
308     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
309 
310     CowReader reader;
311     ASSERT_TRUE(reader.Parse(cow_->fd));
312 
313     const auto& header = reader.header_v3();
314     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
315     ASSERT_EQ(header.prefix.major_version, 3);
316     ASSERT_EQ(header.prefix.minor_version, kCowVersionMinor);
317     ASSERT_EQ(header.block_size, options.block_size);
318 
319     auto iter = reader.GetOpIter();
320     ASSERT_NE(iter, nullptr);
321     ASSERT_FALSE(iter->AtEnd());
322 
323     size_t i = 0;
324     while (!iter->AtEnd()) {
325         auto op = iter->Get();
326         ASSERT_EQ(op->type(), kCowCopyOp);
327         ASSERT_EQ(op->data_length, 0);
328         ASSERT_EQ(op->new_block, 10 + i);
329         ASSERT_EQ(op->source(), 1000 + i);
330         iter->Next();
331         i += 1;
332     }
333 
334     ASSERT_EQ(i, 100);
335 }
336 
TEST_F(CowTestV3,XorOp)337 TEST_F(CowTestV3, XorOp) {
338     CowOptions options;
339     options.op_count_max = 100;
340     auto writer = CreateCowWriter(3, options, GetCowFd());
341 
342     std::string data = "This is test data-1. Testing xor";
343     data.resize(options.block_size, '\0');
344     ASSERT_TRUE(writer->AddXorBlocks(50, data.data(), data.size(), 24, 10));
345     ASSERT_TRUE(writer->Finalize());
346 
347     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
348 
349     CowReader reader;
350 
351     ASSERT_TRUE(reader.Parse(cow_->fd));
352 
353     const auto& header = reader.header_v3();
354     ASSERT_EQ(header.op_count, 1);
355 
356     auto iter = reader.GetOpIter();
357     ASSERT_NE(iter, nullptr);
358     ASSERT_FALSE(iter->AtEnd());
359     auto op = iter->Get();
360     std::string sink(data.size(), '\0');
361 
362     ASSERT_EQ(op->type(), kCowXorOp);
363     ASSERT_EQ(op->data_length, 4096);
364     ASSERT_EQ(op->new_block, 50);
365     ASSERT_EQ(op->source(), 98314);  // 4096 * 24 + 10
366     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
367     ASSERT_EQ(sink, data);
368 }
369 
TEST_F(CowTestV3,ConsecutiveXorOp)370 TEST_F(CowTestV3, ConsecutiveXorOp) {
371     CowOptions options;
372     options.op_count_max = 100;
373     auto writer = CreateCowWriter(3, options, GetCowFd());
374 
375     std::string data;
376     data.resize(options.block_size * 5);
377     for (int i = 0; i < data.size(); i++) {
378         data[i] = char(rand() % 256);
379     }
380 
381     ASSERT_TRUE(writer->AddXorBlocks(50, data.data(), data.size(), 24, 10));
382     ASSERT_TRUE(writer->Finalize());
383 
384     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
385 
386     CowReader reader;
387 
388     ASSERT_TRUE(reader.Parse(cow_->fd));
389 
390     const auto& header = reader.header_v3();
391     ASSERT_EQ(header.op_count, 5);
392 
393     auto iter = reader.GetOpIter();
394     ASSERT_NE(iter, nullptr);
395     ASSERT_FALSE(iter->AtEnd());
396 
397     std::string sink(data.size(), '\0');
398     size_t i = 0;
399 
400     while (!iter->AtEnd()) {
401         auto op = iter->Get();
402         ASSERT_EQ(op->type(), kCowXorOp);
403         ASSERT_EQ(op->data_length, 4096);
404         ASSERT_EQ(op->new_block, 50 + i);
405         ASSERT_EQ(op->source(), 98314 + (i * options.block_size));  // 4096 * 24 + 10
406         ASSERT_TRUE(
407                 ReadData(reader, op, sink.data() + (i * options.block_size), options.block_size));
408         iter->Next();
409         i++;
410     }
411     ASSERT_EQ(sink, data);
412 
413     ASSERT_EQ(i, 5);
414 }
415 
TEST_F(CowTestV3,AllOpsWithCompression)416 TEST_F(CowTestV3, AllOpsWithCompression) {
417     CowOptions options;
418     options.compression = "gz";
419     options.op_count_max = 100;
420     auto writer = CreateCowWriter(3, options, GetCowFd());
421 
422     std::string data;
423     data.resize(options.block_size * 5);
424     for (int i = 0; i < data.size(); i++) {
425         data[i] = char(rand() % 4);
426     }
427 
428     ASSERT_TRUE(writer->AddZeroBlocks(10, 5));
429     ASSERT_TRUE(writer->AddCopy(15, 3, 5));
430     ASSERT_TRUE(writer->AddRawBlocks(18, data.data(), data.size()));
431     ASSERT_TRUE(writer->AddXorBlocks(50, data.data(), data.size(), 24, 10));
432     ASSERT_TRUE(writer->Finalize());
433 
434     CowReader reader;
435 
436     ASSERT_TRUE(reader.Parse(cow_->fd));
437 
438     const auto& header = reader.header_v3();
439     ASSERT_EQ(header.prefix.magic, kCowMagicNumber);
440     ASSERT_EQ(header.prefix.major_version, 3);
441     ASSERT_EQ(header.prefix.minor_version, kCowVersionMinor);
442     ASSERT_EQ(header.block_size, options.block_size);
443     ASSERT_EQ(header.buffer_size, BUFFER_REGION_DEFAULT_SIZE);
444     ASSERT_EQ(header.op_count, 20);
445     ASSERT_EQ(header.op_count_max, 100);
446 
447     auto iter = reader.GetOpIter();
448     ASSERT_NE(iter, nullptr);
449     ASSERT_FALSE(iter->AtEnd());
450 
451     for (size_t i = 0; i < 5; i++) {
452         auto op = iter->Get();
453         ASSERT_EQ(op->type(), kCowZeroOp);
454         ASSERT_EQ(op->new_block, 10 + i);
455         iter->Next();
456     }
457     for (size_t i = 0; i < 5; i++) {
458         auto op = iter->Get();
459         ASSERT_EQ(op->type(), kCowCopyOp);
460         ASSERT_EQ(op->new_block, 15 + i);
461         ASSERT_EQ(op->source(), 3 + i);
462         iter->Next();
463     }
464     std::string sink(data.size(), '\0');
465 
466     for (size_t i = 0; i < 5; i++) {
467         auto op = iter->Get();
468         ASSERT_EQ(op->type(), kCowReplaceOp);
469         ASSERT_EQ(op->new_block, 18 + i);
470         ASSERT_EQ(reader.ReadData(op, sink.data() + (i * options.block_size), options.block_size),
471                   options.block_size);
472         iter->Next();
473     }
474     ASSERT_EQ(sink, data);
475 
476     std::fill(sink.begin(), sink.end(), '\0');
477     for (size_t i = 0; i < 5; i++) {
478         auto op = iter->Get();
479         ASSERT_EQ(op->type(), kCowXorOp);
480         ASSERT_EQ(op->new_block, 50 + i);
481         ASSERT_EQ(op->source(), 98314 + (i * options.block_size));  // 4096 * 24 + 10
482         ASSERT_TRUE(
483                 ReadData(reader, op, sink.data() + (i * options.block_size), options.block_size));
484         iter->Next();
485     }
486     ASSERT_EQ(sink, data);
487 }
488 
TEST_F(CowTestV3,GzCompression)489 TEST_F(CowTestV3, GzCompression) {
490     CowOptions options;
491     options.op_count_max = 100;
492     options.compression = "gz";
493     auto writer = CreateCowWriter(3, options, GetCowFd());
494 
495     std::string data = "This is some data, believe it";
496     data.resize(options.block_size, '\0');
497 
498     ASSERT_TRUE(writer->AddRawBlocks(50, data.data(), data.size()));
499     ASSERT_TRUE(writer->Finalize());
500 
501     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
502 
503     CowReader reader;
504     ASSERT_TRUE(reader.Parse(cow_->fd));
505 
506     auto header = reader.header_v3();
507     ASSERT_EQ(header.compression_algorithm, kCowCompressGz);
508 
509     auto iter = reader.GetOpIter();
510     ASSERT_NE(iter, nullptr);
511     ASSERT_FALSE(iter->AtEnd());
512     auto op = iter->Get();
513 
514     std::string sink(data.size(), '\0');
515 
516     ASSERT_EQ(op->type(), kCowReplaceOp);
517     ASSERT_EQ(op->data_length, 56);  // compressed!
518     ASSERT_EQ(op->new_block, 50);
519     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
520     ASSERT_EQ(sink, data);
521 
522     iter->Next();
523     ASSERT_TRUE(iter->AtEnd());
524 }
525 
TEST_F(CowTestV3,ResumePointTest)526 TEST_F(CowTestV3, ResumePointTest) {
527     CowOptions options;
528     options.op_count_max = 100;
529     auto writer = CreateCowWriter(3, options, GetCowFd());
530 
531     ASSERT_TRUE(writer->AddZeroBlocks(0, 15));
532     ASSERT_TRUE(writer->AddLabel(0));
533     ASSERT_TRUE(writer->AddZeroBlocks(15, 15));
534     ASSERT_TRUE(writer->Finalize());
535 
536     CowReader reader;
537     ASSERT_TRUE(reader.Parse(cow_->fd));
538 
539     auto header = reader.header_v3();
540     ASSERT_EQ(header.op_count, 30);
541 
542     CowWriterV3 second_writer(options, GetCowFd());
543     ASSERT_TRUE(second_writer.Initialize(0));
544     ASSERT_TRUE(second_writer.Finalize());
545 
546     ASSERT_TRUE(reader.Parse(cow_->fd));
547     header = reader.header_v3();
548     ASSERT_EQ(header.op_count, 15);
549 }
550 
TEST_F(CowTestV3,BufferMetadataSyncTest)551 TEST_F(CowTestV3, BufferMetadataSyncTest) {
552     CowOptions options;
553     options.op_count_max = 100;
554     auto writer = CreateCowWriter(3, options, GetCowFd());
555     /*
556     Header metadafields
557     sequence_data_count = 0;
558     resume_point_count = 0;
559     resume_point_max = 4;
560     */
561     ASSERT_TRUE(writer->Finalize());
562 
563     CowReader reader;
564     ASSERT_TRUE(reader.Parse(cow_->fd));
565 
566     auto header = reader.header_v3();
567     ASSERT_EQ(header.sequence_data_count, static_cast<uint64_t>(0));
568     ASSERT_EQ(header.resume_point_count, 0);
569     ASSERT_EQ(header.resume_point_max, 4);
570 
571     writer->AddLabel(0);
572     ASSERT_TRUE(reader.Parse(cow_->fd));
573     header = reader.header_v3();
574     ASSERT_EQ(header.sequence_data_count, static_cast<uint64_t>(0));
575     ASSERT_EQ(header.resume_point_count, 1);
576     ASSERT_EQ(header.resume_point_max, 4);
577 
578     ASSERT_TRUE(reader.Parse(cow_->fd));
579     header = reader.header_v3();
580 
581     /*
582     Header metadafields
583     sequence_data_count = 1;
584     resume_point_count = 0;
585     resume_point_max = 4;
586     */
587 }
588 
TEST_F(CowTestV3,SequenceTest)589 TEST_F(CowTestV3, SequenceTest) {
590     CowOptions options;
591     constexpr int seq_len = std::numeric_limits<uint16_t>::max() / sizeof(uint32_t) + 1;
592     options.op_count_max = seq_len;
593     auto writer = CreateCowWriter(3, options, GetCowFd());
594     // sequence data. This just an arbitrary set of integers that specify the merge order. The
595     // actual calculation is done by update_engine and passed to writer. All we care about here is
596     // writing that data correctly
597     uint32_t sequence[seq_len];
598     for (int i = 0; i < seq_len; i++) {
599         sequence[i] = i + 1;
600     }
601 
602     ASSERT_TRUE(writer->AddSequenceData(seq_len, sequence));
603     ASSERT_TRUE(writer->AddZeroBlocks(1, seq_len - 1));
604     std::vector<uint8_t> data(writer->GetBlockSize());
605     for (size_t i = 0; i < data.size(); i++) {
606         data[i] = static_cast<uint8_t>(i & 0xFF);
607     }
608     ASSERT_TRUE(writer->AddRawBlocks(seq_len, data.data(), data.size()));
609     ASSERT_TRUE(writer->Finalize());
610 
611     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
612 
613     CowReader reader;
614     ASSERT_TRUE(reader.Parse(cow_->fd));
615     auto iter = reader.GetRevMergeOpIter();
616 
617     for (int i = 0; i < seq_len; i++) {
618         ASSERT_TRUE(!iter->AtEnd());
619         const auto& op = iter->Get();
620 
621         ASSERT_EQ(op->new_block, seq_len - i);
622         if (op->new_block == seq_len) {
623             std::vector<uint8_t> read_back(writer->GetBlockSize());
624             ASSERT_EQ(reader.ReadData(op, read_back.data(), read_back.size()),
625                       static_cast<ssize_t>(read_back.size()));
626             ASSERT_EQ(read_back, data);
627         }
628 
629         iter->Next();
630     }
631     ASSERT_TRUE(iter->AtEnd());
632 }
633 
TEST_F(CowTestV3,MissingSeqOp)634 TEST_F(CowTestV3, MissingSeqOp) {
635     CowOptions options;
636     options.op_count_max = std::numeric_limits<uint32_t>::max();
637     auto writer = CreateCowWriter(3, options, GetCowFd());
638     const int seq_len = 10;
639     uint32_t sequence[seq_len];
640     for (int i = 0; i < seq_len; i++) {
641         sequence[i] = i + 1;
642     }
643     ASSERT_TRUE(writer->AddSequenceData(seq_len, sequence));
644     ASSERT_TRUE(writer->AddZeroBlocks(1, seq_len - 1));
645     ASSERT_TRUE(writer->Finalize());
646 
647     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
648 
649     CowReader reader;
650     ASSERT_FALSE(reader.Parse(cow_->fd));
651 }
652 
TEST_F(CowTestV3,ResumeSeqOp)653 TEST_F(CowTestV3, ResumeSeqOp) {
654     CowOptions options;
655     options.op_count_max = std::numeric_limits<uint32_t>::max();
656     auto writer = std::make_unique<CowWriterV3>(options, GetCowFd());
657     const int seq_len = 10;
658     uint32_t sequence[seq_len];
659     for (int i = 0; i < seq_len; i++) {
660         sequence[i] = i + 1;
661     }
662     ASSERT_TRUE(writer->Initialize());
663 
664     ASSERT_TRUE(writer->AddSequenceData(seq_len, sequence));
665     ASSERT_TRUE(writer->AddZeroBlocks(1, seq_len / 2));
666     ASSERT_TRUE(writer->AddLabel(1));
667     ASSERT_TRUE(writer->AddZeroBlocks(1 + seq_len / 2, 1));
668 
669     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
670     auto reader = std::make_unique<CowReader>();
671     ASSERT_TRUE(reader->Parse(cow_->fd, 1));
672     auto itr = reader->GetRevMergeOpIter();
673     ASSERT_TRUE(itr->AtEnd());
674 
675     writer = std::make_unique<CowWriterV3>(options, GetCowFd());
676     ASSERT_TRUE(writer->Initialize({1}));
677     ASSERT_TRUE(writer->AddZeroBlocks(1 + seq_len / 2, seq_len / 2));
678     ASSERT_TRUE(writer->Finalize());
679 
680     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
681 
682     reader = std::make_unique<CowReader>();
683     ASSERT_TRUE(reader->Parse(cow_->fd));
684 
685     auto iter = reader->GetRevMergeOpIter();
686 
687     uint64_t expected_block = 10;
688     while (!iter->AtEnd() && expected_block > 0) {
689         ASSERT_FALSE(iter->AtEnd());
690         const auto& op = iter->Get();
691 
692         ASSERT_EQ(op->new_block, expected_block);
693 
694         iter->Next();
695         expected_block--;
696     }
697     ASSERT_EQ(expected_block, 0);
698     ASSERT_TRUE(iter->AtEnd());
699 }
700 
TEST_F(CowTestV3,SetSourceManyTimes)701 TEST_F(CowTestV3, SetSourceManyTimes) {
702     CowOperationV3 op{};
703     op.set_source(1);
704     ASSERT_EQ(op.source(), 1);
705     op.set_source(2);
706     ASSERT_EQ(op.source(), 2);
707     op.set_source(4);
708     ASSERT_EQ(op.source(), 4);
709     op.set_source(8);
710     ASSERT_EQ(op.source(), 8);
711 }
712 
TEST_F(CowTestV3,SetTypeManyTimes)713 TEST_F(CowTestV3, SetTypeManyTimes) {
714     CowOperationV3 op{};
715     op.set_type(kCowCopyOp);
716     ASSERT_EQ(op.type(), kCowCopyOp);
717     op.set_type(kCowReplaceOp);
718     ASSERT_EQ(op.type(), kCowReplaceOp);
719     op.set_type(kCowZeroOp);
720     ASSERT_EQ(op.type(), kCowZeroOp);
721     op.set_type(kCowXorOp);
722     ASSERT_EQ(op.type(), kCowXorOp);
723 }
724 
TEST_F(CowTestV3,SetTypeSourceInverleave)725 TEST_F(CowTestV3, SetTypeSourceInverleave) {
726     CowOperationV3 op{};
727     op.set_type(kCowCopyOp);
728     ASSERT_EQ(op.type(), kCowCopyOp);
729     op.set_source(0x010203040506);
730     ASSERT_EQ(op.source(), 0x010203040506);
731     ASSERT_EQ(op.type(), kCowCopyOp);
732     op.set_type(kCowReplaceOp);
733     ASSERT_EQ(op.source(), 0x010203040506);
734     ASSERT_EQ(op.type(), kCowReplaceOp);
735 }
736 
TEST_F(CowTestV3,CowSizeEstimate)737 TEST_F(CowTestV3, CowSizeEstimate) {
738     CowOptions options{};
739     options.compression = "none";
740     auto estimator = android::snapshot::CreateCowEstimator(3, options);
741     ASSERT_TRUE(estimator->AddZeroBlocks(0, 1024 * 1024));
742     const auto cow_size = estimator->GetCowSizeInfo().cow_size;
743     options.op_count_max = 1024 * 1024;
744     options.max_blocks = 1024 * 1024;
745     CowWriterV3 writer(options, GetCowFd());
746     ASSERT_TRUE(writer.Initialize());
747     ASSERT_TRUE(writer.AddZeroBlocks(0, 1024 * 1024));
748 
749     ASSERT_LE(writer.GetCowSizeInfo().cow_size, cow_size);
750 }
751 
TEST_F(CowTestV3,CopyOpMany)752 TEST_F(CowTestV3, CopyOpMany) {
753     CowOptions options;
754     options.op_count_max = 100;
755     CowWriterV3 writer(options, GetCowFd());
756     writer.Initialize();
757     ASSERT_TRUE(writer.AddCopy(100, 50, 50));
758     ASSERT_TRUE(writer.AddCopy(150, 100, 50));
759     ASSERT_TRUE(writer.Finalize());
760     CowReader reader;
761     ASSERT_TRUE(reader.Parse(GetCowFd()));
762     auto it = reader.GetOpIter();
763     for (size_t i = 0; i < 100; i++) {
764         ASSERT_FALSE(it->AtEnd()) << " op iterator ended at " << i;
765         const auto op = *it->Get();
766         ASSERT_EQ(op.type(), kCowCopyOp);
767         ASSERT_EQ(op.new_block, 100 + i);
768         it->Next();
769     }
770 }
771 
TEST_F(CowTestV3,CheckOpCount)772 TEST_F(CowTestV3, CheckOpCount) {
773     CowOptions options;
774     options.op_count_max = 20;
775     options.batch_write = true;
776     options.cluster_ops = 200;
777     auto writer = CreateCowWriter(3, options, GetCowFd());
778     ASSERT_TRUE(writer->AddZeroBlocks(0, 19));
779     ASSERT_FALSE(writer->AddZeroBlocks(0, 19));
780 }
781 
782 struct TestParam {
783     std::string compression;
784     int block_size;
785     int num_threads;
786     size_t cluster_ops;
787 };
788 
789 class VariableBlockTest : public ::testing::TestWithParam<TestParam> {
790   protected:
SetUp()791     virtual void SetUp() override {
792         cow_ = std::make_unique<TemporaryFile>();
793         ASSERT_GE(cow_->fd, 0) << strerror(errno);
794     }
795 
TearDown()796     virtual void TearDown() override { cow_ = nullptr; }
797 
GetCowFd()798     unique_fd GetCowFd() { return unique_fd{dup(cow_->fd)}; }
799 
800     std::unique_ptr<TemporaryFile> cow_;
801 };
802 
803 // Helper to check read sizes.
ReadBlockData(CowReader & reader,const CowOperation * op,void * buffer,size_t size)804 static inline void ReadBlockData(CowReader& reader, const CowOperation* op, void* buffer,
805                                  size_t size) {
806     size_t block_size = CowOpCompressionSize(op, 4096);
807     std::string data(block_size, '\0');
808     size_t value = reader.ReadData(op, data.data(), block_size);
809     ASSERT_TRUE(value == block_size);
810     std::memcpy(buffer, data.data(), size);
811 }
812 
TEST_P(VariableBlockTest,VariableBlockCompressionTest)813 TEST_P(VariableBlockTest, VariableBlockCompressionTest) {
814     const TestParam params = GetParam();
815 
816     CowOptions options;
817     options.op_count_max = 100000;
818     options.compression = params.compression;
819     options.num_compress_threads = params.num_threads;
820     options.batch_write = true;
821     options.compression_factor = params.block_size;
822     options.cluster_ops = params.cluster_ops;
823 
824     CowWriterV3 writer(options, GetCowFd());
825 
826     ASSERT_TRUE(writer.Initialize());
827 
828     std::string xor_data = "This is test data-1. Testing xor";
829     xor_data.resize(options.block_size, '\0');
830     ASSERT_TRUE(writer.AddXorBlocks(50, xor_data.data(), xor_data.size(), 24, 10));
831 
832     // Large number of blocks
833     std::string data = "This is test data-2. Testing replace ops";
834     data.resize(options.block_size * 2048, '\0');
835     ASSERT_TRUE(writer.AddRawBlocks(100, data.data(), data.size()));
836 
837     std::string data2 = "This is test data-3. Testing replace ops";
838     data2.resize(options.block_size * 259, '\0');
839     ASSERT_TRUE(writer.AddRawBlocks(6000, data2.data(), data2.size()));
840 
841     // Test data size is smaller than block size
842 
843     // 4k block
844     std::string data3 = "This is test data-4. Testing replace ops";
845     data3.resize(options.block_size, '\0');
846     ASSERT_TRUE(writer.AddRawBlocks(9000, data3.data(), data3.size()));
847 
848     // 8k block
849     std::string data4;
850     data4.resize(options.block_size * 2, '\0');
851     for (size_t i = 0; i < data4.size(); i++) {
852         data4[i] = static_cast<char>('A' + i / options.block_size);
853     }
854     ASSERT_TRUE(writer.AddRawBlocks(10000, data4.data(), data4.size()));
855 
856     // 16k block
857     std::string data5;
858     data.resize(options.block_size * 4, '\0');
859     for (int i = 0; i < data5.size(); i++) {
860         data5[i] = static_cast<char>('C' + i / options.block_size);
861     }
862     ASSERT_TRUE(writer.AddRawBlocks(11000, data5.data(), data5.size()));
863 
864     // 64k Random buffer which cannot be compressed
865     unique_fd rnd_fd(open("/dev/random", O_RDONLY));
866     ASSERT_GE(rnd_fd, 0);
867     std::string random_buffer;
868     random_buffer.resize(65536, '\0');
869     ASSERT_EQ(android::base::ReadFullyAtOffset(rnd_fd, random_buffer.data(), 65536, 0), true);
870     ASSERT_TRUE(writer.AddRawBlocks(12000, random_buffer.data(), 65536));
871 
872     ASSERT_TRUE(writer.Finalize());
873 
874     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
875 
876     CowReader reader;
877     ASSERT_TRUE(reader.Parse(cow_->fd));
878 
879     auto iter = reader.GetOpIter();
880     ASSERT_NE(iter, nullptr);
881 
882     while (!iter->AtEnd()) {
883         auto op = iter->Get();
884 
885         if (op->type() == kCowXorOp) {
886             std::string sink(xor_data.size(), '\0');
887             ASSERT_EQ(op->new_block, 50);
888             ASSERT_EQ(op->source(), 98314);  // 4096 * 24 + 10
889             ReadBlockData(reader, op, sink.data(), sink.size());
890             ASSERT_EQ(sink, xor_data);
891         }
892         if (op->type() == kCowReplaceOp) {
893             if (op->new_block == 100) {
894                 data.resize(options.block_size);
895                 std::string sink(data.size(), '\0');
896                 ReadBlockData(reader, op, sink.data(), sink.size());
897                 ASSERT_EQ(sink.size(), data.size());
898                 ASSERT_EQ(sink, data);
899             }
900             if (op->new_block == 6000) {
901                 data2.resize(options.block_size);
902                 std::string sink(data2.size(), '\0');
903                 ReadBlockData(reader, op, sink.data(), sink.size());
904                 ASSERT_EQ(sink, data2);
905             }
906             if (op->new_block == 9000) {
907                 std::string sink(data3.size(), '\0');
908                 ReadBlockData(reader, op, sink.data(), sink.size());
909                 ASSERT_EQ(sink, data3);
910             }
911             if (op->new_block == 10000) {
912                 data4.resize(options.block_size);
913                 std::string sink(options.block_size, '\0');
914                 ReadBlockData(reader, op, sink.data(), sink.size());
915                 ASSERT_EQ(sink, data4);
916             }
917             if (op->new_block == 11000) {
918                 data5.resize(options.block_size);
919                 std::string sink(options.block_size, '\0');
920                 ReadBlockData(reader, op, sink.data(), sink.size());
921                 ASSERT_EQ(sink, data5);
922             }
923             if (op->new_block == 12000) {
924                 random_buffer.resize(options.block_size);
925                 std::string sink(options.block_size, '\0');
926                 ReadBlockData(reader, op, sink.data(), sink.size());
927                 ASSERT_EQ(sink, random_buffer);
928             }
929         }
930 
931         iter->Next();
932     }
933 }
934 
GetTestConfigs()935 std::vector<TestParam> GetTestConfigs() {
936     std::vector<TestParam> testParams;
937 
938     std::vector<int> block_sizes = {4_KiB, 8_KiB, 16_KiB, 32_KiB, 64_KiB, 128_KiB, 256_KiB};
939     std::vector<std::string> compression_algo = {"none", "lz4", "zstd", "gz"};
940     std::vector<int> threads = {1, 2};
941     // This will also test batch size
942     std::vector<size_t> cluster_ops = {1, 256};
943 
944     // This should test 112 combination
945     for (auto block : block_sizes) {
946         for (auto compression : compression_algo) {
947             for (auto thread : threads) {
948                 for (auto cluster : cluster_ops) {
949                     TestParam param;
950                     param.block_size = block;
951                     param.compression = compression;
952                     param.num_threads = thread;
953                     param.cluster_ops = cluster;
954                     testParams.push_back(std::move(param));
955                 }
956             }
957         }
958     }
959 
960     return testParams;
961 }
962 
963 INSTANTIATE_TEST_SUITE_P(CompressorsWithVariableBlocks, VariableBlockTest,
964                          ::testing::ValuesIn(GetTestConfigs()));
965 
966 }  // namespace snapshot
967 }  // namespace android
968