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 if (DCHECK_IS_ON()) {
36 EXPECT_NOTREACHED_DEATH(AssertComInitialized());
37 }
38 }
39
TEST(TestComInitBalancer,UnbalancedPairsWithComBalancerEnabled)40 TEST(TestComInitBalancer, UnbalancedPairsWithComBalancerEnabled) {
41 {
42 // Assert COM has initialized correctly.
43 ScopedCOMInitializer com_initializer(
44 ScopedCOMInitializer::Uninitialization::kBlockPremature);
45 ASSERT_TRUE(com_initializer.Succeeded());
46
47 // Attempt to prematurely uninitialize the COM library.
48 ::CoUninitialize();
49 ::CoUninitialize();
50
51 // Assert COM is still initialized.
52 AssertComInitialized();
53
54 // Create COM object successfully.
55 ComPtr<IUnknown> shell_link;
56 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
57 IID_PPV_ARGS(&shell_link));
58 EXPECT_TRUE(SUCCEEDED(hr));
59 }
60
61 // ScopedCOMInitializer has gone out of scope and COM has been uninitialized.
62 if (DCHECK_IS_ON()) {
63 EXPECT_NOTREACHED_DEATH(AssertComInitialized());
64 }
65 }
66
TEST(TestComInitBalancer,BalancedPairsWithComBalancerDisabled)67 TEST(TestComInitBalancer, BalancedPairsWithComBalancerDisabled) {
68 {
69 // Assert COM has initialized correctly.
70 ScopedCOMInitializer com_initializer(
71 ScopedCOMInitializer::Uninitialization::kAllow);
72 ASSERT_TRUE(com_initializer.Succeeded());
73
74 // Create COM object successfully.
75 ComPtr<IUnknown> shell_link;
76 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
77 IID_PPV_ARGS(&shell_link));
78 EXPECT_TRUE(SUCCEEDED(hr));
79 }
80
81 // ScopedCOMInitializer has gone out of scope and COM has been uninitialized.
82 if (DCHECK_IS_ON()) {
83 EXPECT_NOTREACHED_DEATH(AssertComInitialized());
84 }
85 }
86
TEST(TestComInitBalancer,UnbalancedPairsWithComBalancerDisabled)87 TEST(TestComInitBalancer, UnbalancedPairsWithComBalancerDisabled) {
88 // Assert COM has initialized correctly.
89 ScopedCOMInitializer com_initializer(
90 ScopedCOMInitializer::Uninitialization::kAllow);
91 ASSERT_TRUE(com_initializer.Succeeded());
92
93 // Attempt to prematurely uninitialize the COM library.
94 ::CoUninitialize();
95 ::CoUninitialize();
96
97 // Assert COM is not initialized.
98 if (DCHECK_IS_ON()) {
99 EXPECT_NOTREACHED_DEATH(AssertComInitialized());
100 }
101
102 // Create COM object unsuccessfully.
103 ComPtr<IUnknown> shell_link;
104 HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_ALL,
105 IID_PPV_ARGS(&shell_link));
106 EXPECT_TRUE(FAILED(hr));
107 EXPECT_EQ(CO_E_NOTINITIALIZED, hr);
108 }
109
TEST(TestComInitBalancer,OneRegisteredSpyRefCount)110 TEST(TestComInitBalancer, OneRegisteredSpyRefCount) {
111 ScopedCOMInitializer com_initializer(
112 ScopedCOMInitializer::Uninitialization::kBlockPremature);
113 ASSERT_TRUE(com_initializer.Succeeded());
114
115 // Reference count should be 1 after initialization.
116 EXPECT_EQ(DWORD(1), com_initializer.GetCOMBalancerReferenceCountForTesting());
117
118 // Attempt to prematurely uninitialize the COM library.
119 ::CoUninitialize();
120
121 // Expect reference count to remain at 1.
122 EXPECT_EQ(DWORD(1), com_initializer.GetCOMBalancerReferenceCountForTesting());
123 }
124
TEST(TestComInitBalancer,ThreeRegisteredSpiesRefCount)125 TEST(TestComInitBalancer, ThreeRegisteredSpiesRefCount) {
126 ScopedCOMInitializer com_initializer_1(
127 ScopedCOMInitializer::Uninitialization::kBlockPremature);
128 ScopedCOMInitializer com_initializer_2(
129 ScopedCOMInitializer::Uninitialization::kBlockPremature);
130 ScopedCOMInitializer com_initializer_3(
131 ScopedCOMInitializer::Uninitialization::kBlockPremature);
132 ASSERT_TRUE(com_initializer_1.Succeeded());
133 ASSERT_TRUE(com_initializer_2.Succeeded());
134 ASSERT_TRUE(com_initializer_3.Succeeded());
135
136 // Reference count should be 3 after initialization.
137 EXPECT_EQ(DWORD(3),
138 com_initializer_1.GetCOMBalancerReferenceCountForTesting());
139 EXPECT_EQ(DWORD(3),
140 com_initializer_2.GetCOMBalancerReferenceCountForTesting());
141 EXPECT_EQ(DWORD(3),
142 com_initializer_3.GetCOMBalancerReferenceCountForTesting());
143
144 // Attempt to prematurely uninitialize the COM library.
145 ::CoUninitialize(); // Reference count -> 2.
146 ::CoUninitialize(); // Reference count -> 1.
147 ::CoUninitialize();
148
149 // Expect reference count to remain at 1.
150 EXPECT_EQ(DWORD(1),
151 com_initializer_1.GetCOMBalancerReferenceCountForTesting());
152 EXPECT_EQ(DWORD(1),
153 com_initializer_2.GetCOMBalancerReferenceCountForTesting());
154 EXPECT_EQ(DWORD(1),
155 com_initializer_3.GetCOMBalancerReferenceCountForTesting());
156 }
157
158 } // namespace win
159 } // namespace base
160