1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 // This file contains unit tests for PEImage.
6
7 #include "testing/gtest/include/gtest/gtest.h"
8 #include "base/win/pe_image.h"
9 #include "base/win/windows_version.h"
10
11 namespace base {
12 namespace win {
13
14 // Just counts the number of invocations.
ExportsCallback(const PEImage & image,DWORD ordinal,DWORD hint,LPCSTR name,PVOID function,LPCSTR forward,PVOID cookie)15 bool ExportsCallback(const PEImage &image,
16 DWORD ordinal,
17 DWORD hint,
18 LPCSTR name,
19 PVOID function,
20 LPCSTR forward,
21 PVOID cookie) {
22 int* count = reinterpret_cast<int*>(cookie);
23 (*count)++;
24 return true;
25 }
26
27 // Just counts the number of invocations.
ImportsCallback(const PEImage & image,LPCSTR module,DWORD ordinal,LPCSTR name,DWORD hint,PIMAGE_THUNK_DATA iat,PVOID cookie)28 bool ImportsCallback(const PEImage &image,
29 LPCSTR module,
30 DWORD ordinal,
31 LPCSTR name,
32 DWORD hint,
33 PIMAGE_THUNK_DATA iat,
34 PVOID cookie) {
35 int* count = reinterpret_cast<int*>(cookie);
36 (*count)++;
37 return true;
38 }
39
40 // Just counts the number of invocations.
SectionsCallback(const PEImage & image,PIMAGE_SECTION_HEADER header,PVOID section_start,DWORD section_size,PVOID cookie)41 bool SectionsCallback(const PEImage &image,
42 PIMAGE_SECTION_HEADER header,
43 PVOID section_start,
44 DWORD section_size,
45 PVOID cookie) {
46 int* count = reinterpret_cast<int*>(cookie);
47 (*count)++;
48 return true;
49 }
50
51 // Just counts the number of invocations.
RelocsCallback(const PEImage & image,WORD type,PVOID address,PVOID cookie)52 bool RelocsCallback(const PEImage &image,
53 WORD type,
54 PVOID address,
55 PVOID cookie) {
56 int* count = reinterpret_cast<int*>(cookie);
57 (*count)++;
58 return true;
59 }
60
61 // Just counts the number of invocations.
ImportChunksCallback(const PEImage & image,LPCSTR module,PIMAGE_THUNK_DATA name_table,PIMAGE_THUNK_DATA iat,PVOID cookie)62 bool ImportChunksCallback(const PEImage &image,
63 LPCSTR module,
64 PIMAGE_THUNK_DATA name_table,
65 PIMAGE_THUNK_DATA iat,
66 PVOID cookie) {
67 int* count = reinterpret_cast<int*>(cookie);
68 (*count)++;
69 return true;
70 }
71
72 // Just counts the number of invocations.
DelayImportChunksCallback(const PEImage & image,PImgDelayDescr delay_descriptor,LPCSTR module,PIMAGE_THUNK_DATA name_table,PIMAGE_THUNK_DATA iat,PIMAGE_THUNK_DATA bound_iat,PIMAGE_THUNK_DATA unload_iat,PVOID cookie)73 bool DelayImportChunksCallback(const PEImage &image,
74 PImgDelayDescr delay_descriptor,
75 LPCSTR module,
76 PIMAGE_THUNK_DATA name_table,
77 PIMAGE_THUNK_DATA iat,
78 PIMAGE_THUNK_DATA bound_iat,
79 PIMAGE_THUNK_DATA unload_iat,
80 PVOID cookie) {
81 int* count = reinterpret_cast<int*>(cookie);
82 (*count)++;
83 return true;
84 }
85
86 // Identifiers for the set of supported expectations.
87 enum ExpectationSet {
88 WIN_2K_SET,
89 WIN_XP_SET,
90 WIN_VISTA_SET,
91 WIN_7_SET,
92 WIN_8_SET,
93 UNSUPPORTED_SET,
94 };
95
96 // We'll be using some known values for the tests.
97 enum Value {
98 sections = 0,
99 imports_dlls,
100 delay_dlls,
101 exports,
102 imports,
103 delay_imports,
104 relocs
105 };
106
GetExpectationSet(DWORD os)107 ExpectationSet GetExpectationSet(DWORD os) {
108 if (os == 50)
109 return WIN_2K_SET;
110 if (os == 51)
111 return WIN_XP_SET;
112 if (os == 60)
113 return WIN_VISTA_SET;
114 if (os == 61)
115 return WIN_7_SET;
116 if (os >= 62)
117 return WIN_8_SET;
118 return UNSUPPORTED_SET;
119 }
120
121 // Retrieves the expected value from advapi32.dll based on the OS.
GetExpectedValue(Value value,DWORD os)122 int GetExpectedValue(Value value, DWORD os) {
123 const int xp_delay_dlls = 2;
124 const int xp_exports = 675;
125 const int xp_imports = 422;
126 const int xp_delay_imports = 8;
127 const int xp_relocs = 9180;
128 const int vista_delay_dlls = 4;
129 const int vista_exports = 799;
130 const int vista_imports = 476;
131 const int vista_delay_imports = 24;
132 const int vista_relocs = 10188;
133 const int w2k_delay_dlls = 0;
134 const int w2k_exports = 566;
135 const int w2k_imports = 357;
136 const int w2k_delay_imports = 0;
137 const int w2k_relocs = 7388;
138 const int win7_delay_dlls = 7;
139 const int win7_exports = 806;
140 const int win7_imports = 568;
141 const int win7_delay_imports = 71;
142 int win7_relocs = 7812;
143 int win7_sections = 4;
144 const int win8_delay_dlls = 9;
145 const int win8_exports = 806;
146 const int win8_imports = 568;
147 const int win8_delay_imports = 113;
148 const int win8_relocs = 9478;
149 int win8_sections = 4;
150 int win8_import_dlls = 17;
151
152 base::win::OSInfo* os_info = base::win::OSInfo::GetInstance();
153 // 32-bit process on a 32-bit system.
154 if (os_info->architecture() == base::win::OSInfo::X86_ARCHITECTURE) {
155 win8_sections = 5;
156 win8_import_dlls = 19;
157
158 // 64-bit process on a 64-bit system.
159 } else if (os_info->wow64_status() == base::win::OSInfo::WOW64_DISABLED) {
160 win7_sections = 6;
161 win7_relocs = 2712;
162 }
163
164 // Contains the expected value, for each enumerated property (Value), and the
165 // OS version: [Value][os_version]
166 const int expected[][5] = {
167 {4, 4, 4, win7_sections, win8_sections},
168 {3, 3, 3, 13, win8_import_dlls},
169 {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls,
170 win8_delay_dlls},
171 {w2k_exports, xp_exports, vista_exports, win7_exports, win8_exports},
172 {w2k_imports, xp_imports, vista_imports, win7_imports, win8_imports},
173 {w2k_delay_imports, xp_delay_imports,
174 vista_delay_imports, win7_delay_imports, win8_delay_imports},
175 {w2k_relocs, xp_relocs, vista_relocs, win7_relocs, win8_relocs}
176 };
177 COMPILE_ASSERT(arraysize(expected[0]) == UNSUPPORTED_SET,
178 expected_value_set_mismatch);
179
180 if (value > relocs)
181 return 0;
182 ExpectationSet expected_set = GetExpectationSet(os);
183 if (expected_set >= arraysize(expected)) {
184 // This should never happen. Log a failure if it does.
185 EXPECT_NE(UNSUPPORTED_SET, expected_set);
186 expected_set = WIN_2K_SET;
187 }
188
189 return expected[value][expected_set];
190 }
191
192
193 // TODO(jschuh): crbug.com/167707 Need to fix test on Win64 bots
194 #if defined(OS_WIN) && defined(ARCH_CPU_X86_64)
195 #define MAYBE_EnumeratesPE DISABLED_EnumeratesPE
196 #else
197 #define MAYBE_EnumeratesPE EnumeratesPE
198 #endif
199
200 // Tests that we are able to enumerate stuff from a PE file, and that
201 // the actual number of items found is within the expected range.
TEST(PEImageTest,MAYBE_EnumeratesPE)202 TEST(PEImageTest, MAYBE_EnumeratesPE) {
203 HMODULE module = LoadLibrary(L"advapi32.dll");
204 ASSERT_TRUE(NULL != module);
205
206 PEImage pe(module);
207 int count = 0;
208 EXPECT_TRUE(pe.VerifyMagic());
209
210 DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion;
211 os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion;
212
213 // Skip this test for unsupported OS versions.
214 if (GetExpectationSet(os) == UNSUPPORTED_SET)
215 return;
216
217 pe.EnumSections(SectionsCallback, &count);
218 EXPECT_EQ(GetExpectedValue(sections, os), count);
219
220 count = 0;
221 pe.EnumImportChunks(ImportChunksCallback, &count);
222 EXPECT_EQ(GetExpectedValue(imports_dlls, os), count);
223
224 count = 0;
225 pe.EnumDelayImportChunks(DelayImportChunksCallback, &count);
226 EXPECT_EQ(GetExpectedValue(delay_dlls, os), count);
227
228 count = 0;
229 pe.EnumExports(ExportsCallback, &count);
230 EXPECT_GT(count, GetExpectedValue(exports, os) - 20);
231 EXPECT_LT(count, GetExpectedValue(exports, os) + 100);
232
233 count = 0;
234 pe.EnumAllImports(ImportsCallback, &count);
235 EXPECT_GT(count, GetExpectedValue(imports, os) - 20);
236 EXPECT_LT(count, GetExpectedValue(imports, os) + 100);
237
238 count = 0;
239 pe.EnumAllDelayImports(ImportsCallback, &count);
240 EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2);
241 EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8);
242
243 count = 0;
244 pe.EnumRelocs(RelocsCallback, &count);
245 EXPECT_GT(count, GetExpectedValue(relocs, os) - 150);
246 EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500);
247
248 FreeLibrary(module);
249 }
250
251 // Tests that we can locate an specific exported symbol, by name and by ordinal.
TEST(PEImageTest,RetrievesExports)252 TEST(PEImageTest, RetrievesExports) {
253 HMODULE module = LoadLibrary(L"advapi32.dll");
254 ASSERT_TRUE(NULL != module);
255
256 PEImage pe(module);
257 WORD ordinal;
258
259 EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal));
260
261 FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW");
262 FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal));
263 EXPECT_TRUE(address1 != NULL);
264 EXPECT_TRUE(address2 != NULL);
265 EXPECT_TRUE(address1 == address2);
266
267 FreeLibrary(module);
268 }
269
270 } // namespace win
271 } // namespace base
272