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
17namespace courgette {
18
19class 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_
151