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