sousa-gecko/browser/components/shell/test/gtest/WindowsDefaultBrowserTests.cpp

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "WindowsDefaultBrowserTests.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <string_view>
#include <string>
#include <windows.h>
#include <ole2.h>
#include <uiautomation.h>
#include <tlhelp32.h>

#include "mozilla/Maybe.h"
#include "mozilla/UniquePtr.h"
#include "WindowsDefaultBrowser.h"

static bool RegDeleteKey(HKEY aRoot, std::wstring_view aKey) {
  LSTATUS ls{::RegDeleteTreeW(HKEY_CURRENT_USER, aKey.data())};

  return (ls == ERROR_SUCCESS);
}

static bool RegDeleteValue(HKEY aRoot, std::wstring_view aKey,
                           std::wstring_view aName) {
  HKEY key;
  LSTATUS ls{
      ::RegOpenKeyExW(HKEY_CURRENT_USER, aKey.data(), 0, KEY_SET_VALUE, &key)};

  if (ls == ERROR_SUCCESS) {
    ::RegDeleteValueW(key, aName.data());
    ::RegCloseKey(key);
    return true;
  }

  return false;
}

static bool RegWriteStringValue(HKEY aRoot, std::wstring_view aKey,
                                const mozilla::Maybe<std::wstring>& aName,
                                std::wstring_view aValue) {
  HKEY key;
  LSTATUS ls{::RegCreateKeyExW(aRoot, aKey.data(), 0, nullptr,
                               REG_OPTION_NON_VOLATILE, KEY_SET_VALUE, nullptr,
                               &key, nullptr)};

  if (ls != ERROR_SUCCESS) {
    return false;
  }

  const BYTE* value{reinterpret_cast<const BYTE*>(aValue.data())};
  const DWORD size{static_cast<DWORD>((aValue.size() + 1) * sizeof(WCHAR))};
  ls = ::RegSetValueExW(key, aName ? aName->c_str() : nullptr, 0, REG_SZ, value,
                        size);

  ::RegCloseKey(key);

  return (ls == ERROR_SUCCESS);
}

static void TerminateProcessById(DWORD aProcessId) {
  HANDLE process{
      ::OpenProcess(PROCESS_TERMINATE | SYNCHRONIZE, FALSE, aProcessId)};
  if (process) {
    ::TerminateProcess(process, 0);
    ::WaitForSingleObject(process, INFINITE);
    ::CloseHandle(process);
  }
}

static void TerminateProcessByName(LPCWSTR aProcessName) {
  HANDLE processes{::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)};
  if (processes == INVALID_HANDLE_VALUE) {
    return;
  }

  PROCESSENTRY32W process{};
  process.dwSize = sizeof(PROCESSENTRY32W);
  if (::Process32FirstW(processes, &process)) {
    do {
      if (::_wcsicmp(process.szExeFile, aProcessName) == 0) {
        TerminateProcessById(process.th32ProcessID);
      }
    } while (::Process32NextW(processes, &process));
  }
  ::CloseHandle(processes);
}

void FindSetDefaultBrowserButtonTests::SetUpTestSuite() {
  ASSERT_TRUE(GetAppRegName(sAppRegName));
  RegisterAsBrowser();
}

void FindSetDefaultBrowserButtonTests::TearDownTestSuite() {
  UnregisterAsBrowser();
}

void FindSetDefaultBrowserButtonTests::TearDown() { TerminateSystemSettings(); }

void FindSetDefaultBrowserButtonTests::TerminateSystemSettings() {
  LPCWSTR processName{L"SystemSettings.exe"};
  TerminateProcessByName(processName);
}

UIWindowElement
FindSetDefaultBrowserButtonTests::WaitForSetDefaultBrowserButton() {
  const int kMaxAttempts{40};
  const DWORD kRetryDelayMs{500};
  for (int i{0}; i < kMaxAttempts; ++i) {
    auto [window, button]{FindSetDefaultBrowserButton()};
    if (window && button) {
      return {window, button};
    }
    Sleep(kRetryDelayMs);
  }
  return {};
}

void FindSetDefaultBrowserButtonTests::RegisterAsBrowser() {
  const std::wstring appRegName{sAppRegName.get()};

  WCHAR exePath[MAX_PATH];
  ::GetModuleFileNameW(nullptr, exePath, MAX_PATH);

  const std::wstring appName{L"Firefox Test"};

  // clang-format off
  // HKEY_CURRENT_USER\Software\RegisteredApplications
  {
    const std::wstring key{L"Software\\RegisteredApplications"};
    const std::wstring value{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(appRegName), value));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Nothing(), appName));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ\Capabilities
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities"};
    const std::wstring name{L"ApplicationDescription"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(name), appName));
  }
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities"};
    const std::wstring name{L"ApplicationIcon"};
    const std::wstring value{std::wstring(exePath) + L",0"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(name), value));
  }
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities"};
    const std::wstring name{L"ApplicationName"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(name), appName));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ\Capabilities\URLAssociations
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities\\URLAssociations"};
    const std::wstring name{L"http"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(name), appRegName));
  }
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\Capabilities\\URLAssociations"};
    const std::wstring name{L"https"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Some(name), appRegName));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ\DefaultIcon
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\DefaultIcon"};
    const std::wstring value{std::wstring(exePath) + L",0"};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Nothing(), value));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ\shell\open\command
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName + L"\\shell\\open\\command"};
    const std::wstring value{L"\"" + std::wstring(exePath) + L"\""};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Nothing(), value));
  }
  // HKEY_CURRENT_USER\Software\Classes\Firefox-XYZ\shell\open\command
  {
    const std::wstring key{L"Software\\Classes\\" + appRegName + L"\\shell\\open\\command"};
    const std::wstring value{L"\"" + std::wstring(exePath) + L"\" -osint -url \"%1\""};
    ASSERT_TRUE(RegWriteStringValue(HKEY_CURRENT_USER, key, mozilla::Nothing(), value));
  }
  // clang-format on
}

void FindSetDefaultBrowserButtonTests::UnregisterAsBrowser() {
  const std::wstring appRegName{sAppRegName.get()};

  // clang-format off
  // HKEY_CURRENT_USER\Software\RegisteredApplications
  {
    const std::wstring key{L"Software\\RegisteredApplications"};
    ASSERT_TRUE(RegDeleteValue(HKEY_CURRENT_USER, key, appRegName));
  }
  // HKEY_CURRENT_USER\Software\Clients\StartMenuInternet\Firefox-XYZ
  {
    const std::wstring key{L"Software\\Clients\\StartMenuInternet\\" + appRegName};
    ASSERT_TRUE(RegDeleteKey(HKEY_CURRENT_USER, key));
  }
  // HKEY_CURRENT_USER\Software\Classes\Firefox-XYZ
  {
    const std::wstring key{L"Software\\Classes\\" + appRegName};
    ASSERT_TRUE(RegDeleteKey(HKEY_CURRENT_USER, key));
  }
  // clang-format on
}

mozilla::UniquePtr<WCHAR[]> FindSetDefaultBrowserButtonTests::sAppRegName;

TEST_F(FindSetDefaultBrowserButtonTests, ButtonFound) {
  ASSERT_TRUE(LaunchModernSettingsDialogDefaultApps());

  auto [window, button]{WaitForSetDefaultBrowserButton()};
  ASSERT_THAT(window, testing::NotNull());
  ASSERT_THAT(button, testing::NotNull());
}