1 files changed, 533 insertions, 0 deletions
diff --git a/toolkit/xre/LauncherRegistryInfo.cpp b/toolkit/xre/LauncherRegistryInfo.cpp
new file mode 100644
index 0000000000..6778827fec
--- /dev/null
+++ b/toolkit/xre/LauncherRegistryInfo.cpp
@@ -0,0 +1,533 @@
+/* -*- 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 https://mozilla.org/MPL/2.0/. */
+
+#include "LauncherRegistryInfo.h"
+
+#include "mozilla/UniquePtr.h"
+#include "mozilla/ArrayUtils.h"
+#include "mozilla/Assertions.h"
+#include "mozilla/NativeNt.h"
+
+#include <string>
+
+#define EXPAND_STRING_MACRO2(t) t
+#define EXPAND_STRING_MACRO(t) EXPAND_STRING_MACRO2(t)
+
+// This function is copied from Chromium base/time/time_win.cc
+// Returns the current value of the performance counter.
+static uint64_t QPCNowRaw() {
+  LARGE_INTEGER perf_counter_now = {};
+  // According to the MSDN documentation for QueryPerformanceCounter(), this
+  // will never fail on systems that run XP or later.
+  // https://docs.microsoft.com/en-us/windows/win32/api/profileapi/nf-profileapi-queryperformancecounter
+  ::QueryPerformanceCounter(&perf_counter_now);
+  return perf_counter_now.QuadPart;
+}
+
+static mozilla::LauncherResult<DWORD> GetCurrentImageTimestamp() {
+  mozilla::nt::PEHeaders headers(::GetModuleHandleW(nullptr));
+  if (!headers) {
+    return LAUNCHER_ERROR_FROM_WIN32(ERROR_BAD_EXE_FORMAT);
+  }
+
+  DWORD timestamp;
+  if (!headers.GetTimeStamp(timestamp)) {
+    return LAUNCHER_ERROR_FROM_WIN32(ERROR_INVALID_DATA);
+  }
+
+  return timestamp;
+}
+
+template <typename T>
+static mozilla::LauncherResult<mozilla::Maybe<T>> ReadRegistryValueData(
+    const nsAutoRegKey& key, const std::wstring& name, DWORD expectedType) {
+  static_assert(mozilla::IsPod<T>::value,
+                "Registry value type must be primitive.");
+  T data;
+  DWORD dataLen = sizeof(data);
+  DWORD type;
+  LSTATUS status = ::RegQueryValueExW(key.get(), name.c_str(), nullptr, &type,
+                                      reinterpret_cast<PBYTE>(&data), &dataLen);
+  if (status == ERROR_FILE_NOT_FOUND) {
+    return mozilla::Maybe<T>();
+  }
+
+  if (status != ERROR_SUCCESS) {
+    return LAUNCHER_ERROR_FROM_WIN32(status);
+  }
+
+  if (type != expectedType) {
+    return LAUNCHER_ERROR_FROM_WIN32(ERROR_DATATYPE_MISMATCH);
+  }
+
+  return mozilla::Some(data);
+}
+
+template <typename T>
+static mozilla::LauncherVoidResult WriteRegistryValueData(
+    const nsAutoRegKey& key, const std::wstring& name, DWORD type, T data) {
+  static_assert(mozilla::IsPod<T>::value,
+                "Registry value type must be primitive.");
+  LSTATUS status =
+      ::RegSetValueExW(key.get(), name.c_str(), 0, type,
+                       reinterpret_cast<PBYTE>(&data), sizeof(data));
+  if (status != ERROR_SUCCESS) {
+    return LAUNCHER_ERROR_FROM_WIN32(status);
+  }
+
+  return mozilla::Ok();
+}
+
+static mozilla::LauncherResult<bool> DeleteRegistryValueData(
+    const nsAutoRegKey& key, const std::wstring& name) {
+  LSTATUS status = ::RegDeleteValueW(key, name.c_str());
+  if (status == ERROR_FILE_NOT_FOUND) {
+    return false;
+  }
+
+  if (status != ERROR_SUCCESS) {
+    return LAUNCHER_ERROR_FROM_WIN32(status);
+  }
+
+  return true;
+}
+
+namespace mozilla {
+
+const wchar_t LauncherRegistryInfo::kLauncherSubKeyPath[] =
+    L"SOFTWARE\\" EXPAND_STRING_MACRO(MOZ_APP_VENDOR) L"\\" EXPAND_STRING_MACRO(
+        MOZ_APP_BASENAME) L"\\Launcher";
+const wchar_t LauncherRegistryInfo::kLauncherSuffix[] = L"|Launcher";
+const wchar_t LauncherRegistryInfo::kBrowserSuffix[] = L"|Browser";
+const wchar_t LauncherRegistryInfo::kImageTimestampSuffix[] = L"|Image";
+const wchar_t LauncherRegistryInfo::kTelemetrySuffix[] = L"|Telemetry";
+
+bool LauncherRegistryInfo::sAllowCommit = true;
+
+LauncherResult<LauncherRegistryInfo::Disposition> LauncherRegistryInfo::Open() {
+  if (!!mRegKey) {
+    return Disposition::OpenedExisting;
+  }
+
+  DWORD disposition;
+  HKEY rawKey;
+  LSTATUS result = ::RegCreateKeyExW(
+      HKEY_CURRENT_USER, kLauncherSubKeyPath, 0, nullptr,
+      REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, nullptr, &rawKey, &disposition);
+  if (result != ERROR_SUCCESS) {
+    return LAUNCHER_ERROR_FROM_WIN32(result);
+  }
+
+  mRegKey.own(rawKey);
+
+  switch (disposition) {
+    case REG_CREATED_NEW_KEY:
+      return Disposition::CreatedNew;
+    case REG_OPENED_EXISTING_KEY:
+      return Disposition::OpenedExisting;
+    default:
+      break;
+  }
+
+  MOZ_ASSERT_UNREACHABLE("Invalid disposition from RegCreateKeyExW");
+  return LAUNCHER_ERROR_GENERIC();
+}
+
+LauncherVoidResult LauncherRegistryInfo::ReflectPrefToRegistry(
+    const bool aEnable) {
+  LauncherResult<EnabledState> curEnabledState = IsEnabled();
+  if (curEnabledState.isErr()) {
+    return curEnabledState.propagateErr();
+  }
+
+  bool isCurrentlyEnabled =
+      curEnabledState.inspect() != EnabledState::ForceDisabled;
+  if (isCurrentlyEnabled == aEnable) {
+    // Don't reflect to the registry unless the new enabled state is actually
+    // changing with respect to the current enabled state.
+    return Ok();
+  }
+
+  // Always delete the launcher timestamp
+  LauncherResult<bool> clearedLauncherTimestamp = ClearLauncherStartTimestamp();
+  MOZ_ASSERT(clearedLauncherTimestamp.isOk());
+  if (clearedLauncherTimestamp.isErr()) {
+    return clearedLauncherTimestamp.propagateErr();
+  }
+
+  // Allow commit when we enable the launcher, otherwise block.
+  sAllowCommit = aEnable;
+
+  if (!aEnable) {
+    // Set the browser timestamp to 0 to indicate force-disabled
+    return WriteBrowserStartTimestamp(0ULL);
+  }
+
+  // Otherwise we delete the browser timestamp to start over fresh
+  LauncherResult<bool> clearedBrowserTimestamp = ClearBrowserStartTimestamp();
+  MOZ_ASSERT(clearedBrowserTimestamp.isOk());
+  if (clearedBrowserTimestamp.isErr()) {
+    return clearedBrowserTimestamp.propagateErr();
+  }
+
+  return Ok();
+}
+
+LauncherVoidResult LauncherRegistryInfo::ReflectTelemetryPrefToRegistry(
+    const bool aEnable) {
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+
+  return WriteRegistryValueData(mRegKey, ResolveTelemetryValueName(), REG_DWORD,
+                                aEnable ? 1UL : 0UL);
+}
+
+LauncherResult<LauncherRegistryInfo::ProcessType> LauncherRegistryInfo::Check(
+    const ProcessType aDesiredType, const CheckOption aOption) {
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+
+  LauncherResult<DWORD> ourImageTimestamp = GetCurrentImageTimestamp();
+  if (ourImageTimestamp.isErr()) {
+    return ourImageTimestamp.propagateErr();
+  }
+
+  LauncherResult<Maybe<DWORD>> savedImageTimestamp = GetSavedImageTimestamp();
+  if (savedImageTimestamp.isErr()) {
+    return savedImageTimestamp.propagateErr();
+  }
+
+  // If we don't have a saved timestamp, or we do but it doesn't match with
+  // our current timestamp, clear previous values unless we're force-disabled.
+  if (savedImageTimestamp.inspect().isNothing() ||
+      savedImageTimestamp.inspect().value() != ourImageTimestamp.inspect()) {
+    LauncherVoidResult clearResult = ClearStartTimestamps();
+    if (clearResult.isErr()) {
+      return clearResult.propagateErr();
+    }
+
+    LauncherVoidResult writeResult =
+        WriteImageTimestamp(ourImageTimestamp.inspect());
+    if (writeResult.isErr()) {
+      return writeResult.propagateErr();
+    }
+  }
+
+  // If we're going to be running as the browser process, or there is no
+  // existing values to check, just write our timestamp and return.
+  if (aDesiredType == ProcessType::Browser) {
+    mBrowserTimestampToWrite = Some(QPCNowRaw());
+    return ProcessType::Browser;
+  }
+
+  if (disposition.inspect() == Disposition::CreatedNew) {
+    mLauncherTimestampToWrite = Some(QPCNowRaw());
+    return ProcessType::Launcher;
+  }
+
+  if (disposition.inspect() != Disposition::OpenedExisting) {
+    MOZ_ASSERT_UNREACHABLE("Invalid |disposition|");
+    return LAUNCHER_ERROR_GENERIC();
+  }
+
+  LauncherResult<Maybe<uint64_t>> lastLauncherTimestampResult =
+      GetLauncherStartTimestamp();
+  if (lastLauncherTimestampResult.isErr()) {
+    return lastLauncherTimestampResult.propagateErr();
+  }
+
+  LauncherResult<Maybe<uint64_t>> lastBrowserTimestampResult =
+      GetBrowserStartTimestamp();
+  if (lastBrowserTimestampResult.isErr()) {
+    return lastBrowserTimestampResult.propagateErr();
+  }
+
+  const Maybe<uint64_t>& lastLauncherTimestamp =
+      lastLauncherTimestampResult.inspect();
+  const Maybe<uint64_t>& lastBrowserTimestamp =
+      lastBrowserTimestampResult.inspect();
+
+  ProcessType typeToRunAs = aDesiredType;
+
+  if (lastLauncherTimestamp.isSome() != lastBrowserTimestamp.isSome()) {
+    // If we have a launcher timestamp but no browser timestamp (or vice versa),
+    // that's bad because it is indicating that the browser can't run with
+    // the launcher process.
+    typeToRunAs = ProcessType::Browser;
+  } else if (lastLauncherTimestamp.isSome()) {
+    // if we have both timestamps, we want to ensure that the launcher timestamp
+    // is earlier than the browser timestamp.
+    if (aDesiredType == ProcessType::Launcher) {
+      bool areTimestampsOk =
+          lastLauncherTimestamp.value() < lastBrowserTimestamp.value();
+      if (!areTimestampsOk) {
+        typeToRunAs = ProcessType::Browser;
+      }
+    }
+  } else {
+    // If we have neither timestamp, then we should try running as suggested
+    // by |aDesiredType|.
+    // We shouldn't really have this scenario unless we're going to be running
+    // as the launcher process.
+    MOZ_ASSERT(typeToRunAs == ProcessType::Launcher);
+    // No change to typeToRunAs
+  }
+
+  // Debugging setting that forces the desired type regardless of the various
+  // tests that have been performed.
+  if (aOption == CheckOption::Force) {
+    typeToRunAs = aDesiredType;
+  }
+
+  switch (typeToRunAs) {
+    case ProcessType::Browser:
+      if (aDesiredType != typeToRunAs) {
+        // We were hoping to run as the launcher, but some failure has caused
+        // us to run as the browser. Set the browser timestamp to zero as an
+        // indicator.
+        mBrowserTimestampToWrite = Some(0ULL);
+      } else {
+        mBrowserTimestampToWrite = Some(QPCNowRaw());
+      }
+      break;
+    case ProcessType::Launcher:
+      mLauncherTimestampToWrite = Some(QPCNowRaw());
+      break;
+    default:
+      MOZ_ASSERT_UNREACHABLE("Invalid |typeToRunAs|");
+      return LAUNCHER_ERROR_GENERIC();
+  }
+
+  return typeToRunAs;
+}
+
+LauncherVoidResult LauncherRegistryInfo::DisableDueToFailure() {
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+  LauncherVoidResult result = WriteBrowserStartTimestamp(0ULL);
+  if (result.isOk()) {
+    // Block commit when we disable the launcher.  It could be allowed
+    // when the image timestamp is updated.
+    sAllowCommit = false;
+  }
+  return result;
+}
+
+LauncherVoidResult LauncherRegistryInfo::Commit() {
+  if (!sAllowCommit) {
+    Abort();
+    return Ok();
+  }
+
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+
+  if (mLauncherTimestampToWrite.isSome()) {
+    LauncherVoidResult writeResult =
+        WriteLauncherStartTimestamp(mLauncherTimestampToWrite.value());
+    if (writeResult.isErr()) {
+      return writeResult.propagateErr();
+    }
+    mLauncherTimestampToWrite = Nothing();
+  }
+
+  if (mBrowserTimestampToWrite.isSome()) {
+    LauncherVoidResult writeResult =
+        WriteBrowserStartTimestamp(mBrowserTimestampToWrite.value());
+    if (writeResult.isErr()) {
+      return writeResult.propagateErr();
+    }
+    mBrowserTimestampToWrite = Nothing();
+  }
+
+  return Ok();
+}
+
+void LauncherRegistryInfo::Abort() {
+  mLauncherTimestampToWrite = mBrowserTimestampToWrite = Nothing();
+}
+
+LauncherRegistryInfo::EnabledState LauncherRegistryInfo::GetEnabledState(
+    const Maybe<uint64_t>& aLauncherTs, const Maybe<uint64_t>& aBrowserTs) {
+  if (aBrowserTs.isSome()) {
+    if (aLauncherTs.isSome()) {
+      if (aLauncherTs.value() < aBrowserTs.value()) {
+        // Both timestamps exist and the browser's timestamp is later.
+        return EnabledState::Enabled;
+      }
+    } else if (aBrowserTs.value() == 0ULL) {
+      // Only browser's timestamp exists and its value is 0.
+      return EnabledState::ForceDisabled;
+    }
+  } else if (aLauncherTs.isNothing()) {
+    // Neither timestamps exist.
+    return EnabledState::Enabled;
+  }
+
+  // Everything else is FailDisabled.
+  return EnabledState::FailDisabled;
+}
+
+LauncherResult<LauncherRegistryInfo::EnabledState>
+LauncherRegistryInfo::IsEnabled() {
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+
+  LauncherResult<Maybe<uint64_t>> lastLauncherTimestamp =
+      GetLauncherStartTimestamp();
+  if (lastLauncherTimestamp.isErr()) {
+    return lastLauncherTimestamp.propagateErr();
+  }
+
+  LauncherResult<Maybe<uint64_t>> lastBrowserTimestamp =
+      GetBrowserStartTimestamp();
+  if (lastBrowserTimestamp.isErr()) {
+    return lastBrowserTimestamp.propagateErr();
+  }
+
+  return GetEnabledState(lastLauncherTimestamp.inspect(),
+                         lastBrowserTimestamp.inspect());
+}
+
+LauncherResult<bool> LauncherRegistryInfo::IsTelemetryEnabled() {
+  LauncherResult<Disposition> disposition = Open();
+  if (disposition.isErr()) {
+    return disposition.propagateErr();
+  }
+
+  LauncherResult<Maybe<DWORD>> result = ReadRegistryValueData<DWORD>(
+      mRegKey, ResolveTelemetryValueName(), REG_DWORD);
+  if (result.isErr()) {
+    return result.propagateErr();
+  }
+
+  if (result.inspect().isNothing()) {
+    // Value does not exist, treat as false
+    return false;
+  }
+
+  return result.inspect().value() != 0;
+}
+
+const std::wstring& LauncherRegistryInfo::ResolveLauncherValueName() {
+  if (mLauncherValueName.empty()) {
+    mLauncherValueName.assign(mBinPath);
+    mLauncherValueName.append(kLauncherSuffix,
+                              ArrayLength(kLauncherSuffix) - 1);
+  }
+
+  return mLauncherValueName;
+}
+
+const std::wstring& LauncherRegistryInfo::ResolveBrowserValueName() {
+  if (mBrowserValueName.empty()) {
+    mBrowserValueName.assign(mBinPath);
+    mBrowserValueName.append(kBrowserSuffix, ArrayLength(kBrowserSuffix) - 1);
+  }
+
+  return mBrowserValueName;
+}
+
+const std::wstring& LauncherRegistryInfo::ResolveImageTimestampValueName() {
+  if (mImageValueName.empty()) {
+    mImageValueName.assign(mBinPath);
+    mImageValueName.append(kImageTimestampSuffix,
+                           ArrayLength(kImageTimestampSuffix) - 1);
+  }
+
+  return mImageValueName;
+}
+
+const std::wstring& LauncherRegistryInfo::ResolveTelemetryValueName() {
+  if (mTelemetryValueName.empty()) {
+    mTelemetryValueName.assign(mBinPath);
+    mTelemetryValueName.append(kTelemetrySuffix,
+                               ArrayLength(kTelemetrySuffix) - 1);
+  }
+
+  return mTelemetryValueName;
+}
+
+LauncherVoidResult LauncherRegistryInfo::WriteLauncherStartTimestamp(
+    uint64_t aValue) {
+  return WriteRegistryValueData(mRegKey, ResolveLauncherValueName(), REG_QWORD,
+                                aValue);
+}
+
+LauncherVoidResult LauncherRegistryInfo::WriteBrowserStartTimestamp(
+    uint64_t aValue) {
+  return WriteRegistryValueData(mRegKey, ResolveBrowserValueName(), REG_QWORD,
+                                aValue);
+}
+
+LauncherVoidResult LauncherRegistryInfo::WriteImageTimestamp(DWORD aTimestamp) {
+  return WriteRegistryValueData(mRegKey, ResolveImageTimestampValueName(),
+                                REG_DWORD, aTimestamp);
+}
+
+LauncherResult<bool> LauncherRegistryInfo::ClearLauncherStartTimestamp() {
+  return DeleteRegistryValueData(mRegKey, ResolveLauncherValueName());
+}
+
+LauncherResult<bool> LauncherRegistryInfo::ClearBrowserStartTimestamp() {
+  return DeleteRegistryValueData(mRegKey, ResolveBrowserValueName());
+}
+
+LauncherVoidResult LauncherRegistryInfo::ClearStartTimestamps() {
+  LauncherResult<EnabledState> enabled = IsEnabled();
+  if (enabled.isOk() && enabled.inspect() == EnabledState::ForceDisabled) {
+    // We don't clear anything when we're force disabled - we need to maintain
+    // the current registry state in this case.
+    return Ok();
+  }
+
+  LauncherResult<bool> clearedLauncherTimestamp = ClearLauncherStartTimestamp();
+  if (clearedLauncherTimestamp.isErr()) {
+    return clearedLauncherTimestamp.propagateErr();
+  }
+
+  LauncherResult<bool> clearedBrowserTimestamp = ClearBrowserStartTimestamp();
+  if (clearedBrowserTimestamp.isErr()) {
+    return clearedBrowserTimestamp.propagateErr();
+  }
+
+  // Reset both timestamps to align with registry deletion
+  mLauncherTimestampToWrite = mBrowserTimestampToWrite = Nothing();
+
+  // Disablement is gone.  Let's allow commit.
+  sAllowCommit = true;
+
+  return Ok();
+}
+
+LauncherResult<Maybe<DWORD>> LauncherRegistryInfo::GetSavedImageTimestamp() {
+  return ReadRegistryValueData<DWORD>(mRegKey, ResolveImageTimestampValueName(),
+                                      REG_DWORD);
+}
+
+LauncherResult<Maybe<uint64_t>>
+LauncherRegistryInfo::GetLauncherStartTimestamp() {
+  return ReadRegistryValueData<uint64_t>(mRegKey, ResolveLauncherValueName(),
+                                         REG_QWORD);
+}
+
+LauncherResult<Maybe<uint64_t>>
+LauncherRegistryInfo::GetBrowserStartTimestamp() {
+  return ReadRegistryValueData<uint64_t>(mRegKey, ResolveBrowserValueName(),
+                                         REG_QWORD);
+}
+
+}  // namespace mozilla