1 // Copyright (c) 2010 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 // We'll be using some known values for the tests.
87 enum Value {
88 sections = 0,
89 imports_dlls,
90 delay_dlls,
91 exports,
92 imports,
93 delay_imports,
94 relocs
95 };
96
97 // Retrieves the expected value from advapi32.dll based on the OS.
GetExpectedValue(Value value,DWORD os)98 int GetExpectedValue(Value value, DWORD os) {
99 const int xp_delay_dlls = 2;
100 const int xp_exports = 675;
101 const int xp_imports = 422;
102 const int xp_delay_imports = 8;
103 const int xp_relocs = 9180;
104 const int vista_delay_dlls = 4;
105 const int vista_exports = 799;
106 const int vista_imports = 476;
107 const int vista_delay_imports = 24;
108 const int vista_relocs = 10188;
109 const int w2k_delay_dlls = 0;
110 const int w2k_exports = 566;
111 const int w2k_imports = 357;
112 const int w2k_delay_imports = 0;
113 const int w2k_relocs = 7388;
114 const int win7_delay_dlls = 7;
115 const int win7_exports = 806;
116 const int win7_imports = 568;
117 const int win7_delay_imports = 71;
118 const int win7_relocs = 7812;
119
120 // Contains the expected value, for each enumerated property (Value), and the
121 // OS version: [Value][os_version]
122 const int expected[][4] = {
123 {4, 4, 4, 4},
124 {3, 3, 3, 13},
125 {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls},
126 {w2k_exports, xp_exports, vista_exports, win7_exports},
127 {w2k_imports, xp_imports, vista_imports, win7_imports},
128 {w2k_delay_imports, xp_delay_imports,
129 vista_delay_imports, win7_delay_imports},
130 {w2k_relocs, xp_relocs, vista_relocs, win7_relocs}
131 };
132
133 if (value > relocs)
134 return 0;
135 if (50 == os)
136 os = 0; // 5.0
137 else if (51 == os || 52 == os)
138 os = 1;
139 else if (os == 60)
140 os = 2; // 6.x
141 else if (os >= 61)
142 os = 3;
143 else
144 return 0;
145
146 return expected[value][os];
147 }
148
149 // Tests that we are able to enumerate stuff from a PE file, and that
150 // the actual number of items found is within the expected range.
TEST(PEImageTest,EnumeratesPE)151 TEST(PEImageTest, EnumeratesPE) {
152 // Windows Server 2003 is not supported as a test environment for this test.
153 if (base::win::GetVersion() == base::win::VERSION_SERVER_2003)
154 return;
155 HMODULE module = LoadLibrary(L"advapi32.dll");
156 ASSERT_TRUE(NULL != module);
157
158 PEImage pe(module);
159 int count = 0;
160 EXPECT_TRUE(pe.VerifyMagic());
161
162 DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion;
163 os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion;
164
165 pe.EnumSections(SectionsCallback, &count);
166 EXPECT_EQ(GetExpectedValue(sections, os), count);
167
168 count = 0;
169 pe.EnumImportChunks(ImportChunksCallback, &count);
170 EXPECT_EQ(GetExpectedValue(imports_dlls, os), count);
171
172 count = 0;
173 pe.EnumDelayImportChunks(DelayImportChunksCallback, &count);
174 EXPECT_EQ(GetExpectedValue(delay_dlls, os), count);
175
176 count = 0;
177 pe.EnumExports(ExportsCallback, &count);
178 EXPECT_GT(count, GetExpectedValue(exports, os) - 20);
179 EXPECT_LT(count, GetExpectedValue(exports, os) + 100);
180
181 count = 0;
182 pe.EnumAllImports(ImportsCallback, &count);
183 EXPECT_GT(count, GetExpectedValue(imports, os) - 20);
184 EXPECT_LT(count, GetExpectedValue(imports, os) + 100);
185
186 count = 0;
187 pe.EnumAllDelayImports(ImportsCallback, &count);
188 EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2);
189 EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8);
190
191 count = 0;
192 pe.EnumRelocs(RelocsCallback, &count);
193 EXPECT_GT(count, GetExpectedValue(relocs, os) - 150);
194 EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500);
195
196 FreeLibrary(module);
197 }
198
199 // Tests that we can locate an specific exported symbol, by name and by ordinal.
TEST(PEImageTest,RetrievesExports)200 TEST(PEImageTest, RetrievesExports) {
201 HMODULE module = LoadLibrary(L"advapi32.dll");
202 ASSERT_TRUE(NULL != module);
203
204 PEImage pe(module);
205 WORD ordinal;
206
207 EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal));
208
209 FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW");
210 FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal));
211 EXPECT_TRUE(address1 != NULL);
212 EXPECT_TRUE(address2 != NULL);
213 EXPECT_TRUE(address1 == address2);
214
215 FreeLibrary(module);
216 }
217
218 } // namespace win
219 } // namespace base
220