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 #include <algorithm>
   7 #include <vector>
   8
   9 #include "base/files/file_path.h"
  10 #include "base/path_service.h"
  11 #include "base/win/pe_image.h"
  12 #include "testing/gtest/include/gtest/gtest.h"
  13
  14 namespace base {
  15 namespace win {
  16
  17 namespace {
  18
  19 // Just counts the number of invocations.
  20 bool ImportsCallback(const PEImage& image,
  21                      LPCSTR module,
  22                      DWORD ordinal,
  23                      LPCSTR name,
  24                      DWORD hint,
  25                      PIMAGE_THUNK_DATA iat,
  26                      PVOID cookie) {
  27   int* count = reinterpret_cast<int*>(cookie);
  28   (*count)++;
  29   return true;
  30 }
  31
  32 // Just counts the number of invocations.
  33 bool SectionsCallback(const PEImage& image,
  34                       PIMAGE_SECTION_HEADER header,
  35                       PVOID section_start,
  36                       DWORD section_size,
  37                       PVOID cookie) {
  38   int* count = reinterpret_cast<int*>(cookie);
  39   (*count)++;
  40   return true;
  41 }
  42
  43 // Just counts the number of invocations.
  44 bool RelocsCallback(const PEImage& image,
  45                     WORD type,
  46                     PVOID address,
  47                     PVOID cookie) {
  48   int* count = reinterpret_cast<int*>(cookie);
  49   (*count)++;
  50   return true;
  51 }
  52
  53 // Just counts the number of invocations.
  54 bool ImportChunksCallback(const PEImage& image,
  55                           LPCSTR module,
  56                           PIMAGE_THUNK_DATA name_table,
  57                           PIMAGE_THUNK_DATA iat,
  58                           PVOID cookie) {
  59   int* count = reinterpret_cast<int*>(cookie);
  60   (*count)++;
  61   return true;
  62 }
  63
  64 // Just counts the number of invocations.
  65 bool DelayImportChunksCallback(const PEImage& image,
  66                                PImgDelayDescr delay_descriptor,
  67                                LPCSTR module,
  68                                PIMAGE_THUNK_DATA name_table,
  69                                PIMAGE_THUNK_DATA iat,
  70                                PIMAGE_THUNK_DATA bound_iat,
  71                                PIMAGE_THUNK_DATA unload_iat,
  72                                PVOID cookie) {
  73   int* count = reinterpret_cast<int*>(cookie);
  74   (*count)++;
  75   return true;
  76 }
  77
  78 // Just counts the number of invocations.
  79 bool ExportsCallback(const PEImage& image,
  80                      DWORD ordinal,
  81                      DWORD hint,
  82                      LPCSTR name,
  83                      PVOID function,
  84                      LPCSTR forward,
  85                      PVOID cookie) {
  86   int* count = reinterpret_cast<int*>(cookie);
  87   (*count)++;
  88   return true;
  89 }
  90
  91 }  // namespace
  92
  93 // Tests that we are able to enumerate stuff from a PE file, and that
  94 // the actual number of items found matches an expected value.
  95 TEST(PEImageTest, EnumeratesPE) {
  96   base::FilePath pe_image_test_path;
  97   ASSERT_TRUE(PathService::Get(DIR_TEST_DATA, &pe_image_test_path));
  98   pe_image_test_path = pe_image_test_path.Append(FILE_PATH_LITERAL("pe_image"));
  99
 100 #if defined(ARCH_CPU_64_BITS)
 101   pe_image_test_path =
 102       pe_image_test_path.Append(FILE_PATH_LITERAL("pe_image_test_64.dll"));
 103   const int sections = 6;
 104   const int imports_dlls = 2;
 105   const int delay_dlls = 2;
 106   const int exports = 2;
 107   const int imports = 69;
 108   const int delay_imports = 2;
 109   const int relocs = 632;
 110 #else
 111   pe_image_test_path =
 112       pe_image_test_path.Append(FILE_PATH_LITERAL("pe_image_test_32.dll"));
 113   const int sections = 5;
 114   const int imports_dlls = 2;
 115   const int delay_dlls = 2;
 116   const int exports = 2;
 117   const int imports = 66;
 118   const int delay_imports = 2;
 119   const int relocs = 1586;
 120 #endif
 121
 122   HMODULE module = LoadLibrary(pe_image_test_path.value().c_str());
 123   ASSERT_TRUE(NULL != module);
 124
 125   PEImage pe(module);
 126   int count = 0;
 127   EXPECT_TRUE(pe.VerifyMagic());
 128
 129   pe.EnumSections(SectionsCallback, &count);
 130   EXPECT_EQ(sections, count);
 131
 132   count = 0;
 133   pe.EnumImportChunks(ImportChunksCallback, &count);
 134   EXPECT_EQ(imports_dlls, count);
 135
 136   count = 0;
 137   pe.EnumDelayImportChunks(DelayImportChunksCallback, &count);
 138   EXPECT_EQ(delay_dlls, count);
 139
 140   count = 0;
 141   pe.EnumExports(ExportsCallback, &count);
 142   EXPECT_EQ(exports, count);
 143
 144   count = 0;
 145   pe.EnumAllImports(ImportsCallback, &count);
 146   EXPECT_EQ(imports, count);
 147
 148   count = 0;
 149   pe.EnumAllDelayImports(ImportsCallback, &count);
 150   EXPECT_EQ(delay_imports, count);
 151
 152   count = 0;
 153   pe.EnumRelocs(RelocsCallback, &count);
 154   EXPECT_EQ(relocs, count);
 155
 156   FreeLibrary(module);
 157 }
 158
 159 // Tests that we can locate an specific exported symbol, by name and by ordinal.
 160 TEST(PEImageTest, RetrievesExports) {
 161   HMODULE module = LoadLibrary(L"advapi32.dll");
 162   ASSERT_TRUE(NULL != module);
 163
 164   PEImage pe(module);
 165   WORD ordinal;
 166
 167   EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal));
 168
 169   FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW");
 170   FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal));
 171   EXPECT_TRUE(address1 != NULL);
 172   EXPECT_TRUE(address2 != NULL);
 173   EXPECT_TRUE(address1 == address2);
 174
 175   FreeLibrary(module);
 176 }
 177
 178 // Test that we can get debug id out of a module.
 179 TEST(PEImageTest, GetDebugId) {
 180   HMODULE module = LoadLibrary(L"advapi32.dll");
 181   ASSERT_TRUE(NULL != module);
 182
 183   PEImage pe(module);
 184   GUID guid = {0};
 185   DWORD age = 0;
 186   EXPECT_TRUE(pe.GetDebugId(&guid, &age));
 187
 188   GUID empty_guid = {0};
 189   EXPECT_TRUE(!IsEqualGUID(empty_guid, guid));
 190   EXPECT_NE(0U, age);
 191   FreeLibrary(module);
 192 }
 193
 194 }  // namespace win
 195 }  // namespace base