4 files changed, 155 insertions, 0 deletions
diff --git a/toolkit/xre/dllservices/tests/gtest/rust/Cargo.toml b/toolkit/xre/dllservices/tests/gtest/rust/Cargo.toml
new file mode 100644
index 0000000000..a342fb8bf2
--- /dev/null
+++ b/toolkit/xre/dllservices/tests/gtest/rust/Cargo.toml
@@ -0,0 +1,12 @@
+[package]
+name = "dllservices-gtest"
+version = "0.1.0"
+authors = ["nobody@mozilla.com"]
+license = "MPL-2.0"
+description = "Tests for dllservices"
+
+[dependencies]
+uuid = { version = "1.0", features = ["v4"] }
+
+[lib]
+path = "test.rs"
diff --git a/toolkit/xre/dllservices/tests/gtest/rust/TestBCryptFallback.cpp b/toolkit/xre/dllservices/tests/gtest/rust/TestBCryptFallback.cpp
new file mode 100644
index 0000000000..ca4dbde3ab
--- /dev/null
+++ b/toolkit/xre/dllservices/tests/gtest/rust/TestBCryptFallback.cpp
@@ -0,0 +1,113 @@
+/* -*- 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 <windows.h>
+#include <ntstatus.h>
+
+#include <bcrypt.h>
+#pragma comment(lib, "bcrypt.lib")
+
+#include "gtest/gtest.h"
+
+#include "nsWindowsDllInterceptor.h"
+
+#define RtlGenRandom SystemFunction036
+extern "C" BOOLEAN NTAPI RtlGenRandom(PVOID aRandomBuffer,
+                                      ULONG aRandomBufferLength);
+
+static mozilla::WindowsDllInterceptor BCryptIntercept;
+static mozilla::WindowsDllInterceptor::FuncHookType<
+    decltype(&::BCryptGenRandom)>
+    stub_BCryptGenRandom;
+
+static mozilla::WindowsDllInterceptor AdvApiIntercept;
+static mozilla::WindowsDllInterceptor::FuncHookType<decltype(&RtlGenRandom)>
+    stub_RtlGenRandom;
+
+volatile bool gAreHooksActive = false;
+volatile bool gHasPanicked = false;
+volatile bool gHasReachedBCryptGenRandom = false;
+volatile bool gHasReachedRtlGenRandom = false;
+
+NTSTATUS WINAPI patched_BCryptGenRandom(BCRYPT_ALG_HANDLE aAlgorithm,
+                                        PUCHAR aBuffer, ULONG aSize,
+                                        ULONG aFlags) {
+  if (gAreHooksActive) {
+    gHasReachedBCryptGenRandom = true;
+    // Force BCryptGenRandom failures when the hook is active.
+    return STATUS_UNSUCCESSFUL;
+  }
+  return stub_BCryptGenRandom(aAlgorithm, aBuffer, aSize, aFlags);
+}
+
+BOOLEAN NTAPI patched_RtlGenRandom(PVOID aRandomBuffer,
+                                   ULONG aRandomBufferLength) {
+  if (gAreHooksActive) {
+    gHasReachedRtlGenRandom = true;
+  }
+  return stub_RtlGenRandom(aRandomBuffer, aRandomBufferLength);
+}
+
+bool InitInterception() {
+  static bool sSuccess = []() {
+    BCryptIntercept.Init(L"bcrypt.dll");
+    AdvApiIntercept.Init(L"advapi32.dll");
+    return stub_BCryptGenRandom.SetDetour(BCryptIntercept, "BCryptGenRandom",
+                                          patched_BCryptGenRandom) &&
+           stub_RtlGenRandom.SetDetour(AdvApiIntercept, "SystemFunction036",
+                                       patched_RtlGenRandom);
+  }();
+  gAreHooksActive = true;
+  return sSuccess;
+}
+
+void ExitInterception() { gAreHooksActive = false; }
+
+DWORD WINAPI TestIsFallbackTriggeredThreadProc(LPVOID aParameter) {
+  auto testedFunction = reinterpret_cast<void (*)()>(aParameter);
+  EXPECT_TRUE(InitInterception());
+  MOZ_SEH_TRY { testedFunction(); }
+  MOZ_SEH_EXCEPT(EXCEPTION_EXECUTE_HANDLER) {
+    // Catch a potential Rust panic
+    gHasPanicked = true;
+  }
+  ExitInterception();
+  return 0;
+}
+
+// This function hooks BCryptGenRandom to make it fail, and hooks RtlGenRandom
+// to allow us to ensure that it gets visited as a fallback for
+// BCryptGenRandom.
+void TestIsFallbackTriggered(void (*aTestedFunction)()) {
+  gHasPanicked = false;
+  gHasReachedBCryptGenRandom = false;
+  gHasReachedRtlGenRandom = false;
+
+  // The HashMap test must run on a new thread, because some random bytes have
+  // already been collected but not used on the current thread by previous
+  // calls to HashMap::new in various locations of the code base. These bytes
+  // would be recycled instead of calling into BCryptGenRandom and RtlGenRandom
+  // if running the HashMap test on the current thread.
+  auto thread =
+      ::CreateThread(nullptr, 0, TestIsFallbackTriggeredThreadProc,
+                     reinterpret_cast<void*>(aTestedFunction), 0, nullptr);
+  EXPECT_TRUE(bool(thread));
+  EXPECT_EQ(::WaitForSingleObject(thread, 5000),
+            static_cast<DWORD>(WAIT_OBJECT_0));
+
+  EXPECT_FALSE(gHasPanicked);
+  EXPECT_TRUE(gHasReachedBCryptGenRandom);
+  EXPECT_TRUE(gHasReachedRtlGenRandom);
+}
+
+extern "C" void Rust_TriggerGenRandomFromHashMap();
+extern "C" void Rust_TriggerGenRandomFromUuid();
+
+TEST(TestBCryptFallback, HashMapTriggersFallback)
+{ TestIsFallbackTriggered(Rust_TriggerGenRandomFromHashMap); }
+
+TEST(TestBCryptFallback, UuidTriggersFallback)
+{ TestIsFallbackTriggered(Rust_TriggerGenRandomFromUuid); }
diff --git a/toolkit/xre/dllservices/tests/gtest/rust/moz.build b/toolkit/xre/dllservices/tests/gtest/rust/moz.build
new file mode 100644
index 0000000000..eea8b9e978
--- /dev/null
+++ b/toolkit/xre/dllservices/tests/gtest/rust/moz.build
@@ -0,0 +1,11 @@
+# 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/.
+
+Library("dllservicestest")
+
+UNIFIED_SOURCES += [
+    "TestBCryptFallback.cpp",
+]
+
+FINAL_LIBRARY = "xul-gtest"
diff --git a/toolkit/xre/dllservices/tests/gtest/rust/test.rs b/toolkit/xre/dllservices/tests/gtest/rust/test.rs
new file mode 100644
index 0000000000..51c2a6e156
--- /dev/null
+++ b/toolkit/xre/dllservices/tests/gtest/rust/test.rs
@@ -0,0 +1,19 @@
+/* -*- Mode: rust; rust-indent-offset: 2 -*- */
+/* 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/. */
+
+use std::collections::HashMap;
+
+extern crate uuid;
+use uuid::Uuid;
+
+#[no_mangle]
+pub extern "C" fn Rust_TriggerGenRandomFromHashMap() -> () {
+    let _: HashMap<u32, u32> = HashMap::new();
+}
+
+#[no_mangle]
+pub extern "C" fn Rust_TriggerGenRandomFromUuid() -> () {
+    let _: Uuid = Uuid::new_v4();
+}