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 "base/win/com_init_balancer.h"
6
7 #include <shlobj.h>
8 #include <wrl/client.h>
9
10 #include "base/test/gtest_util.h"
11 #include "base/win/com_init_util.h"
12 #include "base/win/scoped_com_initializer.h"
13 #include "testing/gtest/include/gtest/gtest.h"
14
15 namespace base {
16 namespace win {
17
18 using Microsoft::WRL::ComPtr;
19
TEST(TestComInitBalancer,BalancedPairsWithComBalancerEnabled)20 TEST(TestComInitBalancer, BalancedPairsWithComBalancerEnabled) {
21 {
22 // Assert COM has initialized correctly.
23 ScopedCOMInitializer com_initializer(
24 ScopedCOMInitializer::Uninitialization::kBlockPremature);
25 ASSERT_TRUE(com_initializer.Succeeded());
26
27 // Create COM object successfully.
28 ComPtr<IUnknown> shell_link;
29 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
30 IID_PPV_ARGS(&shell_link));
31 EXPECT_TRUE(SUCCEEDED(hr));
32 }
33
34 // ScopedCOMInitializer has gone out of scope and COM has been uninitialized.
35 EXPECT_DCHECK_DEATH(AssertComInitialized());
36 }
37
TEST(TestComInitBalancer,UnbalancedPairsWithComBalancerEnabled)38 TEST(TestComInitBalancer, UnbalancedPairsWithComBalancerEnabled) {
39 {
40 // Assert COM has initialized correctly.
41 ScopedCOMInitializer com_initializer(
42 ScopedCOMInitializer::Uninitialization::kBlockPremature);
43 ASSERT_TRUE(com_initializer.Succeeded());
44
45 // Attempt to prematurely uninitialize the COM library.
46 ::CoUninitialize();
47 ::CoUninitialize();
48
49 // Assert COM is still initialized.
50 AssertComInitialized();
51
52 // Create COM object successfully.
53 ComPtr<IUnknown> shell_link;
54 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
55 IID_PPV_ARGS(&shell_link));
56 EXPECT_TRUE(SUCCEEDED(hr));
57 }
58
59 // ScopedCOMInitializer has gone out of scope and COM has been uninitialized.
60 EXPECT_DCHECK_DEATH(AssertComInitialized());
61 }
62
TEST(TestComInitBalancer,BalancedPairsWithComBalancerDisabled)63 TEST(TestComInitBalancer, BalancedPairsWithComBalancerDisabled) {
64 {
65 // Assert COM has initialized correctly.
66 ScopedCOMInitializer com_initializer(
67 ScopedCOMInitializer::Uninitialization::kAllow);
68 ASSERT_TRUE(com_initializer.Succeeded());
69
70 // Create COM object successfully.
71 ComPtr<IUnknown> shell_link;
72 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
73 IID_PPV_ARGS(&shell_link));
74 EXPECT_TRUE(SUCCEEDED(hr));
75 }
76
77 // ScopedCOMInitializer has gone out of scope and COM has been uninitialized.
78 EXPECT_DCHECK_DEATH(AssertComInitialized());
79 }
80
TEST(TestComInitBalancer,UnbalancedPairsWithComBalancerDisabled)81 TEST(TestComInitBalancer, UnbalancedPairsWithComBalancerDisabled) {
82 // Assert COM has initialized correctly.
83 ScopedCOMInitializer com_initializer(
84 ScopedCOMInitializer::Uninitialization::kAllow);
85 ASSERT_TRUE(com_initializer.Succeeded());
86
87 // Attempt to prematurely uninitialize the COM library.
88 ::CoUninitialize();
89 ::CoUninitialize();
90
91 // Assert COM is not initialized.
92 EXPECT_DCHECK_DEATH(AssertComInitialized());
93
94 // Create COM object unsuccessfully.
95 ComPtr<IUnknown> shell_link;
96 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
97 IID_PPV_ARGS(&shell_link));
98 EXPECT_TRUE(FAILED(hr));
99 EXPECT_EQ(CO_E_NOTINITIALIZED, hr);
100 }
101
TEST(TestComInitBalancer,OneRegisteredSpyRefCount)102 TEST(TestComInitBalancer, OneRegisteredSpyRefCount) {
103 ScopedCOMInitializer com_initializer(
104 ScopedCOMInitializer::Uninitialization::kBlockPremature);
105 ASSERT_TRUE(com_initializer.Succeeded());
106
107 // Reference count should be 1 after initialization.
108 EXPECT_EQ(DWORD(1), com_initializer.GetCOMBalancerReferenceCountForTesting());
109
110 // Attempt to prematurely uninitialize the COM library.
111 ::CoUninitialize();
112
113 // Expect reference count to remain at 1.
114 EXPECT_EQ(DWORD(1), com_initializer.GetCOMBalancerReferenceCountForTesting());
115 }
116
TEST(TestComInitBalancer,ThreeRegisteredSpiesRefCount)117 TEST(TestComInitBalancer, ThreeRegisteredSpiesRefCount) {
118 ScopedCOMInitializer com_initializer_1(
119 ScopedCOMInitializer::Uninitialization::kBlockPremature);
120 ScopedCOMInitializer com_initializer_2(
121 ScopedCOMInitializer::Uninitialization::kBlockPremature);
122 ScopedCOMInitializer com_initializer_3(
123 ScopedCOMInitializer::Uninitialization::kBlockPremature);
124 ASSERT_TRUE(com_initializer_1.Succeeded());
125 ASSERT_TRUE(com_initializer_2.Succeeded());
126 ASSERT_TRUE(com_initializer_3.Succeeded());
127
128 // Reference count should be 3 after initialization.
129 EXPECT_EQ(DWORD(3),
130 com_initializer_1.GetCOMBalancerReferenceCountForTesting());
131 EXPECT_EQ(DWORD(3),
132 com_initializer_2.GetCOMBalancerReferenceCountForTesting());
133 EXPECT_EQ(DWORD(3),
134 com_initializer_3.GetCOMBalancerReferenceCountForTesting());
135
136 // Attempt to prematurely uninitialize the COM library.
137 ::CoUninitialize(); // Reference count -> 2.
138 ::CoUninitialize(); // Reference count -> 1.
139 ::CoUninitialize();
140
141 // Expect reference count to remain at 1.
142 EXPECT_EQ(DWORD(1),
143 com_initializer_1.GetCOMBalancerReferenceCountForTesting());
144 EXPECT_EQ(DWORD(1),
145 com_initializer_2.GetCOMBalancerReferenceCountForTesting());
146 EXPECT_EQ(DWORD(1),
147 com_initializer_3.GetCOMBalancerReferenceCountForTesting());
148 }
149
150 } // namespace win
151 } // namespace base
152