1 : // Copyright 2013 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/pe/coff_relinker.h"
16 :
17 : #include "base/file_util.h"
18 : #include "syzygy/pe/coff_decomposer.h"
19 : #include "syzygy/pe/coff_file_writer.h"
20 : #include "syzygy/pe/coff_image_layout_builder.h"
21 : #include "syzygy/pe/pe_utils.h"
22 : #include "syzygy/pe/transforms/coff_convert_legacy_code_references_transform.h"
23 : #include "syzygy/pe/transforms/coff_prepare_headers_transform.h"
24 :
25 : namespace pe {
26 :
27 : namespace {
28 :
29 : using block_graph::ApplyBlockGraphTransform;
30 : using block_graph::BlockGraph;
31 : using block_graph::OrderedBlockGraph;
32 : using core::RelativeAddress;
33 :
34 : typedef block_graph::BlockGraphOrdererInterface Orderer;
35 : typedef block_graph::BlockGraphTransformInterface Transform;
36 :
37 : // Decompose @p image_file into @p image_layout.
38 : //
39 : // @param image_file the COFF file to decompose.
40 : // @param image_layout the layout to fill with the results.
41 : // @param headers_block where to place the headers block pointer.
42 : // @returns true on success, or false on failure.
43 : bool Decompose(const CoffFile& image_file,
44 : ImageLayout* image_layout,
45 E : BlockGraph::Block** headers_block) {
46 E : DCHECK(image_layout != NULL);
47 E : DCHECK(headers_block != NULL);
48 :
49 E : LOG(INFO) << "Decomposing module: " << image_file.path().value() << ".";
50 :
51 E : BlockGraph* block_graph = image_layout->blocks.graph();
52 :
53 : // Decompose the input image.
54 E : CoffDecomposer decomposer(image_file);
55 E : if (!decomposer.Decompose(image_layout)) {
56 i : LOG(ERROR) << "Unable to decompose module: "
57 : << image_file.path().value() << ".";
58 i : return false;
59 : }
60 :
61 : // Get the headers block.
62 E : *headers_block = image_layout->blocks.GetBlockByAddress(RelativeAddress(0));
63 E : if (*headers_block == NULL) {
64 i : LOG(ERROR) << "Unable to find the headers block.";
65 i : return false;
66 : }
67 :
68 E : return true;
69 E : }
70 :
71 : // Build an image layout from an ordered block graph.
72 : //
73 : // @param ordered_graph the ordered block graph to lay out.
74 : // @param headers_block the headers block in @p ordered_graph.
75 : // @param image_layout the image layout to fill with the results.
76 : // @returns true on success, or false on failure.
77 : bool BuildImageLayout(const OrderedBlockGraph& ordered_graph,
78 : BlockGraph::Block* headers_block,
79 E : ImageLayout* image_layout) {
80 E : DCHECK(headers_block != NULL);
81 E : DCHECK(image_layout != NULL);
82 :
83 E : LOG(INFO) << "Building image layout.";
84 :
85 E : CoffImageLayoutBuilder builder(image_layout);
86 E : if (!builder.LayoutImage(ordered_graph)) {
87 i : LOG(ERROR) << "Image layout failed.";
88 i : return false;
89 : }
90 :
91 E : return true;
92 E : }
93 :
94 : // Write an image layout to disk.
95 : //
96 : // @param image_layout the image layout.
97 : // @param output_path the path to which the result is written.
98 : // @returns true on success, or false on failure.
99 : bool WriteImage(const ImageLayout& image_layout,
100 E : const base::FilePath& output_path) {
101 E : CoffFileWriter writer(&image_layout);
102 :
103 E : LOG(INFO) << "Writing image to file: " << output_path.value() << ".";
104 E : if (!writer.WriteImage(output_path)) {
105 i : LOG(ERROR) << "Failed to write image: " << output_path.value() << ".";
106 i : return false;
107 : }
108 :
109 E : return true;
110 E : }
111 :
112 : } // namespace
113 :
114 : CoffRelinker::CoffRelinker(const CoffTransformPolicy* transform_policy)
115 E : : PECoffRelinker(transform_policy) {
116 E : }
117 :
118 E : bool CoffRelinker::Init() {
119 E : DCHECK(inited_ == false);
120 :
121 : // Initialize the paths.
122 E : if (!CheckPaths())
123 E : return false;
124 :
125 E : LOG(INFO) << "Input module: " << input_path_.value() << ".";
126 E : LOG(INFO) << "Output module: " << output_path_.value() << ".";
127 :
128 : // Open the input PE file.
129 E : if (!input_image_file_.Init(input_path_)) {
130 i : LOG(ERROR) << "Unable to load input image: " << input_path_.value() << ".";
131 i : return false;
132 : }
133 :
134 : // Decompose the image.
135 E : if (!Decompose(input_image_file_, &input_image_layout_, &headers_block_))
136 i : return false;
137 :
138 E : inited_ = true;
139 :
140 E : return true;
141 E : }
142 :
143 E : bool CoffRelinker::Relink() {
144 E : if (!inited_) {
145 i : LOG(ERROR) << "Init() has not been successfully called.";
146 i : return false;
147 : }
148 :
149 E : if (!ApplyUserTransforms())
150 E : return false;
151 :
152 : // We apply the extra prepare headers transform. This ensures that the
153 : // header block is properly sized to receive layout information
154 : // post-ordering.
155 : //
156 : // TODO(chrisha): Remove CoffConvertLegacyCodeReferencesTransform when the
157 : // basic block assembler is made fully COFF-compatible.
158 E : pe::transforms::CoffConvertLegacyCodeReferencesTransform fix_refs_tx;
159 E : pe::transforms::CoffPrepareHeadersTransform prep_headers_tx;
160 E : std::vector<Transform*> post_transforms;
161 E : post_transforms.push_back(&fix_refs_tx);
162 E : post_transforms.push_back(&prep_headers_tx);
163 : if (!block_graph::ApplyBlockGraphTransforms(
164 E : post_transforms, transform_policy_, &block_graph_, headers_block_)) {
165 i : return false;
166 : }
167 :
168 E : OrderedBlockGraph ordered_graph(&block_graph_);
169 E : if (!ApplyUserOrderers(&ordered_graph))
170 E : return false;
171 :
172 : // Lay it out.
173 E : ImageLayout output_image_layout(&block_graph_);
174 : if (!BuildImageLayout(ordered_graph, headers_block_,
175 E : &output_image_layout)) {
176 i : return false;
177 : }
178 :
179 : // Write the image.
180 E : if (!WriteImage(output_image_layout, output_path_))
181 i : return false;
182 :
183 E : return true;
184 E : }
185 :
186 E : bool CoffRelinker::CheckPaths() {
187 : // At a very minimum we have to specify input and output.
188 E : if (input_path_.empty() || output_path_.empty()) {
189 E : LOG(ERROR) << "Input path and output path must be set and non-empty.";
190 E : return false;
191 : }
192 :
193 E : if (!base::PathExists(input_path_)) {
194 E : LOG(ERROR) << "Input file not found: " << input_path_.value() << ".";
195 E : return false;
196 : }
197 :
198 : // Ensure we aren't about to overwrite anything we don't want to. We do this
199 : // early on so that we abort before decomposition, transformation, etc.
200 E : if (!allow_overwrite_) {
201 E : if (base::PathExists(output_path_)) {
202 E : LOG(ERROR) << "Output file already exists: "
203 : << output_path_.value() << ".";
204 E : return false;
205 : }
206 : }
207 :
208 E : return true;
209 E : }
210 :
211 : } // namespace pe
|