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/command_line.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/core/unittest_util.h"
21 : #include "syzygy/pe/unittest_util.h"
22 :
23 : namespace instrument {
24 : namespace instrumenters {
25 :
26 : namespace {
27 :
28 : using testing::StrictMock;
29 : using testing::Return;
30 : using testing::_;
31 :
32 : const char kTestAgentDllName[] = "test_agent_dll.dll";
33 :
34 : class MockPERelinker : public pe::PERelinker {
35 : public:
36 E : MockPERelinker() : pe::PERelinker(&policy_) {
37 E : }
38 :
39 E : MOCK_METHOD0(Init, bool());
40 E : MOCK_METHOD0(Relink, bool());
41 :
42 : private:
43 : pe::PETransformPolicy policy_;
44 : };
45 :
46 : class MockCoffRelinker : public pe::CoffRelinker {
47 : public:
48 E : MockCoffRelinker() : pe::CoffRelinker(&policy_) {
49 E : }
50 :
51 E : MOCK_METHOD0(Init, bool());
52 E : MOCK_METHOD0(Relink, bool());
53 :
54 : private:
55 : pe::CoffTransformPolicy policy_;
56 : };
57 :
58 : class TestInstrumenterWithAgent : public InstrumenterWithAgent {
59 : public:
60 : using InstrumenterWithAgent::input_image_path_;
61 : using InstrumenterWithAgent::input_pdb_path_;
62 : using InstrumenterWithAgent::output_image_path_;
63 : using InstrumenterWithAgent::output_pdb_path_;
64 : using InstrumenterWithAgent::allow_overwrite_;
65 : using InstrumenterWithAgent::old_decomposer_;
66 : using InstrumenterWithAgent::no_augment_pdb_;
67 : using InstrumenterWithAgent::no_strip_strings_;
68 :
69 E : TestInstrumenterWithAgent() {
70 E : agent_dll_ = kTestAgentDllName;
71 E : }
72 :
73 : // For the purposes of testing, our instrumenter supports all image formats.
74 E : virtual bool ImageFormatIsSupported(pe::ImageFormat image_format) OVERRIDE {
75 E : return true;
76 E : }
77 :
78 E : MOCK_METHOD0(InstrumentImpl, bool());
79 :
80 E : pe::PERelinker* GetPERelinker() OVERRIDE {
81 E : return &mock_pe_relinker_;
82 E : }
83 :
84 E : pe::CoffRelinker* GetCoffRelinker() OVERRIDE {
85 E : return &mock_coff_relinker_;
86 E : }
87 :
88 i : virtual const char* InstrumentationMode() OVERRIDE { return "test"; }
89 :
90 : StrictMock<MockPERelinker> mock_pe_relinker_;
91 : StrictMock<MockCoffRelinker> mock_coff_relinker_;
92 : };
93 :
94 : class InstrumenterWithAgentTest : public testing::PELibUnitTest {
95 : public:
96 : typedef testing::PELibUnitTest Super;
97 :
98 E : InstrumenterWithAgentTest()
99 : : cmd_line_(base::FilePath(L"instrument.exe")) {
100 E : }
101 :
102 E : virtual void SetUp() OVERRIDE {
103 E : testing::Test::SetUp();
104 :
105 : // Several of the tests generate progress and (deliberate) error messages
106 : // that would otherwise clutter the unittest output.
107 E : logging::SetMinLogLevel(logging::LOG_FATAL);
108 :
109 : // Setup the IO streams.
110 E : CreateTemporaryDir(&temp_dir_);
111 E : stdin_path_ = temp_dir_.Append(L"NUL");
112 E : stdout_path_ = temp_dir_.Append(L"stdout.txt");
113 E : stderr_path_ = temp_dir_.Append(L"stderr.txt");
114 E : InitStreams(stdin_path_, stdout_path_, stderr_path_);
115 :
116 : // Initialize the (potential) input and output path values.
117 : abs_input_pe_image_path_ = testing::GetExeRelativePath(
118 E : testing::kTestDllName);
119 E : input_pe_image_path_ = testing::GetRelativePath(abs_input_pe_image_path_);
120 E : abs_input_pdb_path_ = testing::GetExeRelativePath(testing::kTestDllPdbName);
121 E : input_pdb_path_ = testing::GetRelativePath(abs_input_pdb_path_);
122 E : output_pe_image_path_ = temp_dir_.Append(input_pe_image_path_.BaseName());
123 E : output_pdb_path_ = temp_dir_.Append(input_pdb_path_.BaseName());
124 :
125 : abs_input_coff_image_path_ = testing::GetExeTestDataRelativePath(
126 E : testing::kTestDllCoffObjName);
127 : input_coff_image_path_ = testing::GetRelativePath(
128 E : abs_input_coff_image_path_);
129 E : output_coff_image_path_ = temp_dir_.Append(testing::kTestDllCoffObjName);
130 E : }
131 :
132 E : void SetUpValidCommandLinePE() {
133 E : cmd_line_.AppendSwitchPath("input-image", input_pe_image_path_);
134 E : cmd_line_.AppendSwitchPath("output-image", output_pe_image_path_);
135 E : }
136 :
137 E : void SetUpValidCommandLineCoff() {
138 E : cmd_line_.AppendSwitchPath("input-image", input_coff_image_path_);
139 E : cmd_line_.AppendSwitchPath("output-image", output_coff_image_path_);
140 E : }
141 :
142 : protected:
143 : base::FilePath temp_dir_;
144 :
145 : // @name The redirected streams paths.
146 : // @{
147 : base::FilePath stdin_path_;
148 : base::FilePath stdout_path_;
149 : base::FilePath stderr_path_;
150 : // @}
151 :
152 : // @name Command-line and parameters.
153 : // @{
154 : CommandLine cmd_line_;
155 : base::FilePath input_pe_image_path_;
156 : base::FilePath input_pdb_path_;
157 : base::FilePath output_pe_image_path_;
158 : base::FilePath output_pdb_path_;
159 : base::FilePath input_coff_image_path_;
160 : base::FilePath output_coff_image_path_;
161 : // @}
162 :
163 : // @name Expected final values of input parameters.
164 : // @{
165 : base::FilePath abs_input_pe_image_path_;
166 : base::FilePath abs_input_pdb_path_;
167 : base::FilePath abs_input_coff_image_path_;
168 : // @}
169 : };
170 :
171 : } // namespace
172 :
173 E : TEST_F(InstrumenterWithAgentTest, EmptyCommandLineFails) {
174 E : TestInstrumenterWithAgent instrumenter;
175 E : ASSERT_FALSE(instrumenter.ParseCommandLine(&cmd_line_));
176 E : }
177 :
178 E : TEST_F(InstrumenterWithAgentTest, ParseWithNoInputImageFails) {
179 E : cmd_line_.AppendSwitchPath("output-image", output_pe_image_path_);
180 :
181 E : TestInstrumenterWithAgent instrumenter;
182 E : ASSERT_FALSE(instrumenter.ParseCommandLine(&cmd_line_));
183 E : }
184 :
185 E : TEST_F(InstrumenterWithAgentTest, ParseWithNoOutputImageFails) {
186 E : cmd_line_.AppendSwitchPath("input-image", input_pe_image_path_);
187 :
188 E : TestInstrumenterWithAgent instrumenter;
189 E : ASSERT_FALSE(instrumenter.ParseCommandLine(&cmd_line_));
190 E : }
191 :
192 E : TEST_F(InstrumenterWithAgentTest, ParseInputImages) {
193 E : cmd_line_.AppendSwitchPath("input-image", input_pe_image_path_);
194 E : cmd_line_.AppendSwitchPath("output-image", output_pe_image_path_);
195 :
196 E : TestInstrumenterWithAgent instrumenter;
197 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
198 E : EXPECT_EQ(abs_input_pe_image_path_, instrumenter.input_image_path_);
199 E : EXPECT_EQ(output_pe_image_path_, instrumenter.output_image_path_);
200 :
201 E : EXPECT_FALSE(instrumenter.allow_overwrite_);
202 E : EXPECT_FALSE(instrumenter.old_decomposer_);
203 E : EXPECT_FALSE(instrumenter.no_augment_pdb_);
204 E : EXPECT_FALSE(instrumenter.no_strip_strings_);
205 E : }
206 :
207 E : TEST_F(InstrumenterWithAgentTest, agent_dll) {
208 E : TestInstrumenterWithAgent instrumenter;
209 E : EXPECT_STREQ(kTestAgentDllName, instrumenter.agent_dll().c_str());
210 E : }
211 :
212 E : TEST_F(InstrumenterWithAgentTest, InstrumentPE) {
213 E : SetUpValidCommandLinePE();
214 :
215 E : TestInstrumenterWithAgent instrumenter;
216 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
217 E : EXPECT_CALL(instrumenter.mock_pe_relinker_, Init()).WillOnce(Return(true));
218 E : EXPECT_CALL(instrumenter.mock_pe_relinker_, Relink()).WillOnce(Return(true));
219 E : EXPECT_CALL(instrumenter, InstrumentImpl()).WillOnce(Return(true));
220 :
221 E : EXPECT_TRUE(instrumenter.Instrument());
222 E : }
223 :
224 E : TEST_F(InstrumenterWithAgentTest, InstrumentCoff) {
225 E : SetUpValidCommandLineCoff();
226 :
227 E : TestInstrumenterWithAgent instrumenter;
228 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
229 E : EXPECT_CALL(instrumenter.mock_coff_relinker_, Init()).WillOnce(Return(true));
230 : EXPECT_CALL(instrumenter.mock_coff_relinker_, Relink()).WillOnce(
231 E : Return(true));
232 E : EXPECT_CALL(instrumenter, InstrumentImpl()).WillOnce(Return(true));
233 :
234 E : EXPECT_TRUE(instrumenter.Instrument());
235 E : }
236 :
237 E : TEST_F(InstrumenterWithAgentTest, InstrumentFailsInitPE) {
238 E : TestInstrumenterWithAgent instrumenter;
239 E : SetUpValidCommandLinePE();
240 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
241 :
242 E : EXPECT_CALL(instrumenter.mock_pe_relinker_, Init()).WillOnce(Return(false));
243 :
244 E : EXPECT_FALSE(instrumenter.Instrument());
245 E : }
246 :
247 E : TEST_F(InstrumenterWithAgentTest, InstrumentFailsInitCoff) {
248 E : TestInstrumenterWithAgent instrumenter;
249 E : SetUpValidCommandLineCoff();
250 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
251 :
252 E : EXPECT_CALL(instrumenter.mock_coff_relinker_, Init()).WillOnce(Return(false));
253 :
254 E : EXPECT_FALSE(instrumenter.Instrument());
255 E : }
256 :
257 E : TEST_F(InstrumenterWithAgentTest, InstrumentFailsRelinkPE) {
258 E : TestInstrumenterWithAgent instrumenter;
259 E : SetUpValidCommandLinePE();
260 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
261 :
262 E : EXPECT_CALL(instrumenter.mock_pe_relinker_, Init()).WillOnce(Return(true));
263 : EXPECT_CALL(instrumenter.mock_pe_relinker_, Relink()).WillOnce(
264 E : Return(false));
265 E : EXPECT_CALL(instrumenter, InstrumentImpl()).WillOnce(Return(true));
266 :
267 E : EXPECT_FALSE(instrumenter.Instrument());
268 E : }
269 :
270 E : TEST_F(InstrumenterWithAgentTest, InstrumentFailsRelinkCoff) {
271 E : TestInstrumenterWithAgent instrumenter;
272 E : SetUpValidCommandLineCoff();
273 E : EXPECT_TRUE(instrumenter.ParseCommandLine(&cmd_line_));
274 :
275 E : EXPECT_CALL(instrumenter.mock_coff_relinker_, Init()).WillOnce(Return(true));
276 : EXPECT_CALL(instrumenter.mock_coff_relinker_, Relink()).WillOnce(
277 E : Return(false));
278 E : EXPECT_CALL(instrumenter, InstrumentImpl()).WillOnce(Return(true));
279 :
280 E : EXPECT_FALSE(instrumenter.Instrument());
281 E : }
282 :
283 : } // namespace instrumenters
284 : } // namespace instrument
|