com_init_util.cc 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  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_util.h"
  5. #include <windows.h>
  6. #include <winternl.h>
  7. #include "base/logging.h"
  8. #include "base/notreached.h"
  9. namespace base {
  10. namespace win {
  11. namespace {
  12. #if DCHECK_IS_ON()
  13. const char kComNotInitialized[] = "COM is not initialized on this thread.";
  14. #endif // DCHECK_IS_ON()
  15. // Derived from combase.dll.
  16. struct OleTlsData {
  17. enum ApartmentFlags {
  18. LOGICAL_THREAD_REGISTERED = 0x2,
  19. STA = 0x80,
  20. MTA = 0x140,
  21. };
  22. void* thread_base;
  23. void* sm_allocator;
  24. DWORD apartment_id;
  25. DWORD apartment_flags;
  26. // There are many more fields than this, but for our purposes, we only care
  27. // about |apartment_flags|. Correctly declaring the previous types allows this
  28. // to work between x86 and x64 builds.
  29. };
  30. OleTlsData* GetOleTlsData() {
  31. TEB* teb = NtCurrentTeb();
  32. return reinterpret_cast<OleTlsData*>(teb->ReservedForOle);
  33. }
  34. } // namespace
  35. ComApartmentType GetComApartmentTypeForThread() {
  36. OleTlsData* ole_tls_data = GetOleTlsData();
  37. if (!ole_tls_data)
  38. return ComApartmentType::NONE;
  39. if (ole_tls_data->apartment_flags & OleTlsData::ApartmentFlags::STA)
  40. return ComApartmentType::STA;
  41. if ((ole_tls_data->apartment_flags & OleTlsData::ApartmentFlags::MTA) ==
  42. OleTlsData::ApartmentFlags::MTA) {
  43. return ComApartmentType::MTA;
  44. }
  45. return ComApartmentType::NONE;
  46. }
  47. #if DCHECK_IS_ON()
  48. void AssertComInitialized(const char* message) {
  49. if (GetComApartmentTypeForThread() != ComApartmentType::NONE)
  50. return;
  51. // COM worker threads don't always set up the apartment, but they do perform
  52. // some thread registration, so we allow those.
  53. OleTlsData* ole_tls_data = GetOleTlsData();
  54. if (ole_tls_data && (ole_tls_data->apartment_flags &
  55. OleTlsData::ApartmentFlags::LOGICAL_THREAD_REGISTERED)) {
  56. return;
  57. }
  58. NOTREACHED() << (message ? message : kComNotInitialized);
  59. }
  60. void AssertComApartmentType(ComApartmentType apartment_type) {
  61. DCHECK_EQ(apartment_type, GetComApartmentTypeForThread());
  62. }
  63. #endif // DCHECK_IS_ON()
  64. } // namespace win
  65. } // namespace base