1 files changed, 278 insertions, 0 deletions
diff --git a/toolkit/profile/ProfileUnlockerWin.cpp b/toolkit/profile/ProfileUnlockerWin.cpp
new file mode 100644
index 000000000..0e0139188
--- /dev/null
+++ b/toolkit/profile/ProfileUnlockerWin.cpp
@@ -0,0 +1,278 @@
+/* -*- 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)
+  , mQueryFullProcessImageName(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;
+  }
+
+  mQueryFullProcessImageName =
+    reinterpret_cast<QUERYFULLPROCESSIMAGENAME>(::GetProcAddress(
+                                  ::GetModuleHandleW(L"kernel32.dll"),
+                                  "QueryFullProcessImageNameW"));
+  if (!mQueryFullProcessImageName) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  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 (!mQueryFullProcessImageName(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 (!mQueryFullProcessImageName(::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
+