base/win/pe_image_unittest.cc

   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.
  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.
  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.
  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.
  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.
  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.
  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
 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.
 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.
 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.
 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