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/entry_thunk_instrumenter.h"
16 :
17 : #include "base/command_line.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/core/unittest_util.h"
20 : #include "syzygy/pe/unittest_util.h"
21 :
22 : namespace instrument {
23 : namespace instrumenters {
24 :
25 : namespace {
26 :
27 : class TestEntryThunkInstrumenter : public EntryThunkInstrumenter {
28 : public:
29 : using EntryThunkInstrumenter::agent_dll_;
30 : using EntryThunkInstrumenter::input_image_path_;
31 : using EntryThunkInstrumenter::input_pdb_path_;
32 : using EntryThunkInstrumenter::output_image_path_;
33 : using EntryThunkInstrumenter::output_pdb_path_;
34 : using EntryThunkInstrumenter::allow_overwrite_;
35 : using EntryThunkInstrumenter::old_decomposer_;
36 : using EntryThunkInstrumenter::no_augment_pdb_;
37 : using EntryThunkInstrumenter::no_strip_strings_;
38 : using EntryThunkInstrumenter::instrument_unsafe_references_;
39 : using EntryThunkInstrumenter::module_entry_only_;
40 : using EntryThunkInstrumenter::thunk_imports_;
41 : using EntryThunkInstrumenter::debug_friendly_;
42 : using EntryThunkInstrumenter::instrumentation_mode_;
43 : using EntryThunkInstrumenter::kAgentDllProfile;
44 : using EntryThunkInstrumenter::kAgentDllRpc;
45 : using EntryThunkInstrumenter::InstrumentImpl;
46 : using InstrumenterWithAgent::CreateRelinker;
47 :
48 E : explicit TestEntryThunkInstrumenter(Mode instrumentation_mode)
49 : : EntryThunkInstrumenter(instrumentation_mode) {
50 : // Call the GetPERelinker function to initialize it.
51 E : EXPECT_TRUE(GetPERelinker() != NULL);
52 E : }
53 : };
54 :
55 : class EntryThunkInstrumenterTest : public testing::PELibUnitTest {
56 : public:
57 : typedef testing::PELibUnitTest Super;
58 :
59 E : EntryThunkInstrumenterTest()
60 : : cmd_line_(base::FilePath(L"instrument.exe")) {
61 E : }
62 :
63 E : virtual void SetUp() OVERRIDE {
64 E : testing::Test::SetUp();
65 :
66 : // Several of the tests generate progress and (deliberate) error messages
67 : // that would otherwise clutter the unittest output.
68 E : logging::SetMinLogLevel(logging::LOG_FATAL);
69 :
70 : // Setup the IO streams.
71 E : CreateTemporaryDir(&temp_dir_);
72 E : stdin_path_ = temp_dir_.Append(L"NUL");
73 E : stdout_path_ = temp_dir_.Append(L"stdout.txt");
74 E : stderr_path_ = temp_dir_.Append(L"stderr.txt");
75 E : InitStreams(stdin_path_, stdout_path_, stderr_path_);
76 :
77 : // Initialize the (potential) input and output path values.
78 E : abs_input_image_path_ = testing::GetExeRelativePath(testing::kTestDllName);
79 E : input_image_path_ = testing::GetRelativePath(abs_input_image_path_);
80 E : abs_input_pdb_path_ = testing::GetExeRelativePath(testing::kTestDllPdbName);
81 E : input_pdb_path_ = testing::GetRelativePath(abs_input_pdb_path_);
82 E : output_image_path_ = temp_dir_.Append(input_image_path_.BaseName());
83 E : output_pdb_path_ = temp_dir_.Append(input_pdb_path_.BaseName());
84 E : }
85 :
86 E : void SetUpValidCommandLine() {
87 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
88 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
89 E : }
90 :
91 : protected:
92 : base::FilePath temp_dir_;
93 :
94 : // @name The redirected streams paths.
95 : // @{
96 : base::FilePath stdin_path_;
97 : base::FilePath stdout_path_;
98 : base::FilePath stderr_path_;
99 : // @}
100 :
101 : // @name Command-line and parameters.
102 : // @{
103 : CommandLine cmd_line_;
104 : base::FilePath input_image_path_;
105 : base::FilePath input_pdb_path_;
106 : base::FilePath output_image_path_;
107 : base::FilePath output_pdb_path_;
108 : // @}
109 :
110 : // @name Expected final values of input parameters.
111 : // @{
112 : base::FilePath abs_input_image_path_;
113 : base::FilePath abs_input_pdb_path_;
114 : // @}
115 :
116 : // The fake instrumenter we delegate to.
117 : scoped_ptr<TestEntryThunkInstrumenter> instrumenter_;
118 : };
119 :
120 : } // namespace
121 :
122 E : TEST_F(EntryThunkInstrumenterTest, ParseMinimalCallTrace) {
123 E : SetUpValidCommandLine();
124 : instrumenter_.reset(
125 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::CALL_TRACE));
126 :
127 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
128 :
129 : EXPECT_EQ(EntryThunkInstrumenter::CALL_TRACE,
130 E : instrumenter_->instrumentation_mode_);
131 E : EXPECT_EQ(abs_input_image_path_, instrumenter_->input_image_path_);
132 E : EXPECT_EQ(output_image_path_, instrumenter_->output_image_path_);
133 : EXPECT_EQ(std::string(TestEntryThunkInstrumenter::kAgentDllRpc),
134 E : instrumenter_->agent_dll_);
135 E : EXPECT_FALSE(instrumenter_->allow_overwrite_);
136 E : EXPECT_FALSE(instrumenter_->old_decomposer_);
137 E : EXPECT_FALSE(instrumenter_->no_augment_pdb_);
138 E : EXPECT_FALSE(instrumenter_->no_strip_strings_);
139 E : EXPECT_FALSE(instrumenter_->debug_friendly_);
140 E : EXPECT_FALSE(instrumenter_->thunk_imports_);
141 E : EXPECT_TRUE(instrumenter_->instrument_unsafe_references_);
142 E : EXPECT_FALSE(instrumenter_->module_entry_only_);
143 E : }
144 :
145 E : TEST_F(EntryThunkInstrumenterTest, ParseFullCallTrace) {
146 E : SetUpValidCommandLine();
147 : instrumenter_.reset(
148 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::CALL_TRACE));
149 :
150 E : cmd_line_.AppendSwitchASCII("agent", "foo.dll");
151 E : cmd_line_.AppendSwitch("debug-friendly");
152 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
153 E : cmd_line_.AppendSwitch("old-decomposer");
154 E : cmd_line_.AppendSwitch("no-augment-pdb");
155 E : cmd_line_.AppendSwitch("no-strip-strings");
156 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
157 E : cmd_line_.AppendSwitch("overwrite");
158 E : cmd_line_.AppendSwitch("instrument-imports");
159 E : cmd_line_.AppendSwitch("module-entry-only");
160 E : cmd_line_.AppendSwitch("no-unsafe-refs");
161 :
162 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
163 :
164 : EXPECT_EQ(EntryThunkInstrumenter::CALL_TRACE,
165 E : instrumenter_->instrumentation_mode_);
166 E : EXPECT_EQ(abs_input_image_path_, instrumenter_->input_image_path_);
167 E : EXPECT_EQ(output_image_path_, instrumenter_->output_image_path_);
168 E : EXPECT_EQ(abs_input_pdb_path_, instrumenter_->input_pdb_path_);
169 E : EXPECT_EQ(output_pdb_path_, instrumenter_->output_pdb_path_);
170 E : EXPECT_EQ(std::string("foo.dll"), instrumenter_->agent_dll_);
171 E : EXPECT_TRUE(instrumenter_->allow_overwrite_);
172 E : EXPECT_TRUE(instrumenter_->old_decomposer_);
173 E : EXPECT_TRUE(instrumenter_->no_augment_pdb_);
174 E : EXPECT_TRUE(instrumenter_->no_strip_strings_);
175 E : EXPECT_TRUE(instrumenter_->debug_friendly_);
176 E : EXPECT_TRUE(instrumenter_->thunk_imports_);
177 E : EXPECT_FALSE(instrumenter_->instrument_unsafe_references_);
178 E : EXPECT_TRUE(instrumenter_->module_entry_only_);
179 E : }
180 :
181 E : TEST_F(EntryThunkInstrumenterTest, ParseMinimalProfile) {
182 E : SetUpValidCommandLine();
183 : instrumenter_.reset(
184 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::PROFILE));
185 :
186 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
187 :
188 : EXPECT_EQ(EntryThunkInstrumenter::PROFILE,
189 E : instrumenter_->instrumentation_mode_);
190 E : EXPECT_EQ(abs_input_image_path_, instrumenter_->input_image_path_);
191 E : EXPECT_EQ(output_image_path_, instrumenter_->output_image_path_);
192 :
193 : EXPECT_EQ(std::string(TestEntryThunkInstrumenter::kAgentDllProfile),
194 E : instrumenter_->agent_dll_);
195 :
196 E : EXPECT_FALSE(instrumenter_->allow_overwrite_);
197 E : EXPECT_FALSE(instrumenter_->old_decomposer_);
198 E : EXPECT_FALSE(instrumenter_->no_augment_pdb_);
199 E : EXPECT_FALSE(instrumenter_->no_strip_strings_);
200 E : EXPECT_FALSE(instrumenter_->debug_friendly_);
201 E : EXPECT_FALSE(instrumenter_->thunk_imports_);
202 E : EXPECT_FALSE(instrumenter_->instrument_unsafe_references_);
203 E : EXPECT_FALSE(instrumenter_->module_entry_only_);
204 E : }
205 :
206 E : TEST_F(EntryThunkInstrumenterTest, ParseFullProfile) {
207 E : SetUpValidCommandLine();
208 : instrumenter_.reset(
209 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::PROFILE));
210 E : cmd_line_.AppendSwitchASCII("agent", "foo.dll");
211 E : cmd_line_.AppendSwitch("debug-friendly");
212 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
213 E : cmd_line_.AppendSwitch("old-decomposer");
214 E : cmd_line_.AppendSwitch("no-augment-pdb");
215 E : cmd_line_.AppendSwitch("no-strip-strings");
216 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
217 E : cmd_line_.AppendSwitch("overwrite");
218 E : cmd_line_.AppendSwitch("instrument-imports");
219 :
220 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
221 :
222 : EXPECT_EQ(EntryThunkInstrumenter::PROFILE,
223 E : instrumenter_->instrumentation_mode_);
224 E : EXPECT_EQ(abs_input_image_path_, instrumenter_->input_image_path_);
225 E : EXPECT_EQ(output_image_path_, instrumenter_->output_image_path_);
226 E : EXPECT_EQ(abs_input_pdb_path_, instrumenter_->input_pdb_path_);
227 E : EXPECT_EQ(output_pdb_path_, instrumenter_->output_pdb_path_);
228 E : EXPECT_EQ(std::string("foo.dll"), instrumenter_->agent_dll_);
229 E : EXPECT_TRUE(instrumenter_->allow_overwrite_);
230 E : EXPECT_TRUE(instrumenter_->old_decomposer_);
231 E : EXPECT_TRUE(instrumenter_->no_augment_pdb_);
232 E : EXPECT_TRUE(instrumenter_->no_strip_strings_);
233 E : EXPECT_TRUE(instrumenter_->debug_friendly_);
234 E : EXPECT_TRUE(instrumenter_->thunk_imports_);
235 E : }
236 :
237 E : TEST_F(EntryThunkInstrumenterTest, InstrumentImplCallTrace) {
238 E : SetUpValidCommandLine();
239 : instrumenter_.reset(
240 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::CALL_TRACE));
241 :
242 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
243 E : EXPECT_TRUE(instrumenter_->CreateRelinker());
244 E : EXPECT_TRUE(instrumenter_->InstrumentImpl());
245 E : }
246 :
247 E : TEST_F(EntryThunkInstrumenterTest, InstrumentImplProfile) {
248 E : SetUpValidCommandLine();
249 : instrumenter_.reset(
250 E : new TestEntryThunkInstrumenter(EntryThunkInstrumenter::PROFILE));
251 :
252 E : EXPECT_TRUE(instrumenter_->ParseCommandLine(&cmd_line_));
253 E : EXPECT_TRUE(instrumenter_->CreateRelinker());
254 E : EXPECT_TRUE(instrumenter_->InstrumentImpl());
255 E : }
256 :
257 : } // namespace instrumenters
258 : } // namespace instrument
|