ab_generator.cc revision 3b2c7d0e6d859e6fc75c82b3417f87cf5968a49d 
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    int data_file_fd,
26    off_t* data_file_size,
27    vector<AnnotatedOperation>* rootfs_ops,
28    vector<AnnotatedOperation>* kernel_ops) {
29
30  off_t chunk_blocks = (config.chunk_size == -1 ? -1 :
31                        config.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      chunk_blocks,
39      data_file_fd,
40      data_file_size,
41      false,  // skip_block_0
42      true));  // src_ops_allowed
43  LOG(INFO) << "done reading normal files";
44
45  // Read kernel partition
46  TEST_AND_RETURN_FALSE(diff_utils::DeltaReadPartition(
47      kernel_ops,
48      config.source.kernel,
49      config.target.kernel,
50      chunk_blocks,
51      data_file_fd,
52      data_file_size,
53      false,  // skip_block_0
54      true));  // src_ops_allowed
55  LOG(INFO) << "done reading kernel";
56
57  TEST_AND_RETURN_FALSE(FragmentOperations(rootfs_ops,
58                                           config.target.rootfs.path,
59                                           data_file_fd,
60                                           data_file_size));
61  TEST_AND_RETURN_FALSE(FragmentOperations(kernel_ops,
62                                           config.target.kernel.path,
63                                           data_file_fd,
64                                           data_file_size));
65  SortOperationsByDestination(rootfs_ops);
66  SortOperationsByDestination(kernel_ops);
67  // TODO(alliewood): Change merge operations to use config.chunk_size once
68  // specifying chunk_size on the command line works. crbug/485397.
69  TEST_AND_RETURN_FALSE(MergeOperations(rootfs_ops,
70                                        kDefaultChunkSize,
71                                        config.target.rootfs.path,
72                                        data_file_fd,
73                                        data_file_size));
74  TEST_AND_RETURN_FALSE(MergeOperations(kernel_ops,
75                                        kDefaultChunkSize,
76                                        config.target.kernel.path,
77                                        data_file_fd,
78                                        data_file_size));
79  return true;
80}
81
82void ABGenerator::SortOperationsByDestination(
83    vector<AnnotatedOperation>* aops) {
84  sort(aops->begin(), aops->end(), diff_utils::CompareAopsByDestination);
85}
86
87bool ABGenerator::FragmentOperations(
88    vector<AnnotatedOperation>* aops,
89    const string& target_part_path,
90    int data_fd,
91    off_t* data_file_size) {
92  vector<AnnotatedOperation> fragmented_aops;
93  for (const AnnotatedOperation& aop : *aops) {
94    if (aop.op.type() ==
95        DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY) {
96      TEST_AND_RETURN_FALSE(SplitSourceCopy(aop, &fragmented_aops));
97    } else if ((aop.op.type() ==
98                DeltaArchiveManifest_InstallOperation_Type_REPLACE) ||
99               (aop.op.type() ==
100                DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ)) {
101      TEST_AND_RETURN_FALSE(SplitReplaceOrReplaceBz(aop, &fragmented_aops,
102                                                    target_part_path, data_fd,
103                                                    data_file_size));
104    } else {
105      fragmented_aops.push_back(aop);
106    }
107  }
108  *aops = fragmented_aops;
109  return true;
110}
111
112bool ABGenerator::SplitSourceCopy(
113    const AnnotatedOperation& original_aop,
114    vector<AnnotatedOperation>* result_aops) {
115  DeltaArchiveManifest_InstallOperation original_op = original_aop.op;
116  TEST_AND_RETURN_FALSE(original_op.type() ==
117                        DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY);
118  // Keeps track of the index of curr_src_ext.
119  int curr_src_ext_index = 0;
120  Extent curr_src_ext = original_op.src_extents(curr_src_ext_index);
121  for (int i = 0; i < original_op.dst_extents_size(); i++) {
122    Extent dst_ext = original_op.dst_extents(i);
123    // The new operation which will have only one dst extent.
124    DeltaArchiveManifest_InstallOperation new_op;
125    uint64_t blocks_left = dst_ext.num_blocks();
126    while (blocks_left > 0) {
127      if (curr_src_ext.num_blocks() <= blocks_left) {
128        // If the curr_src_ext is smaller than dst_ext, add it.
129        blocks_left -= curr_src_ext.num_blocks();
130        *(new_op.add_src_extents()) = curr_src_ext;
131        if (curr_src_ext_index + 1 < original_op.src_extents().size()) {
132          curr_src_ext = original_op.src_extents(++curr_src_ext_index);
133        } else {
134          break;
135        }
136      } else {
137        // Split src_exts that are bigger than the dst_ext we're dealing with.
138        Extent first_ext;
139        first_ext.set_num_blocks(blocks_left);
140        first_ext.set_start_block(curr_src_ext.start_block());
141        *(new_op.add_src_extents()) = first_ext;
142        // Keep the second half of the split op.
143        curr_src_ext.set_num_blocks(curr_src_ext.num_blocks() - blocks_left);
144        curr_src_ext.set_start_block(curr_src_ext.start_block() + blocks_left);
145        blocks_left -= first_ext.num_blocks();
146      }
147    }
148    // Fix up our new operation and add it to the results.
149    new_op.set_type(DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY);
150    *(new_op.add_dst_extents()) = dst_ext;
151    new_op.set_src_length(dst_ext.num_blocks() * kBlockSize);
152    new_op.set_dst_length(dst_ext.num_blocks() * kBlockSize);
153
154    AnnotatedOperation new_aop;
155    new_aop.op = new_op;
156    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
157    result_aops->push_back(new_aop);
158  }
159  if (curr_src_ext_index != original_op.src_extents().size() - 1) {
160    LOG(FATAL) << "Incorrectly split SOURCE_COPY operation. Did not use all "
161               << "source extents.";
162  }
163  return true;
164}
165
166bool ABGenerator::SplitReplaceOrReplaceBz(
167    const AnnotatedOperation& original_aop,
168    vector<AnnotatedOperation>* result_aops,
169    const string& target_part_path,
170    int data_fd,
171    off_t* data_file_size) {
172  DeltaArchiveManifest_InstallOperation original_op = original_aop.op;
173  const bool is_replace =
174      original_op.type() == DeltaArchiveManifest_InstallOperation_Type_REPLACE;
175  TEST_AND_RETURN_FALSE(
176      is_replace ||
177      (original_op.type() ==
178       DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ));
179
180  uint32_t data_offset = original_op.data_offset();
181  for (int i = 0; i < original_op.dst_extents_size(); i++) {
182    Extent dst_ext = original_op.dst_extents(i);
183    // Make a new operation with only one dst extent.
184    DeltaArchiveManifest_InstallOperation new_op;
185    *(new_op.add_dst_extents()) = dst_ext;
186    uint32_t data_size = dst_ext.num_blocks() * kBlockSize;
187    new_op.set_dst_length(data_size);
188    // If this is a REPLACE, attempt to reuse portions of the existing blob.
189    if (is_replace) {
190      new_op.set_type(DeltaArchiveManifest_InstallOperation_Type_REPLACE);
191      new_op.set_data_length(data_size);
192      new_op.set_data_offset(data_offset);
193      data_offset += data_size;
194    }
195
196    AnnotatedOperation new_aop;
197    new_aop.op = new_op;
198    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
199    TEST_AND_RETURN_FALSE(AddDataAndSetType(&new_aop, target_part_path, data_fd,
200                                            data_file_size));
201
202    result_aops->push_back(new_aop);
203  }
204  return true;
205}
206
207bool ABGenerator::MergeOperations(vector<AnnotatedOperation>* aops,
208                                  off_t chunk_size,
209                                  const string& target_part_path,
210                                  int data_fd,
211                                  off_t* data_file_size) {
212  vector<AnnotatedOperation> new_aops;
213  for (const AnnotatedOperation& curr_aop : *aops) {
214    if (new_aops.empty()) {
215      new_aops.push_back(curr_aop);
216      continue;
217    }
218    AnnotatedOperation& last_aop = new_aops.back();
219
220    if (last_aop.op.dst_extents_size() <= 0 ||
221        curr_aop.op.dst_extents_size() <= 0) {
222      new_aops.push_back(curr_aop);
223      continue;
224    }
225    uint32_t last_dst_idx = last_aop.op.dst_extents_size() - 1;
226    uint32_t last_end_block =
227        last_aop.op.dst_extents(last_dst_idx).start_block() +
228        last_aop.op.dst_extents(last_dst_idx).num_blocks();
229    uint32_t curr_start_block = curr_aop.op.dst_extents(0).start_block();
230    uint32_t combined_block_count =
231        last_aop.op.dst_extents(last_dst_idx).num_blocks() +
232        curr_aop.op.dst_extents(0).num_blocks();
233    bool good_op_type =
234        curr_aop.op.type() ==
235            DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY ||
236        curr_aop.op.type() ==
237            DeltaArchiveManifest_InstallOperation_Type_REPLACE ||
238        curr_aop.op.type() ==
239            DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ;
240    if (good_op_type &&
241        last_aop.op.type() == curr_aop.op.type() &&
242        last_end_block == curr_start_block &&
243        static_cast<off_t>(combined_block_count * kBlockSize) <= chunk_size) {
244      // If the operations have the same type (which is a type that we can
245      // merge), are contiguous, are fragmented to have one destination extent,
246      // and their combined block count would be less than chunk size, merge
247      // them.
248      last_aop.name = base::StringPrintf("%s,%s",
249                                         last_aop.name.c_str(),
250                                         curr_aop.name.c_str());
251
252      ExtendExtents(last_aop.op.mutable_src_extents(),
253                    curr_aop.op.src_extents());
254      if (curr_aop.op.src_length() > 0)
255        last_aop.op.set_src_length(last_aop.op.src_length() +
256                                   curr_aop.op.src_length());
257      ExtendExtents(last_aop.op.mutable_dst_extents(),
258                    curr_aop.op.dst_extents());
259      if (curr_aop.op.dst_length() > 0)
260        last_aop.op.set_dst_length(last_aop.op.dst_length() +
261                                   curr_aop.op.dst_length());
262      // Set the data length to zero so we know to add the blob later.
263      if (curr_aop.op.type() ==
264          DeltaArchiveManifest_InstallOperation_Type_REPLACE ||
265          curr_aop.op.type() ==
266          DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ) {
267        last_aop.op.set_data_length(0);
268      }
269    } else {
270      // Otherwise just include the extent as is.
271      new_aops.push_back(curr_aop);
272    }
273  }
274
275  // Set the blobs for REPLACE/REPLACE_BZ operations that have been merged.
276  for (AnnotatedOperation& curr_aop : new_aops) {
277    if (curr_aop.op.data_length() == 0 &&
278        (curr_aop.op.type() ==
279            DeltaArchiveManifest_InstallOperation_Type_REPLACE ||
280         curr_aop.op.type() ==
281            DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ)) {
282      TEST_AND_RETURN_FALSE(AddDataAndSetType(&curr_aop, target_part_path,
283                                              data_fd, data_file_size));
284    }
285  }
286
287  *aops = new_aops;
288  return true;
289}
290
291bool ABGenerator::AddDataAndSetType(AnnotatedOperation* aop,
292                                    const string& target_part_path,
293                                    int data_fd,
294                                    off_t* data_file_size) {
295  TEST_AND_RETURN_FALSE(
296      aop->op.type() == DeltaArchiveManifest_InstallOperation_Type_REPLACE ||
297      aop->op.type() == DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ);
298
299  chromeos::Blob data(aop->op.dst_length());
300  vector<Extent> dst_extents;
301  ExtentsToVector(aop->op.dst_extents(), &dst_extents);
302  TEST_AND_RETURN_FALSE(utils::ReadExtents(target_part_path,
303                                           dst_extents,
304                                           &data,
305                                           data.size(),
306                                           kBlockSize));
307
308  chromeos::Blob data_bz;
309  TEST_AND_RETURN_FALSE(BzipCompress(data, &data_bz));
310  CHECK(!data_bz.empty());
311
312  chromeos::Blob* data_p = nullptr;
313  DeltaArchiveManifest_InstallOperation_Type new_op_type;
314  if (data_bz.size() < data.size()) {
315    new_op_type = DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ;
316    data_p = &data_bz;
317  } else {
318    new_op_type = DeltaArchiveManifest_InstallOperation_Type_REPLACE;
319    data_p = &data;
320  }
321
322  // If the operation already points to a data blob, check whether it's
323  // identical to the new one, in which case don't add it.
324  if (aop->op.type() == new_op_type &&
325      aop->op.data_length() == data_p->size()) {
326    chromeos::Blob current_data(data_p->size());
327    ssize_t bytes_read;
328    TEST_AND_RETURN_FALSE(utils::PReadAll(data_fd,
329                                          current_data.data(),
330                                          aop->op.data_length(),
331                                          aop->op.data_offset(),
332                                          &bytes_read));
333    TEST_AND_RETURN_FALSE(bytes_read ==
334                          static_cast<ssize_t>(aop->op.data_length()));
335    if (current_data == *data_p)
336      data_p = nullptr;
337  }
338
339  if (data_p) {
340    aop->op.set_type(new_op_type);
341    aop->SetOperationBlob(data_p, data_fd, data_file_size);
342  }
343
344  return true;
345}
346
347}  // namespace chromeos_update_engine
348