// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
// This file contains the code to apply a Courgette patch.
#include "courgette/ensemble.h"
#include "base/basictypes.h"
#include "base/file_util.h"
#include "base/files/memory_mapped_file.h"
#include "base/logging.h"
#include "courgette/crc.h"
#include "courgette/region.h"
#include "courgette/streams.h"
#include "courgette/simple_delta.h"
#include "courgette/patcher_x86_32.h"
namespace courgette {
// EnsemblePatchApplication is all the logic and data required to apply the
// multi-stage patch.
class EnsemblePatchApplication {
public:
EnsemblePatchApplication();
~EnsemblePatchApplication();
Status ReadHeader(SourceStream* header_stream);
Status InitBase(const Region& region);
Status ValidateBase();
Status ReadInitialParameters(SourceStream* initial_parameters);
Status PredictTransformParameters(SinkStreamSet* predicted_parameters);
Status SubpatchTransformParameters(SinkStreamSet* prediction,
SourceStream* correction,
SourceStreamSet* corrected_parameters);
Status TransformUp(SourceStreamSet* parameters,
SinkStreamSet* transformed_elements);
Status SubpatchTransformedElements(SinkStreamSet* elements,
SourceStream* correction,
SourceStreamSet* corrected_elements);
Status TransformDown(SourceStreamSet* transformed_elements,
SinkStream* basic_elements);
Status SubpatchFinalOutput(SourceStream* original,
SourceStream* correction,
SinkStream* corrected_ensemble);
private:
Status SubpatchStreamSets(SinkStreamSet* predicted_items,
SourceStream* correction,
SourceStreamSet* corrected_items,
SinkStream* corrected_items_storage);
Region base_region_; // Location of in-memory copy of 'old' version.
uint32 source_checksum_;
uint32 target_checksum_;
uint32 final_patch_input_size_prediction_;
std::vector<TransformationPatcher*> patchers_;
SinkStream corrected_parameters_storage_;
SinkStream corrected_elements_storage_;
DISALLOW_COPY_AND_ASSIGN(EnsemblePatchApplication);
};
EnsemblePatchApplication::EnsemblePatchApplication()
: source_checksum_(0), target_checksum_(0),
final_patch_input_size_prediction_(0) {
}
EnsemblePatchApplication::~EnsemblePatchApplication() {
for (size_t i = 0; i < patchers_.size(); ++i) {
delete patchers_[i];
}
}
Status EnsemblePatchApplication::ReadHeader(SourceStream* header_stream) {
uint32 magic;
if (!header_stream->ReadVarint32(&magic))
return C_BAD_ENSEMBLE_MAGIC;
if (magic != CourgettePatchFile::kMagic)
return C_BAD_ENSEMBLE_MAGIC;
uint32 version;
if (!header_stream->ReadVarint32(&version))
return C_BAD_ENSEMBLE_VERSION;
if (version != CourgettePatchFile::kVersion)
return C_BAD_ENSEMBLE_VERSION;
if (!header_stream->ReadVarint32(&source_checksum_))
return C_BAD_ENSEMBLE_HEADER;
if (!header_stream->ReadVarint32(&target_checksum_))
return C_BAD_ENSEMBLE_HEADER;
if (!header_stream->ReadVarint32(&final_patch_input_size_prediction_))
return C_BAD_ENSEMBLE_HEADER;
return C_OK;
}
Status EnsemblePatchApplication::InitBase(const Region& region) {
base_region_.assign(region);
return C_OK;
}
Status EnsemblePatchApplication::ValidateBase() {
uint32 checksum = CalculateCrc(base_region_.start(), base_region_.length());
if (source_checksum_ != checksum)
return C_BAD_ENSEMBLE_CRC;
return C_OK;
}
Status EnsemblePatchApplication::ReadInitialParameters(
SourceStream* transformation_parameters) {
uint32 number_of_transformations = 0;
if (!transformation_parameters->ReadVarint32(&number_of_transformations))
return C_BAD_ENSEMBLE_HEADER;
for (size_t i = 0; i < number_of_transformations; ++i) {
uint32 kind;
if (!transformation_parameters->ReadVarint32(&kind))
return C_BAD_ENSEMBLE_HEADER;
TransformationPatcher* patcher = NULL;
switch (kind)
{
case EXE_WIN_32_X86:
patcher = new PatcherX86_32(base_region_);
break;
case EXE_ELF_32_X86:
patcher = new PatcherX86_32(base_region_);
break;
case EXE_ELF_32_ARM:
patcher = new PatcherX86_32(base_region_);
break;
case EXE_WIN_32_X64:
patcher = new PatcherX86_32(base_region_);
break;
}
if (patcher)
patchers_.push_back(patcher);
else
return C_BAD_ENSEMBLE_HEADER;
}
for (size_t i = 0; i < patchers_.size(); ++i) {
Status status = patchers_[i]->Init(transformation_parameters);
if (status != C_OK)
return status;
}
// All transformation_parameters should have been consumed by the above loop.
if (!transformation_parameters->Empty())
return C_BAD_ENSEMBLE_HEADER;
return C_OK;
}
Status EnsemblePatchApplication::PredictTransformParameters(
SinkStreamSet* all_predicted_parameters) {
for (size_t i = 0; i < patchers_.size(); ++i) {
SinkStreamSet single_predicted_parameters;
Status status =
patchers_[i]->PredictTransformParameters(&single_predicted_parameters);
if (status != C_OK)
return status;
if (!all_predicted_parameters->WriteSet(&single_predicted_parameters))
return C_STREAM_ERROR;
}
return C_OK;
}
Status EnsemblePatchApplication::SubpatchTransformParameters(
SinkStreamSet* predicted_parameters,
SourceStream* correction,
SourceStreamSet* corrected_parameters) {
return SubpatchStreamSets(predicted_parameters,
correction,
corrected_parameters,
&corrected_parameters_storage_);
}
Status EnsemblePatchApplication::TransformUp(
SourceStreamSet* parameters,
SinkStreamSet* transformed_elements) {
for (size_t i = 0; i < patchers_.size(); ++i) {
SourceStreamSet single_parameters;
if (!parameters->ReadSet(&single_parameters))
return C_STREAM_ERROR;
SinkStreamSet single_transformed_element;
Status status = patchers_[i]->Transform(&single_parameters,
&single_transformed_element);
if (status != C_OK)
return status;
if (!single_parameters.Empty())
return C_STREAM_NOT_CONSUMED;
if (!transformed_elements->WriteSet(&single_transformed_element))
return C_STREAM_ERROR;
}
if (!parameters->Empty())
return C_STREAM_NOT_CONSUMED;
return C_OK;
}
Status EnsemblePatchApplication::SubpatchTransformedElements(
SinkStreamSet* predicted_elements,
SourceStream* correction,
SourceStreamSet* corrected_elements) {
return SubpatchStreamSets(predicted_elements,
correction,
corrected_elements,
&corrected_elements_storage_);
}
Status EnsemblePatchApplication::TransformDown(
SourceStreamSet* transformed_elements,
SinkStream* basic_elements) {
// Construct blob of original input followed by reformed elements.
if (!basic_elements->Reserve(final_patch_input_size_prediction_)) {
return C_STREAM_ERROR;
}
// The original input:
if (!basic_elements->Write(base_region_.start(), base_region_.length()))
return C_STREAM_ERROR;
for (size_t i = 0; i < patchers_.size(); ++i) {
SourceStreamSet single_corrected_element;
if (!transformed_elements->ReadSet(&single_corrected_element))
return C_STREAM_ERROR;
Status status = patchers_[i]->Reform(&single_corrected_element,
basic_elements);
if (status != C_OK)
return status;
if (!single_corrected_element.Empty())
return C_STREAM_NOT_CONSUMED;
}
if (!transformed_elements->Empty())
return C_STREAM_NOT_CONSUMED;
// We have totally consumed transformed_elements, so can free the
// storage to which it referred.
corrected_elements_storage_.Retire();
return C_OK;
}
Status EnsemblePatchApplication::SubpatchFinalOutput(
SourceStream* original,
SourceStream* correction,
SinkStream* corrected_ensemble) {
Status delta_status = ApplySimpleDelta(original, correction,
corrected_ensemble);
if (delta_status != C_OK)
return delta_status;
if (CalculateCrc(corrected_ensemble->Buffer(),
corrected_ensemble->Length()) != target_checksum_)
return C_BAD_ENSEMBLE_CRC;
return C_OK;
}
Status EnsemblePatchApplication::SubpatchStreamSets(
SinkStreamSet* predicted_items,
SourceStream* correction,
SourceStreamSet* corrected_items,
SinkStream* corrected_items_storage) {
SinkStream linearized_predicted_items;
if (!predicted_items->CopyTo(&linearized_predicted_items))
return C_STREAM_ERROR;
SourceStream prediction;
prediction.Init(linearized_predicted_items);
Status status = ApplySimpleDelta(&prediction,
correction,
corrected_items_storage);
if (status != C_OK)
return status;
if (!corrected_items->Init(corrected_items_storage->Buffer(),
corrected_items_storage->Length()))
return C_STREAM_ERROR;
return C_OK;
}
Status ApplyEnsemblePatch(SourceStream* base,
SourceStream* patch,
SinkStream* output) {
Status status;
EnsemblePatchApplication patch_process;
status = patch_process.ReadHeader(patch);
if (status != C_OK)
return status;
status = patch_process.InitBase(Region(base->Buffer(), base->Remaining()));
if (status != C_OK)
return status;
status = patch_process.ValidateBase();
if (status != C_OK)
return status;
// The rest of the patch stream is a StreamSet.
SourceStreamSet patch_streams;
patch_streams.Init(patch);
SourceStream* transformation_descriptions = patch_streams.stream(0);
SourceStream* parameter_correction = patch_streams.stream(1);
SourceStream* transformed_elements_correction = patch_streams.stream(2);
SourceStream* ensemble_correction = patch_streams.stream(3);
status = patch_process.ReadInitialParameters(transformation_descriptions);
if (status != C_OK)
return status;
SinkStreamSet predicted_parameters;
status = patch_process.PredictTransformParameters(&predicted_parameters);
if (status != C_OK)
return status;
SourceStreamSet corrected_parameters;
status = patch_process.SubpatchTransformParameters(&predicted_parameters,
parameter_correction,
&corrected_parameters);
if (status != C_OK)
return status;
SinkStreamSet transformed_elements;
status = patch_process.TransformUp(&corrected_parameters,
&transformed_elements);
if (status != C_OK)
return status;
SourceStreamSet corrected_transformed_elements;
status = patch_process.SubpatchTransformedElements(
&transformed_elements,
transformed_elements_correction,
&corrected_transformed_elements);
if (status != C_OK)
return status;
SinkStream original_ensemble_and_corrected_base_elements;
status = patch_process.TransformDown(
&corrected_transformed_elements,
&original_ensemble_and_corrected_base_elements);
if (status != C_OK)
return status;
SourceStream final_patch_prediction;
final_patch_prediction.Init(original_ensemble_and_corrected_base_elements);
status = patch_process.SubpatchFinalOutput(&final_patch_prediction,
ensemble_correction, output);
if (status != C_OK)
return status;
return C_OK;
}
Status ApplyEnsemblePatch(const base::FilePath::CharType* old_file_name,
const base::FilePath::CharType* patch_file_name,
const base::FilePath::CharType* new_file_name) {
// First read enough of the patch file to validate the header is well-formed.
// A few varint32 numbers should fit in 100.
base::FilePath patch_file_path(patch_file_name);
base::MemoryMappedFile patch_file;
if (!patch_file.Initialize(patch_file_path))
return C_READ_OPEN_ERROR;
// 'Dry-run' the first step of the patch process to validate format of header.
SourceStream patch_header_stream;
patch_header_stream.Init(patch_file.data(), patch_file.length());
EnsemblePatchApplication patch_process;
Status status = patch_process.ReadHeader(&patch_header_stream);
if (status != C_OK)
return status;
// Read the old_file.
base::FilePath old_file_path(old_file_name);
base::MemoryMappedFile old_file;
if (!old_file.Initialize(old_file_path))
return C_READ_ERROR;
// Apply patch on streams.
SourceStream old_source_stream;
SourceStream patch_source_stream;
old_source_stream.Init(old_file.data(), old_file.length());
patch_source_stream.Init(patch_file.data(), patch_file.length());
SinkStream new_sink_stream;
status = ApplyEnsemblePatch(&old_source_stream, &patch_source_stream,
&new_sink_stream);
if (status != C_OK)
return status;
// Write the patched data to |new_file_name|.
base::FilePath new_file_path(new_file_name);
int written =
file_util::WriteFile(
new_file_path,
reinterpret_cast<const char*>(new_sink_stream.Buffer()),
static_cast<int>(new_sink_stream.Length()));
if (written == -1)
return C_WRITE_OPEN_ERROR;
if (static_cast<size_t>(written) != new_sink_stream.Length())
return C_WRITE_ERROR;
return C_OK;
}
} // namespace