1 : // Copyright 2015 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/instrument/instrumenters/instrumenter_with_relinker.h"
16 :
17 : #include "base/logging.h"
18 : #include "base/files/file_util.h"
19 : #include "syzygy/application/application.h"
20 : #include "syzygy/core/file_util.h"
21 :
22 : namespace instrument {
23 : namespace instrumenters {
24 :
25 : namespace {
26 :
27 : using block_graph::BlockGraph;
28 :
29 : bool GetImageFormat(const base::FilePath& path,
30 E : BlockGraph::ImageFormat* image_format) {
31 E : DCHECK(image_format != nullptr);
32 :
33 : // Determine the type of the input.
34 E : core::FileType file_type = core::kUnknownFileType;
35 E : if (!core::GuessFileType(path, &file_type)) {
36 i : LOG(ERROR) << "Failed to determine file type of \""
37 : << path.value() << "\".";
38 i : return false;
39 : }
40 :
41 E : if (file_type == core::kCoffFileType) {
42 E : *image_format = BlockGraph::COFF_IMAGE;
43 E : return true;
44 : }
45 :
46 E : if (file_type == core::kPeFileType) {
47 E : *image_format = BlockGraph::PE_IMAGE;
48 E : return true;
49 : }
50 :
51 i : LOG(ERROR) << "File is not a PE or COFF image: " << path.value();
52 i : return false;
53 E : }
54 :
55 : } // namespace
56 :
57 : bool InstrumenterWithRelinker::ParseCommandLine(
58 E : const base::CommandLine* command_line) {
59 : return DoCommandLineParse(command_line) &&
60 E : CheckCommandLineParse(command_line);
61 E : }
62 :
63 : bool InstrumenterWithRelinker::DoCommandLineParse(
64 E : const base::CommandLine* command_line) {
65 E : DCHECK(command_line != nullptr);
66 : // No super class.
67 :
68 : // TODO(chrisha): Simplify the input/output image parsing once external
69 : // tools have been updated.
70 :
71 : // Parse the input image.
72 E : if (command_line->HasSwitch("input-dll")) {
73 E : LOG(WARNING) << "DEPRECATED: Using --input-dll.";
74 : input_image_path_ = application::AppImplBase::AbsolutePath(
75 E : command_line->GetSwitchValuePath("input-dll"));
76 E : } else {
77 : input_image_path_ = application::AppImplBase::AbsolutePath(
78 E : command_line->GetSwitchValuePath("input-image"));
79 : }
80 :
81 : // Parse the output image.
82 E : if (command_line->HasSwitch("output-dll")) {
83 E : LOG(WARNING) << "DEPRECATED: Using --output-dll.";
84 : output_image_path_ = application::AppImplBase::AbsolutePath(
85 E : command_line->GetSwitchValuePath("output-dll"));
86 E : } else {
87 : output_image_path_ = application::AppImplBase::AbsolutePath(
88 E : command_line->GetSwitchValuePath("output-image"));
89 : }
90 :
91 : // Ensure that both input and output have been specified.
92 E : if (input_image_path_.empty() || output_image_path_.empty()) {
93 E : LOG(ERROR) << "You must provide input and output file names.";
94 E : return false;
95 : }
96 :
97 : // Parse the remaining command line arguments.
98 : input_pdb_path_ = application::AppImplBase::AbsolutePath(
99 E : command_line->GetSwitchValuePath("input-pdb"));
100 : output_pdb_path_ = application::AppImplBase::AbsolutePath(
101 E : command_line->GetSwitchValuePath("output-pdb"));
102 E : allow_overwrite_ = command_line->HasSwitch("overwrite");
103 E : debug_friendly_ = command_line->HasSwitch("debug-friendly");
104 E : no_augment_pdb_ = command_line->HasSwitch("no-augment-pdb");
105 E : no_strip_strings_ = command_line->HasSwitch("no-strip-strings");
106 :
107 E : return true;
108 E : }
109 :
110 : bool InstrumenterWithRelinker::CheckCommandLineParse(
111 E : const base::CommandLine* command_line) {
112 E : return true; // No super class.
113 E : }
114 :
115 E : bool InstrumenterWithRelinker::Instrument() {
116 E : if (!InstrumentPrepare())
117 i : return false;
118 :
119 E : if (!CreateRelinker())
120 i : return false;
121 :
122 : // Initialize the relinker. This does the decomposition, etc.
123 E : if (!relinker_->Init()) {
124 E : LOG(ERROR) << "Failed to initialize relinker.";
125 E : return false;
126 : }
127 :
128 : // Let the instrumenter implementation set up the relinker and anything else
129 : // that is required.
130 E : if (!InstrumentImpl())
131 i : return false;
132 :
133 : // Do the actual instrumentation by running the relinker.
134 E : if (!relinker_->Relink()) {
135 E : LOG(ERROR) << "Unable to relink input image.";
136 E : return false;
137 : }
138 :
139 E : return true;
140 E : }
141 :
142 : bool InstrumenterWithRelinker::ImageFormatIsSupported(
143 E : BlockGraph::ImageFormat image_format) {
144 : // By default we only support PE images.
145 E : if (image_format == BlockGraph::PE_IMAGE)
146 E : return true;
147 i : return false;
148 E : }
149 :
150 E : pe::PETransformPolicy* InstrumenterWithRelinker::GetPETransformPolicy() {
151 E : DCHECK_EQ(BlockGraph::PE_IMAGE, image_format_);
152 E : DCHECK(policy_object_.get() == nullptr);
153 E : policy_object_.reset(new pe::PETransformPolicy());
154 E : return static_cast<pe::PETransformPolicy*>(policy_object_.get());
155 E : }
156 :
157 E : pe::CoffTransformPolicy* InstrumenterWithRelinker::GetCoffTransformPolicy() {
158 E : DCHECK_EQ(BlockGraph::COFF_IMAGE, image_format_);
159 E : DCHECK(policy_object_.get() == nullptr);
160 E : policy_object_.reset(new pe::CoffTransformPolicy());
161 E : return static_cast<pe::CoffTransformPolicy*>(policy_object_.get());
162 E : }
163 :
164 E : pe::PERelinker* InstrumenterWithRelinker::GetPERelinker() {
165 E : DCHECK_EQ(BlockGraph::PE_IMAGE, image_format_);
166 E : DCHECK(relinker_object_.get() == nullptr);
167 E : relinker_object_.reset(new pe::PERelinker(GetPETransformPolicy()));
168 E : return static_cast<pe::PERelinker*>(relinker_object_.get());
169 E : }
170 :
171 E : pe::CoffRelinker* InstrumenterWithRelinker::GetCoffRelinker() {
172 E : DCHECK_EQ(BlockGraph::COFF_IMAGE, image_format_);
173 E : relinker_object_.reset(new pe::CoffRelinker(GetCoffTransformPolicy()));
174 E : return static_cast<pe::CoffRelinker*>(relinker_object_.get());
175 E : }
176 :
177 E : bool InstrumenterWithRelinker::CreateRelinker() {
178 : // Get the image format by quickly inspecting the image. This logs verbosely
179 : // on failure.
180 E : if (!GetImageFormat(input_image_path_, &image_format_))
181 i : return false;
182 :
183 : // Check if the format is supported and bail if it isn't.
184 E : if (!ImageFormatIsSupported(image_format_)) {
185 i : LOG(ERROR) << "Instrumenter \"" << InstrumentationMode()
186 : << "\" does not support input image format.";
187 i : return false;
188 : }
189 :
190 : // Create and setup an image format specific relinker.
191 E : if (image_format_ == BlockGraph::COFF_IMAGE) {
192 E : pe::CoffRelinker* relinker = GetCoffRelinker();
193 E : DCHECK_NE(reinterpret_cast<pe::CoffRelinker*>(nullptr), relinker);
194 E : relinker_ = relinker;
195 E : relinker->set_input_path(input_image_path_);
196 E : relinker->set_output_path(output_image_path_);
197 E : relinker->set_allow_overwrite(allow_overwrite_);
198 E : } else {
199 E : pe::PERelinker* relinker = GetPERelinker();
200 E : DCHECK_NE(reinterpret_cast<pe::PERelinker*>(nullptr), relinker);
201 E : relinker_ = relinker;
202 E : relinker->set_input_path(input_image_path_);
203 E : relinker->set_input_pdb_path(input_pdb_path_);
204 E : relinker->set_output_path(output_image_path_);
205 E : relinker->set_output_pdb_path(output_pdb_path_);
206 E : relinker->set_allow_overwrite(allow_overwrite_);
207 E : relinker->set_augment_pdb(!no_augment_pdb_);
208 E : relinker->set_strip_strings(!no_strip_strings_);
209 : }
210 :
211 E : DCHECK_EQ(image_format_, relinker_->image_format());
212 :
213 E : return true;
214 E : }
215 :
216 : } // namespace instrumenters
217 : } // namespace instrument
|