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/optimize/optimize_app.h"
16 :
17 : #include "syzygy/block_graph/block_graph.h"
18 : #include "syzygy/block_graph/transforms/fuzzing_transform.h"
19 : #include "syzygy/block_graph/transforms/named_transform.h"
20 : #include "syzygy/common/indexed_frequency_data.h"
21 : #include "syzygy/grinder/basic_block_util.h"
22 : #include "syzygy/optimize/application_profile.h"
23 : #include "syzygy/optimize/transforms/basic_block_reordering_transform.h"
24 : #include "syzygy/optimize/transforms/block_alignment_transform.h"
25 : #include "syzygy/optimize/transforms/chained_subgraph_transforms.h"
26 : #include "syzygy/optimize/transforms/inlining_transform.h"
27 : #include "syzygy/optimize/transforms/peephole_transform.h"
28 : #include "syzygy/optimize/transforms/unreachable_block_transform.h"
29 : #include "syzygy/pe/pe_relinker.h"
30 : #include "syzygy/pe/pe_transform_policy.h"
31 :
32 : namespace optimize {
33 :
34 : namespace {
35 :
36 : using block_graph::transforms::FuzzingTransform;
37 : using common::IndexedFrequencyData;
38 : using grinder::basic_block_util::IndexedFrequencyMap;
39 : using grinder::basic_block_util::LoadBranchStatisticsFromFile;
40 : using optimize::transforms::BasicBlockReorderingTransform;
41 : using optimize::transforms::BlockAlignmentTransform;
42 : using optimize::transforms::ChainedSubgraphTransforms;
43 : using optimize::transforms::InliningTransform;
44 : using optimize::transforms::PeepholeTransform;
45 : using optimize::transforms::UnreachableBlockTransform;
46 :
47 : const char kUsageFormatStr[] =
48 : "Usage: %ls [options]\n"
49 : " Required Options:\n"
50 : " --input-image=<path> The input image file to optimize.\n"
51 : " --output-image=<path> Output path for the rewritten image file.\n"
52 : "\n"
53 : " Options:\n"
54 : " --branch-file=<path> Branch statistics in JSON format.\n"
55 : " --input-pdb=<path> The PDB file associated with the input DLL.\n"
56 : " Default is inferred from input-image.\n"
57 : " --output-pdb=<path> Output path for the rewritten PDB file.\n"
58 : " Default is inferred from output-image.\n"
59 : " --overwrite Allow output files to be overwritten.\n"
60 : "\n"
61 : " Optimization Options:\n"
62 : " --all Enable all optimizations.\n"
63 : " --allow-inline-assembly\n"
64 : " Enable the decomposition of inline assembly\n"
65 : " blocks.\n"
66 : " --basic-block-reorder Enable basic block reodering.\n"
67 : " --block-alignment Enable block realignment.\n"
68 : " --inlining Enable function inlining.\n"
69 : " --peephole Enable peephole optimization.\n"
70 : " --unreachable-block Enable unreachable block optimization.\n"
71 : "\n"
72 : " Unreachable block options:\n"
73 : " --dump-unreachable-graph=<path>\n"
74 : " Dump the unreachable graph.\n"
75 : "\n"
76 : " Testing Options:\n"
77 : " --fuzz Fuzz the binary.\n"
78 : "\n";
79 :
80 : } // namespace
81 :
82 E : bool OptimizeApp::ParseCommandLine(const base::CommandLine* cmd_line) {
83 E : if (cmd_line->HasSwitch("help"))
84 E : return Usage(cmd_line, "");
85 :
86 E : input_image_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-image"));
87 E : output_image_path_ = cmd_line->GetSwitchValuePath("output-image");
88 E : input_pdb_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-pdb"));
89 E : output_pdb_path_ = cmd_line->GetSwitchValuePath("output-pdb");
90 E : branch_file_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("branch-file"));
91 :
92 E : basic_block_reorder_ = cmd_line->HasSwitch("basic-block-reorder");
93 E : block_alignment_ = cmd_line->HasSwitch("block-alignment");
94 E : fuzz_ = cmd_line->HasSwitch("fuzz");
95 E : inlining_ = cmd_line->HasSwitch("inlining");
96 E : allow_inline_assembly_ = cmd_line->HasSwitch("allow-inline-assembly");
97 E : peephole_ = cmd_line->HasSwitch("peephole");
98 E : unreachable_block_ = cmd_line->HasSwitch("unreachable-block");
99 E : overwrite_ = cmd_line->HasSwitch("overwrite");
100 :
101 : // Enable all optimization transforms.
102 E : if (cmd_line->HasSwitch("all")) {
103 E : basic_block_reorder_ = true;
104 E : block_alignment_ = true;
105 E : inlining_ = true;
106 E : peephole_ = true;
107 E : unreachable_block_ = true;
108 : }
109 :
110 : unreachable_graph_path_ =
111 E : AbsolutePath(cmd_line->GetSwitchValuePath("dump-unreachable-graph"));
112 :
113 : // The --input-image argument is required.
114 E : if (input_image_path_.empty())
115 E : return Usage(cmd_line, "You must specify --input-image.");
116 :
117 : // The --output-image argument is required.
118 E : if (output_image_path_.empty())
119 E : return Usage(cmd_line, "You must specify --output-image.");
120 :
121 E : return true;
122 E : }
123 :
124 E : bool OptimizeApp::SetUp() {
125 E : DCHECK(!input_image_path_.empty());
126 E : DCHECK(!output_image_path_.empty());
127 E : return true;
128 E : }
129 :
130 E : int OptimizeApp::Run() {
131 E : pe::PETransformPolicy policy;
132 E : policy.set_allow_inline_assembly(allow_inline_assembly_);
133 E : pe::PERelinker relinker(&policy);
134 E : relinker.set_input_path(input_image_path_);
135 E : relinker.set_input_pdb_path(input_pdb_path_);
136 E : relinker.set_output_path(output_image_path_);
137 E : relinker.set_output_pdb_path(output_pdb_path_);
138 E : relinker.set_allow_overwrite(overwrite_);
139 :
140 : // Initialize the relinker. This does the decomposition, etc.
141 E : if (!relinker.Init()) {
142 i : LOG(ERROR) << "Failed to initialize relinker.";
143 i : return 1;
144 : }
145 :
146 : // Get module signature and layout.
147 E : pe::PEFile::Signature signature;
148 E : relinker.input_pe_file().GetSignature(&signature);
149 E : const pe::ImageLayout& image_layout = relinker.input_image_layout();
150 :
151 : // Load profile information from file.
152 E : ApplicationProfile profile(&image_layout);
153 E : if (!branch_file_path_.empty()) {
154 i : IndexedFrequencyMap frequencies;
155 : if (!LoadBranchStatisticsFromFile(branch_file_path_,
156 : signature,
157 i : &frequencies)) {
158 i : LOG(ERROR) << "Unable to load profile information.";
159 i : return 1;
160 : }
161 i : if (!profile.ImportFrequencies(frequencies)) {
162 i : LOG(ERROR) << "Could not import metrics for '"
163 : << branch_file_path_.value() << "'.";
164 i : return false;
165 : }
166 i : }
167 :
168 : // Compute global profile information for the current block graph.
169 E : if (!profile.ComputeGlobalProfile()) {
170 i : LOG(ERROR) << "Unable to build profile information.";
171 i : return 1;
172 : }
173 :
174 : // Construct a chain of basic block transforms.
175 E : ChainedSubgraphTransforms chains(&profile);
176 :
177 : // Declare transforms we may apply.
178 E : scoped_ptr<BasicBlockReorderingTransform> basic_block_reordering_transform;
179 E : scoped_ptr<BlockAlignmentTransform> block_alignment_transform;
180 E : scoped_ptr<FuzzingTransform> fuzzing_transform;
181 E : scoped_ptr<InliningTransform> inlining_transform;
182 E : scoped_ptr<PeepholeTransform> peephole_transform;
183 E : scoped_ptr<UnreachableBlockTransform> unreachable_block_transform;
184 :
185 : // If block block reordering is enabled, add it to the chain.
186 E : if (peephole_) {
187 i : peephole_transform.reset(new PeepholeTransform());
188 i : chains.AppendTransform(peephole_transform.get());
189 : }
190 :
191 : // If inlining is enabled, add it to the chain.
192 E : if (inlining_) {
193 i : inlining_transform.reset(new InliningTransform());
194 i : chains.AppendTransform(inlining_transform.get());
195 : }
196 :
197 : // If block block reordering is enabled, add it to the chain.
198 E : if (basic_block_reorder_) {
199 i : basic_block_reordering_transform.reset(new BasicBlockReorderingTransform());
200 i : chains.AppendTransform(basic_block_reordering_transform.get());
201 : }
202 :
203 : // If block alignment is enabled, add it to the chain.
204 E : if (block_alignment_) {
205 i : block_alignment_transform.reset(new BlockAlignmentTransform());
206 i : chains.AppendTransform(block_alignment_transform.get());
207 : }
208 :
209 : // Append the chain to the relinker.
210 E : if (!relinker.AppendTransform(&chains))
211 i : return false;
212 :
213 : // If unreachable-block is enabled, add it to the relinker.
214 E : if (unreachable_block_) {
215 i : unreachable_block_transform.reset(new UnreachableBlockTransform());
216 : unreachable_block_transform->set_unreachable_graph_path(
217 i : unreachable_graph_path_);
218 i : relinker.AppendTransform(unreachable_block_transform.get());
219 : }
220 :
221 : // If fuzzing is enabled, add it to the relinker.
222 E : if (fuzz_) {
223 i : fuzzing_transform.reset(new block_graph::transforms::FuzzingTransform);
224 i : relinker.AppendTransform(fuzzing_transform.get());
225 : }
226 :
227 : // Perform the actual relink.
228 E : if (!relinker.Relink()) {
229 i : LOG(ERROR) << "Unable to relink input image.";
230 i : return 1;
231 : }
232 :
233 E : return 0;
234 E : }
235 :
236 : bool OptimizeApp::Usage(const base::CommandLine* cmd_line,
237 E : const base::StringPiece& message) const {
238 E : if (!message.empty()) {
239 E : ::fwrite(message.data(), 1, message.length(), err());
240 E : ::fprintf(err(), "\n\n");
241 : }
242 :
243 : ::fprintf(err(),
244 : kUsageFormatStr,
245 E : cmd_line->GetProgram().BaseName().value().c_str());
246 :
247 E : return false;
248 E : }
249 :
250 : } // namespace optimize
|