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