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      case EXE_WIN_32_X64:
152        patcher = new PatcherX86_32(base_region_);
153        break;
154    }
155
156    if (patcher)
157      patchers_.push_back(patcher);
158    else
159      return C_BAD_ENSEMBLE_HEADER;
160  }
161
162  for (size_t i = 0;  i < patchers_.size();  ++i) {
163    Status status = patchers_[i]->Init(transformation_parameters);
164    if (status != C_OK)
165      return status;
166  }
167
168  // All transformation_parameters should have been consumed by the above loop.
169  if (!transformation_parameters->Empty())
170    return C_BAD_ENSEMBLE_HEADER;
171
172  return C_OK;
173}
174
175Status EnsemblePatchApplication::PredictTransformParameters(
176    SinkStreamSet* all_predicted_parameters) {
177  for (size_t i = 0;  i < patchers_.size();  ++i) {
178    SinkStreamSet single_predicted_parameters;
179    Status status =
180        patchers_[i]->PredictTransformParameters(&single_predicted_parameters);
181    if (status != C_OK)
182      return status;
183    if (!all_predicted_parameters->WriteSet(&single_predicted_parameters))
184      return C_STREAM_ERROR;
185  }
186  return C_OK;
187}
188
189Status EnsemblePatchApplication::SubpatchTransformParameters(
190    SinkStreamSet* predicted_parameters,
191    SourceStream* correction,
192    SourceStreamSet* corrected_parameters) {
193  return SubpatchStreamSets(predicted_parameters,
194                            correction,
195                            corrected_parameters,
196                            &corrected_parameters_storage_);
197}
198
199Status EnsemblePatchApplication::TransformUp(
200    SourceStreamSet* parameters,
201    SinkStreamSet* transformed_elements) {
202  for (size_t i = 0;  i < patchers_.size();  ++i) {
203    SourceStreamSet single_parameters;
204    if (!parameters->ReadSet(&single_parameters))
205      return C_STREAM_ERROR;
206    SinkStreamSet single_transformed_element;
207    Status status = patchers_[i]->Transform(&single_parameters,
208                                            &single_transformed_element);
209    if (status != C_OK)
210      return status;
211    if (!single_parameters.Empty())
212      return C_STREAM_NOT_CONSUMED;
213    if (!transformed_elements->WriteSet(&single_transformed_element))
214      return C_STREAM_ERROR;
215  }
216
217  if (!parameters->Empty())
218    return C_STREAM_NOT_CONSUMED;
219  return C_OK;
220}
221
222Status EnsemblePatchApplication::SubpatchTransformedElements(
223    SinkStreamSet* predicted_elements,
224    SourceStream* correction,
225    SourceStreamSet* corrected_elements) {
226  return SubpatchStreamSets(predicted_elements,
227                            correction,
228                            corrected_elements,
229                            &corrected_elements_storage_);
230}
231
232Status EnsemblePatchApplication::TransformDown(
233    SourceStreamSet* transformed_elements,
234    SinkStream* basic_elements) {
235  // Construct blob of original input followed by reformed elements.
236
237  if (!basic_elements->Reserve(final_patch_input_size_prediction_)) {
238    return C_STREAM_ERROR;
239  }
240
241  // The original input:
242  if (!basic_elements->Write(base_region_.start(), base_region_.length()))
243    return C_STREAM_ERROR;
244
245  for (size_t i = 0;  i < patchers_.size();  ++i) {
246    SourceStreamSet single_corrected_element;
247    if (!transformed_elements->ReadSet(&single_corrected_element))
248      return C_STREAM_ERROR;
249    Status status = patchers_[i]->Reform(&single_corrected_element,
250                                         basic_elements);
251    if (status != C_OK)
252      return status;
253    if (!single_corrected_element.Empty())
254      return C_STREAM_NOT_CONSUMED;
255  }
256
257  if (!transformed_elements->Empty())
258    return C_STREAM_NOT_CONSUMED;
259  // We have totally consumed transformed_elements, so can free the
260  // storage to which it referred.
261  corrected_elements_storage_.Retire();
262
263  return C_OK;
264}
265
266Status EnsemblePatchApplication::SubpatchFinalOutput(
267    SourceStream* original,
268    SourceStream* correction,
269    SinkStream* corrected_ensemble) {
270  Status delta_status = ApplySimpleDelta(original, correction,
271                                         corrected_ensemble);
272  if (delta_status != C_OK)
273    return delta_status;
274
275  if (CalculateCrc(corrected_ensemble->Buffer(),
276                   corrected_ensemble->Length()) != target_checksum_)
277    return C_BAD_ENSEMBLE_CRC;
278
279  return C_OK;
280}
281
282Status EnsemblePatchApplication::SubpatchStreamSets(
283    SinkStreamSet* predicted_items,
284    SourceStream* correction,
285    SourceStreamSet* corrected_items,
286    SinkStream* corrected_items_storage) {
287  SinkStream linearized_predicted_items;
288  if (!predicted_items->CopyTo(&linearized_predicted_items))
289    return C_STREAM_ERROR;
290
291  SourceStream prediction;
292  prediction.Init(linearized_predicted_items);
293
294  Status status = ApplySimpleDelta(&prediction,
295                                   correction,
296                                   corrected_items_storage);
297  if (status != C_OK)
298    return status;
299
300  if (!corrected_items->Init(corrected_items_storage->Buffer(),
301                             corrected_items_storage->Length()))
302    return C_STREAM_ERROR;
303
304  return C_OK;
305}
306
307Status ApplyEnsemblePatch(SourceStream* base,
308                          SourceStream* patch,
309                          SinkStream* output) {
310  Status status;
311  EnsemblePatchApplication patch_process;
312
313  status = patch_process.ReadHeader(patch);
314  if (status != C_OK)
315    return status;
316
317  status = patch_process.InitBase(Region(base->Buffer(), base->Remaining()));
318  if (status != C_OK)
319    return status;
320
321  status = patch_process.ValidateBase();
322  if (status != C_OK)
323    return status;
324
325  // The rest of the patch stream is a StreamSet.
326  SourceStreamSet patch_streams;
327  patch_streams.Init(patch);
328
329  SourceStream* transformation_descriptions     = patch_streams.stream(0);
330  SourceStream* parameter_correction            = patch_streams.stream(1);
331  SourceStream* transformed_elements_correction = patch_streams.stream(2);
332  SourceStream* ensemble_correction             = patch_streams.stream(3);
333
334  status = patch_process.ReadInitialParameters(transformation_descriptions);
335  if (status != C_OK)
336    return status;
337
338  SinkStreamSet predicted_parameters;
339  status = patch_process.PredictTransformParameters(&predicted_parameters);
340  if (status != C_OK)
341    return status;
342
343  SourceStreamSet corrected_parameters;
344  status = patch_process.SubpatchTransformParameters(&predicted_parameters,
345                                                     parameter_correction,
346                                                     &corrected_parameters);
347  if (status != C_OK)
348    return status;
349
350  SinkStreamSet transformed_elements;
351  status = patch_process.TransformUp(&corrected_parameters,
352                                     &transformed_elements);
353  if (status != C_OK)
354    return status;
355
356  SourceStreamSet corrected_transformed_elements;
357  status = patch_process.SubpatchTransformedElements(
358          &transformed_elements,
359          transformed_elements_correction,
360          &corrected_transformed_elements);
361  if (status != C_OK)
362    return status;
363
364  SinkStream original_ensemble_and_corrected_base_elements;
365  status = patch_process.TransformDown(
366      &corrected_transformed_elements,
367      &original_ensemble_and_corrected_base_elements);
368  if (status != C_OK)
369    return status;
370
371  SourceStream final_patch_prediction;
372  final_patch_prediction.Init(original_ensemble_and_corrected_base_elements);
373  status = patch_process.SubpatchFinalOutput(&final_patch_prediction,
374                                             ensemble_correction, output);
375  if (status != C_OK)
376    return status;
377
378  return C_OK;
379}
380
381Status ApplyEnsemblePatch(const base::FilePath::CharType* old_file_name,
382                          const base::FilePath::CharType* patch_file_name,
383                          const base::FilePath::CharType* new_file_name) {
384  // First read enough of the patch file to validate the header is well-formed.
385  // A few varint32 numbers should fit in 100.
386  base::FilePath patch_file_path(patch_file_name);
387  base::MemoryMappedFile patch_file;
388  if (!patch_file.Initialize(patch_file_path))
389    return C_READ_OPEN_ERROR;
390
391  // 'Dry-run' the first step of the patch process to validate format of header.
392  SourceStream patch_header_stream;
393  patch_header_stream.Init(patch_file.data(), patch_file.length());
394  EnsemblePatchApplication patch_process;
395  Status status = patch_process.ReadHeader(&patch_header_stream);
396  if (status != C_OK)
397    return status;
398
399  // Read the old_file.
400  base::FilePath old_file_path(old_file_name);
401  base::MemoryMappedFile old_file;
402  if (!old_file.Initialize(old_file_path))
403    return C_READ_ERROR;
404
405  // Apply patch on streams.
406  SourceStream old_source_stream;
407  SourceStream patch_source_stream;
408  old_source_stream.Init(old_file.data(), old_file.length());
409  patch_source_stream.Init(patch_file.data(), patch_file.length());
410  SinkStream new_sink_stream;
411  status = ApplyEnsemblePatch(&old_source_stream, &patch_source_stream,
412                              &new_sink_stream);
413  if (status != C_OK)
414    return status;
415
416  // Write the patched data to |new_file_name|.
417  base::FilePath new_file_path(new_file_name);
418  int written =
419      file_util::WriteFile(
420          new_file_path,
421          reinterpret_cast<const char*>(new_sink_stream.Buffer()),
422          static_cast<int>(new_sink_stream.Length()));
423  if (written == -1)
424    return C_WRITE_OPEN_ERROR;
425  if (static_cast<size_t>(written) != new_sink_stream.Length())
426    return C_WRITE_ERROR;
427
428  return C_OK;
429}
430
431}  // namespace
432