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/crc.h"
  14 #include "courgette/patcher_x86_32.h"
  15 #include "courgette/region.h"
  16 #include "courgette/simple_delta.h"
  17 #include "courgette/streams.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       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
 175 Status 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
 189 Status 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
 199 Status 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
 222 Status 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
 232 Status 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
 266 Status 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
 282 Status 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
 307 Status 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
 381 Status 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   base::FilePath patch_file_path(patch_file_name);
 385   base::MemoryMappedFile patch_file;
 386   if (!patch_file.Initialize(patch_file_path))
 387     return C_READ_OPEN_ERROR;
 388
 389   // 'Dry-run' the first step of the patch process to validate format of header.
 390   SourceStream patch_header_stream;
 391   patch_header_stream.Init(patch_file.data(), patch_file.length());
 392   EnsemblePatchApplication patch_process;
 393   Status status = patch_process.ReadHeader(&patch_header_stream);
 394   if (status != C_OK)
 395     return status;
 396
 397   // Read the old_file.
 398   base::FilePath old_file_path(old_file_name);
 399   base::MemoryMappedFile old_file;
 400   if (!old_file.Initialize(old_file_path))
 401     return C_READ_ERROR;
 402
 403   // Apply patch on streams.
 404   SourceStream old_source_stream;
 405   SourceStream patch_source_stream;
 406   old_source_stream.Init(old_file.data(), old_file.length());
 407   patch_source_stream.Init(patch_file.data(), patch_file.length());
 408   SinkStream new_sink_stream;
 409   status = ApplyEnsemblePatch(&old_source_stream, &patch_source_stream,
 410                               &new_sink_stream);
 411   if (status != C_OK)
 412     return status;
 413
 414   // Write the patched data to |new_file_name|.
 415   base::FilePath new_file_path(new_file_name);
 416   int written =
 417       base::WriteFile(
 418           new_file_path,
 419           reinterpret_cast<const char*>(new_sink_stream.Buffer()),
 420           static_cast<int>(new_sink_stream.Length()));
 421   if (written == -1)
 422     return C_WRITE_OPEN_ERROR;
 423   if (static_cast<size_t>(written) != new_sink_stream.Length())
 424     return C_WRITE_ERROR;
 425
 426   return C_OK;
 427 }
 428
 429 }  // namespace