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