ensemble_apply.cc revision eb525c5499e34cc9c4b825d6d9e75bb07cc06ace 
1// Copyright (c) 2011 The Chromium 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// This file contains the code to apply a Courgette patch.
6
7#include "courgette/ensemble.h"
8
9#include "base/basictypes.h"
10#include "base/file_util.h"
11#include "base/files/memory_mapped_file.h"
12#include "base/logging.h"
13#include "courgette/crc.h"
14#include "courgette/region.h"
15#include "courgette/streams.h"
16#include "courgette/simple_delta.h"
17#include "courgette/patcher_x86_32.h"
18
19namespace courgette {
20
21// EnsemblePatchApplication is all the logic and data required to apply the
22// multi-stage patch.
23class EnsemblePatchApplication {
24 public:
25  EnsemblePatchApplication();
26  ~EnsemblePatchApplication();
27
28  Status ReadHeader(SourceStream* header_stream);
29
30  Status InitBase(const Region& region);
31
32  Status ValidateBase();
33
34  Status ReadInitialParameters(SourceStream* initial_parameters);
35
36  Status PredictTransformParameters(SinkStreamSet* predicted_parameters);
37
38  Status SubpatchTransformParameters(SinkStreamSet* prediction,
39                                     SourceStream* correction,
40                                     SourceStreamSet* corrected_parameters);
41
42  Status TransformUp(SourceStreamSet* parameters,
43                     SinkStreamSet* transformed_elements);
44
45  Status SubpatchTransformedElements(SinkStreamSet* elements,
46                                     SourceStream* correction,
47                                     SourceStreamSet* corrected_elements);
48
49  Status TransformDown(SourceStreamSet* transformed_elements,
50                       SinkStream* basic_elements);
51
52  Status SubpatchFinalOutput(SourceStream* original,
53                             SourceStream* correction,
54                             SinkStream* corrected_ensemble);
55
56 private:
57  Status SubpatchStreamSets(SinkStreamSet* predicted_items,
58                            SourceStream* correction,
59                            SourceStreamSet* corrected_items,
60                            SinkStream* corrected_items_storage);
61
62  Region base_region_;       // Location of in-memory copy of 'old' version.
63
64  uint32 source_checksum_;
65  uint32 target_checksum_;
66  uint32 final_patch_input_size_prediction_;
67
68  std::vector<TransformationPatcher*> patchers_;
69
70  SinkStream corrected_parameters_storage_;
71  SinkStream corrected_elements_storage_;
72
73  DISALLOW_COPY_AND_ASSIGN(EnsemblePatchApplication);
74};
75
76EnsemblePatchApplication::EnsemblePatchApplication()
77    : source_checksum_(0), target_checksum_(0),
78      final_patch_input_size_prediction_(0) {
79}
80
81EnsemblePatchApplication::~EnsemblePatchApplication() {
82  for (size_t i = 0;  i < patchers_.size();  ++i) {
83    delete patchers_[i];
84  }
85}
86
87Status EnsemblePatchApplication::ReadHeader(SourceStream* header_stream) {
88  uint32 magic;
89  if (!header_stream->ReadVarint32(&magic))
90    return C_BAD_ENSEMBLE_MAGIC;
91
92  if (magic != CourgettePatchFile::kMagic)
93    return C_BAD_ENSEMBLE_MAGIC;
94
95  uint32 version;
96  if (!header_stream->ReadVarint32(&version))
97    return C_BAD_ENSEMBLE_VERSION;
98
99  if (version != CourgettePatchFile::kVersion)
100    return C_BAD_ENSEMBLE_VERSION;
101
102  if (!header_stream->ReadVarint32(&source_checksum_))
103    return C_BAD_ENSEMBLE_HEADER;
104
105  if (!header_stream->ReadVarint32(&target_checksum_))
106    return C_BAD_ENSEMBLE_HEADER;
107
108  if (!header_stream->ReadVarint32(&final_patch_input_size_prediction_))
109    return C_BAD_ENSEMBLE_HEADER;
110
111  return C_OK;
112}
113
114Status EnsemblePatchApplication::InitBase(const Region& region) {
115  base_region_.assign(region);
116  return C_OK;
117}
118
119Status EnsemblePatchApplication::ValidateBase() {
120  uint32 checksum = CalculateCrc(base_region_.start(), base_region_.length());
121  if (source_checksum_ != checksum)
122    return C_BAD_ENSEMBLE_CRC;
123
124  return C_OK;
125}
126
127Status EnsemblePatchApplication::ReadInitialParameters(
128    SourceStream* transformation_parameters) {
129  uint32 number_of_transformations = 0;
130  if (!transformation_parameters->ReadVarint32(&number_of_transformations))
131    return C_BAD_ENSEMBLE_HEADER;
132
133  for (size_t i = 0;  i < number_of_transformations;  ++i) {
134    uint32 kind;
135    if (!transformation_parameters->ReadVarint32(&kind))
136      return C_BAD_ENSEMBLE_HEADER;
137
138    TransformationPatcher* patcher = NULL;
139
140    switch (kind)
141    {
142      case EXE_WIN_32_X86:
143        patcher = new PatcherX86_32(base_region_);
144        break;
145      case EXE_ELF_32_X86:
146        patcher = new PatcherX86_32(base_region_);
147        break;
148      case EXE_ELF_32_ARM:
149        patcher = new PatcherX86_32(base_region_);
150        break;
151    }
152
153    if (patcher)
154      patchers_.push_back(patcher);
155    else
156      return C_BAD_ENSEMBLE_HEADER;
157  }
158
159  for (size_t i = 0;  i < patchers_.size();  ++i) {
160    Status status = patchers_[i]->Init(transformation_parameters);
161    if (status != C_OK)
162      return status;
163  }
164
165  // All transformation_parameters should have been consumed by the above loop.
166  if (!transformation_parameters->Empty())
167    return C_BAD_ENSEMBLE_HEADER;
168
169  return C_OK;
170}
171
172Status EnsemblePatchApplication::PredictTransformParameters(
173    SinkStreamSet* all_predicted_parameters) {
174  for (size_t i = 0;  i < patchers_.size();  ++i) {
175    SinkStreamSet single_predicted_parameters;
176    Status status =
177        patchers_[i]->PredictTransformParameters(&single_predicted_parameters);
178    if (status != C_OK)
179      return status;
180    if (!all_predicted_parameters->WriteSet(&single_predicted_parameters))
181      return C_STREAM_ERROR;
182  }
183  return C_OK;
184}
185
186Status EnsemblePatchApplication::SubpatchTransformParameters(
187    SinkStreamSet* predicted_parameters,
188    SourceStream* correction,
189    SourceStreamSet* corrected_parameters) {
190  return SubpatchStreamSets(predicted_parameters,
191                            correction,
192                            corrected_parameters,
193                            &corrected_parameters_storage_);
194}
195
196Status EnsemblePatchApplication::TransformUp(
197    SourceStreamSet* parameters,
198    SinkStreamSet* transformed_elements) {
199  for (size_t i = 0;  i < patchers_.size();  ++i) {
200    SourceStreamSet single_parameters;
201    if (!parameters->ReadSet(&single_parameters))
202      return C_STREAM_ERROR;
203    SinkStreamSet single_transformed_element;
204    Status status = patchers_[i]->Transform(&single_parameters,
205                                            &single_transformed_element);
206    if (status != C_OK)
207      return status;
208    if (!single_parameters.Empty())
209      return C_STREAM_NOT_CONSUMED;
210    if (!transformed_elements->WriteSet(&single_transformed_element))
211      return C_STREAM_ERROR;
212  }
213
214  if (!parameters->Empty())
215    return C_STREAM_NOT_CONSUMED;
216  return C_OK;
217}
218
219Status EnsemblePatchApplication::SubpatchTransformedElements(
220    SinkStreamSet* predicted_elements,
221    SourceStream* correction,
222    SourceStreamSet* corrected_elements) {
223  return SubpatchStreamSets(predicted_elements,
224                            correction,
225                            corrected_elements,
226                            &corrected_elements_storage_);
227}
228
229Status EnsemblePatchApplication::TransformDown(
230    SourceStreamSet* transformed_elements,
231    SinkStream* basic_elements) {
232  // Construct blob of original input followed by reformed elements.
233
234  if (!basic_elements->Reserve(final_patch_input_size_prediction_)) {
235    return C_STREAM_ERROR;
236  }
237
238  // The original input:
239  if (!basic_elements->Write(base_region_.start(), base_region_.length()))
240    return C_STREAM_ERROR;
241
242  for (size_t i = 0;  i < patchers_.size();  ++i) {
243    SourceStreamSet single_corrected_element;
244    if (!transformed_elements->ReadSet(&single_corrected_element))
245      return C_STREAM_ERROR;
246    Status status = patchers_[i]->Reform(&single_corrected_element,
247                                         basic_elements);
248    if (status != C_OK)
249      return status;
250    if (!single_corrected_element.Empty())
251      return C_STREAM_NOT_CONSUMED;
252  }
253
254  if (!transformed_elements->Empty())
255    return C_STREAM_NOT_CONSUMED;
256  // We have totally consumed transformed_elements, so can free the
257  // storage to which it referred.
258  corrected_elements_storage_.Retire();
259
260  return C_OK;
261}
262
263Status EnsemblePatchApplication::SubpatchFinalOutput(
264    SourceStream* original,
265    SourceStream* correction,
266    SinkStream* corrected_ensemble) {
267  Status delta_status = ApplySimpleDelta(original, correction,
268                                         corrected_ensemble);
269  if (delta_status != C_OK)
270    return delta_status;
271
272  if (CalculateCrc(corrected_ensemble->Buffer(),
273                   corrected_ensemble->Length()) != target_checksum_)
274    return C_BAD_ENSEMBLE_CRC;
275
276  return C_OK;
277}
278
279Status EnsemblePatchApplication::SubpatchStreamSets(
280    SinkStreamSet* predicted_items,
281    SourceStream* correction,
282    SourceStreamSet* corrected_items,
283    SinkStream* corrected_items_storage) {
284  SinkStream linearized_predicted_items;
285  if (!predicted_items->CopyTo(&linearized_predicted_items))
286    return C_STREAM_ERROR;
287
288  SourceStream prediction;
289  prediction.Init(linearized_predicted_items);
290
291  Status status = ApplySimpleDelta(&prediction,
292                                   correction,
293                                   corrected_items_storage);
294  if (status != C_OK)
295    return status;
296
297  if (!corrected_items->Init(corrected_items_storage->Buffer(),
298                             corrected_items_storage->Length()))
299    return C_STREAM_ERROR;
300
301  return C_OK;
302}
303
304Status ApplyEnsemblePatch(SourceStream* base,
305                          SourceStream* patch,
306                          SinkStream* output) {
307  Status status;
308  EnsemblePatchApplication patch_process;
309
310  status = patch_process.ReadHeader(patch);
311  if (status != C_OK)
312    return status;
313
314  status = patch_process.InitBase(Region(base->Buffer(), base->Remaining()));
315  if (status != C_OK)
316    return status;
317
318  status = patch_process.ValidateBase();
319  if (status != C_OK)
320    return status;
321
322  // The rest of the patch stream is a StreamSet.
323  SourceStreamSet patch_streams;
324  patch_streams.Init(patch);
325
326  SourceStream* transformation_descriptions     = patch_streams.stream(0);
327  SourceStream* parameter_correction            = patch_streams.stream(1);
328  SourceStream* transformed_elements_correction = patch_streams.stream(2);
329  SourceStream* ensemble_correction             = patch_streams.stream(3);
330
331  status = patch_process.ReadInitialParameters(transformation_descriptions);
332  if (status != C_OK)
333    return status;
334
335  SinkStreamSet predicted_parameters;
336  status = patch_process.PredictTransformParameters(&predicted_parameters);
337  if (status != C_OK)
338    return status;
339
340  SourceStreamSet corrected_parameters;
341  status = patch_process.SubpatchTransformParameters(&predicted_parameters,
342                                                     parameter_correction,
343                                                     &corrected_parameters);
344  if (status != C_OK)
345    return status;
346
347  SinkStreamSet transformed_elements;
348  status = patch_process.TransformUp(&corrected_parameters,
349                                     &transformed_elements);
350  if (status != C_OK)
351    return status;
352
353  SourceStreamSet corrected_transformed_elements;
354  status = patch_process.SubpatchTransformedElements(
355          &transformed_elements,
356          transformed_elements_correction,
357          &corrected_transformed_elements);
358  if (status != C_OK)
359    return status;
360
361  SinkStream original_ensemble_and_corrected_base_elements;
362  status = patch_process.TransformDown(
363      &corrected_transformed_elements,
364      &original_ensemble_and_corrected_base_elements);
365  if (status != C_OK)
366    return status;
367
368  SourceStream final_patch_prediction;
369  final_patch_prediction.Init(original_ensemble_and_corrected_base_elements);
370  status = patch_process.SubpatchFinalOutput(&final_patch_prediction,
371                                             ensemble_correction, output);
372  if (status != C_OK)
373    return status;
374
375  return C_OK;
376}
377
378Status ApplyEnsemblePatch(const base::FilePath::CharType* old_file_name,
379                          const base::FilePath::CharType* patch_file_name,
380                          const base::FilePath::CharType* new_file_name) {
381  // First read enough of the patch file to validate the header is well-formed.
382  // A few varint32 numbers should fit in 100.
383  base::FilePath patch_file_path(patch_file_name);
384  base::MemoryMappedFile patch_file;
385  if (!patch_file.Initialize(patch_file_path))
386    return C_READ_OPEN_ERROR;
387
388  // 'Dry-run' the first step of the patch process to validate format of header.
389  SourceStream patch_header_stream;
390  patch_header_stream.Init(patch_file.data(), patch_file.length());
391  EnsemblePatchApplication patch_process;
392  Status status = patch_process.ReadHeader(&patch_header_stream);
393  if (status != C_OK)
394    return status;
395
396  // Read the old_file.
397  base::FilePath old_file_path(old_file_name);
398  base::MemoryMappedFile old_file;
399  if (!old_file.Initialize(old_file_path))
400    return C_READ_ERROR;
401
402  // Apply patch on streams.
403  SourceStream old_source_stream;
404  SourceStream patch_source_stream;
405  old_source_stream.Init(old_file.data(), old_file.length());
406  patch_source_stream.Init(patch_file.data(), patch_file.length());
407  SinkStream new_sink_stream;
408  status = ApplyEnsemblePatch(&old_source_stream, &patch_source_stream,
409                              &new_sink_stream);
410  if (status != C_OK)
411    return status;
412
413  // Write the patched data to |new_file_name|.
414  base::FilePath new_file_path(new_file_name);
415  int written =
416      file_util::WriteFile(
417          new_file_path,
418          reinterpret_cast<const char*>(new_sink_stream.Buffer()),
419          static_cast<int>(new_sink_stream.Length()));
420  if (written == -1)
421    return C_WRITE_OPEN_ERROR;
422  if (static_cast<size_t>(written) != new_sink_stream.Length())
423    return C_WRITE_ERROR;
424
425  return C_OK;
426}
427
428}  // namespace
429