courgette/ensemble_apply.cc

   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
  19 namespace courgette {
  20
  21 // EnsemblePatchApplication is all the logic and data required to apply the
  22 // multi-stage patch.
  23 class 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
  76 EnsemblePatchApplication::EnsemblePatchApplication()
  77     : source_checksum_(0), target_checksum_(0),
  78       final_patch_input_size_prediction_(0) {
  79 }
  80
  81 EnsemblePatchApplication::~EnsemblePatchApplication() {
  82   for (size_t i = 0;  i < patchers_.size();  ++i) {
  83     delete patchers_[i];
  84   }
  85 }
  86
  87 Status 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
 114 Status EnsemblePatchApplication::InitBase(const Region& region) {
 115   base_region_.assign(region);
 116   return C_OK;
 117 }
 118
 119 Status 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
 127 Status 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
 172 Status 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
 186 Status 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
 196 Status 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
 219 Status 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
 229 Status 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
 263 Status 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
 279 Status 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
 304 Status 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
 378 Status 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