Rework PE image tests to use a compiled DLL instead of a system one.

base/win/pe_image_unittest.cc

 // 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 base {
 namespace win {
 
+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_X86_ON_X64,
+    ARCH_X64_ON_X64,
+    ARCH_X86_ON_X86,
+    ARCH_ALL
+  };
+
+  int GetExpectation(Value value);
+  void SetOverride(Value value, Version version, Arch arch, int count);
+  void SetOverride(Value value, Version version, int count);
+  void SetOverride(Value value, Arch arch, int count);
+
+  void SetDefault(Value value, int count) {
+    SetOverride(value, VERSION_WIN_LAST, ARCH_ALL, count);
+  }
+
+  Expectations() {
+    my_version_ = GetVersion();
+#if defined(ARCH_CPU_64_BITS)
+    my_arch_ = ARCH_X64_ON_X64;
+#else
+    OSInfo* os_info = OSInfo::GetInstance();
+    if (os_info->wow64_status() == OSInfo::WOW64_DISABLED)
+      my_arch_ = ARCH_X86_ON_X86;
+    else
+      my_arch_ = ARCH_X86_ON_X64;
+#endif
+  }
+
+ private:
+  class Override {
+   public:
+    enum MatchType { MATCH_VERSION, MATCH_ARCH, MATCH_BOTH, MATCH_NONE };
+    explicit Override(Value value, Version version, Arch arch, int count)
+        : value_(value), version_(version), arch_(arch), count_(count){};
+
+    bool Matches(Value value, Version version, Arch arch, MatchType type) {
+      if (value_ != value)
+        return false;
+      switch (type) {
+        case MATCH_BOTH:
+          return MatchesVersion(version) && MatchesArch(arch);
+          break;
+        case MATCH_ARCH:
+          return MatchesArch(arch);
+          break;
+        case MATCH_VERSION:
+          return MatchesVersion(version);
+          break;
+        case MATCH_NONE:
+          return true;
+          break;
+      }
+      return false;
+    }
+
+    int GetCount() { return count_; }
+
+   private:
+    bool MatchesArch(Arch arch) {
+      switch (arch) {
+        case ARCH_X64_ON_X64:
+          return (arch_ == ARCH_X64_ON_X64 || arch_ == ARCH_X64);
+          break;
+        case ARCH_X86_ON_X86:
+          return (arch_ == ARCH_X86_ON_X86 || arch_ == ARCH_X86);
+          break;
+        case ARCH_X86_ON_X64:
+          return (arch_ == ARCH_X86_ON_X64 || arch_ == ARCH_X86);
+          break;
+        default:
+          return false;
+      }
+    }
+
+    bool MatchesVersion(Version version) { return (version == version_); }
+
+    Value value_;
+    Version version_;
+    Arch arch_;
+    int count_;
+  };
+
+  bool MatchesMyArch(Arch arch);
+
+  std::vector<Override> overrides_;
+  Arch my_arch_;
+  Version my_version_;
+};
+
+int Expectations::GetExpectation(Value value) {
+  // Prefer OS version specificity over Arch specificity.
+  for (auto type : {Expectations::Override::MATCH_BOTH,
+                    Expectations::Override::MATCH_VERSION,
+                    Expectations::Override::MATCH_ARCH,
+                    Expectations::Override::MATCH_NONE})
+    for (auto or : overrides_)
+      if (or.Matches(value, my_version_, my_arch_, type))
+        return or.GetCount();
+  return 0;
+}
+
+void Expectations::SetOverride(Value value,
+                               Version version,
+                               Arch arch,
+                               int count) {
+  overrides_.push_back(Override(value, version, arch, count));
+}
+
+void Expectations::SetOverride(Value value, Version version, int count) {
+  SetOverride(value, version, ARCH_ALL, count);
+}
+
+void Expectations::SetOverride(Value value, Arch arch, int count) {
+  SetOverride(value, VERSION_WIN_LAST, arch, count);
+}
+
 // 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)++;
@@ skipping 52 matching lines @@
                                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;
 }
 
-// Identifiers for the set of supported expectations.
-enum ExpectationSet {
-  WIN_2K_SET,
-  WIN_XP_SET,
-  WIN_VISTA_SET,
-  WIN_7_SET,
-  WIN_8_SET,
-  UNSUPPORTED_SET,
-};
-
-// We'll be using some known values for the tests.
-enum Value {
-  sections = 0,
-  imports_dlls,
-  delay_dlls,
-  exports,
-  imports,
-  delay_imports,
-  relocs
-};
-
-ExpectationSet GetExpectationSet(DWORD os) {
-  if (os == 50)
-    return WIN_2K_SET;
-  if (os == 51)
-    return WIN_XP_SET;
-  if (os == 60)
-    return WIN_VISTA_SET;
-  if (os == 61)
-    return WIN_7_SET;
-  if (os >= 62)
-    return WIN_8_SET;
-  return UNSUPPORTED_SET;
-}
-
-// Retrieves the expected value from advapi32.dll based on the OS.
-int GetExpectedValue(Value value, DWORD os) {
-  const int xp_delay_dlls = 2;
-  const int xp_exports = 675;
-  const int xp_imports = 422;
-  const int xp_delay_imports = 8;
-  const int xp_relocs = 9180;
-  const int vista_delay_dlls = 4;
-  const int vista_exports = 799;
-  const int vista_imports = 476;
-  const int vista_delay_imports = 24;
-  const int vista_relocs = 10188;
-  const int w2k_delay_dlls = 0;
-  const int w2k_exports = 566;
-  const int w2k_imports = 357;
-  const int w2k_delay_imports = 0;
-  const int w2k_relocs = 7388;
-  const int win7_delay_dlls = 7;
-  const int win7_exports = 806;
-  const int win7_imports = 568;
-  const int win7_delay_imports = 71;
-  int win7_relocs = 7812;
-  int win7_sections = 4;
-  const int win8_delay_dlls = 9;
-  const int win8_exports = 806;
-  const int win8_imports = 568;
-  const int win8_delay_imports = 113;
-  const int win8_relocs = 9478;
-  int win8_sections = 4;
-  int win8_import_dlls = 17;
-
-  base::win::OSInfo* os_info = base::win::OSInfo::GetInstance();
-  // 32-bit process on a 32-bit system.
-  if (os_info->architecture() == base::win::OSInfo::X86_ARCHITECTURE) {
-    win8_sections = 5;
-    win8_import_dlls = 19;
-
-  // 64-bit process on a 64-bit system.
-  } else if (os_info->wow64_status() == base::win::OSInfo::WOW64_DISABLED) {
-    win7_sections = 6;
-    win7_relocs = 2712;
-  }
-
-  // Contains the expected value, for each enumerated property (Value), and the
-  // OS version: [Value][os_version]
-  const int expected[][5] = {
-    {4, 4, 4, win7_sections, win8_sections},
-    {3, 3, 3, 13, win8_import_dlls},
-    {w2k_delay_dlls, xp_delay_dlls, vista_delay_dlls, win7_delay_dlls,
-     win8_delay_dlls},
-    {w2k_exports, xp_exports, vista_exports, win7_exports, win8_exports},
-    {w2k_imports, xp_imports, vista_imports, win7_imports, win8_imports},
-    {w2k_delay_imports, xp_delay_imports,
-     vista_delay_imports, win7_delay_imports, win8_delay_imports},
-    {w2k_relocs, xp_relocs, vista_relocs, win7_relocs, win8_relocs}
-  };
-  COMPILE_ASSERT(arraysize(expected[0]) == UNSUPPORTED_SET,
-                 expected_value_set_mismatch);
-
-  if (value > relocs)
-    return 0;
-  ExpectationSet expected_set = GetExpectationSet(os);
-  if (expected_set >= arraysize(expected)) {
-    // This should never happen.  Log a failure if it does.
-    EXPECT_NE(UNSUPPORTED_SET, expected_set);
-    expected_set = WIN_2K_SET;
-  }
-
-  return expected[value][expected_set];
-}
-
-
-// TODO(jschuh): crbug.com/167707 Need to fix test on Win64 bots
-#if defined(OS_WIN) && defined(ARCH_CPU_X86_64)
-#define MAYBE_EnumeratesPE DISABLED_EnumeratesPE
-#else
-#define MAYBE_EnumeratesPE EnumeratesPE
-#endif
+}  // 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, MAYBE_EnumeratesPE) {
-  HMODULE module = LoadLibrary(L"advapi32.dll");
+TEST(PEImageTest, EnumeratesPE) {
+  Expectations expectations;
+
+  // Default expectations.
+  expectations.SetDefault(Expectations::sections, 5);
+  expectations.SetDefault(Expectations::imports_dlls, 1);
+  expectations.SetDefault(Expectations::delay_dlls, 1);
+  expectations.SetDefault(Expectations::exports, 1);
+  expectations.SetDefault(Expectations::imports, 63);
+  expectations.SetDefault(Expectations::delay_imports, 2);
+  expectations.SetDefault(Expectations::relocs, 1584);
+
+  // win_chromium_x64_rel_ng (Win7 x64 on x64)
+  expectations.SetOverride(Expectations::sections, VERSION_WIN7,
+                           Expectations::ARCH_X64, 6);
+  expectations.SetOverride(Expectations::imports, VERSION_WIN7,
+                           Expectations::ARCH_X64, 66);
+  expectations.SetOverride(Expectations::relocs, VERSION_WIN7,
+                           Expectations::ARCH_X64, 630);
+
+  HMODULE module = LoadLibrary(L"pe_image_test.dll");
   ASSERT_TRUE(NULL != module);
 
   PEImage pe(module);
   int count = 0;
   EXPECT_TRUE(pe.VerifyMagic());
 
-  DWORD os = pe.GetNTHeaders()->OptionalHeader.MajorOperatingSystemVersion;
-  os = os * 10 + pe.GetNTHeaders()->OptionalHeader.MinorOperatingSystemVersion;
-
-  // Skip this test for unsupported OS versions.
-  if (GetExpectationSet(os) == UNSUPPORTED_SET)
-    return;
-
   pe.EnumSections(SectionsCallback, &count);
-  EXPECT_EQ(GetExpectedValue(sections, os), count);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::sections), count);
 
   count = 0;
   pe.EnumImportChunks(ImportChunksCallback, &count);
-  EXPECT_EQ(GetExpectedValue(imports_dlls, os), count);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::imports_dlls), count);
 
   count = 0;
   pe.EnumDelayImportChunks(DelayImportChunksCallback, &count);
-  EXPECT_EQ(GetExpectedValue(delay_dlls, os), count);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::delay_dlls), count);
 
   count = 0;
   pe.EnumExports(ExportsCallback, &count);
-  EXPECT_GT(count, GetExpectedValue(exports, os) - 20);
-  EXPECT_LT(count, GetExpectedValue(exports, os) + 100);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::exports), count);
 
   count = 0;
   pe.EnumAllImports(ImportsCallback, &count);
-  EXPECT_GT(count, GetExpectedValue(imports, os) - 20);
-  EXPECT_LT(count, GetExpectedValue(imports, os) + 100);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::imports), count);
 
   count = 0;
   pe.EnumAllDelayImports(ImportsCallback, &count);
-  EXPECT_GT(count, GetExpectedValue(delay_imports, os) - 2);
-  EXPECT_LT(count, GetExpectedValue(delay_imports, os) + 8);
+  EXPECT_EQ(expectations.GetExpectation(Expectations::delay_imports), count);
 
   count = 0;
   pe.EnumRelocs(RelocsCallback, &count);
-  EXPECT_GT(count, GetExpectedValue(relocs, os) - 150);
-  EXPECT_LT(count, GetExpectedValue(relocs, os) + 1500);
+  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);
@@ skipping 21 matching lines @@
   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