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   const int win7_relocs = 7812;
 143   const int win8_delay_dlls = 9;
 144   const int win8_exports = 806;
 145   const int win8_imports = 568;
 146   const int win8_delay_imports = 113;
 147   const int win8_relocs = 9478;
 148   int win8_sections = 4;
 149   int win8_import_dlls = 17;
 150
 151   base::win::OSInfo* os_info = base::win::OSInfo::GetInstance();
 152   if (os_info->architecture() == base::win::OSInfo::X86_ARCHITECTURE) {
 153     win8_sections = 5;
 154     win8_import_dlls = 19;
 155   }
 156
 157   // Contains the expected value, for each enumerated property (Value), and the
 158   // OS version: [Value][os_version]
 159   const int expected[][5] = {
 160     {4, 4, 4, 4, win8_sections},
 161     {3, 3, 3, 13, win8_import_dlls},
 162     {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls,
 163      win8_delay_dlls},
 164     {w2k_exports, xp_exports, vista_exports, win7_exports, win8_exports},
 165     {w2k_imports, xp_imports, vista_imports, win7_imports, win8_imports},
 166     {w2k_delay_imports, xp_delay_imports,
 167      vista_delay_imports, win7_delay_imports, win8_delay_imports},
 168     {w2k_relocs, xp_relocs, vista_relocs, win7_relocs, win8_relocs}
 169   };
 170   COMPILE_ASSERT(arraysize(expected[0]) == UNSUPPORTED_SET,
 171                  expected_value_set_mismatch);
 172
 173   if (value > relocs)
 174     return 0;
 175   ExpectationSet expected_set = GetExpectationSet(os);
 176   if (expected_set >= arraysize(expected)) {
 177     // This should never happen.  Log a failure if it does.
 178     EXPECT_NE(UNSUPPORTED_SET, expected_set);
 179     expected_set = WIN_2K_SET;
 180   }
 181
 182   return expected[value][expected_set];
 183 }
 184
 185 // Tests that we are able to enumerate stuff from a PE file, and that
 186 // the actual number of items found is within the expected range.
 187 TEST(PEImageTest, EnumeratesPE) {
 188   HMODULE module = LoadLibrary(L"advapi32.dll");
 189   ASSERT_TRUE(NULL != module);
 190
 191   PEImage pe(module);
 192   int count = 0;
 193   EXPECT_TRUE(pe.VerifyMagic());
 194
 195   DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion;
 196   os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion;
 197
 198   // Skip this test for unsupported OS versions.
 199   if (GetExpectationSet(os) == UNSUPPORTED_SET)
 200     return;
 201
 202   pe.EnumSections(SectionsCallback, &count);
 203   EXPECT_EQ(GetExpectedValue(sections, os), count);
 204
 205   count = 0;
 206   pe.EnumImportChunks(ImportChunksCallback, &count);
 207   EXPECT_EQ(GetExpectedValue(imports_dlls, os), count);
 208
 209   count = 0;
 210   pe.EnumDelayImportChunks(DelayImportChunksCallback, &count);
 211   EXPECT_EQ(GetExpectedValue(delay_dlls, os), count);
 212
 213   count = 0;
 214   pe.EnumExports(ExportsCallback, &count);
 215   EXPECT_GT(count, GetExpectedValue(exports, os) - 20);
 216   EXPECT_LT(count, GetExpectedValue(exports, os) + 100);
 217
 218   count = 0;
 219   pe.EnumAllImports(ImportsCallback, &count);
 220   EXPECT_GT(count, GetExpectedValue(imports, os) - 20);
 221   EXPECT_LT(count, GetExpectedValue(imports, os) + 100);
 222
 223   count = 0;
 224   pe.EnumAllDelayImports(ImportsCallback, &count);
 225   EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2);
 226   EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8);
 227
 228   count = 0;
 229   pe.EnumRelocs(RelocsCallback, &count);
 230   EXPECT_GT(count, GetExpectedValue(relocs, os) - 150);
 231   EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500);
 232
 233   FreeLibrary(module);
 234 }
 235
 236 // Tests that we can locate an specific exported symbol, by name and by ordinal.
 237 TEST(PEImageTest, RetrievesExports) {
 238   HMODULE module = LoadLibrary(L"advapi32.dll");
 239   ASSERT_TRUE(NULL != module);
 240
 241   PEImage pe(module);
 242   WORD ordinal;
 243
 244   EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal));
 245
 246   FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW");
 247   FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal));
 248   EXPECT_TRUE(address1 != NULL);
 249   EXPECT_TRUE(address2 != NULL);
 250   EXPECT_TRUE(address1 == address2);
 251
 252   FreeLibrary(module);
 253 }
 254
 255 }  // namespace win
 256 }  // namespace base