| // Copyright (c) 2012 The Chromium Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| // This file contains unit tests for PEImage. |
| #include <algorithm> |
| #include <vector> |
| |
| #include "testing/gtest/include/gtest/gtest.h" |
| #include "base/win/pe_image.h" |
| #include "base/win/windows_version.h" |
| |
| namespace { |
| |
| class Expectations { |
| public: |
| enum Value { |
| SECTIONS = 0, |
| IMPORTS_DLLS, |
| DELAY_DLLS, |
| EXPORTS, |
| IMPORTS, |
| DELAY_IMPORTS, |
| RELOCS |
| }; |
| |
| enum Arch { |
| ARCH_X86 = 0, |
| ARCH_X64, |
| ARCH_ALL |
| }; |
| |
| Expectations(); |
| |
| void SetDefault(Value value, int count); |
| void SetOverride(Value value, base::win::Version version, |
| Arch arch, int count); |
| void SetOverride(Value value, base::win::Version version, int count); |
| void SetOverride(Value value, Arch arch, int count); |
| |
| // returns -1 on failure. |
| int GetExpectation(Value value); |
| |
| private: |
| class Override { |
| public: |
| enum MatchType { MATCH_VERSION, MATCH_ARCH, MATCH_BOTH, MATCH_NONE }; |
| |
| Override(Value value, base::win::Version version, Arch arch, int count) |
| : value_(value), version_(version), arch_(arch), count_(count) { |
| }; |
| |
| bool Matches(Value value, base::win::Version version, |
| Arch arch, MatchType type) { |
| if (value_ != value) |
| return false; |
| |
| switch (type) { |
| case MATCH_BOTH: |
| return (arch == arch_ && version == version_); |
| case MATCH_ARCH: |
| return (arch == arch_ && version_ == base::win::VERSION_WIN_LAST); |
| case MATCH_VERSION: |
| return (arch_ == ARCH_ALL && version == version_); |
| case MATCH_NONE: |
| return (arch_ == ARCH_ALL && version_ == base::win::VERSION_WIN_LAST); |
| } |
| return false; |
| } |
| |
| int GetCount() { return count_; } |
| |
| private: |
| Value value_; |
| base::win::Version version_; |
| Arch arch_; |
| int count_; |
| }; |
| |
| bool MatchesMyArch(Arch arch); |
| |
| std::vector<Override> overrides_; |
| Arch my_arch_; |
| base::win::Version my_version_; |
| }; |
| |
| Expectations::Expectations() { |
| my_version_ = base::win::GetVersion(); |
| #if defined(ARCH_CPU_64_BITS) |
| my_arch_ = ARCH_X64; |
| #else |
| my_arch_ = ARCH_X86; |
| #endif |
| } |
| |
| int Expectations::GetExpectation(Value value) { |
| // Prefer OS version specificity over Arch specificity. |
| for (auto type : { Override::MATCH_BOTH, |
| Override::MATCH_VERSION, |
| Override::MATCH_ARCH, |
| Override::MATCH_NONE }) { |
| for (auto override : overrides_) { |
| if (override.Matches(value, my_version_, my_arch_, type)) |
| return override.GetCount(); |
| } |
| } |
| return -1; |
| } |
| |
| void Expectations::SetDefault(Value value, int count) { |
| SetOverride(value, base::win::VERSION_WIN_LAST, ARCH_ALL, count); |
| } |
| |
| void Expectations::SetOverride(Value value, |
| base::win::Version version, |
| Arch arch, |
| int count) { |
| overrides_.push_back(Override(value, version, arch, count)); |
| } |
| |
| void Expectations::SetOverride(Value value, |
| base::win::Version version, |
| int count) { |
| SetOverride(value, version, ARCH_ALL, count); |
| } |
| |
| void Expectations::SetOverride(Value value, Arch arch, int count) { |
| SetOverride(value, base::win::VERSION_WIN_LAST, arch, count); |
| } |
| |
| } // namespace |
| |
| namespace base { |
| namespace win { |
| |
| namespace { |
| |
| // Just counts the number of invocations. |
| bool ImportsCallback(const PEImage& image, |
| LPCSTR module, |
| DWORD ordinal, |
| LPCSTR name, |
| DWORD hint, |
| PIMAGE_THUNK_DATA iat, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| // Just counts the number of invocations. |
| bool SectionsCallback(const PEImage& image, |
| PIMAGE_SECTION_HEADER header, |
| PVOID section_start, |
| DWORD section_size, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| // Just counts the number of invocations. |
| bool RelocsCallback(const PEImage& image, |
| WORD type, |
| PVOID address, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| // Just counts the number of invocations. |
| bool ImportChunksCallback(const PEImage& image, |
| LPCSTR module, |
| PIMAGE_THUNK_DATA name_table, |
| PIMAGE_THUNK_DATA iat, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| // Just counts the number of invocations. |
| bool DelayImportChunksCallback(const PEImage& image, |
| PImgDelayDescr delay_descriptor, |
| LPCSTR module, |
| PIMAGE_THUNK_DATA name_table, |
| PIMAGE_THUNK_DATA iat, |
| PIMAGE_THUNK_DATA bound_iat, |
| PIMAGE_THUNK_DATA unload_iat, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| // Just counts the number of invocations. |
| bool ExportsCallback(const PEImage& image, |
| DWORD ordinal, |
| DWORD hint, |
| LPCSTR name, |
| PVOID function, |
| LPCSTR forward, |
| PVOID cookie) { |
| int* count = reinterpret_cast<int*>(cookie); |
| (*count)++; |
| return true; |
| } |
| |
| } // namespace |
| |
| // Tests that we are able to enumerate stuff from a PE file, and that |
| // the actual number of items found is within the expected range. |
| TEST(PEImageTest, EnumeratesPE) { |
| Expectations expectations; |
| |
| #ifndef NDEBUG |
| // Default Debug expectations. |
| expectations.SetDefault(Expectations::SECTIONS, 7); |
| expectations.SetDefault(Expectations::IMPORTS_DLLS, 3); |
| expectations.SetDefault(Expectations::DELAY_DLLS, 2); |
| expectations.SetDefault(Expectations::EXPORTS, 2); |
| expectations.SetDefault(Expectations::IMPORTS, 49); |
| expectations.SetDefault(Expectations::DELAY_IMPORTS, 2); |
| expectations.SetDefault(Expectations::RELOCS, 438); |
| |
| // 64-bit Debug expectations. |
| expectations.SetOverride(Expectations::SECTIONS, Expectations::ARCH_X64, 8); |
| expectations.SetOverride(Expectations::IMPORTS, Expectations::ARCH_X64, 69); |
| expectations.SetOverride(Expectations::RELOCS, Expectations::ARCH_X64, 632); |
| #else |
| // Default Release expectations. |
| expectations.SetDefault(Expectations::SECTIONS, 5); |
| expectations.SetDefault(Expectations::IMPORTS_DLLS, 2); |
| expectations.SetDefault(Expectations::DELAY_DLLS, 2); |
| expectations.SetDefault(Expectations::EXPORTS, 2); |
| expectations.SetDefault(Expectations::IMPORTS, 66); |
| expectations.SetDefault(Expectations::DELAY_IMPORTS, 2); |
| expectations.SetDefault(Expectations::RELOCS, 1586); |
| |
| // 64-bit Release expectations. |
| expectations.SetOverride(Expectations::SECTIONS, Expectations::ARCH_X64, 6); |
| expectations.SetOverride(Expectations::IMPORTS, Expectations::ARCH_X64, 69); |
| expectations.SetOverride(Expectations::RELOCS, Expectations::ARCH_X64, 632); |
| #endif |
| |
| HMODULE module = LoadLibrary(L"pe_image_test.dll"); |
| ASSERT_TRUE(NULL != module); |
| |
| PEImage pe(module); |
| int count = 0; |
| EXPECT_TRUE(pe.VerifyMagic()); |
| |
| pe.EnumSections(SectionsCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::SECTIONS), count); |
| |
| count = 0; |
| pe.EnumImportChunks(ImportChunksCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::IMPORTS_DLLS), count); |
| |
| count = 0; |
| pe.EnumDelayImportChunks(DelayImportChunksCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::DELAY_DLLS), count); |
| |
| count = 0; |
| pe.EnumExports(ExportsCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::EXPORTS), count); |
| |
| count = 0; |
| pe.EnumAllImports(ImportsCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::IMPORTS), count); |
| |
| count = 0; |
| pe.EnumAllDelayImports(ImportsCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::DELAY_IMPORTS), count); |
| |
| count = 0; |
| pe.EnumRelocs(RelocsCallback, &count); |
| EXPECT_EQ(expectations.GetExpectation(Expectations::RELOCS), count); |
| |
| FreeLibrary(module); |
| } |
| |
| // Tests that we can locate an specific exported symbol, by name and by ordinal. |
| TEST(PEImageTest, RetrievesExports) { |
| HMODULE module = LoadLibrary(L"advapi32.dll"); |
| ASSERT_TRUE(NULL != module); |
| |
| PEImage pe(module); |
| WORD ordinal; |
| |
| EXPECT_TRUE(pe.GetProcOrdinal("RegEnumKeyExW", &ordinal)); |
| |
| FARPROC address1 = pe.GetProcAddress("RegEnumKeyExW"); |
| FARPROC address2 = pe.GetProcAddress(reinterpret_cast<char*>(ordinal)); |
| EXPECT_TRUE(address1 != NULL); |
| EXPECT_TRUE(address2 != NULL); |
| EXPECT_TRUE(address1 == address2); |
| |
| FreeLibrary(module); |
| } |
| |
| // Test that we can get debug id out of a module. |
| TEST(PEImageTest, GetDebugId) { |
| HMODULE module = LoadLibrary(L"advapi32.dll"); |
| ASSERT_TRUE(NULL != module); |
| |
| PEImage pe(module); |
| GUID guid = {0}; |
| DWORD age = 0; |
| EXPECT_TRUE(pe.GetDebugId(&guid, &age)); |
| |
| GUID empty_guid = {0}; |
| EXPECT_TRUE(!IsEqualGUID(empty_guid, guid)); |
| EXPECT_NE(0U, age); |
| FreeLibrary(module); |
| } |
| |
| } // namespace win |
| } // namespace base |