scoped_mock_context.cc 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. // Copyright 2021 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 "chrome/elevation_service/scoped_mock_context.h"
  5. #include <objbase.h>
  6. #include <objidl.h>
  7. #include <unknwn.h>
  8. #include <wrl/implements.h>
  9. #include "base/win/com_init_util.h"
  10. #include "testing/gtest/include/gtest/gtest.h"
  11. namespace {
  12. // A mock implementation of IServerSecurity that allows for the production code
  13. // that calls ::CoImpersonateClient() to work.
  14. class MockServerSecurity
  15. : public Microsoft::WRL::RuntimeClass<
  16. Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
  17. IServerSecurity> {
  18. public:
  19. MockServerSecurity() = default;
  20. MockServerSecurity(const MockServerSecurity&) = delete;
  21. MockServerSecurity& operator=(const MockServerSecurity&) = delete;
  22. IFACEMETHODIMP QueryBlanket(DWORD* authentication_service,
  23. DWORD* authorization_service,
  24. OLECHAR** server_principal_name,
  25. DWORD* authentication_level,
  26. DWORD* impersonation_level,
  27. void** privilege,
  28. DWORD* capabilities) override {
  29. return E_NOTIMPL;
  30. }
  31. IFACEMETHODIMP ImpersonateClient() override {
  32. is_impersonating_ = true;
  33. return S_OK;
  34. }
  35. IFACEMETHODIMP RevertToSelf() override {
  36. is_impersonating_ = false;
  37. return S_OK;
  38. }
  39. IFACEMETHODIMP_(BOOL) IsImpersonating() override { return is_impersonating_; }
  40. private:
  41. ~MockServerSecurity() override { EXPECT_FALSE(is_impersonating_); }
  42. bool is_impersonating_ = false;
  43. };
  44. } // namespace
  45. namespace elevation_service {
  46. ScopedMockContext::ScopedMockContext() {
  47. base::win::AssertComInitialized();
  48. auto mock_call_context = Microsoft::WRL::Make<MockServerSecurity>();
  49. // We set the call context to a mock object that implements IServerSecurity.
  50. // This allows for the production code that calls ::CoImpersonateClient() to
  51. // succeed.
  52. auto hresult =
  53. ::CoSwitchCallContext(mock_call_context.Get(), &original_call_context_);
  54. EXPECT_HRESULT_SUCCEEDED(hresult);
  55. if (FAILED(hresult))
  56. return;
  57. mock_call_context_ = std::move(mock_call_context);
  58. EXPECT_EQ(original_call_context_, nullptr);
  59. }
  60. ScopedMockContext::~ScopedMockContext() {
  61. base::win::AssertComInitialized();
  62. if (!Succeeded())
  63. return;
  64. IUnknown* this_call_context = nullptr;
  65. EXPECT_HRESULT_SUCCEEDED(
  66. ::CoSwitchCallContext(original_call_context_, &this_call_context));
  67. EXPECT_EQ(this_call_context, mock_call_context_.Get())
  68. << "CoSwitchCallContext switched out someone else's context.";
  69. }
  70. } // namespace elevation_service