1 // Copyright (c) 2010 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 // We'll be using some known values for the tests. 87 enum Value { 88 sections = 0, 89 imports_dlls, 90 delay_dlls, 91 exports, 92 imports, 93 delay_imports, 94 relocs 95 }; 96 97 // Retrieves the expected value from advapi32.dll based on the OS. 98 int GetExpectedValue(Value value, DWORD os) { 99 const int xp_delay_dlls = 2; 100 const int xp_exports = 675; 101 const int xp_imports = 422; 102 const int xp_delay_imports = 8; 103 const int xp_relocs = 9180; 104 const int vista_delay_dlls = 4; 105 const int vista_exports = 799; 106 const int vista_imports = 476; 107 const int vista_delay_imports = 24; 108 const int vista_relocs = 10188; 109 const int w2k_delay_dlls = 0; 110 const int w2k_exports = 566; 111 const int w2k_imports = 357; 112 const int w2k_delay_imports = 0; 113 const int w2k_relocs = 7388; 114 const int win7_delay_dlls = 7; 115 const int win7_exports = 806; 116 const int win7_imports = 568; 117 const int win7_delay_imports = 71; 118 const int win7_relocs = 7812; 119 120 // Contains the expected value, for each enumerated property (Value), and the 121 // OS version: [Value][os_version] 122 const int expected[][4] = { 123 {4, 4, 4, 4}, 124 {3, 3, 3, 13}, 125 {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls}, 126 {w2k_exports, xp_exports, vista_exports, win7_exports}, 127 {w2k_imports, xp_imports, vista_imports, win7_imports}, 128 {w2k_delay_imports, xp_delay_imports, 129 vista_delay_imports, win7_delay_imports}, 130 {w2k_relocs, xp_relocs, vista_relocs, win7_relocs} 131 }; 132 133 if (value > relocs) 134 return 0; 135 if (50 == os) 136 os = 0; // 5.0 137 else if (51 == os || 52 == os) 138 os = 1; 139 else if (os == 60) 140 os = 2; // 6.x 141 else if (os >= 61) 142 os = 3; 143 else 144 return 0; 145 146 return expected[value][os]; 147 } 148 149 // Tests that we are able to enumerate stuff from a PE file, and that 150 // the actual number of items found is within the expected range. 151 TEST(PEImageTest, EnumeratesPE) { 152 // Windows Server 2003 is not supported as a test environment for this test. 153 if (base::win::GetVersion() == base::win::VERSION_SERVER_2003) 154 return; 155 HMODULE module = LoadLibrary(L"advapi32.dll"); 156 ASSERT_TRUE(NULL != module); 157 158 PEImage pe(module); 159 int count = 0; 160 EXPECT_TRUE(pe.VerifyMagic()); 161 162 DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion; 163 os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion; 164 165 pe.EnumSections(SectionsCallback, &count); 166 EXPECT_EQ(GetExpectedValue(sections, os), count); 167 168 count = 0; 169 pe.EnumImportChunks(ImportChunksCallback, &count); 170 EXPECT_EQ(GetExpectedValue(imports_dlls, os), count); 171 172 count = 0; 173 pe.EnumDelayImportChunks(DelayImportChunksCallback, &count); 174 EXPECT_EQ(GetExpectedValue(delay_dlls, os), count); 175 176 count = 0; 177 pe.EnumExports(ExportsCallback, &count); 178 EXPECT_GT(count, GetExpectedValue(exports, os) - 20); 179 EXPECT_LT(count, GetExpectedValue(exports, os) + 100); 180 181 count = 0; 182 pe.EnumAllImports(ImportsCallback, &count); 183 EXPECT_GT(count, GetExpectedValue(imports, os) - 20); 184 EXPECT_LT(count, GetExpectedValue(imports, os) + 100); 185 186 count = 0; 187 pe.EnumAllDelayImports(ImportsCallback, &count); 188 EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2); 189 EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8); 190 191 count = 0; 192 pe.EnumRelocs(RelocsCallback, &count); 193 EXPECT_GT(count, GetExpectedValue(relocs, os) - 150); 194 EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500); 195 196 FreeLibrary(module); 197 } 198 199 // Tests that we can locate an specific exported symbol, by name and by ordinal. 200 TEST(PEImageTest, RetrievesExports) { 201 HMODULE module = LoadLibrary(L"advapi32.dll"); 202 ASSERT_TRUE(NULL != module); 203 204 PEImage pe(module); 205 WORD ordinal; 206 207 EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal)); 208 209 FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW"); 210 FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal)); 211 EXPECT_TRUE(address1 != NULL); 212 EXPECT_TRUE(address2 != NULL); 213 EXPECT_TRUE(address1 == address2); 214 215 FreeLibrary(module); 216 } 217 218 } // namespace win 219 } // namespace base 220
