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