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