ab_generator.cc revision 55c4f9ba7f5c59e3345f2c1869464433ffa8dc2b 
1//
2// Copyright (C) 2015 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/ab_generator.h"
18
19#include <algorithm>
20
21#include <base/strings/stringprintf.h>
22
23#include "update_engine/common/hash_calculator.h"
24#include "update_engine/common/utils.h"
25#include "update_engine/payload_consumer/payload_constants.h"
26#include "update_engine/payload_generator/annotated_operation.h"
27#include "update_engine/payload_generator/bzip.h"
28#include "update_engine/payload_generator/delta_diff_generator.h"
29#include "update_engine/payload_generator/delta_diff_utils.h"
30
31using std::string;
32using std::vector;
33
34namespace chromeos_update_engine {
35
36bool ABGenerator::GenerateOperations(
37    const PayloadGenerationConfig& config,
38    const PartitionConfig& old_part,
39    const PartitionConfig& new_part,
40    BlobFileWriter* blob_file,
41    vector<AnnotatedOperation>* aops) {
42  TEST_AND_RETURN_FALSE(old_part.name == new_part.name);
43
44  ssize_t hard_chunk_blocks = (config.hard_chunk_size == -1 ? -1 :
45                               config.hard_chunk_size / config.block_size);
46  size_t soft_chunk_blocks = config.soft_chunk_size / config.block_size;
47
48  aops->clear();
49  TEST_AND_RETURN_FALSE(diff_utils::DeltaReadPartition(
50      aops,
51      old_part,
52      new_part,
53      hard_chunk_blocks,
54      soft_chunk_blocks,
55      blob_file,
56      config.imgdiff_allowed,
57      true));  // src_ops_allowed
58  LOG(INFO) << "done reading " << new_part.name;
59
60  TEST_AND_RETURN_FALSE(FragmentOperations(aops,
61                                           new_part.path,
62                                           blob_file));
63  SortOperationsByDestination(aops);
64
65  // Use the soft_chunk_size when merging operations to prevent merging all
66  // the operations into a huge one if there's no hard limit.
67  size_t merge_chunk_blocks = soft_chunk_blocks;
68  if (hard_chunk_blocks != -1 &&
69      static_cast<size_t>(hard_chunk_blocks) < soft_chunk_blocks) {
70    merge_chunk_blocks = hard_chunk_blocks;
71  }
72
73  TEST_AND_RETURN_FALSE(MergeOperations(aops,
74                                        merge_chunk_blocks,
75                                        new_part.path,
76                                        blob_file));
77
78  if (config.minor_version >= kOpSrcHashMinorPayloadVersion)
79    TEST_AND_RETURN_FALSE(AddSourceHash(aops, old_part.path));
80
81  return true;
82}
83
84void ABGenerator::SortOperationsByDestination(
85    vector<AnnotatedOperation>* aops) {
86  sort(aops->begin(), aops->end(), diff_utils::CompareAopsByDestination);
87}
88
89bool ABGenerator::FragmentOperations(
90    vector<AnnotatedOperation>* aops,
91    const string& target_part_path,
92    BlobFileWriter* blob_file) {
93  vector<AnnotatedOperation> fragmented_aops;
94  for (const AnnotatedOperation& aop : *aops) {
95    if (aop.op.type() == InstallOperation::SOURCE_COPY) {
96      TEST_AND_RETURN_FALSE(SplitSourceCopy(aop, &fragmented_aops));
97    } else if ((aop.op.type() == InstallOperation::REPLACE) ||
98               (aop.op.type() == InstallOperation::REPLACE_BZ)) {
99      TEST_AND_RETURN_FALSE(SplitReplaceOrReplaceBz(aop, &fragmented_aops,
100                                                    target_part_path,
101                                                    blob_file));
102    } else {
103      fragmented_aops.push_back(aop);
104    }
105  }
106  *aops = fragmented_aops;
107  return true;
108}
109
110bool ABGenerator::SplitSourceCopy(
111    const AnnotatedOperation& original_aop,
112    vector<AnnotatedOperation>* result_aops) {
113  InstallOperation original_op = original_aop.op;
114  TEST_AND_RETURN_FALSE(original_op.type() == InstallOperation::SOURCE_COPY);
115  // Keeps track of the index of curr_src_ext.
116  int curr_src_ext_index = 0;
117  Extent curr_src_ext = original_op.src_extents(curr_src_ext_index);
118  for (int i = 0; i < original_op.dst_extents_size(); i++) {
119    Extent dst_ext = original_op.dst_extents(i);
120    // The new operation which will have only one dst extent.
121    InstallOperation new_op;
122    uint64_t blocks_left = dst_ext.num_blocks();
123    while (blocks_left > 0) {
124      if (curr_src_ext.num_blocks() <= blocks_left) {
125        // If the curr_src_ext is smaller than dst_ext, add it.
126        blocks_left -= curr_src_ext.num_blocks();
127        *(new_op.add_src_extents()) = curr_src_ext;
128        if (curr_src_ext_index + 1 < original_op.src_extents().size()) {
129          curr_src_ext = original_op.src_extents(++curr_src_ext_index);
130        } else {
131          break;
132        }
133      } else {
134        // Split src_exts that are bigger than the dst_ext we're dealing with.
135        Extent first_ext;
136        first_ext.set_num_blocks(blocks_left);
137        first_ext.set_start_block(curr_src_ext.start_block());
138        *(new_op.add_src_extents()) = first_ext;
139        // Keep the second half of the split op.
140        curr_src_ext.set_num_blocks(curr_src_ext.num_blocks() - blocks_left);
141        curr_src_ext.set_start_block(curr_src_ext.start_block() + blocks_left);
142        blocks_left -= first_ext.num_blocks();
143      }
144    }
145    // Fix up our new operation and add it to the results.
146    new_op.set_type(InstallOperation::SOURCE_COPY);
147    *(new_op.add_dst_extents()) = dst_ext;
148    new_op.set_src_length(dst_ext.num_blocks() * kBlockSize);
149    new_op.set_dst_length(dst_ext.num_blocks() * kBlockSize);
150
151    AnnotatedOperation new_aop;
152    new_aop.op = new_op;
153    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
154    result_aops->push_back(new_aop);
155  }
156  if (curr_src_ext_index != original_op.src_extents().size() - 1) {
157    LOG(FATAL) << "Incorrectly split SOURCE_COPY operation. Did not use all "
158               << "source extents.";
159  }
160  return true;
161}
162
163bool ABGenerator::SplitReplaceOrReplaceBz(
164    const AnnotatedOperation& original_aop,
165    vector<AnnotatedOperation>* result_aops,
166    const string& target_part_path,
167    BlobFileWriter* blob_file) {
168  InstallOperation original_op = original_aop.op;
169  const bool is_replace = original_op.type() == InstallOperation::REPLACE;
170  TEST_AND_RETURN_FALSE(is_replace ||
171                        original_op.type() == InstallOperation::REPLACE_BZ);
172
173  uint32_t data_offset = original_op.data_offset();
174  for (int i = 0; i < original_op.dst_extents_size(); i++) {
175    Extent dst_ext = original_op.dst_extents(i);
176    // Make a new operation with only one dst extent.
177    InstallOperation new_op;
178    *(new_op.add_dst_extents()) = dst_ext;
179    uint32_t data_size = dst_ext.num_blocks() * kBlockSize;
180    new_op.set_dst_length(data_size);
181    // If this is a REPLACE, attempt to reuse portions of the existing blob.
182    if (is_replace) {
183      new_op.set_type(InstallOperation::REPLACE);
184      new_op.set_data_length(data_size);
185      new_op.set_data_offset(data_offset);
186      data_offset += data_size;
187    }
188
189    AnnotatedOperation new_aop;
190    new_aop.op = new_op;
191    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
192    TEST_AND_RETURN_FALSE(AddDataAndSetType(&new_aop, target_part_path,
193                                            blob_file));
194
195    result_aops->push_back(new_aop);
196  }
197  return true;
198}
199
200bool ABGenerator::MergeOperations(vector<AnnotatedOperation>* aops,
201                                  size_t chunk_blocks,
202                                  const string& target_part_path,
203                                  BlobFileWriter* blob_file) {
204  vector<AnnotatedOperation> new_aops;
205  for (const AnnotatedOperation& curr_aop : *aops) {
206    if (new_aops.empty()) {
207      new_aops.push_back(curr_aop);
208      continue;
209    }
210    AnnotatedOperation& last_aop = new_aops.back();
211
212    if (last_aop.op.dst_extents_size() <= 0 ||
213        curr_aop.op.dst_extents_size() <= 0) {
214      new_aops.push_back(curr_aop);
215      continue;
216    }
217    uint32_t last_dst_idx = last_aop.op.dst_extents_size() - 1;
218    uint32_t last_end_block =
219        last_aop.op.dst_extents(last_dst_idx).start_block() +
220        last_aop.op.dst_extents(last_dst_idx).num_blocks();
221    uint32_t curr_start_block = curr_aop.op.dst_extents(0).start_block();
222    uint32_t combined_block_count =
223        last_aop.op.dst_extents(last_dst_idx).num_blocks() +
224        curr_aop.op.dst_extents(0).num_blocks();
225    bool good_op_type = curr_aop.op.type() == InstallOperation::SOURCE_COPY ||
226                        curr_aop.op.type() == InstallOperation::REPLACE ||
227                        curr_aop.op.type() == InstallOperation::REPLACE_BZ;
228    if (good_op_type &&
229        last_aop.op.type() == curr_aop.op.type() &&
230        last_end_block == curr_start_block &&
231        combined_block_count <= chunk_blocks) {
232      // If the operations have the same type (which is a type that we can
233      // merge), are contiguous, are fragmented to have one destination extent,
234      // and their combined block count would be less than chunk size, merge
235      // them.
236      last_aop.name = base::StringPrintf("%s,%s",
237                                         last_aop.name.c_str(),
238                                         curr_aop.name.c_str());
239
240      ExtendExtents(last_aop.op.mutable_src_extents(),
241                    curr_aop.op.src_extents());
242      if (curr_aop.op.src_length() > 0)
243        last_aop.op.set_src_length(last_aop.op.src_length() +
244                                   curr_aop.op.src_length());
245      ExtendExtents(last_aop.op.mutable_dst_extents(),
246                    curr_aop.op.dst_extents());
247      if (curr_aop.op.dst_length() > 0)
248        last_aop.op.set_dst_length(last_aop.op.dst_length() +
249                                   curr_aop.op.dst_length());
250      // Set the data length to zero so we know to add the blob later.
251      if (curr_aop.op.type() == InstallOperation::REPLACE ||
252          curr_aop.op.type() == InstallOperation::REPLACE_BZ) {
253        last_aop.op.set_data_length(0);
254      }
255    } else {
256      // Otherwise just include the extent as is.
257      new_aops.push_back(curr_aop);
258    }
259  }
260
261  // Set the blobs for REPLACE/REPLACE_BZ operations that have been merged.
262  for (AnnotatedOperation& curr_aop : new_aops) {
263    if (curr_aop.op.data_length() == 0 &&
264        (curr_aop.op.type() == InstallOperation::REPLACE ||
265         curr_aop.op.type() == InstallOperation::REPLACE_BZ)) {
266      TEST_AND_RETURN_FALSE(AddDataAndSetType(&curr_aop, target_part_path,
267                                              blob_file));
268    }
269  }
270
271  *aops = new_aops;
272  return true;
273}
274
275bool ABGenerator::AddDataAndSetType(AnnotatedOperation* aop,
276                                    const string& target_part_path,
277                                    BlobFileWriter* blob_file) {
278  TEST_AND_RETURN_FALSE(aop->op.type() == InstallOperation::REPLACE ||
279                        aop->op.type() == InstallOperation::REPLACE_BZ);
280
281  brillo::Blob data(aop->op.dst_length());
282  vector<Extent> dst_extents;
283  ExtentsToVector(aop->op.dst_extents(), &dst_extents);
284  TEST_AND_RETURN_FALSE(utils::ReadExtents(target_part_path,
285                                           dst_extents,
286                                           &data,
287                                           data.size(),
288                                           kBlockSize));
289
290  brillo::Blob data_bz;
291  TEST_AND_RETURN_FALSE(BzipCompress(data, &data_bz));
292  CHECK(!data_bz.empty());
293
294  brillo::Blob* data_p = nullptr;
295  InstallOperation_Type new_op_type;
296  if (data_bz.size() < data.size()) {
297    new_op_type = InstallOperation::REPLACE_BZ;
298    data_p = &data_bz;
299  } else {
300    new_op_type = InstallOperation::REPLACE;
301    data_p = &data;
302  }
303
304  // If the operation doesn't point to a data blob, then we add it.
305  if (aop->op.type() != new_op_type ||
306      aop->op.data_length() != data_p->size()) {
307    aop->op.set_type(new_op_type);
308    aop->SetOperationBlob(data_p, blob_file);
309  }
310
311  return true;
312}
313
314bool ABGenerator::AddSourceHash(vector<AnnotatedOperation>* aops,
315                                const string& source_part_path) {
316  for (AnnotatedOperation& aop : *aops) {
317    if (aop.op.src_extents_size() == 0)
318      continue;
319
320    vector<Extent> src_extents;
321    ExtentsToVector(aop.op.src_extents(), &src_extents);
322    brillo::Blob src_data, src_hash;
323    uint64_t src_length =
324        aop.op.has_src_length()
325            ? aop.op.src_length()
326            : BlocksInExtents(aop.op.src_extents()) * kBlockSize;
327    TEST_AND_RETURN_FALSE(utils::ReadExtents(
328        source_part_path, src_extents, &src_data, src_length, kBlockSize));
329    TEST_AND_RETURN_FALSE(HashCalculator::RawHashOfData(src_data, &src_hash));
330    aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
331  }
332  return true;
333}
334
335}  // namespace chromeos_update_engine
336