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 272