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
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// We'll be using some known values for the tests.
87enum 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.
98int 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.
151TEST(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.
200TEST(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