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/instrument/instrumenters/instrumenter_with_agent.h"
16 :
17 : #include "base/file_util.h"
18 : #include "base/logging.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 != NULL);
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 E : LOG(ERROR) << "Failed to determine file type of \""
37 : << path.value() << "\".";
38 E : 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 E : bool InstrumenterWithAgent::ParseCommandLine(const CommandLine* command_line) {
58 E : DCHECK(command_line != NULL);
59 :
60 : // TODO(chrisha): Simplify the input/output image parsing once external
61 : // tools have been updated.
62 :
63 : // Parse the input image.
64 E : if (command_line->HasSwitch("input-dll")) {
65 E : LOG(WARNING) << "DEPRECATED: Using --input-dll.";
66 : input_image_path_ = application::AppImplBase::AbsolutePath(
67 E : command_line->GetSwitchValuePath("input-dll"));
68 E : } else {
69 : input_image_path_ = application::AppImplBase::AbsolutePath(
70 E : command_line->GetSwitchValuePath("input-image"));
71 : }
72 :
73 : // Parse the output image.
74 E : if (command_line->HasSwitch("output-dll")) {
75 E : LOG(WARNING) << "DEPRECATED: Using --output-dll.";
76 : output_image_path_ = application::AppImplBase::AbsolutePath(
77 E : command_line->GetSwitchValuePath("output-dll"));
78 E : } else {
79 : output_image_path_ = application::AppImplBase::AbsolutePath(
80 E : command_line->GetSwitchValuePath("output-image"));
81 : }
82 :
83 : // Ensure that both input and output have been specified.
84 E : if (input_image_path_.empty() || output_image_path_.empty()) {
85 E : LOG(ERROR) << "You must provide input and output file names.";
86 E : return false;
87 : }
88 :
89 : // Parse the remaining command line arguments.
90 : input_pdb_path_ = application::AppImplBase::AbsolutePath(
91 E : command_line->GetSwitchValuePath("input-pdb"));
92 : output_pdb_path_ = application::AppImplBase::AbsolutePath(
93 E : command_line->GetSwitchValuePath("output-pdb"));
94 E : allow_overwrite_ = command_line->HasSwitch("overwrite");
95 E : debug_friendly_ = command_line->HasSwitch("debug-friendly");
96 E : no_augment_pdb_ = command_line->HasSwitch("no-augment-pdb");
97 E : no_strip_strings_ = command_line->HasSwitch("no-strip-strings");
98 :
99 E : if (!agent_dll_.empty()) {
100 E : LOG(INFO) << "Default agent DLL for " << InstrumentationMode() << " mode "
101 : << "is \"" << agent_dll_ << "\".";
102 : }
103 :
104 : // Parse the custom agent if one is specified.
105 E : if (command_line->HasSwitch("agent")) {
106 E : std::string new_agent_dll = command_line->GetSwitchValueASCII("agent");
107 E : if (new_agent_dll != agent_dll_) {
108 E : agent_dll_ = new_agent_dll;
109 E : LOG(INFO) << "Using custom agent DLL \"" << agent_dll_ << "\".";
110 : }
111 E : }
112 :
113 E : if (agent_dll_.empty()) {
114 i : LOG(ERROR) << "No agent DLL has been specified.";
115 i : return false;
116 : }
117 :
118 E : if (!ParseAdditionalCommandLineArguments(command_line)) {
119 E : LOG(ERROR) << "Unable to parse the additional arguments from the command "
120 : << "line.";
121 E : return false;
122 : }
123 :
124 E : return true;
125 E : }
126 :
127 E : bool InstrumenterWithAgent::Instrument() {
128 E : if (!CreateRelinker())
129 E : return false;
130 :
131 : // Initialize the relinker. This does the decomposition, etc.
132 E : if (!relinker_->Init()) {
133 E : LOG(ERROR) << "Failed to initialize relinker.";
134 E : return false;
135 : }
136 :
137 : // Let the instrumenter implementation set up the relinker and anything else
138 : // that is required.
139 E : if (!InstrumentImpl())
140 i : return false;
141 :
142 : // Do the actual instrumentation by running the relinker.
143 E : if (!relinker_->Relink()) {
144 E : LOG(ERROR) << "Unable to relink input image.";
145 E : return false;
146 : }
147 :
148 E : return true;
149 E : }
150 :
151 : bool InstrumenterWithAgent::ImageFormatIsSupported(
152 E : BlockGraph::ImageFormat image_format) {
153 : // By default we only support PE images.
154 E : if (image_format == BlockGraph::PE_IMAGE)
155 E : return true;
156 i : return false;
157 E : }
158 :
159 E : pe::PETransformPolicy* InstrumenterWithAgent::GetPETransformPolicy() {
160 E : DCHECK_EQ(BlockGraph::PE_IMAGE, image_format_);
161 E : DCHECK(policy_object_.get() == NULL);
162 E : policy_object_.reset(new pe::PETransformPolicy());
163 E : return static_cast<pe::PETransformPolicy*>(policy_object_.get());
164 E : }
165 :
166 E : pe::CoffTransformPolicy* InstrumenterWithAgent::GetCoffTransformPolicy() {
167 E : DCHECK_EQ(BlockGraph::COFF_IMAGE, image_format_);
168 E : DCHECK(policy_object_.get() == NULL);
169 E : policy_object_.reset(new pe::CoffTransformPolicy());
170 E : return static_cast<pe::CoffTransformPolicy*>(policy_object_.get());
171 E : }
172 :
173 E : pe::PERelinker* InstrumenterWithAgent::GetPERelinker() {
174 E : DCHECK_EQ(BlockGraph::PE_IMAGE, image_format_);
175 E : DCHECK(relinker_object_.get() == NULL);
176 E : relinker_object_.reset(new pe::PERelinker(GetPETransformPolicy()));
177 E : return static_cast<pe::PERelinker*>(relinker_object_.get());
178 E : }
179 :
180 E : pe::CoffRelinker* InstrumenterWithAgent::GetCoffRelinker() {
181 E : DCHECK_EQ(BlockGraph::COFF_IMAGE, image_format_);
182 E : relinker_object_.reset(new pe::CoffRelinker(GetCoffTransformPolicy()));
183 E : return static_cast<pe::CoffRelinker*>(relinker_object_.get());
184 E : }
185 :
186 E : bool InstrumenterWithAgent::CreateRelinker() {
187 : // Get the image format by quickly inspecting the image. This logs verbosely
188 : // on failure.
189 E : if (!GetImageFormat(input_image_path_, &image_format_))
190 E : return false;
191 :
192 : // Check if the format is supported and bail if it isn't.
193 E : if (!ImageFormatIsSupported(image_format_)) {
194 i : LOG(ERROR) << "Instrumenter \"" << InstrumentationMode()
195 : << "\" does not support input image format.";
196 i : return false;
197 : }
198 :
199 : // Create and setup an image format specific relinker.
200 E : if (image_format_ == BlockGraph::COFF_IMAGE) {
201 E : pe::CoffRelinker* relinker = GetCoffRelinker();
202 E : DCHECK_NE(reinterpret_cast<pe::CoffRelinker*>(NULL), relinker);
203 E : relinker_ = relinker;
204 E : relinker->set_input_path(input_image_path_);
205 E : relinker->set_output_path(output_image_path_);
206 E : relinker->set_allow_overwrite(allow_overwrite_);
207 E : } else {
208 E : pe::PERelinker* relinker = GetPERelinker();
209 E : DCHECK_NE(reinterpret_cast<pe::PERelinker*>(NULL), relinker);
210 E : relinker_ = relinker;
211 E : relinker->set_input_path(input_image_path_);
212 E : relinker->set_input_pdb_path(input_pdb_path_);
213 E : relinker->set_output_path(output_image_path_);
214 E : relinker->set_output_pdb_path(output_pdb_path_);
215 E : relinker->set_allow_overwrite(allow_overwrite_);
216 E : relinker->set_augment_pdb(!no_augment_pdb_);
217 E : relinker->set_strip_strings(!no_strip_strings_);
218 : }
219 :
220 E : DCHECK_EQ(image_format_, relinker_->image_format());
221 :
222 E : return true;
223 E : }
224 :
225 : } // namespace instrumenters
226 : } // namespace instrument
|