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
11namespace base {
12namespace win {
13
14// Just counts the number of invocations.
15bool 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.
28bool 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.
41bool 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.
52bool 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.
62bool 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.
73bool 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.
87enum 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.
97enum Value {
98  sections = 0,
99  imports_dlls,
100  delay_dlls,
101  exports,
102  imports,
103  delay_imports,
104  relocs
105};
106
107ExpectationSet 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.
122int 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.
202TEST(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.
252TEST(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