com_init_check_hook_unittest.cc 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. // Copyright 2017 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. #include "base/win/com_init_check_hook.h"
  5. #include <objbase.h>
  6. #include <shlobj.h>
  7. #include <wrl/client.h>
  8. #include "base/test/gtest_util.h"
  9. #include "base/win/com_init_util.h"
  10. #include "base/win/patch_util.h"
  11. #include "base/win/scoped_com_initializer.h"
  12. #include "testing/gtest/include/gtest/gtest.h"
  13. namespace base {
  14. namespace win {
  15. using Microsoft::WRL::ComPtr;
  16. TEST(ComInitCheckHook, AssertNotInitialized) {
  17. ComInitCheckHook com_check_hook;
  18. AssertComApartmentType(ComApartmentType::NONE);
  19. ComPtr<IUnknown> shell_link;
  20. #if defined(COM_INIT_CHECK_HOOK_ENABLED)
  21. EXPECT_DCHECK_DEATH(::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  22. IID_PPV_ARGS(&shell_link)));
  23. #else
  24. EXPECT_EQ(CO_E_NOTINITIALIZED,
  25. ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  26. IID_PPV_ARGS(&shell_link)));
  27. #endif
  28. }
  29. TEST(ComInitCheckHook, HookRemoval) {
  30. AssertComApartmentType(ComApartmentType::NONE);
  31. { ComInitCheckHook com_check_hook; }
  32. ComPtr<IUnknown> shell_link;
  33. EXPECT_EQ(CO_E_NOTINITIALIZED,
  34. ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  35. IID_PPV_ARGS(&shell_link)));
  36. }
  37. TEST(ComInitCheckHook, NoAssertComInitialized) {
  38. ComInitCheckHook com_check_hook;
  39. ScopedCOMInitializer com_initializer;
  40. ComPtr<IUnknown> shell_link;
  41. EXPECT_TRUE(SUCCEEDED(::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  42. IID_PPV_ARGS(&shell_link))));
  43. }
  44. TEST(ComInitCheckHook, MultipleHooks) {
  45. ComInitCheckHook com_check_hook_1;
  46. ComInitCheckHook com_check_hook_2;
  47. AssertComApartmentType(ComApartmentType::NONE);
  48. ComPtr<IUnknown> shell_link;
  49. #if defined(COM_INIT_CHECK_HOOK_ENABLED)
  50. EXPECT_DCHECK_DEATH(::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  51. IID_PPV_ARGS(&shell_link)));
  52. #else
  53. EXPECT_EQ(CO_E_NOTINITIALIZED,
  54. ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
  55. IID_PPV_ARGS(&shell_link)));
  56. #endif
  57. }
  58. TEST(ComInitCheckHook, UnexpectedHook) {
  59. #if defined(COM_INIT_CHECK_HOOK_ENABLED)
  60. HMODULE ole32_library = ::LoadLibrary(L"ole32.dll");
  61. ASSERT_TRUE(ole32_library);
  62. uint32_t co_create_instance_padded_address =
  63. reinterpret_cast<uint32_t>(
  64. GetProcAddress(ole32_library, "CoCreateInstance")) -
  65. 5;
  66. const unsigned char* co_create_instance_bytes =
  67. reinterpret_cast<const unsigned char*>(co_create_instance_padded_address);
  68. const unsigned char original_byte = co_create_instance_bytes[0];
  69. const unsigned char unexpected_byte = 0xdb;
  70. ASSERT_EQ(static_cast<DWORD>(NO_ERROR),
  71. internal::ModifyCode(
  72. reinterpret_cast<void*>(co_create_instance_padded_address),
  73. reinterpret_cast<const void*>(&unexpected_byte),
  74. sizeof(unexpected_byte)));
  75. EXPECT_DCHECK_DEATH({ ComInitCheckHook com_check_hook; });
  76. // If this call fails, really bad things are going to happen to other tests
  77. // so CHECK here.
  78. CHECK_EQ(static_cast<DWORD>(NO_ERROR),
  79. internal::ModifyCode(
  80. reinterpret_cast<void*>(co_create_instance_padded_address),
  81. reinterpret_cast<const void*>(&original_byte),
  82. sizeof(original_byte)));
  83. ::FreeLibrary(ole32_library);
  84. ole32_library = nullptr;
  85. #endif
  86. }
  87. TEST(ComInitCheckHook, ExternallyHooked) {
  88. #if defined(COM_INIT_CHECK_HOOK_ENABLED)
  89. HMODULE ole32_library = ::LoadLibrary(L"ole32.dll");
  90. ASSERT_TRUE(ole32_library);
  91. uint32_t co_create_instance_address = reinterpret_cast<uint32_t>(
  92. GetProcAddress(ole32_library, "CoCreateInstance"));
  93. const unsigned char* co_create_instance_bytes =
  94. reinterpret_cast<const unsigned char*>(co_create_instance_address);
  95. const unsigned char original_byte = co_create_instance_bytes[0];
  96. const unsigned char jmp_byte = 0xe9;
  97. ASSERT_EQ(static_cast<DWORD>(NO_ERROR),
  98. internal::ModifyCode(
  99. reinterpret_cast<void*>(co_create_instance_address),
  100. reinterpret_cast<const void*>(&jmp_byte), sizeof(jmp_byte)));
  101. // Externally patched instances should crash so we catch these cases on bots.
  102. EXPECT_DCHECK_DEATH({ ComInitCheckHook com_check_hook; });
  103. // If this call fails, really bad things are going to happen to other tests
  104. // so CHECK here.
  105. CHECK_EQ(
  106. static_cast<DWORD>(NO_ERROR),
  107. internal::ModifyCode(reinterpret_cast<void*>(co_create_instance_address),
  108. reinterpret_cast<const void*>(&original_byte),
  109. sizeof(original_byte)));
  110. ::FreeLibrary(ole32_library);
  111. ole32_library = nullptr;
  112. #endif
  113. }
  114. TEST(ComInitCheckHook, UnexpectedChangeDuringHook) {
  115. #if defined(COM_INIT_CHECK_HOOK_ENABLED)
  116. HMODULE ole32_library = ::LoadLibrary(L"ole32.dll");
  117. ASSERT_TRUE(ole32_library);
  118. uint32_t co_create_instance_padded_address =
  119. reinterpret_cast<uint32_t>(
  120. GetProcAddress(ole32_library, "CoCreateInstance")) -
  121. 5;
  122. const unsigned char* co_create_instance_bytes =
  123. reinterpret_cast<const unsigned char*>(co_create_instance_padded_address);
  124. const unsigned char original_byte = co_create_instance_bytes[0];
  125. const unsigned char unexpected_byte = 0xdb;
  126. ASSERT_EQ(static_cast<DWORD>(NO_ERROR),
  127. internal::ModifyCode(
  128. reinterpret_cast<void*>(co_create_instance_padded_address),
  129. reinterpret_cast<const void*>(&unexpected_byte),
  130. sizeof(unexpected_byte)));
  131. EXPECT_DCHECK_DEATH({
  132. ComInitCheckHook com_check_hook;
  133. internal::ModifyCode(
  134. reinterpret_cast<void*>(co_create_instance_padded_address),
  135. reinterpret_cast<const void*>(&unexpected_byte),
  136. sizeof(unexpected_byte));
  137. });
  138. // If this call fails, really bad things are going to happen to other tests
  139. // so CHECK here.
  140. CHECK_EQ(static_cast<DWORD>(NO_ERROR),
  141. internal::ModifyCode(
  142. reinterpret_cast<void*>(co_create_instance_padded_address),
  143. reinterpret_cast<const void*>(&original_byte),
  144. sizeof(original_byte)));
  145. ::FreeLibrary(ole32_library);
  146. ole32_library = nullptr;
  147. #endif
  148. }
  149. } // namespace win
  150. } // namespace base