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 "update_engine/payload_generator/cow_size_estimator.h"
18 
19 #include <algorithm>
20 #include <string>
21 #include <utility>
22 #include <vector>
23 
24 #include <android-base/unique_fd.h>
25 #include <libsnapshot/cow_writer.h>
26 #include <libsnapshot/cow_format.h>
27 
28 #include "update_engine/payload_consumer/block_extent_writer.h"
29 #include "update_engine/payload_consumer/snapshot_extent_writer.h"
30 #include "update_engine/payload_consumer/xor_extent_writer.h"
31 #include "update_engine/common/utils.h"
32 #include "update_engine/payload_consumer/vabc_partition_writer.h"
33 #include "update_engine/payload_generator/extent_ranges.h"
34 #include "update_engine/payload_generator/extent_utils.h"
35 #include "update_engine/update_metadata.pb.h"
36 
37 namespace chromeos_update_engine {
38 using android::snapshot::CreateCowEstimator;
39 using android::snapshot::ICowWriter;
40 // Compute XOR map, a map from dst extent to corresponding merge operation
ComputeXorMap(const google::protobuf::RepeatedPtrField<CowMergeOperation> & merge_ops)41 static ExtentMap<const CowMergeOperation*, ExtentLess> ComputeXorMap(
42     const google::protobuf::RepeatedPtrField<CowMergeOperation>& merge_ops) {
43   ExtentMap<const CowMergeOperation*, ExtentLess> xor_map;
44   for (const auto& merge_op : merge_ops) {
45     if (merge_op.type() == CowMergeOperation::COW_XOR) {
46       xor_map.AddExtent(merge_op.dst_extent(), &merge_op);
47     }
48   }
49   return xor_map;
50 }
51 
CowDryRun(FileDescriptorPtr source_fd,FileDescriptorPtr target_fd,const google::protobuf::RepeatedPtrField<InstallOperation> & operations,const google::protobuf::RepeatedPtrField<CowMergeOperation> & merge_operations,const size_t block_size,android::snapshot::ICowWriter * cow_writer,const size_t new_partition_size,const size_t old_partition_size,const bool xor_enabled)52 bool CowDryRun(
53     FileDescriptorPtr source_fd,
54     FileDescriptorPtr target_fd,
55     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
56     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
57         merge_operations,
58     const size_t block_size,
59     android::snapshot::ICowWriter* cow_writer,
60     const size_t new_partition_size,
61     const size_t old_partition_size,
62     const bool xor_enabled) {
63   CHECK_NE(target_fd, nullptr);
64   CHECK(target_fd->IsOpen());
65   VABCPartitionWriter::WriteMergeSequence(merge_operations, cow_writer);
66   ExtentRanges visited;
67   SnapshotExtentWriter extent_writer(cow_writer);
68   ExtentMap<const CowMergeOperation*, ExtentLess> xor_map =
69       ComputeXorMap(merge_operations);
70   ExtentRanges copy_blocks;
71   for (const auto& cow_op : merge_operations) {
72     if (cow_op.type() != CowMergeOperation::COW_COPY) {
73       continue;
74     }
75     copy_blocks.AddExtent(cow_op.dst_extent());
76   }
77   for (const auto& op : operations) {
78     switch (op.type()) {
79       case InstallOperation::SOURCE_BSDIFF:
80       case InstallOperation::BROTLI_BSDIFF:
81       case InstallOperation::PUFFDIFF:
82       case InstallOperation::ZUCCHINI:
83       case InstallOperation::LZ4DIFF_PUFFDIFF:
84       case InstallOperation::LZ4DIFF_BSDIFF: {
85         if (xor_enabled) {
86           std::unique_ptr<XORExtentWriter> writer =
87               std::make_unique<XORExtentWriter>(
88                   op, source_fd, cow_writer, xor_map, old_partition_size);
89           TEST_AND_RETURN_FALSE(writer->Init(op.dst_extents(), block_size));
90           for (const auto& ext : op.dst_extents()) {
91             visited.AddExtent(ext);
92             ssize_t bytes_read = 0;
93             std::vector<unsigned char> new_data(ext.num_blocks() * block_size);
94             if (!utils::PReadAll(target_fd,
95                                  new_data.data(),
96                                  new_data.size(),
97                                  ext.start_block() * block_size,
98                                  &bytes_read)) {
99               PLOG(ERROR) << "Failed to read target data at " << ext;
100               return false;
101             }
102             if (!writer->Write(new_data.data(),
103                                ext.num_blocks() * block_size)) {
104               LOG(ERROR) << "Failed to write XOR operation for extent: "
105                          << ext.start_block();
106               return false;
107             }
108           }
109           cow_writer->AddLabel(0);
110           break;
111         }
112         [[fallthrough]];
113       }
114       case InstallOperation::REPLACE:
115       case InstallOperation::REPLACE_BZ:
116       case InstallOperation::REPLACE_XZ: {
117         TEST_AND_RETURN_FALSE(extent_writer.Init(op.dst_extents(), block_size));
118         for (const auto& ext : op.dst_extents()) {
119           visited.AddExtent(ext);
120           std::vector<unsigned char> data(ext.num_blocks() * block_size);
121           ssize_t bytes_read = 0;
122           if (!utils::PReadAll(target_fd,
123                                data.data(),
124                                data.size(),
125                                ext.start_block() * block_size,
126                                &bytes_read)) {
127             PLOG(ERROR) << "Failed to read new block data at " << ext;
128             return false;
129           }
130           if (!extent_writer.Write(data.data(), data.size())) {
131             LOG(ERROR) << "Failed to write REPLACE op for extent: "
132                        << ext.start_block();
133             return false;
134           }
135         }
136         cow_writer->AddLabel(0);
137         break;
138       }
139       case InstallOperation::ZERO:
140       case InstallOperation::DISCARD: {
141         for (const auto& ext : op.dst_extents()) {
142           visited.AddExtent(ext);
143           cow_writer->AddZeroBlocks(ext.start_block(), ext.num_blocks());
144         }
145         cow_writer->AddLabel(0);
146         break;
147       }
148       case InstallOperation::SOURCE_COPY: {
149         for (const auto& ext : op.dst_extents()) {
150           visited.AddExtent(ext);
151         }
152         if (!VABCPartitionWriter::ProcessSourceCopyOperation(
153                 op, block_size, copy_blocks, source_fd, cow_writer, true)) {
154           LOG(ERROR) << "Failed to process source copy operation: " << op.type()
155                      << "\nsource extents: " << op.src_extents()
156                      << "\ndestination extents: " << op.dst_extents();
157           return false;
158         }
159         break;
160       }
161       default:
162         LOG(ERROR) << "unknown op: " << op.type();
163     }
164   }
165 
166   const size_t last_block = new_partition_size / block_size;
167   const auto unvisited_extents =
168       FilterExtentRanges({ExtentForRange(0, last_block)}, visited);
169   for (const auto& ext : unvisited_extents) {
170     std::vector<unsigned char> data(ext.num_blocks() * block_size);
171     ssize_t bytes_read = 0;
172     if (!utils::PReadAll(target_fd,
173                          data.data(),
174                          data.size(),
175                          ext.start_block() * block_size,
176                          &bytes_read)) {
177       PLOG(ERROR) << "Failed to read new block data at " << ext;
178       return false;
179     }
180     auto to_write = data.size();
181     // FEC data written on device is chunked to 1mb. We want to mirror that here
182     while (to_write) {
183       auto curr_write = std::min(block_size, to_write);
184       cow_writer->AddRawBlocks(
185           ext.start_block() + ((data.size() - to_write) / block_size),
186           data.data() + (data.size() - to_write),
187           curr_write);
188       to_write -= curr_write;
189     }
190     CHECK_EQ(to_write, 0ULL);
191     cow_writer->AddLabel(0);
192   }
193 
194   TEST_AND_RETURN_FALSE(cow_writer->Finalize());
195 
196   return true;
197 }
198 
EstimateCowSizeInfo(FileDescriptorPtr source_fd,FileDescriptorPtr target_fd,const google::protobuf::RepeatedPtrField<InstallOperation> & operations,const google::protobuf::RepeatedPtrField<CowMergeOperation> & merge_operations,const size_t block_size,std::string compression,const size_t new_partition_size,const size_t old_partition_size,const bool xor_enabled,uint32_t cow_version,uint64_t compression_factor)199 android::snapshot::CowSizeInfo EstimateCowSizeInfo(
200     FileDescriptorPtr source_fd,
201     FileDescriptorPtr target_fd,
202     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
203     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
204         merge_operations,
205     const size_t block_size,
206     std::string compression,
207     const size_t new_partition_size,
208     const size_t old_partition_size,
209     const bool xor_enabled,
210     uint32_t cow_version,
211     uint64_t compression_factor) {
212   android::snapshot::CowOptions options{
213       .block_size = static_cast<uint32_t>(block_size),
214       .compression = std::move(compression),
215       .max_blocks = (new_partition_size / block_size),
216       .compression_factor = compression_factor};
217   auto cow_writer = CreateCowEstimator(cow_version, options);
218   CHECK_NE(cow_writer, nullptr) << "Could not create cow estimator";
219   CHECK(CowDryRun(source_fd,
220                   target_fd,
221                   operations,
222                   merge_operations,
223                   block_size,
224                   cow_writer.get(),
225                   new_partition_size,
226                   old_partition_size,
227                   xor_enabled));
228   return cow_writer->GetCowSizeInfo();
229 }
230 
231 }  // namespace chromeos_update_engine
232