1 : // Copyright 2014 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/swapimport/swapimport_app.h"
16 :
17 : #include "base/files/file_util.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/core/unittest_util.h"
20 : #include "syzygy/pe/unittest_util.h"
21 :
22 : namespace swapimport {
23 :
24 : namespace {
25 :
26 : class TestSwapImportApp : public SwapImportApp {
27 : public:
28 : using SwapImportApp::import_name_;
29 : using SwapImportApp::input_image_;
30 : using SwapImportApp::output_image_;
31 : using SwapImportApp::overwrite_;
32 : using SwapImportApp::verbose_;
33 : };
34 :
35 : typedef application::Application<TestSwapImportApp> TestApp;
36 :
37 : class SwapImportAppTest : public testing::PELibUnitTest {
38 : public:
39 : typedef testing::PELibUnitTest Super;
40 :
41 : SwapImportAppTest()
42 : : cmd_line_(base::FilePath(L"swapimport.exe")),
43 E : test_impl_(test_app_.implementation()) {
44 E : }
45 :
46 E : void SetUp() {
47 E : Super::SetUp();
48 :
49 E : logging::SetMinLogLevel(logging::LOG_ERROR);
50 :
51 : // Setup the IO streams.
52 E : CreateTemporaryDir(&temp_dir_);
53 E : stdin_path_ = temp_dir_.Append(L"NUL");
54 E : stdout_path_ = temp_dir_.Append(L"stdout.txt");
55 E : stderr_path_ = temp_dir_.Append(L"stderr.txt");
56 E : InitStreams(stdin_path_, stdout_path_, stderr_path_);
57 :
58 E : ASSERT_NO_FATAL_FAILURE(ConfigureTestApp(&test_app_));
59 :
60 E : input_image_ = testing::GetOutputRelativePath(testing::kTestDllName);
61 E : output_image_ = temp_dir_.Append(testing::kTestDllName);
62 :
63 E : input_image_64_ = testing::GetOutputRelativePath(testing::kTestDllName64);
64 E : output_image_64_ = temp_dir_.Append(testing::kTestDllName64);
65 E : }
66 :
67 : template <class PEFileType>
68 : void ValidateFirstImport(const base::FilePath& image,
69 E : const char* import_name) {
70 E : PEFileType pe_file;
71 E : ASSERT_TRUE(pe_file.Init(image));
72 :
73 : // Get the absolute address of the first import entry.
74 : PEFileType::AbsoluteAddress iid_addr(
75 : pe_file.nt_headers()->OptionalHeader.ImageBase +
76 : pe_file.nt_headers()->OptionalHeader.DataDirectory[
77 E : IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
78 :
79 : // Read the import entry.
80 E : IMAGE_IMPORT_DESCRIPTOR iid = {};
81 E : ASSERT_TRUE(pe_file.ReadImage(iid_addr, &iid, sizeof(iid)));
82 :
83 : // Read the name of the import entry.
84 : PEFileType::AbsoluteAddress name_addr(
85 E : pe_file.nt_headers()->OptionalHeader.ImageBase + iid.Name);
86 E : std::string name;
87 E : ASSERT_TRUE(pe_file.ReadImageString(name_addr, &name));
88 E : ASSERT_EQ(0, base::CompareCaseInsensitiveASCII(name.c_str(), import_name));
89 E : }
90 :
91 E : void ValidateImportsSwapped() {
92 E : ASSERT_NO_FATAL_FAILURE(ValidateFirstImport<pe::PEFile>(input_image_,
93 : "export_dll.dll"));
94 E : ASSERT_NO_FATAL_FAILURE(ValidateFirstImport<pe::PEFile>(output_image_,
95 : "kernel32.dll"));
96 E : }
97 :
98 E : void ValidateImportsSwapped64() {
99 E : ASSERT_NO_FATAL_FAILURE(ValidateFirstImport<pe::PEFile64>(input_image_64_,
100 : "user32.dll"));
101 E : ASSERT_NO_FATAL_FAILURE(ValidateFirstImport<pe::PEFile64>(output_image_64_,
102 : "kernel32.dll"));
103 E : }
104 :
105 : // Points the application at the fixture's command-line and IO streams.
106 : template<typename TestAppType>
107 E : void ConfigureTestApp(TestAppType* test_app) {
108 E : test_app->set_command_line(&cmd_line_);
109 E : test_app->set_in(in());
110 E : test_app->set_out(out());
111 E : test_app->set_err(err());
112 E : }
113 :
114 : // Stashes the current log-level before each test instance and restores it
115 : // after each test completes.
116 : testing::ScopedLogLevelSaver log_level_saver;
117 :
118 : // @name The application under test.
119 : // @{
120 : TestApp test_app_;
121 : TestApp::Implementation& test_impl_;
122 : base::FilePath temp_dir_;
123 : base::FilePath stdin_path_;
124 : base::FilePath stdout_path_;
125 : base::FilePath stderr_path_;
126 : // @}
127 :
128 : base::CommandLine cmd_line_;
129 : base::FilePath input_image_;
130 : base::FilePath output_image_;
131 : base::FilePath input_image_64_;
132 : base::FilePath output_image_64_;
133 : };
134 :
135 : } // namespace
136 :
137 E : TEST_F(SwapImportAppTest, GetHelp) {
138 E : cmd_line_.AppendSwitch("help");
139 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
140 E : }
141 :
142 E : TEST_F(SwapImportAppTest, ParseEmptyCommandLineFails) {
143 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
144 E : }
145 :
146 E : TEST_F(SwapImportAppTest, ParseEmptyCommandFailsNoInputImage) {
147 E : cmd_line_.AppendSwitchPath("output-image", input_image_);
148 E : cmd_line_.AppendArg("kernel32.dll");
149 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
150 E : }
151 :
152 E : TEST_F(SwapImportAppTest, ParseEmptyCommandFailsNoOutputImage) {
153 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
154 E : cmd_line_.AppendArg("kernel32.dll");
155 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
156 E : }
157 :
158 E : TEST_F(SwapImportAppTest, ParseEmptyCommandFailsNoImportName) {
159 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
160 E : cmd_line_.AppendSwitchPath("output-image", input_image_);
161 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
162 E : }
163 :
164 E : TEST_F(SwapImportAppTest, ParseEmptyCommandMinimal) {
165 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
166 E : cmd_line_.AppendSwitchPath("output-image", input_image_);
167 E : cmd_line_.AppendArg("kernel32.dll");
168 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
169 E : }
170 :
171 E : TEST_F(SwapImportAppTest, ParseEmptyCommandMaximal) {
172 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
173 E : cmd_line_.AppendSwitchPath("output-image", input_image_);
174 E : cmd_line_.AppendSwitch("overwrite");
175 E : cmd_line_.AppendSwitch("verbose");
176 E : cmd_line_.AppendArg("kernel32.dll");
177 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
178 E : }
179 :
180 E : TEST_F(SwapImportAppTest, RunFailsInputAndOutputSame) {
181 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
182 E : cmd_line_.AppendSwitchPath("output-image", input_image_);
183 E : cmd_line_.AppendArg("kernel32.dll");
184 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
185 E : EXPECT_NE(0, test_impl_.Run());
186 E : }
187 :
188 E : TEST_F(SwapImportAppTest, RunFailsOutputExists) {
189 E : base::WriteFile(output_image_, "a", 1);
190 :
191 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
192 E : cmd_line_.AppendSwitchPath("output-image", output_image_);
193 E : cmd_line_.AppendArg("kernel32.dll");
194 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
195 E : EXPECT_NE(0, test_impl_.Run());
196 E : }
197 :
198 E : TEST_F(SwapImportAppTest, RunFailsImportNameNotMatched) {
199 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
200 E : cmd_line_.AppendSwitchPath("output-image", output_image_);
201 E : cmd_line_.AppendArg("nosuchimport.dll");
202 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
203 E : EXPECT_NE(0, test_impl_.Run());
204 E : }
205 :
206 E : TEST_F(SwapImportAppTest, RunSucceeds) {
207 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
208 E : cmd_line_.AppendSwitchPath("output-image", output_image_);
209 E : cmd_line_.AppendArg("kernel32.dll");
210 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
211 E : EXPECT_EQ(0, test_impl_.Run());
212 :
213 E : ASSERT_NO_FATAL_FAILURE(ValidateImportsSwapped());
214 E : }
215 :
216 E : TEST_F(SwapImportAppTest, RunSucceedsOverwrite) {
217 E : base::WriteFile(output_image_, "a", 1);
218 :
219 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
220 E : cmd_line_.AppendSwitchPath("output-image", output_image_);
221 E : cmd_line_.AppendSwitch("overwrite");
222 E : cmd_line_.AppendArg("kernel32.dll");
223 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
224 E : EXPECT_EQ(0, test_impl_.Run());
225 :
226 E : ASSERT_NO_FATAL_FAILURE(ValidateImportsSwapped());
227 E : }
228 :
229 E : TEST_F(SwapImportAppTest, RunSucceeds64) {
230 E : cmd_line_.AppendSwitch("x64");
231 E : cmd_line_.AppendSwitchPath("input-image", input_image_64_);
232 E : cmd_line_.AppendSwitchPath("output-image", output_image_64_);
233 E : cmd_line_.AppendArg("kernel32.dll");
234 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
235 E : EXPECT_EQ(0, test_impl_.Run());
236 :
237 E : ASSERT_NO_FATAL_FAILURE(ValidateImportsSwapped64());
238 E : }
239 :
240 E : TEST_F(SwapImportAppTest, Run64On32BitBinaryFails) {
241 E : cmd_line_.AppendSwitch("x64");
242 E : cmd_line_.AppendSwitchPath("input-image", input_image_);
243 E : cmd_line_.AppendSwitchPath("output-image", output_image_);
244 E : cmd_line_.AppendArg("kernel32.dll");
245 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
246 E : EXPECT_NE(0, test_impl_.Run());
247 E : EXPECT_FALSE(base::PathExists(output_image_));
248 E : }
249 :
250 E : TEST_F(SwapImportAppTest, Run32On64BitBinaryFails) {
251 E : cmd_line_.AppendSwitchPath("input-image", input_image_64_);
252 E : cmd_line_.AppendSwitchPath("output-image", output_image_64_);
253 E : cmd_line_.AppendArg("kernel32.dll");
254 E : ASSERT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
255 E : EXPECT_NE(0, test_impl_.Run());
256 E : EXPECT_FALSE(base::PathExists(output_image_64_));
257 E : }
258 :
259 : } // namespace swapimport
|