1 // Copyright (c) 2006-2008 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/pe_image.h" 9 #include "base/win_util.h" 10 11 // Just counts the number of invocations. 12 bool ExportsCallback(const PEImage &image, 13 DWORD ordinal, 14 DWORD hint, 15 LPCSTR name, 16 PVOID function, 17 LPCSTR forward, 18 PVOID cookie) { 19 int* count = reinterpret_cast<int*>(cookie); 20 (*count)++; 21 return true; 22 } 23 24 // Just counts the number of invocations. 25 bool ImportsCallback(const PEImage &image, 26 LPCSTR module, 27 DWORD ordinal, 28 LPCSTR name, 29 DWORD hint, 30 PIMAGE_THUNK_DATA iat, 31 PVOID cookie) { 32 int* count = reinterpret_cast<int*>(cookie); 33 (*count)++; 34 return true; 35 } 36 37 // Just counts the number of invocations. 38 bool SectionsCallback(const PEImage &image, 39 PIMAGE_SECTION_HEADER header, 40 PVOID section_start, 41 DWORD section_size, 42 PVOID cookie) { 43 int* count = reinterpret_cast<int*>(cookie); 44 (*count)++; 45 return true; 46 } 47 48 // Just counts the number of invocations. 49 bool RelocsCallback(const PEImage &image, 50 WORD type, 51 PVOID address, 52 PVOID cookie) { 53 int* count = reinterpret_cast<int*>(cookie); 54 (*count)++; 55 return true; 56 } 57 58 // Just counts the number of invocations. 59 bool ImportChunksCallback(const PEImage &image, 60 LPCSTR module, 61 PIMAGE_THUNK_DATA name_table, 62 PIMAGE_THUNK_DATA iat, 63 PVOID cookie) { 64 int* count = reinterpret_cast<int*>(cookie); 65 (*count)++; 66 return true; 67 } 68 69 // Just counts the number of invocations. 70 bool DelayImportChunksCallback(const PEImage &image, 71 PImgDelayDescr delay_descriptor, 72 LPCSTR module, 73 PIMAGE_THUNK_DATA name_table, 74 PIMAGE_THUNK_DATA iat, 75 PIMAGE_THUNK_DATA bound_iat, 76 PIMAGE_THUNK_DATA unload_iat, 77 PVOID cookie) { 78 int* count = reinterpret_cast<int*>(cookie); 79 (*count)++; 80 return true; 81 } 82 83 // We'll be using some known values for the tests. 84 enum Value { 85 sections = 0, 86 imports_dlls, 87 delay_dlls, 88 exports, 89 imports, 90 delay_imports, 91 relocs 92 }; 93 94 // Retrieves the expected value from advapi32.dll based on the OS. 95 int GetExpectedValue(Value value, DWORD os) { 96 const int xp_delay_dlls = 2; 97 const int xp_exports = 675; 98 const int xp_imports = 422; 99 const int xp_delay_imports = 8; 100 const int xp_relocs = 9180; 101 const int vista_delay_dlls = 4; 102 const int vista_exports = 799; 103 const int vista_imports = 476; 104 const int vista_delay_imports = 24; 105 const int vista_relocs = 10188; 106 const int w2k_delay_dlls = 0; 107 const int w2k_exports = 566; 108 const int w2k_imports = 357; 109 const int w2k_delay_imports = 0; 110 const int w2k_relocs = 7388; 111 const int win7_delay_dlls = 7; 112 const int win7_exports = 806; 113 const int win7_imports = 568; 114 const int win7_delay_imports = 71; 115 const int win7_relocs = 7812; 116 117 // Contains the expected value, for each enumerated property (Value), and the 118 // OS version: [Value][os_version] 119 const int expected[][4] = { 120 {4, 4, 4, 4}, 121 {3, 3, 3, 13}, 122 {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls}, 123 {w2k_exports, xp_exports, vista_exports, win7_exports}, 124 {w2k_imports, xp_imports, vista_imports, win7_imports}, 125 {w2k_delay_imports, xp_delay_imports, 126 vista_delay_imports, win7_delay_imports}, 127 {w2k_relocs, xp_relocs, vista_relocs, win7_relocs} 128 }; 129 130 if (value > relocs) 131 return 0; 132 if (50 == os) 133 os = 0; // 5.0 134 else if (51 == os || 52 == os) 135 os = 1; 136 else if (os == 60) 137 os = 2; // 6.x 138 else if (os >= 61) 139 os = 3; 140 else 141 return 0; 142 143 return expected[value][os]; 144 } 145 146 // Tests that we are able to enumerate stuff from a PE file, and that 147 // the actual number of items found is within the expected range. 148 TEST(PEImageTest, EnumeratesPE) { 149 // Windows Server 2003 is not supported as a test environment for this test. 150 if (win_util::GetWinVersion() == win_util::WINVERSION_SERVER_2003) 151 return; 152 HMODULE module = LoadLibrary(L"advapi32.dll"); 153 ASSERT_TRUE(NULL != module); 154 155 PEImage pe(module); 156 int count = 0; 157 EXPECT_TRUE(pe.VerifyMagic()); 158 159 DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion; 160 os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion; 161 162 pe.EnumSections(SectionsCallback, &count); 163 EXPECT_EQ(GetExpectedValue(sections, os), count); 164 165 count = 0; 166 pe.EnumImportChunks(ImportChunksCallback, &count); 167 EXPECT_EQ(GetExpectedValue(imports_dlls, os), count); 168 169 count = 0; 170 pe.EnumDelayImportChunks(DelayImportChunksCallback, &count); 171 EXPECT_EQ(GetExpectedValue(delay_dlls, os), count); 172 173 count = 0; 174 pe.EnumExports(ExportsCallback, &count); 175 EXPECT_GT(count, GetExpectedValue(exports, os) - 20); 176 EXPECT_LT(count, GetExpectedValue(exports, os) + 100); 177 178 count = 0; 179 pe.EnumAllImports(ImportsCallback, &count); 180 EXPECT_GT(count, GetExpectedValue(imports, os) - 20); 181 EXPECT_LT(count, GetExpectedValue(imports, os) + 100); 182 183 count = 0; 184 pe.EnumAllDelayImports(ImportsCallback, &count); 185 EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2); 186 EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8); 187 188 count = 0; 189 pe.EnumRelocs(RelocsCallback, &count); 190 EXPECT_GT(count, GetExpectedValue(relocs, os) - 150); 191 EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500); 192 193 FreeLibrary(module); 194 } 195 196 // Tests that we can locate an specific exported symbol, by name and by ordinal. 197 TEST(PEImageTest, RetrievesExports) { 198 HMODULE module = LoadLibrary(L"advapi32.dll"); 199 ASSERT_TRUE(NULL != module); 200 201 PEImage pe(module); 202 WORD ordinal; 203 204 EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal)); 205 206 FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW"); 207 FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal)); 208 EXPECT_TRUE(address1 != NULL); 209 EXPECT_TRUE(address2 != NULL); 210 EXPECT_TRUE(address1 == address2); 211 212 FreeLibrary(module); 213 } 214
