1 // Copyright 2020 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 <objbase.h> 6 7 #include "base/check_op.h" 8 #include "base/win/com_init_balancer.h" 9 10 namespace base { 11 namespace win { 12 namespace internal { 13 ComInitBalancer(DWORD co_init)14ComInitBalancer::ComInitBalancer(DWORD co_init) : co_init_(co_init) { 15 ULARGE_INTEGER spy_cookie = {}; 16 HRESULT hr = ::CoRegisterInitializeSpy(this, &spy_cookie); 17 if (SUCCEEDED(hr)) 18 spy_cookie_ = spy_cookie; 19 } 20 ~ComInitBalancer()21ComInitBalancer::~ComInitBalancer() { 22 DCHECK(!spy_cookie_.has_value()); 23 } 24 Disable()25void ComInitBalancer::Disable() { 26 if (spy_cookie_.has_value()) { 27 ::CoRevokeInitializeSpy(spy_cookie_.value()); 28 reference_count_ = 0; 29 spy_cookie_.reset(); 30 } 31 } 32 GetReferenceCountForTesting() const33DWORD ComInitBalancer::GetReferenceCountForTesting() const { 34 return reference_count_; 35 } 36 37 IFACEMETHODIMP PreInitialize(DWORD apartment_type,DWORD reference_count)38ComInitBalancer::PreInitialize(DWORD apartment_type, DWORD reference_count) { 39 return S_OK; 40 } 41 42 IFACEMETHODIMP PostInitialize(HRESULT result,DWORD apartment_type,DWORD new_reference_count)43ComInitBalancer::PostInitialize(HRESULT result, 44 DWORD apartment_type, 45 DWORD new_reference_count) { 46 reference_count_ = new_reference_count; 47 return result; 48 } 49 50 IFACEMETHODIMP PreUninitialize(DWORD reference_count)51ComInitBalancer::PreUninitialize(DWORD reference_count) { 52 if (reference_count == 1 && spy_cookie_.has_value()) { 53 // Increase the reference count to prevent premature and unbalanced 54 // uninitalization of the COM library. 55 ::CoInitializeEx(nullptr, co_init_); 56 } 57 return S_OK; 58 } 59 60 IFACEMETHODIMP PostUninitialize(DWORD new_reference_count)61ComInitBalancer::PostUninitialize(DWORD new_reference_count) { 62 reference_count_ = new_reference_count; 63 return S_OK; 64 } 65 66 } // namespace internal 67 } // namespace win 68 } // namespace base 69