1 files changed, 245 insertions, 0 deletions
diff --git a/toolkit/profile/ProfileUnlockerWin.cpp b/toolkit/profile/ProfileUnlockerWin.cpp
new file mode 100644
index 0000000000..d7ab5e0a30
--- /dev/null
+++ b/toolkit/profile/ProfileUnlockerWin.cpp
@@ -0,0 +1,245 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* 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 "ProfileUnlockerWin.h"
+#include "nsCOMPtr.h"
+#include "nsIFile.h"
+#include "nsTArray.h"
+#include "nsXPCOM.h"
+
+namespace mozilla {
+
+/**
+ * RAII class to obtain and manage a handle to a Restart Manager session.
+ * It opens a new handle upon construction and releases it upon destruction.
+ */
+class MOZ_STACK_CLASS ScopedRestartManagerSession {
+ public:
+  explicit ScopedRestartManagerSession(ProfileUnlockerWin& aUnlocker)
+      : mError(ERROR_INVALID_HANDLE),
+        mHandle((DWORD)-1)  // 0 is a valid restart manager handle
+        ,
+        mUnlocker(aUnlocker) {
+    mError = mUnlocker.StartSession(mHandle);
+  }
+
+  ~ScopedRestartManagerSession() {
+    if (mError == ERROR_SUCCESS) {
+      mUnlocker.EndSession(mHandle);
+    }
+  }
+
+  /**
+   * @return true if the handle is a valid Restart Ranager handle.
+   */
+  inline bool ok() { return mError == ERROR_SUCCESS; }
+
+  /**
+   * @return the Restart Manager handle to pass to other Restart Manager APIs.
+   */
+  inline DWORD handle() { return mHandle; }
+
+ private:
+  DWORD mError;
+  DWORD mHandle;
+  ProfileUnlockerWin& mUnlocker;
+};
+
+ProfileUnlockerWin::ProfileUnlockerWin(const nsAString& aFileName)
+    : mRmStartSession(nullptr),
+      mRmRegisterResources(nullptr),
+      mRmGetList(nullptr),
+      mRmEndSession(nullptr),
+      mFileName(aFileName) {}
+
+ProfileUnlockerWin::~ProfileUnlockerWin() {}
+
+NS_IMPL_ISUPPORTS(ProfileUnlockerWin, nsIProfileUnlocker)
+
+nsresult ProfileUnlockerWin::Init() {
+  MOZ_ASSERT(!mRestartMgrModule);
+  if (mFileName.IsEmpty()) {
+    return NS_ERROR_ILLEGAL_VALUE;
+  }
+
+  nsModuleHandle module(::LoadLibraryW(L"Rstrtmgr.dll"));
+  if (!module) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+  mRmStartSession = reinterpret_cast<RMSTARTSESSION>(
+      ::GetProcAddress(module, "RmStartSession"));
+  if (!mRmStartSession) {
+    return NS_ERROR_UNEXPECTED;
+  }
+  mRmRegisterResources = reinterpret_cast<RMREGISTERRESOURCES>(
+      ::GetProcAddress(module, "RmRegisterResources"));
+  if (!mRmRegisterResources) {
+    return NS_ERROR_UNEXPECTED;
+  }
+  mRmGetList =
+      reinterpret_cast<RMGETLIST>(::GetProcAddress(module, "RmGetList"));
+  if (!mRmGetList) {
+    return NS_ERROR_UNEXPECTED;
+  }
+  mRmEndSession =
+      reinterpret_cast<RMENDSESSION>(::GetProcAddress(module, "RmEndSession"));
+  if (!mRmEndSession) {
+    return NS_ERROR_UNEXPECTED;
+  }
+
+  mRestartMgrModule.steal(module);
+  return NS_OK;
+}
+
+DWORD
+ProfileUnlockerWin::StartSession(DWORD& aHandle) {
+  WCHAR sessionKey[CCH_RM_SESSION_KEY + 1] = {0};
+  return mRmStartSession(&aHandle, 0, sessionKey);
+}
+
+void ProfileUnlockerWin::EndSession(DWORD aHandle) { mRmEndSession(aHandle); }
+
+NS_IMETHODIMP
+ProfileUnlockerWin::Unlock(uint32_t aSeverity) {
+  if (!mRestartMgrModule) {
+    return NS_ERROR_NOT_INITIALIZED;
+  }
+
+  if (aSeverity != FORCE_QUIT) {
+    return NS_ERROR_NOT_IMPLEMENTED;
+  }
+
+  ScopedRestartManagerSession session(*this);
+  if (!session.ok()) {
+    return NS_ERROR_FAILURE;
+  }
+
+  LPCWSTR resources[] = {mFileName.get()};
+  DWORD error = mRmRegisterResources(session.handle(), 1, resources, 0, nullptr,
+                                     0, nullptr);
+  if (error != ERROR_SUCCESS) {
+    return NS_ERROR_FAILURE;
+  }
+
+  // Using a AutoTArray here because we expect the required size to be 1.
+  AutoTArray<RM_PROCESS_INFO, 1> info;
+  UINT numEntries;
+  UINT numEntriesNeeded = 1;
+  error = ERROR_MORE_DATA;
+  DWORD reason = RmRebootReasonNone;
+  while (error == ERROR_MORE_DATA) {
+    info.SetLength(numEntriesNeeded);
+    numEntries = numEntriesNeeded;
+    error = mRmGetList(session.handle(), &numEntriesNeeded, &numEntries,
+                       &info[0], &reason);
+  }
+  if (error != ERROR_SUCCESS) {
+    return NS_ERROR_FAILURE;
+  }
+  if (numEntries == 0) {
+    // Nobody else is locking the file; the other process must have terminated
+    return NS_OK;
+  }
+
+  nsresult rv = NS_ERROR_FAILURE;
+  for (UINT i = 0; i < numEntries; ++i) {
+    rv = TryToTerminate(info[i].Process);
+    if (NS_SUCCEEDED(rv)) {
+      return NS_OK;
+    }
+  }
+
+  // If nothing could be unlocked then we return the error code of the last
+  // failure that was returned.
+  return rv;
+}
+
+nsresult ProfileUnlockerWin::TryToTerminate(RM_UNIQUE_PROCESS& aProcess) {
+  // Subtle: If the target process terminated before this call to OpenProcess,
+  // this call will still succeed. This is because the restart manager session
+  // internally retains a handle to the target process. The rules for Windows
+  // PIDs state that the PID of a terminated process remains valid as long as
+  // at least one handle to that process remains open, so when we reach this
+  // point the PID is still valid and the process will open successfully.
+  DWORD accessRights = PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_TERMINATE;
+  nsAutoHandle otherProcess(
+      ::OpenProcess(accessRights, FALSE, aProcess.dwProcessId));
+  if (!otherProcess) {
+    return NS_ERROR_FAILURE;
+  }
+
+  FILETIME creationTime, exitTime, kernelTime, userTime;
+  if (!::GetProcessTimes(otherProcess, &creationTime, &exitTime, &kernelTime,
+                         &userTime)) {
+    return NS_ERROR_FAILURE;
+  }
+  if (::CompareFileTime(&aProcess.ProcessStartTime, &creationTime)) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  WCHAR imageName[MAX_PATH];
+  DWORD imageNameLen = MAX_PATH;
+  if (!::QueryFullProcessImageNameW(otherProcess, 0, imageName,
+                                    &imageNameLen)) {
+    // The error codes for this function are not very descriptive. There are
+    // actually two failure cases here: Either the call failed because the
+    // process is no longer running, or it failed for some other reason. We
+    // need to know which case that is.
+    DWORD otherProcessExitCode;
+    if (!::GetExitCodeProcess(otherProcess, &otherProcessExitCode) ||
+        otherProcessExitCode == STILL_ACTIVE) {
+      // The other process is still running.
+      return NS_ERROR_FAILURE;
+    }
+    // The other process must have terminated. We should return NS_OK so that
+    // this process may proceed with startup.
+    return NS_OK;
+  }
+  nsCOMPtr<nsIFile> otherProcessImageName;
+  if (NS_FAILED(NS_NewLocalFile(nsDependentString(imageName, imageNameLen),
+                                false,
+                                getter_AddRefs(otherProcessImageName)))) {
+    return NS_ERROR_FAILURE;
+  }
+  nsAutoString otherProcessLeafName;
+  if (NS_FAILED(otherProcessImageName->GetLeafName(otherProcessLeafName))) {
+    return NS_ERROR_FAILURE;
+  }
+
+  imageNameLen = MAX_PATH;
+  if (!::QueryFullProcessImageNameW(::GetCurrentProcess(), 0, imageName,
+                                    &imageNameLen)) {
+    return NS_ERROR_FAILURE;
+  }
+  nsCOMPtr<nsIFile> thisProcessImageName;
+  if (NS_FAILED(NS_NewLocalFile(nsDependentString(imageName, imageNameLen),
+                                false, getter_AddRefs(thisProcessImageName)))) {
+    return NS_ERROR_FAILURE;
+  }
+  nsAutoString thisProcessLeafName;
+  if (NS_FAILED(thisProcessImageName->GetLeafName(thisProcessLeafName))) {
+    return NS_ERROR_FAILURE;
+  }
+
+  // Make sure the image leaf names match
+  if (_wcsicmp(otherProcessLeafName.get(), thisProcessLeafName.get())) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  // We know that another process holds the lock and that it shares the same
+  // image name as our process. Let's kill it.
+  // Subtle: TerminateProcess returning ERROR_ACCESS_DENIED is actually an
+  // indicator that the target process managed to shut down on its own. In that
+  // case we should return NS_OK since we may proceed with startup.
+  if (!::TerminateProcess(otherProcess, 1) &&
+      ::GetLastError() != ERROR_ACCESS_DENIED) {
+    return NS_ERROR_FAILURE;
+  }
+
+  return NS_OK;
+}
+
+}  // namespace mozilla