courgette/patch_generator_x86_32.h

   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 is the transformation and adjustment for Windows X86 executables.
   6 // The same code can be used for Windows X64 executables.
   7
   8 #ifndef COURGETTE_WIN32_X86_GENERATOR_H_
   9 #define COURGETTE_WIN32_X86_GENERATOR_H_
  10
  11 #include "base/logging.h"
  12 #include "base/memory/scoped_ptr.h"
  13
  14 #include "courgette/assembly_program.h"
  15 #include "courgette/ensemble.h"
  16
  17 namespace courgette {
  18
  19 class PatchGeneratorX86_32 : public TransformationPatchGenerator {
  20  public:
  21   PatchGeneratorX86_32(Element* old_element,
  22                        Element* new_element,
  23                        PatcherX86_32* patcher,
  24                        ExecutableType kind)
  25       : TransformationPatchGenerator(old_element, new_element, patcher),
  26         kind_(kind) {
  27   }
  28
  29   virtual ExecutableType Kind() { return kind_; }
  30
  31   Status WriteInitialParameters(SinkStream* parameter_stream) {
  32     if (!parameter_stream->WriteSizeVarint32(
  33             old_element_->offset_in_ensemble()) ||
  34         !parameter_stream->WriteSizeVarint32(old_element_->region().length())) {
  35       return C_STREAM_ERROR;
  36     }
  37     return C_OK;
  38     // TODO(sra): Initialize |patcher_| with these parameters.
  39   }
  40
  41   Status PredictTransformParameters(SinkStreamSet* prediction) {
  42     return TransformationPatchGenerator::PredictTransformParameters(prediction);
  43   }
  44
  45   Status CorrectedTransformParameters(SinkStreamSet* parameters) {
  46     // No code needed to write an 'empty' parameter set.
  47     return C_OK;
  48   }
  49
  50   // The format of a transformed_element is a serialized EncodedProgram.  We
  51   // first disassemble the original old and new Elements into AssemblyPrograms.
  52   // Then we adjust the new AssemblyProgram to make it as much like the old one
  53   // as possible, before converting the AssemblyPrograms to EncodedPrograms and
  54   // serializing them.
  55   Status Transform(SourceStreamSet* corrected_parameters,
  56                    SinkStreamSet* old_transformed_element,
  57                    SinkStreamSet* new_transformed_element) {
  58     // Don't expect any corrected parameters.
  59     if (!corrected_parameters->Empty())
  60       return C_GENERAL_ERROR;
  61
  62     // Generate old version of program using |corrected_parameters|.
  63     // TODO(sra): refactor to use same code from patcher_.
  64     AssemblyProgram* old_program = NULL;
  65     Status old_parse_status =
  66         ParseDetectedExecutable(old_element_->region().start(),
  67                                 old_element_->region().length(),
  68                                 &old_program);
  69     if (old_parse_status != C_OK) {
  70       LOG(ERROR) << "Cannot parse as WinPE " << old_element_->Name();
  71       return old_parse_status;
  72     }
  73
  74     AssemblyProgram* new_program = NULL;
  75     Status new_parse_status =
  76         ParseDetectedExecutable(new_element_->region().start(),
  77                                 new_element_->region().length(),
  78                                 &new_program);
  79     if (new_parse_status != C_OK) {
  80       DeleteAssemblyProgram(old_program);
  81       LOG(ERROR) << "Cannot parse as WinPE " << new_element_->Name();
  82       return new_parse_status;
  83     }
  84
  85     // Trim labels below a certain threshold
  86     Status trim_old_status = TrimLabels(old_program);
  87     if (trim_old_status != C_OK) {
  88       DeleteAssemblyProgram(old_program);
  89       return trim_old_status;
  90     }
  91
  92     Status trim_new_status = TrimLabels(new_program);
  93     if (trim_new_status != C_OK) {
  94       DeleteAssemblyProgram(new_program);
  95       return trim_new_status;
  96     }
  97
  98     EncodedProgram* old_encoded = NULL;
  99     Status old_encode_status = Encode(old_program, &old_encoded);
 100     if (old_encode_status != C_OK) {
 101       DeleteAssemblyProgram(old_program);
 102       return old_encode_status;
 103     }
 104
 105     Status old_write_status =
 106         WriteEncodedProgram(old_encoded, old_transformed_element);
 107     DeleteEncodedProgram(old_encoded);
 108     if (old_write_status != C_OK) {
 109       DeleteAssemblyProgram(old_program);
 110       return old_write_status;
 111     }
 112
 113     Status adjust_status = Adjust(*old_program, new_program);
 114     DeleteAssemblyProgram(old_program);
 115     if (adjust_status != C_OK) {
 116       DeleteAssemblyProgram(new_program);
 117       return adjust_status;
 118     }
 119
 120     EncodedProgram* new_encoded = NULL;
 121     Status new_encode_status = Encode(new_program, &new_encoded);
 122     DeleteAssemblyProgram(new_program);
 123     if (new_encode_status != C_OK)
 124       return new_encode_status;
 125
 126     Status new_write_status =
 127         WriteEncodedProgram(new_encoded, new_transformed_element);
 128     DeleteEncodedProgram(new_encoded);
 129     if (new_write_status != C_OK)
 130       return new_write_status;
 131
 132     return C_OK;
 133   }
 134
 135   Status Reform(SourceStreamSet* transformed_element,
 136                 SinkStream* reformed_element) {
 137     return TransformationPatchGenerator::Reform(transformed_element,
 138                                                 reformed_element);
 139   }
 140
 141  private:
 142   virtual ~PatchGeneratorX86_32() { }
 143
 144   ExecutableType kind_;
 145
 146   DISALLOW_COPY_AND_ASSIGN(PatchGeneratorX86_32);
 147 };
 148
 149 }  // namespace courgette
 150 #endif  // COURGETTE_WIN32_X86_GENERATOR_H_