Bug 1805509 - Set default extension handlers via defaultagent r=nalexander

Differential Revision: https://phabricator.services.mozilla.com/D169816

toolkit/mozapps/defaultagent/SetDefaultBrowser.cpp

@@ -16,6 +16,28 @@
 #include "EventLog.h"
 #include "SetDefaultBrowser.h"
 
+/*
+ * The implementation for setting extension handlers by writing UserChoice.
+ *
+ * This is used by both SetDefaultBrowserUserChoice and
+ * SetDefaultExtensionHandlersUserChoice.
+ *
+ * @param aAumi The AUMI of the installation to set as default.
+ *
+ * @param aSid Current user's string SID
+ *
+ * @param aFileExtensions Optional null-terminated list of file association
+ * pairs to set as default, like `{ L".pdf", "FirefoxPDF", nullptr }`.
+ *
+ * @returns S_OK           All associations set and checked successfully.
+ *          MOZ_E_REJECTED UserChoice was set, but checking the default did not
+ *                         return our ProgID.
+ *          E_FAIL         Failed to set at least one association.
+ */
+static HRESULT SetDefaultExtensionHandlersUserChoiceImpl(
+    const wchar_t* aAumi, const wchar_t* const aSid,
+    const wchar_t* const* aFileExtensions);
+
 static bool AddMillisecondsToSystemTime(SYSTEMTIME& aSystemTime,
                                         ULONGLONG aIncrementMS) {
   FILETIME fileTime;
@@ -258,23 +280,15 @@ HRESULT SetDefaultBrowserUserChoice(
     }
   }
 
-  const wchar_t* const* extraFileExtension = aExtraFileExtensions;
-  const wchar_t* const* extraProgIDRoot = aExtraFileExtensions + 1;
-  while (ok && extraFileExtension && *extraFileExtension && extraProgIDRoot &&
-         *extraProgIDRoot) {
-    // Formatting the ProgID here prevents using this helper to target arbitrary
-    // ProgIDs.
-    auto extraProgID = FormatProgID(*extraProgIDRoot, aAumi);
-
-    if (!SetUserChoice(*extraFileExtension, sid.get(), extraProgID.get())) {
+  if (ok) {
+    HRESULT hr = SetDefaultExtensionHandlersUserChoiceImpl(
+        aAumi, sid.get(), aExtraFileExtensions);
+    if (hr == MOZ_E_REJECTED) {
       ok = false;
-    } else if (!VerifyUserDefault(*extraFileExtension, extraProgID.get())) {
       defaultRejected = true;
+    } else if (hr == E_FAIL) {
       ok = false;
     }
-
-    extraFileExtension += 2;
-    extraProgIDRoot += 2;
   }
 
   // Notify shell to refresh icons
@@ -286,7 +300,66 @@ HRESULT SetDefaultBrowserUserChoice(
       return MOZ_E_REJECTED;
     }
     return E_FAIL;
-  } else {
-    return S_OK;
   }
+
+  return S_OK;
+}
+
+HRESULT SetDefaultExtensionHandlersUserChoice(
+    const wchar_t* aAumi, const wchar_t* const* aFileExtensions) {
+  auto sid = GetCurrentUserStringSid();
+  if (!sid) {
+    return E_FAIL;
+  }
+
+  bool ok = true;
+  bool defaultRejected = false;
+
+  HRESULT hr = SetDefaultExtensionHandlersUserChoiceImpl(aAumi, sid.get(),
+                                                         aFileExtensions);
+  if (hr == MOZ_E_REJECTED) {
+    ok = false;
+    defaultRejected = true;
+  } else if (hr == E_FAIL) {
+    ok = false;
+  }
+
+  // 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;
+  }
+
+  return S_OK;
+}
+
+HRESULT SetDefaultExtensionHandlersUserChoiceImpl(
+    const wchar_t* aAumi, const wchar_t* const aSid,
+    const wchar_t* const* aFileExtensions) {
+  const wchar_t* const* extraFileExtension = aFileExtensions;
+  const wchar_t* const* extraProgIDRoot = aFileExtensions + 1;
+  while (extraFileExtension && *extraFileExtension && extraProgIDRoot &&
+         *extraProgIDRoot) {
+    // Formatting the ProgID here prevents using this helper to target arbitrary
+    // ProgIDs.
+    auto extraProgID = FormatProgID(*extraProgIDRoot, aAumi);
+
+    if (!SetUserChoice(*extraFileExtension, aSid, extraProgID.get())) {
+      return E_FAIL;
+    }
+
+    if (!VerifyUserDefault(*extraFileExtension, extraProgID.get())) {
+      return MOZ_E_REJECTED;
+    }
+
+    extraFileExtension += 2;
+    extraProgIDRoot += 2;
+  }
+
+  return S_OK;
 }

toolkit/mozapps/defaultagent/SetDefaultBrowser.h

@@ -34,6 +34,23 @@
 HRESULT SetDefaultBrowserUserChoice(
     const wchar_t* aAumi, const wchar_t* const* aExtraFileExtensions = nullptr);
 
+/*
+ * Set the default extension handlers for the given file extensions by writing
+ * the UserChoice registry keys.
+ *
+ * @param aAumi The AUMI of the installation to set as default.
+ *
+ * @param aFileExtensions Optional null-terminated list of file association
+ * pairs to set as default, like `{ L".pdf", "FirefoxPDF", nullptr }`.
+ *
+ * @returns S_OK           All associations set and checked successfully.
+ *          MOZ_E_REJECTED UserChoice was set, but checking the default did not
+ *                         return our ProgID.
+ *          E_FAIL         Failed to set at least one association.
+ */
+HRESULT SetDefaultExtensionHandlersUserChoice(
+    const wchar_t* aAumi, const wchar_t* const* aFileExtensions = nullptr);
+
 /*
  * Additional HRESULT error codes from SetDefaultBrowserUserChoice
  *

toolkit/mozapps/defaultagent/main.cpp

@@ -408,6 +408,14 @@ int wmain(int argc, wchar_t** argv) {
     // `argv` is itself null-terminated, so we can safely pass the tail of the
     // array here.
     return SetDefaultBrowserUserChoice(argv[2], &argv[3]);
+  } else if (!wcscmp(argv[1], L"set-default-extension-handlers-user-choice")) {
+    if (argc < 3 || !argv[2]) {
+      return E_INVALIDARG;
+    }
+
+    // `argv` is itself null-terminated, so we can safely pass the tail of the
+    // array here.
+    return SetDefaultExtensionHandlersUserChoice(argv[2], &argv[3]);
   } else {
     return E_INVALIDARG;
   }