1 // Copyright 2017 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #include "base/win/com_init_util.h" 6 7 #include <windows.h> 8 9 #include <winternl.h> 10 #include "base/logging.h" 11 #include "base/memory/raw_ptr_exclusion.h" 12 #include "base/notreached.h" 13 14 namespace base { 15 namespace win { 16 17 namespace { 18 19 #if DCHECK_IS_ON() 20 const char kComNotInitialized[] = "COM is not initialized on this thread."; 21 #endif // DCHECK_IS_ON() 22 23 // Derived from combase.dll. 24 struct OleTlsData { 25 enum ApartmentFlags { 26 LOGICAL_THREAD_REGISTERED = 0x2, 27 STA = 0x80, 28 MTA = 0x140, 29 }; 30 31 // This field is not a raw_ptr<> because it was filtered by the rewriter for: 32 // #reinterpret-cast-trivial-type 33 RAW_PTR_EXCLUSION void* thread_base; 34 // This field is not a raw_ptr<> because it was filtered by the rewriter for: 35 // #reinterpret-cast-trivial-type 36 RAW_PTR_EXCLUSION void* sm_allocator; 37 DWORD apartment_id; 38 DWORD apartment_flags; 39 // There are many more fields than this, but for our purposes, we only care 40 // about |apartment_flags|. Correctly declaring the previous types allows this 41 // to work between x86 and x64 builds. 42 }; 43 GetOleTlsData()44OleTlsData* GetOleTlsData() { 45 TEB* teb = NtCurrentTeb(); 46 return reinterpret_cast<OleTlsData*>(teb->ReservedForOle); 47 } 48 49 } // namespace 50 GetComApartmentTypeForThread()51ComApartmentType GetComApartmentTypeForThread() { 52 OleTlsData* ole_tls_data = GetOleTlsData(); 53 if (!ole_tls_data) 54 return ComApartmentType::NONE; 55 56 if (ole_tls_data->apartment_flags & OleTlsData::ApartmentFlags::STA) 57 return ComApartmentType::STA; 58 59 if ((ole_tls_data->apartment_flags & OleTlsData::ApartmentFlags::MTA) == 60 OleTlsData::ApartmentFlags::MTA) { 61 return ComApartmentType::MTA; 62 } 63 64 return ComApartmentType::NONE; 65 } 66 67 #if DCHECK_IS_ON() 68 AssertComInitialized(const char * message)69void AssertComInitialized(const char* message) { 70 if (GetComApartmentTypeForThread() != ComApartmentType::NONE) 71 return; 72 73 // COM worker threads don't always set up the apartment, but they do perform 74 // some thread registration, so we allow those. 75 OleTlsData* ole_tls_data = GetOleTlsData(); 76 if (ole_tls_data && (ole_tls_data->apartment_flags & 77 OleTlsData::ApartmentFlags::LOGICAL_THREAD_REGISTERED)) { 78 return; 79 } 80 81 NOTREACHED() << (message ? message : kComNotInitialized); 82 } 83 AssertComApartmentType(ComApartmentType apartment_type)84void AssertComApartmentType(ComApartmentType apartment_type) { 85 DCHECK_EQ(apartment_type, GetComApartmentTypeForThread()); 86 } 87 88 #endif // DCHECK_IS_ON() 89 90 } // namespace win 91 } // namespace base 92