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_prepare_headers_transform.h"
23 :
24 : namespace pe {
25 :
26 : namespace {
27 :
28 : using block_graph::ApplyBlockGraphTransform;
29 : using block_graph::BlockGraph;
30 : using block_graph::OrderedBlockGraph;
31 : using core::RelativeAddress;
32 :
33 : typedef block_graph::BlockGraphOrdererInterface Orderer;
34 : typedef block_graph::BlockGraphTransformInterface Transform;
35 :
36 : // Decompose @p image_file into @p image_layout.
37 : //
38 : // @param image_file the COFF file to decompose.
39 : // @param image_layout the layout to fill with the results.
40 : // @param headers_block where to place the headers block pointer.
41 : // @returns true on success, or false on failure.
42 : bool Decompose(const CoffFile& image_file,
43 : ImageLayout* image_layout,
44 E : BlockGraph::Block** headers_block) {
45 E : DCHECK(image_layout != NULL);
46 E : DCHECK(headers_block != NULL);
47 :
48 E : LOG(INFO) << "Decomposing module: " << image_file.path().value() << ".";
49 :
50 E : BlockGraph* block_graph = image_layout->blocks.graph();
51 :
52 : // Decompose the input image.
53 E : CoffDecomposer decomposer(image_file);
54 E : if (!decomposer.Decompose(image_layout)) {
55 i : LOG(ERROR) << "Unable to decompose module: "
56 : << image_file.path().value() << ".";
57 i : return false;
58 : }
59 :
60 : // Get the headers block.
61 E : *headers_block = image_layout->blocks.GetBlockByAddress(RelativeAddress(0));
62 E : if (*headers_block == NULL) {
63 i : LOG(ERROR) << "Unable to find the headers block.";
64 i : return false;
65 : }
66 :
67 E : return true;
68 E : }
69 :
70 : // Build an image layout from an ordered block graph.
71 : //
72 : // @param ordered_graph the ordered block graph to lay out.
73 : // @param headers_block the headers block in @p ordered_graph.
74 : // @param image_layout the image layout to fill with the results.
75 : // @returns true on success, or false on failure.
76 : bool BuildImageLayout(const OrderedBlockGraph& ordered_graph,
77 : BlockGraph::Block* headers_block,
78 E : ImageLayout* image_layout) {
79 E : DCHECK(headers_block != NULL);
80 E : DCHECK(image_layout != NULL);
81 :
82 E : LOG(INFO) << "Building image layout.";
83 :
84 E : CoffImageLayoutBuilder builder(image_layout);
85 E : if (!builder.LayoutImage(ordered_graph)) {
86 i : LOG(ERROR) << "Image layout failed.";
87 i : return false;
88 : }
89 :
90 E : return true;
91 E : }
92 :
93 : // Write an image layout to disk.
94 : //
95 : // @param image_layout the image layout.
96 : // @param output_path the path to which the result is written.
97 : // @returns true on success, or false on failure.
98 : bool WriteImage(const ImageLayout& image_layout,
99 E : const base::FilePath& output_path) {
100 E : CoffFileWriter writer(&image_layout);
101 :
102 E : LOG(INFO) << "Writing image to file: " << output_path.value() << ".";
103 E : if (!writer.WriteImage(output_path)) {
104 i : LOG(ERROR) << "Failed to write image: " << output_path.value() << ".";
105 i : return false;
106 : }
107 :
108 E : return true;
109 E : }
110 :
111 : } // namespace
112 :
113 : CoffRelinker::CoffRelinker(const CoffTransformPolicy* transform_policy)
114 E : : PECoffRelinker(transform_policy) {
115 E : }
116 :
117 E : bool CoffRelinker::Init() {
118 E : DCHECK(inited_ == false);
119 :
120 : // Initialize the paths.
121 E : if (!CheckPaths())
122 E : return false;
123 :
124 E : LOG(INFO) << "Input module: " << input_path_.value() << ".";
125 E : LOG(INFO) << "Output module: " << output_path_.value() << ".";
126 :
127 : // Open the input PE file.
128 E : if (!input_image_file_.Init(input_path_)) {
129 i : LOG(ERROR) << "Unable to load input image: " << input_path_.value() << ".";
130 i : return false;
131 : }
132 :
133 : // Decompose the image.
134 E : if (!Decompose(input_image_file_, &input_image_layout_, &headers_block_))
135 i : return false;
136 :
137 E : inited_ = true;
138 :
139 E : return true;
140 E : }
141 :
142 E : bool CoffRelinker::Relink() {
143 E : if (!inited_) {
144 i : LOG(ERROR) << "Init() has not been successfully called.";
145 i : return false;
146 : }
147 :
148 E : if (!ApplyUserTransforms())
149 E : return false;
150 :
151 : // We apply the extra prepare headers transform. This ensures that the
152 : // header block is properly sized to receive layout information
153 : // post-ordering.
154 E : std::vector<Transform*> post_transforms;
155 E : pe::transforms::CoffPrepareHeadersTransform prep_headers_tx;
156 E : post_transforms.push_back(&prep_headers_tx);
157 : if (!block_graph::ApplyBlockGraphTransforms(
158 E : post_transforms, transform_policy_, &block_graph_, headers_block_)) {
159 i : return false;
160 : }
161 :
162 E : OrderedBlockGraph ordered_graph(&block_graph_);
163 E : if (!ApplyUserOrderers(&ordered_graph))
164 E : return false;
165 :
166 : // Lay it out.
167 E : ImageLayout output_image_layout(&block_graph_);
168 : if (!BuildImageLayout(ordered_graph, headers_block_,
169 E : &output_image_layout)) {
170 i : return false;
171 : }
172 :
173 : // Write the image.
174 E : if (!WriteImage(output_image_layout, output_path_))
175 i : return false;
176 :
177 E : return true;
178 E : }
179 :
180 E : bool CoffRelinker::CheckPaths() {
181 : // At a very minimum we have to specify input and output.
182 E : if (input_path_.empty() || output_path_.empty()) {
183 E : LOG(ERROR) << "Input path and output path must be set and non-empty.";
184 E : return false;
185 : }
186 :
187 E : if (!file_util::PathExists(input_path_)) {
188 E : LOG(ERROR) << "Input file not found: " << input_path_.value() << ".";
189 E : return false;
190 : }
191 :
192 : // Ensure we aren't about to overwrite anything we don't want to. We do this
193 : // early on so that we abort before decomposition, transformation, etc.
194 E : if (!allow_overwrite_) {
195 E : if (file_util::PathExists(output_path_)) {
196 E : LOG(ERROR) << "Output file already exists: "
197 : << output_path_.value() << ".";
198 E : return false;
199 : }
200 : }
201 :
202 E : return true;
203 E : }
204 :
205 : } // namespace pe
|