Bug 1674904: Part 2 - New dll registration code in mscom/oop/module; r=Jamie

We add new DLL registration code. This is a rather generic function that
permits the following:

* Registering multiple `CLSID`s for the same DLL;
* Registering an optional `AppID`. Registering an `AppID` allows us to use a
  `DllSurrogate` to host the DLL out-of-process using Windows' built-in
  `dllhost.exe`. I'll be using this feature in a future bug.
* Supporting all available threading modelsl;
* Capable of registering either inproc servers or inproc handlers;
* Using the transaction-based registry API so that we can cleanly rollback
  during registration if any part(s) of it fail.

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

Author: dblohm7

ipc/mscom/Utils.cpp

@@ -16,6 +16,8 @@
 #  endif
 #endif
 
+#include "mozilla/ArrayUtils.h"
+#include "mozilla/DebugOnly.h"
 #include "mozilla/mscom/Objref.h"
 #include "mozilla/mscom/Utils.h"
 #include "mozilla/RefPtr.h"
@@ -139,6 +141,50 @@ uintptr_t GetContainingModuleHandle() {
   return reinterpret_cast<uintptr_t>(thisModule);
 }
 
+namespace detail {
+
+long BuildRegGuidPath(REFGUID aGuid, const GuidType aGuidType, wchar_t* aBuf,
+                      const size_t aBufLen) {
+  constexpr wchar_t kClsid[] = L"CLSID\\";
+  constexpr wchar_t kAppid[] = L"AppID\\";
+  constexpr wchar_t kSubkeyBase[] = L"SOFTWARE\\Classes\\";
+
+  // We exclude null terminators in these length calculations because we include
+  // the stringified GUID's null terminator at the end. Since kClsid and kAppid
+  // have identical lengths, we just choose one to compute this length.
+  constexpr size_t kSubkeyBaseLen = mozilla::ArrayLength(kSubkeyBase) - 1;
+  constexpr size_t kSubkeyLen =
+      kSubkeyBaseLen + mozilla::ArrayLength(kClsid) - 1;
+  // Guid length as formatted for the registry (including curlies and dashes),
+  // but excluding null terminator.
+  constexpr size_t kGuidLen = kGuidRegFormatCharLenInclNul - 1;
+  constexpr size_t kExpectedPathLenInclNul = kSubkeyLen + kGuidLen + 1;
+
+  if (aBufLen < kExpectedPathLenInclNul) {
+    // Buffer is too short
+    return E_INVALIDARG;
+  }
+
+  if (wcscpy_s(aBuf, aBufLen, kSubkeyBase)) {
+    return E_INVALIDARG;
+  }
+
+  const wchar_t* strGuidType = aGuidType == GuidType::CLSID ? kClsid : kAppid;
+  if (wcscat_s(aBuf, aBufLen, strGuidType)) {
+    return E_INVALIDARG;
+  }
+
+  int guidConversionResult =
+      ::StringFromGUID2(aGuid, &aBuf[kSubkeyLen], aBufLen - kSubkeyLen);
+  if (!guidConversionResult) {
+    return E_INVALIDARG;
+  }
+
+  return S_OK;
+}
+
+}  // namespace detail
+
 long CreateStream(const uint8_t* aInitBuf, const uint32_t aInitBufSize,
                   IStream** aOutStream) {
   if (!aInitBufSize || !aOutStream) {
@@ -271,6 +317,15 @@ void GUIDToString(REFGUID aGuid, nsAString& aOutString) {
   }
 }
 
+#else
+
+void GUIDToString(REFGUID aGuid,
+                  wchar_t (&aOutBuf)[kGuidRegFormatCharLenInclNul]) {
+  DebugOnly<int> result =
+      ::StringFromGUID2(aGuid, aOutBuf, ArrayLength(aOutBuf));
+  MOZ_ASSERT(result);
+}
+
 #endif  // defined(MOZILLA_INTERNAL_API)
 
 #if defined(ACCESSIBILITY)

ipc/mscom/Utils.h

@@ -19,6 +19,17 @@ struct IUnknown;
 
 namespace mozilla {
 namespace mscom {
+namespace detail {
+
+enum class GuidType {
+  CLSID,
+  AppID,
+};
+
+long BuildRegGuidPath(REFGUID aGuid, const GuidType aGuidType, wchar_t* aBuf,
+                      const size_t aBufLen);
+
+}  // namespace detail
 
 bool IsCOMInitializedOnCurrentThread();
 bool IsCurrentThreadMTA();
@@ -31,6 +42,16 @@ bool IsProxy(IUnknown* aUnknown);
 bool IsValidGUID(REFGUID aCheckGuid);
 uintptr_t GetContainingModuleHandle();
 
+template <size_t N>
+inline long BuildAppidPath(REFGUID aAppId, wchar_t (&aPath)[N]) {
+  return detail::BuildRegGuidPath(aAppId, detail::GuidType::AppID, aPath, N);
+}
+
+template <size_t N>
+inline long BuildClsidPath(REFCLSID aClsid, wchar_t (&aPath)[N]) {
+  return detail::BuildRegGuidPath(aClsid, detail::GuidType::CLSID, aPath, N);
+}
+
 /**
  * Given a buffer, create a new IStream object.
  * @param aBuf Buffer containing data to initialize the stream. This parameter
@@ -53,6 +74,12 @@ long CreateStream(const uint8_t* aBuf, const uint32_t aBufLen,
  */
 long CopySerializedProxy(IStream* aInStream, IStream** aOutStream);
 
+/**
+ * Length of a stringified GUID as formatted for the registry, i.e. including
+ * curly-braces and dashes.
+ */
+constexpr size_t kGuidRegFormatCharLenInclNul = 39;
+
 #if defined(MOZILLA_INTERNAL_API)
 /**
  * Checks the registry to see if |aClsid| is a thread-aware in-process server.
@@ -78,6 +105,9 @@ long CopySerializedProxy(IStream* aInStream, IStream** aOutStream);
 bool IsClassThreadAwareInprocServer(REFCLSID aClsid);
 
 void GUIDToString(REFGUID aGuid, nsAString& aOutString);
+#else
+void GUIDToString(REFGUID aGuid,
+                  wchar_t (&aOutBuf)[kGuidRegFormatCharLenInclNul]);
 #endif  // defined(MOZILLA_INTERNAL_API)
 
 #if defined(ACCESSIBILITY)

ipc/mscom/oop/Module.cpp

@@ -6,13 +6,263 @@
 
 #include "Module.h"
 
+#include <stdlib.h>
+
+#include <ktmw32.h>
 #include <memory.h>
 #include <rpc.h>
 
+#include "mozilla/ArrayUtils.h"
+#include "mozilla/Assertions.h"
+#include "mozilla/mscom/Utils.h"
+#include "mozilla/Range.h"
+#include "nsWindowsHelpers.h"
+
+template <size_t N>
+static const mozilla::Range<const wchar_t> LiteralToRange(
+    const wchar_t (&aArg)[N]) {
+  return mozilla::Range(aArg, N);
+}
+
 namespace mozilla {
 namespace mscom {
 
 ULONG Module::sRefCount = 0;
 
+static const wchar_t* SubkeyNameFromClassType(
+    const Module::ClassType aClassType) {
+  switch (aClassType) {
+    case Module::ClassType::InprocServer:
+      return L"InprocServer32";
+    case Module::ClassType::InprocHandler:
+      return L"InprocHandler32";
+    default:
+      MOZ_CRASH("Unknown ClassType");
+      return nullptr;
+  }
+}
+
+static const Range<const wchar_t> ThreadingModelAsString(
+    const Module::ThreadingModel aThreadingModel) {
+  switch (aThreadingModel) {
+    case Module::ThreadingModel::DedicatedUiThreadOnly:
+      return LiteralToRange(L"Apartment");
+    case Module::ThreadingModel::MultiThreadedApartmentOnly:
+      return LiteralToRange(L"Free");
+    case Module::ThreadingModel::DedicatedUiThreadXorMultiThreadedApartment:
+      return LiteralToRange(L"Both");
+    case Module::ThreadingModel::AllThreadsAllApartments:
+      return LiteralToRange(L"Neutral");
+    default:
+      MOZ_CRASH("Unknown ThreadingModel");
+      return Range<const wchar_t>();
+  }
+}
+
+/* static */
+HRESULT Module::Register(const CLSID* const* aClsids, const size_t aNumClsids,
+                         const ThreadingModel aThreadingModel,
+                         const ClassType aClassType, const GUID* const aAppId) {
+  MOZ_ASSERT(aClsids && aNumClsids);
+  if (!aClsids || !aNumClsids) {
+    return E_INVALIDARG;
+  }
+
+  const wchar_t* inprocName = SubkeyNameFromClassType(aClassType);
+
+  const Range<const wchar_t> threadingModelStr =
+      ThreadingModelAsString(aThreadingModel);
+  const DWORD threadingModelStrLenBytesInclNul =
+      threadingModelStr.length() * sizeof(wchar_t);
+
+  wchar_t strAppId[kGuidRegFormatCharLenInclNul] = {};
+  if (aAppId) {
+    GUIDToString(*aAppId, strAppId);
+  }
+
+  // Obtain the full path to this DLL
+  HMODULE thisModule;
+  if (!::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                                GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                            reinterpret_cast<LPCWSTR>(&Module::CanUnload),
+                            &thisModule)) {
+    return HRESULT_FROM_WIN32(::GetLastError());
+  }
+
+  wchar_t absThisModulePath[MAX_PATH + 1] = {};
+  DWORD actualPathLenCharsExclNul = ::GetModuleFileNameW(
+      thisModule, absThisModulePath, ArrayLength(absThisModulePath));
+  if (!actualPathLenCharsExclNul ||
+      actualPathLenCharsExclNul == ArrayLength(absThisModulePath)) {
+    return HRESULT_FROM_WIN32(::GetLastError());
+  }
+  const DWORD actualPathLenBytesInclNul =
+      (actualPathLenCharsExclNul + 1) * sizeof(wchar_t);
+
+  // Use the name of this DLL as the name of the transaction
+  wchar_t txnName[_MAX_FNAME] = {};
+  if (_wsplitpath_s(absThisModulePath, nullptr, 0, nullptr, 0, txnName,
+                    ArrayLength(txnName), nullptr, 0)) {
+    return HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND);
+  }
+
+  // Manipulate the registry using a transaction so that any failures are
+  // rolled back.
+  nsAutoHandle txn(::CreateTransaction(
+      nullptr, nullptr, TRANSACTION_DO_NOT_PROMOTE, 0, 0, 0, txnName));
+  if (txn.get() == INVALID_HANDLE_VALUE) {
+    return HRESULT_FROM_WIN32(::GetLastError());
+  }
+
+  HRESULT hr;
+  LSTATUS status;
+
+  // A single DLL may serve multiple components. For each CLSID, we register
+  // this DLL as its server and, when an AppId is specified, set up a reference
+  // from the CLSID to the specified AppId.
+  for (size_t idx = 0; idx < aNumClsids; ++idx) {
+    if (!aClsids[idx]) {
+      return E_INVALIDARG;
+    }
+
+    wchar_t clsidKeyPath[256];
+    hr = BuildClsidPath(*aClsids[idx], clsidKeyPath);
+    if (FAILED(hr)) {
+      return hr;
+    }
+
+    // Create the CLSID key
+    HKEY rawClsidKey;
+    status = ::RegCreateKeyTransactedW(
+        HKEY_LOCAL_MACHINE, clsidKeyPath, 0, nullptr, REG_OPTION_NON_VOLATILE,
+        KEY_ALL_ACCESS, nullptr, &rawClsidKey, nullptr, txn, nullptr);
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+    nsAutoRegKey clsidKey(rawClsidKey);
+
+    if (aAppId) {
+      // This value associates the registered CLSID with the specified AppID
+      status = ::RegSetValueExW(clsidKey, L"AppID", 0, REG_SZ,
+                                reinterpret_cast<const BYTE*>(strAppId),
+                                ArrayLength(strAppId) * sizeof(wchar_t));
+      if (status != ERROR_SUCCESS) {
+        return HRESULT_FROM_WIN32(status);
+      }
+    }
+
+    HKEY rawInprocKey;
+    status = ::RegCreateKeyTransactedW(
+        clsidKey, inprocName, 0, nullptr, REG_OPTION_NON_VOLATILE,
+        KEY_ALL_ACCESS, nullptr, &rawInprocKey, nullptr, txn, nullptr);
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+    nsAutoRegKey inprocKey(rawInprocKey);
+
+    // Set the component's path to this DLL
+    status = ::RegSetValueExW(inprocKey, nullptr, 0, REG_EXPAND_SZ,
+                              reinterpret_cast<const BYTE*>(absThisModulePath),
+                              actualPathLenBytesInclNul);
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+
+    status = ::RegSetValueExW(
+        inprocKey, L"ThreadingModel", 0, REG_SZ,
+        reinterpret_cast<const BYTE*>(threadingModelStr.begin().get()),
+        threadingModelStrLenBytesInclNul);
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+  }
+
+  if (aAppId) {
+    // When specified, we must also create a key for the AppID.
+    wchar_t appidKeyPath[256];
+    hr = BuildAppidPath(*aAppId, appidKeyPath);
+    if (FAILED(hr)) {
+      return hr;
+    }
+
+    HKEY rawAppidKey;
+    status = ::RegCreateKeyTransactedW(
+        HKEY_LOCAL_MACHINE, appidKeyPath, 0, nullptr, REG_OPTION_NON_VOLATILE,
+        KEY_ALL_ACCESS, nullptr, &rawAppidKey, nullptr, txn, nullptr);
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+    nsAutoRegKey appidKey(rawAppidKey);
+
+    // Setting DllSurrogate to a null or empty string indicates to Windows that
+    // we want to use the default surrogate (i.e. dllhost.exe) to load our DLL.
+    status =
+        ::RegSetValueExW(appidKey, L"DllSurrogate", 0, REG_SZ,
+                         reinterpret_cast<const BYTE*>(L""), sizeof(wchar_t));
+    if (status != ERROR_SUCCESS) {
+      return HRESULT_FROM_WIN32(status);
+    }
+  }
+
+  if (!::CommitTransaction(txn)) {
+    return HRESULT_FROM_WIN32(::GetLastError());
+  }
+
+  return S_OK;
+}
+
+/**
+ * Unfortunately the registry transaction APIs are not as well-developed for
+ * deleting things as they are for creating them. We just use RegDeleteTree
+ * for the implementation of this method.
+ */
+HRESULT Module::Deregister(const CLSID* const* aClsids, const size_t aNumClsids,
+                           const GUID* const aAppId) {
+  MOZ_ASSERT(aClsids && aNumClsids);
+  if (!aClsids || !aNumClsids) {
+    return E_INVALIDARG;
+  }
+
+  HRESULT hr;
+  LSTATUS status;
+
+  // Delete the key for each CLSID. This will also delete any references to
+  // the AppId.
+  for (size_t idx = 0; idx < aNumClsids; ++idx) {
+    if (!aClsids[idx]) {
+      return E_INVALIDARG;
+    }
+
+    wchar_t clsidKeyPath[256];
+    hr = BuildClsidPath(*aClsids[idx], clsidKeyPath);
+    if (FAILED(hr)) {
+      return hr;
+    }
+
+    status = ::RegDeleteTreeW(HKEY_LOCAL_MACHINE, clsidKeyPath);
+    // We allow the deletion to succeed if the key was already gone
+    if (status != ERROR_SUCCESS && status != ERROR_FILE_NOT_FOUND) {
+      return HRESULT_FROM_WIN32(status);
+    }
+  }
+
+  // Now delete the AppID key, if desired.
+  if (aAppId) {
+    wchar_t appidKeyPath[256];
+    hr = BuildAppidPath(*aAppId, appidKeyPath);
+    if (FAILED(hr)) {
+      return hr;
+    }
+
+    status = ::RegDeleteTreeW(HKEY_LOCAL_MACHINE, appidKeyPath);
+    // We allow the deletion to succeed if the key was already gone
+    if (status != ERROR_SUCCESS && status != ERROR_FILE_NOT_FOUND) {
+      return HRESULT_FROM_WIN32(status);
+    }
+  }
+
+  return S_OK;
+}
+
 }  // namespace mscom
 }  // namespace mozilla