1 : // Copyright 2012 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/pe/transforms/pe_add_imports_transform.h"
16 :
17 : #include "gtest/gtest.h"
18 : #include "syzygy/block_graph/unittest_util.h"
19 : #include "syzygy/core/unittest_util.h"
20 : #include "syzygy/pe/decomposer.h"
21 : #include "syzygy/pe/pe_utils.h"
22 : #include "syzygy/pe/unittest_util.h"
23 :
24 : namespace pe {
25 : namespace transforms {
26 :
27 : using block_graph::BlockGraph;
28 : using core::RelativeAddress;
29 :
30 : namespace {
31 :
32 : class PEAddImportsTransformTest : public testing::PELibUnitTest {
33 : public:
34 : PEAddImportsTransformTest()
35 : : image_layout_(&block_graph_),
36 E : dos_header_block_(NULL) {
37 E : }
38 :
39 E : virtual void SetUp() override {
40 E : testing::PELibUnitTest::SetUp();
41 :
42 : base::FilePath image_path(
43 E : testing::GetExeRelativePath(testing::kTestDllName));
44 :
45 E : ASSERT_TRUE(pe_file_.Init(image_path));
46 :
47 : // Decompose the test image and look at the result.
48 E : Decomposer decomposer(pe_file_);
49 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
50 :
51 : // Retrieve and validate the DOS header.
52 : dos_header_block_ =
53 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
54 E : ASSERT_TRUE(dos_header_block_ != NULL);
55 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
56 E : }
57 :
58 : PEFile pe_file_;
59 : testing::DummyTransformPolicy policy_;
60 : BlockGraph block_graph_;
61 : ImageLayout image_layout_;
62 : BlockGraph::Block* dos_header_block_;
63 : };
64 :
65 : // Given an ImportedModule tests that all of its symbols have been properly
66 : // processed.
67 E : void TestSymbols(const ImportedModule& module) {
68 E : for (size_t i = 0; i < module.size(); ++i) {
69 E : BlockGraph::Reference ref;
70 E : EXPECT_TRUE(module.GetSymbolReference(i, &ref));
71 E : EXPECT_TRUE(ref.referenced() != NULL);
72 E : EXPECT_GE(ref.offset(), 0);
73 : EXPECT_LT(ref.offset(),
74 E : static_cast<BlockGraph::Offset>(ref.referenced()->size()));
75 E : }
76 E : }
77 :
78 : } // namespace
79 :
80 E : TEST_F(PEAddImportsTransformTest, AddImportsExisting) {
81 E : ImportedModule module("export_dll.dll");
82 : size_t function1 = module.AddSymbol("function1",
83 E : ImportedModule::kAlwaysImport);
84 : size_t function3 = module.AddSymbol("function3",
85 E : ImportedModule::kAlwaysImport);
86 E : EXPECT_EQ("function1", module.GetSymbolName(function1));
87 E : EXPECT_EQ("function3", module.GetSymbolName(function3));
88 E : EXPECT_EQ(ImportedModule::kAlwaysImport, module.mode());
89 : EXPECT_EQ(ImportedModule::kAlwaysImport,
90 E : module.GetSymbolMode(function1));
91 : EXPECT_EQ(ImportedModule::kAlwaysImport,
92 E : module.GetSymbolMode(function3));
93 :
94 E : PEAddImportsTransform transform;
95 E : transform.AddModule(&module);
96 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
97 E : &transform, &policy_, &block_graph_, dos_header_block_));
98 E : EXPECT_EQ(0u, transform.modules_added());
99 E : EXPECT_EQ(0u, transform.symbols_added());
100 :
101 E : EXPECT_TRUE(module.ModuleIsImported());
102 E : EXPECT_TRUE(module.SymbolIsImported(function1));
103 E : EXPECT_TRUE(module.SymbolIsImported(function3));
104 :
105 E : EXPECT_FALSE(module.ModuleWasAdded());
106 E : EXPECT_FALSE(module.SymbolWasAdded(function1));
107 E : EXPECT_FALSE(module.SymbolWasAdded(function3));
108 :
109 E : EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
110 E : }
111 :
112 E : TEST_F(PEAddImportsTransformTest, AddImportsExistingDelayLoad) {
113 E : ImportedModule module("ole32.dll");
114 : size_t co_create_guid = module.AddSymbol("CoCreateGuid",
115 E : ImportedModule::kFindOnly);
116 E : EXPECT_EQ("CoCreateGuid", module.GetSymbolName(co_create_guid));
117 E : EXPECT_EQ(ImportedModule::kFindOnly, module.mode());
118 : EXPECT_EQ(ImportedModule::kFindOnly,
119 E : module.GetSymbolMode(co_create_guid));
120 :
121 E : PEAddImportsTransform transform;
122 E : transform.AddModule(&module);
123 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
124 E : &transform, &policy_, &block_graph_, dos_header_block_));
125 E : EXPECT_EQ(0u, transform.modules_added());
126 E : EXPECT_EQ(0u, transform.symbols_added());
127 :
128 E : EXPECT_TRUE(module.ModuleIsImported());
129 E : EXPECT_TRUE(module.SymbolIsImported(co_create_guid));
130 :
131 E : EXPECT_FALSE(module.ModuleWasAdded());
132 E : EXPECT_FALSE(module.SymbolWasAdded(co_create_guid));
133 :
134 E : EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
135 E : }
136 :
137 E : TEST_F(PEAddImportsTransformTest, AddImportsNewSymbol) {
138 E : ImportedModule module("export_dll.dll");
139 : size_t function1 = module.AddSymbol("function1",
140 E : ImportedModule::kAlwaysImport);
141 : size_t function3 = module.AddSymbol("function3",
142 E : ImportedModule::kAlwaysImport);
143 : size_t function4 = module.AddSymbol("function4",
144 E : ImportedModule::kAlwaysImport);
145 E : EXPECT_EQ("function1", module.GetSymbolName(function1));
146 E : EXPECT_EQ("function3", module.GetSymbolName(function3));
147 E : EXPECT_EQ("function4", module.GetSymbolName(function4));
148 E : EXPECT_EQ(ImportedModule::kAlwaysImport, module.mode());
149 : EXPECT_EQ(ImportedModule::kAlwaysImport,
150 E : module.GetSymbolMode(function1));
151 : EXPECT_EQ(ImportedModule::kAlwaysImport,
152 E : module.GetSymbolMode(function3));
153 : EXPECT_EQ(ImportedModule::kAlwaysImport,
154 E : module.GetSymbolMode(function4));
155 :
156 E : PEAddImportsTransform transform;
157 E : transform.AddModule(&module);
158 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
159 E : &transform, &policy_, &block_graph_, dos_header_block_));
160 E : EXPECT_EQ(0u, transform.modules_added());
161 E : EXPECT_EQ(1u, transform.symbols_added());
162 :
163 E : EXPECT_TRUE(module.ModuleIsImported());
164 E : EXPECT_TRUE(module.SymbolIsImported(function1));
165 E : EXPECT_TRUE(module.SymbolIsImported(function3));
166 E : EXPECT_TRUE(module.SymbolIsImported(function4));
167 :
168 E : EXPECT_FALSE(module.ModuleWasAdded());
169 E : EXPECT_FALSE(module.SymbolWasAdded(function1));
170 E : EXPECT_FALSE(module.SymbolWasAdded(function3));
171 E : EXPECT_TRUE(module.SymbolWasAdded(function4));
172 :
173 E : EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
174 :
175 : // TODO(chrisha): Write the image and try to load it!
176 E : }
177 :
178 E : TEST_F(PEAddImportsTransformTest, AddImportsNewModule) {
179 E : ImportedModule module("call_trace_client_rpc.dll");
180 : size_t indirect_penter = module.AddSymbol(
181 E : "_indirect_penter", ImportedModule::kAlwaysImport);
182 : size_t indirect_penter_dllmain = module.AddSymbol(
183 E : "_indirect_penter_dllmain", ImportedModule::kAlwaysImport);
184 : EXPECT_EQ("_indirect_penter",
185 E : module.GetSymbolName(indirect_penter));
186 : EXPECT_EQ("_indirect_penter_dllmain",
187 E : module.GetSymbolName(indirect_penter_dllmain));
188 E : EXPECT_EQ(ImportedModule::kAlwaysImport, module.mode());
189 : EXPECT_EQ(ImportedModule::kAlwaysImport,
190 E : module.GetSymbolMode(indirect_penter));
191 : EXPECT_EQ(ImportedModule::kAlwaysImport,
192 E : module.GetSymbolMode(indirect_penter_dllmain));
193 :
194 E : PEAddImportsTransform transform;
195 E : transform.AddModule(&module);
196 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
197 E : &transform, &policy_, &block_graph_, dos_header_block_));
198 E : EXPECT_EQ(1u, transform.modules_added());
199 E : EXPECT_EQ(2u, transform.symbols_added());
200 :
201 E : EXPECT_TRUE(module.ModuleIsImported());
202 E : EXPECT_TRUE(module.SymbolIsImported(indirect_penter));
203 E : EXPECT_TRUE(module.SymbolIsImported(indirect_penter_dllmain));
204 :
205 E : EXPECT_TRUE(module.ModuleWasAdded());
206 E : EXPECT_TRUE(module.SymbolWasAdded(indirect_penter));
207 E : EXPECT_TRUE(module.SymbolWasAdded(indirect_penter_dllmain));
208 :
209 E : EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
210 :
211 : // TODO(chrisha): Write the image and try to load it!
212 E : }
213 :
214 E : TEST_F(PEAddImportsTransformTest, FindImportsExisting) {
215 E : ImportedModule module("export_dll.dll");
216 : size_t function1 = module.AddSymbol("function1",
217 E : ImportedModule::kFindOnly);
218 : size_t function3 = module.AddSymbol("function3",
219 E : ImportedModule::kFindOnly);
220 E : EXPECT_EQ("function1", module.GetSymbolName(function1));
221 E : EXPECT_EQ("function3", module.GetSymbolName(function3));
222 E : EXPECT_EQ(ImportedModule::kFindOnly, module.mode());
223 E : EXPECT_EQ(ImportedModule::kFindOnly, module.GetSymbolMode(function1));
224 E : EXPECT_EQ(ImportedModule::kFindOnly, module.GetSymbolMode(function3));
225 :
226 E : PEAddImportsTransform transform;
227 E : transform.AddModule(&module);
228 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
229 E : &transform, &policy_, &block_graph_, dos_header_block_));
230 E : EXPECT_EQ(0u, transform.modules_added());
231 E : EXPECT_EQ(0u, transform.symbols_added());
232 :
233 E : EXPECT_TRUE(module.ModuleIsImported());
234 E : EXPECT_TRUE(module.SymbolIsImported(function1));
235 E : EXPECT_TRUE(module.SymbolIsImported(function3));
236 :
237 E : EXPECT_FALSE(module.ModuleWasAdded());
238 E : EXPECT_FALSE(module.SymbolWasAdded(function1));
239 E : EXPECT_FALSE(module.SymbolWasAdded(function3));
240 E : }
241 :
242 E : TEST_F(PEAddImportsTransformTest, FindImportsNewSymbol) {
243 E : ImportedModule module("export_dll.dll");
244 : size_t function1 = module.AddSymbol("function1",
245 E : ImportedModule::kFindOnly);
246 : size_t function3 = module.AddSymbol("function3",
247 E : ImportedModule::kFindOnly);
248 : size_t function4 = module.AddSymbol("function4",
249 E : ImportedModule::kFindOnly);
250 E : EXPECT_EQ("function1", module.GetSymbolName(function1));
251 E : EXPECT_EQ("function3", module.GetSymbolName(function3));
252 E : EXPECT_EQ("function4", module.GetSymbolName(function4));
253 E : EXPECT_EQ(ImportedModule::kFindOnly, module.mode());
254 E : EXPECT_EQ(ImportedModule::kFindOnly, module.GetSymbolMode(function1));
255 E : EXPECT_EQ(ImportedModule::kFindOnly, module.GetSymbolMode(function3));
256 E : EXPECT_EQ(ImportedModule::kFindOnly, module.GetSymbolMode(function4));
257 :
258 E : PEAddImportsTransform transform;
259 E : transform.AddModule(&module);
260 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
261 E : &transform, &policy_, &block_graph_, dos_header_block_));
262 E : EXPECT_EQ(0u, transform.modules_added());
263 E : EXPECT_EQ(0u, transform.symbols_added());
264 :
265 E : EXPECT_TRUE(module.ModuleIsImported());
266 E : EXPECT_TRUE(module.SymbolIsImported(function1));
267 E : EXPECT_TRUE(module.SymbolIsImported(function3));
268 E : EXPECT_FALSE(module.SymbolIsImported(function4));
269 :
270 E : EXPECT_FALSE(module.ModuleWasAdded());
271 E : EXPECT_FALSE(module.SymbolWasAdded(function1));
272 E : EXPECT_FALSE(module.SymbolWasAdded(function3));
273 E : EXPECT_FALSE(module.SymbolWasAdded(function4));
274 E : }
275 :
276 E : TEST_F(PEAddImportsTransformTest, FindImportsNewModule) {
277 E : ImportedModule module("call_trace_client_rpc.dll");
278 : size_t indirect_penter = module.AddSymbol(
279 E : "_indirect_penter", ImportedModule::kFindOnly);
280 : size_t indirect_penter_dllmain = module.AddSymbol(
281 E : "_indirect_penter_dllmain", ImportedModule::kFindOnly);
282 : EXPECT_EQ("_indirect_penter",
283 E : module.GetSymbolName(indirect_penter));
284 : EXPECT_EQ("_indirect_penter_dllmain",
285 E : module.GetSymbolName(indirect_penter_dllmain));
286 E : EXPECT_EQ(ImportedModule::kFindOnly, module.mode());
287 : EXPECT_EQ(ImportedModule::kFindOnly,
288 E : module.GetSymbolMode(indirect_penter));
289 : EXPECT_EQ(ImportedModule::kFindOnly,
290 E : module.GetSymbolMode(indirect_penter_dllmain));
291 :
292 E : PEAddImportsTransform transform;
293 E : transform.AddModule(&module);
294 : EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
295 E : &transform, &policy_, &block_graph_, dos_header_block_));
296 E : EXPECT_EQ(0u, transform.modules_added());
297 E : EXPECT_EQ(0u, transform.symbols_added());
298 :
299 E : EXPECT_FALSE(module.ModuleIsImported());
300 E : EXPECT_FALSE(module.SymbolIsImported(indirect_penter));
301 E : EXPECT_FALSE(module.SymbolIsImported(indirect_penter_dllmain));
302 :
303 E : EXPECT_FALSE(module.ModuleWasAdded());
304 E : EXPECT_FALSE(module.SymbolWasAdded(indirect_penter));
305 E : EXPECT_FALSE(module.SymbolWasAdded(indirect_penter_dllmain));
306 E : }
307 :
308 : } // namespace transforms
309 : } // namespace pe
|