Bug 1703578 - Part 2: Add WDBA command to set default browser. r=bytesized

agashlin · agashlin · commit f22ce8ffdd8c · 2021-06-17T18:06:34.000Z
Depends on D113426 Differential Revision: https://phabricator.services.mozilla.com/D113427
diff --git a/toolkit/mozapps/defaultagent/SetDefaultBrowser.cpp b/toolkit/mozapps/defaultagent/SetDefaultBrowser.cpp
@@ -0,0 +1,260 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* 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 <windows.h>
+#include <shlobj.h>  // for SHChangeNotify and IApplicationAssociationRegistration
+
+#include "mozilla/ArrayUtils.h"
+#include "mozilla/RefPtr.h"
+#include "mozilla/UniquePtr.h"
+#include "mozilla/WinHeaderOnlyUtils.h"
+#include "WindowsUserChoice.h"
+
+#include "EventLog.h"
+#include "SetDefaultBrowser.h"
+
+static bool AddMillisecondsToSystemTime(SYSTEMTIME& aSystemTime,
+                                        ULONGLONG aIncrementMS) {
+  FILETIME fileTime;
+  ULARGE_INTEGER fileTimeInt;
+  if (!::SystemTimeToFileTime(&aSystemTime, &fileTime)) {
+    return false;
+  }
+  fileTimeInt.LowPart = fileTime.dwLowDateTime;
+  fileTimeInt.HighPart = fileTime.dwHighDateTime;
+
+  // FILETIME is in units of 100ns.
+  fileTimeInt.QuadPart += aIncrementMS * 1000 * 10;
+
+  fileTime.dwLowDateTime = fileTimeInt.LowPart;
+  fileTime.dwHighDateTime = fileTimeInt.HighPart;
+  SYSTEMTIME tmpSystemTime;
+  if (!::FileTimeToSystemTime(&fileTime, &tmpSystemTime)) {
+    return false;
+  }
+
+  aSystemTime = tmpSystemTime;
+  return true;
+}
+
+// Compare two SYSTEMTIMEs as FILETIME after clearing everything
+// below minutes.
+static bool CheckEqualMinutes(SYSTEMTIME aSystemTime1,
+                              SYSTEMTIME aSystemTime2) {
+  aSystemTime1.wSecond = 0;
+  aSystemTime1.wMilliseconds = 0;
+
+  aSystemTime2.wSecond = 0;
+  aSystemTime2.wMilliseconds = 0;
+
+  FILETIME fileTime1;
+  FILETIME fileTime2;
+  if (!::SystemTimeToFileTime(&aSystemTime1, &fileTime1) ||
+      !::SystemTimeToFileTime(&aSystemTime2, &fileTime2)) {
+    return false;
+  }
+
+  return (fileTime1.dwLowDateTime == fileTime2.dwLowDateTime) &&
+         (fileTime1.dwHighDateTime == fileTime2.dwHighDateTime);
+}
+
+/*
+ * Set an association with a UserChoice key
+ *
+ * Removes the old key, creates a new one with ProgID and Hash set to
+ * enable a new asociation.
+ *
+ * @param aExt      File type or protocol to associate
+ * @param aSid      Current user's string SID
+ * @param aProgID   ProgID to use for the asociation
+ *
+ * @return true if successful, false on error.
+ */
+static bool SetUserChoice(const wchar_t* aExt, const wchar_t* aSid,
+                          const wchar_t* aProgID) {
+  SYSTEMTIME hashTimestamp;
+  ::GetSystemTime(&hashTimestamp);
+  auto hash = GenerateUserChoiceHash(aExt, aSid, aProgID, hashTimestamp);
+  if (!hash) {
+    return false;
+  }
+
+  // The hash changes at the end of each minute, so check that the hash should
+  // be the same by the time we're done writing.
+  const ULONGLONG kWriteTimingThresholdMilliseconds = 100;
+  // Generating the hash could have taken some time, so start from now.
+  SYSTEMTIME writeEndTimestamp;
+  ::GetSystemTime(&writeEndTimestamp);
+  if (!AddMillisecondsToSystemTime(writeEndTimestamp,
+                                   kWriteTimingThresholdMilliseconds)) {
+    return false;
+  }
+  if (!CheckEqualMinutes(hashTimestamp, writeEndTimestamp)) {
+    LOG_ERROR_MESSAGE(
+        L"Hash is too close to expiration, sleeping until next hash.");
+    ::Sleep(kWriteTimingThresholdMilliseconds * 2);
+
+    // For consistency, use the current time.
+    ::GetSystemTime(&hashTimestamp);
+    hash = GenerateUserChoiceHash(aExt, aSid, aProgID, hashTimestamp);
+    if (!hash) {
+      return false;
+    }
+  }
+
+  auto assocKeyPath = GetAssociationKeyPath(aExt);
+  if (!assocKeyPath) {
+    return false;
+  }
+
+  LSTATUS ls;
+  HKEY rawAssocKey;
+  ls = ::RegOpenKeyExW(HKEY_CURRENT_USER, assocKeyPath.get(), 0,
+                       KEY_READ | KEY_WRITE, &rawAssocKey);
+  if (ls != ERROR_SUCCESS) {
+    LOG_ERROR(HRESULT_FROM_WIN32(ls));
+    return false;
+  }
+  nsAutoRegKey assocKey(rawAssocKey);
+
+  // When Windows creates this key, it is read-only (Deny Set Value), so we need
+  // to delete it first.
+  // We don't set any similar special permissions.
+  ls = ::RegDeleteKeyW(assocKey.get(), L"UserChoice");
+  if (ls != ERROR_SUCCESS) {
+    LOG_ERROR(HRESULT_FROM_WIN32(ls));
+    return false;
+  }
+
+  HKEY rawUserChoiceKey;
+  ls = ::RegCreateKeyExW(assocKey.get(), L"UserChoice", 0, nullptr,
+                         0 /* options */, KEY_READ | KEY_WRITE,
+                         0 /* security attributes */, &rawUserChoiceKey,
+                         nullptr);
+  if (ls != ERROR_SUCCESS) {
+    LOG_ERROR(HRESULT_FROM_WIN32(ls));
+    return false;
+  }
+  nsAutoRegKey userChoiceKey(rawUserChoiceKey);
+
+  DWORD progIdByteCount = (::lstrlenW(aProgID) + 1) * sizeof(wchar_t);
+  ls = ::RegSetValueExW(userChoiceKey.get(), L"ProgID", 0, REG_SZ,
+                        reinterpret_cast<const unsigned char*>(aProgID),
+                        progIdByteCount);
+  if (ls != ERROR_SUCCESS) {
+    LOG_ERROR(HRESULT_FROM_WIN32(ls));
+    return false;
+  }
+
+  DWORD hashByteCount = (::lstrlenW(hash.get()) + 1) * sizeof(wchar_t);
+  ls = ::RegSetValueExW(userChoiceKey.get(), L"Hash", 0, REG_SZ,
+                        reinterpret_cast<const unsigned char*>(hash.get()),
+                        hashByteCount);
+  if (ls != ERROR_SUCCESS) {
+    LOG_ERROR(HRESULT_FROM_WIN32(ls));
+    return false;
+  }
+
+  return true;
+}
+
+static bool VerifyUserDefault(const wchar_t* aExt, const wchar_t* aProgID) {
+  RefPtr<IApplicationAssociationRegistration> pAAR;
+  HRESULT hr = ::CoCreateInstance(
+      CLSID_ApplicationAssociationRegistration, nullptr, CLSCTX_INPROC,
+      IID_IApplicationAssociationRegistration, getter_AddRefs(pAAR));
+  if (FAILED(hr)) {
+    LOG_ERROR(hr);
+    return false;
+  }
+
+  wchar_t* rawRegisteredApp;
+  bool isProtocol = aExt[0] != L'.';
+  // Note: Checks AL_USER instead of AL_EFFECTIVE.
+  hr = pAAR->QueryCurrentDefault(aExt,
+                                 isProtocol ? AT_URLPROTOCOL : AT_FILEEXTENSION,
+                                 AL_USER, &rawRegisteredApp);
+  if (FAILED(hr)) {
+    if (hr == HRESULT_FROM_WIN32(ERROR_NO_ASSOCIATION)) {
+      LOG_ERROR_MESSAGE(L"UserChoice ProgID %s for %s was rejected", aProgID,
+                        aExt);
+    } else {
+      LOG_ERROR(hr);
+    }
+    return false;
+  }
+  mozilla::UniquePtr<wchar_t, mozilla::CoTaskMemFreeDeleter> registeredApp(
+      rawRegisteredApp);
+
+  if (::CompareStringOrdinal(registeredApp.get(), -1, aProgID, -1, FALSE) !=
+      CSTR_EQUAL) {
+    LOG_ERROR_MESSAGE(
+        L"Default was %s after writing ProgID %s to UserChoice for %s",
+        registeredApp.get(), aProgID, aExt);
+    return false;
+  }
+
+  return true;
+}
+
+HRESULT SetDefaultBrowserUserChoice(const wchar_t* aAumi) {
+  auto urlProgID = FormatProgID(L"FirefoxURL", aAumi);
+  if (!CheckProgIDExists(urlProgID.get())) {
+    LOG_ERROR_MESSAGE(L"ProgID %s not found", urlProgID.get());
+    return MOZ_E_NO_PROGID;
+  }
+
+  auto htmlProgID = FormatProgID(L"FirefoxHTML", aAumi);
+  if (!CheckProgIDExists(htmlProgID.get())) {
+    LOG_ERROR_MESSAGE(L"ProgID %s not found", htmlProgID.get());
+    return MOZ_E_NO_PROGID;
+  }
+
+  if (!CheckBrowserUserChoiceHashes()) {
+    LOG_ERROR_MESSAGE(L"UserChoice Hash mismatch");
+    return MOZ_E_HASH_CHECK;
+  }
+
+  auto sid = GetCurrentUserStringSid();
+  if (!sid) {
+    return E_FAIL;
+  }
+
+  bool ok = true;
+  bool defaultRejected = false;
+
+  struct {
+    const wchar_t* ext;
+    const wchar_t* progID;
+  } associations[] = {{L"https", urlProgID.get()},
+                      {L"http", urlProgID.get()},
+                      {L".html", htmlProgID.get()},
+                      {L".htm", htmlProgID.get()}};
+  for (int i = 0; i < mozilla::ArrayLength(associations); ++i) {
+    if (!SetUserChoice(associations[i].ext, sid.get(),
+                       associations[i].progID)) {
+      ok = false;
+      break;
+    } else if (!VerifyUserDefault(associations[i].ext,
+                                  associations[i].progID)) {
+      defaultRejected = true;
+      ok = false;
+      break;
+    }
+  }
+
+  // Notify shell to refresh icons
+  ::SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, nullptr, nullptr);
+
+  if (!ok) {
+    LOG_ERROR_MESSAGE(L"Failed setting default with %s", aAumi);
+    if (defaultRejected) {
+      return MOZ_E_REJECTED;
+    }
+    return E_FAIL;
+  } else {
+    return S_OK;
+  }
+}
diff --git a/toolkit/mozapps/defaultagent/SetDefaultBrowser.h b/toolkit/mozapps/defaultagent/SetDefaultBrowser.h
@@ -0,0 +1,37 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* 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/. */
+
+#ifndef DEFAULT_BROWSER_SET_DEFAULT_BROWSER_H__
+#define DEFAULT_BROWSER_SET_DEFAULT_BROWSER_H__
+
+/*
+ * Set the default browser by writing the UserChoice registry keys.
+ *
+ * This sets the associations for https, http, .html, and .htm.
+ *
+ * When the agent is run with set-default-browser-user-choice,
+ * the exit code is the result of this function.
+ *
+ * @param aAumi The AUMI of the installation to set as default.
+ *
+ * @return S_OK             All associations set and checked successfully.
+ *         MOZ_E_NO_PROGID  The ProgID classes had not been registered.
+ *         MOZ_E_HASH_CHECK The existing UserChoice Hash could not be verified.
+ *         MOZ_E_REJECTED   UserChoice was set, but checking the default
+ *                          did not return our ProgID.
+ *         E_FAIL           other failure
+ */
+HRESULT SetDefaultBrowserUserChoice(const wchar_t* aAumi);
+
+/*
+ * Additional HRESULT error codes from SetDefaultBrowserUserChoice
+ *
+ * 0x20000000 is set to put these in the customer-defined range.
+ */
+const HRESULT MOZ_E_NO_PROGID = 0xa0000001L;
+const HRESULT MOZ_E_HASH_CHECK = 0xa0000002L;
+const HRESULT MOZ_E_REJECTED = 0xa0000003L;
+
+#endif  // DEFAULT_BROWSER_SET_DEFAULT_BROWSER_H__
diff --git a/toolkit/mozapps/defaultagent/main.cpp b/toolkit/mozapps/defaultagent/main.cpp
@@ -20,6 +20,7 @@
 #include "Registry.h"
 #include "RemoteSettings.h"
 #include "ScheduledTask.h"
+#include "SetDefaultBrowser.h"
 #include "Telemetry.h"
 
 // The AGENT_REGKEY_NAME is dependent on MOZ_APP_VENDOR and MOZ_APP_BASENAME,
@@ -250,6 +251,8 @@ static bool CheckIfAppRanRecently(bool* aResult) {
 //   Actually performs the default agent task, which currently means generating
 //   and sending our telemetry ping and possibly showing a notification to the
 //   user if their browser has switched from Firefox to Edge with Blink.
+// set-default-browser-user-choice [app-user-model-id]
+//   Set the default browser via the UserChoice registry keys.
 int wmain(int argc, wchar_t** argv) {
   if (argc < 2 || !argv[1]) {
     return E_INVALIDARG;
@@ -386,6 +389,12 @@ int wmain(int argc, wchar_t** argv) {
         MaybeShowNotification(browserInfo, argv[2]);
 
     return SendDefaultBrowserPing(browserInfo, activitiesPerformed);
+  } else if (!wcscmp(argv[1], L"set-default-browser-user-choice")) {
+    if (argc < 3 || !argv[2]) {
+      return E_INVALIDARG;
+    }
+
+    return SetDefaultBrowserUserChoice(argv[2]);
   } else {
     return E_INVALIDARG;
   }
diff --git a/toolkit/mozapps/defaultagent/moz.build b/toolkit/mozapps/defaultagent/moz.build
@@ -23,11 +23,13 @@ UNIFIED_SOURCES += [
     "Registry.cpp",
     "RemoteSettings.cpp",
     "ScheduledTask.cpp",
+    "SetDefaultBrowser.cpp",
     "Telemetry.cpp",
 ]
 
 SOURCES += [
     "/browser/components/shell/WindowsDefaultBrowser.cpp",
+    "/browser/components/shell/WindowsUserChoice.cpp",
     "/other-licenses/nsis/Contrib/CityHash/cityhash/city.cpp",
     "/third_party/WinToast/wintoastlib.cpp",
     "/toolkit/mozapps/update/common/readstrings.cpp",
@@ -50,7 +52,9 @@ LOCAL_INCLUDES += [
 
 OS_LIBS += [
     "advapi32",
+    "bcrypt",
     "comsupp",
+    "crypt32",
     "kernel32",
     "netapi32",
     "ole32",
@@ -70,9 +74,9 @@ DEFINES["IMPL_MFBT"] = True
 DEFINES["UNICODE"] = True
 DEFINES["_UNICODE"] = True
 
-# If defines are added to this list that are required by the Cache or its
-# dependencies (Registry, EventLog, common), tests/gtest/moz.build will need to
-# be updated as well.
+# If defines are added to this list that are required by the Cache,
+# SetDefaultBrowser, or their dependencies (Registry, EventLog, common),
+# tests/gtest/moz.build will need to be updated as well.
 for var in ("MOZ_APP_BASENAME", "MOZ_APP_DISPLAYNAME", "MOZ_APP_VENDOR"):
     DEFINES[var] = '"%s"' % CONFIG[var]
 
diff --git a/toolkit/mozapps/defaultagent/tests/gtest/SetDefaultBrowserTest.cpp b/toolkit/mozapps/defaultagent/tests/gtest/SetDefaultBrowserTest.cpp
diff --git a/toolkit/mozapps/defaultagent/tests/gtest/moz.build b/toolkit/mozapps/defaultagent/tests/gtest/moz.build
