From c21ea0b35bd661ff73fa83d52d1b895d1f1c5e40 Mon Sep 17 00:00:00 2001
From: Mike Hommey <mh+mozilla@glandium.org>
Date: Tue, 25 Apr 2023 22:18:52 +0000
Subject: [PATCH] Bug 1829780 - Adjust our OOM hook to the new API from rustc
 1.71. r=emilio

The new API landed in https://github.com/rust-lang/rust/pull/109507

Differential Revision: https://phabricator.services.mozilla.com/D176383
---
 js/src/vm/Initialization.cpp |  6 ++----
 mozglue/static/rust/build.rs |  7 ++++--
 mozglue/static/rust/lib.rs   | 42 +++++++++++++++++++++++-------------
 toolkit/xre/nsAppRunner.cpp  |  8 ++-----
 4 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/js/src/vm/Initialization.cpp b/js/src/vm/Initialization.cpp
index d03f95af32794..a452190e3b737 100644
--- a/js/src/vm/Initialization.cpp
+++ b/js/src/vm/Initialization.cpp
@@ -105,8 +105,7 @@ static void SetupCanonicalNaN() {
     if (!code) return #code " failed"; \
   } while (0)
 
-extern "C" void install_rust_panic_hook();
-extern "C" void install_rust_oom_hook();
+extern "C" void install_rust_hooks();
 
 JS_PUBLIC_API const char* JS::detail::InitWithFailureDiagnostic(
     bool isDebugBuild) {
@@ -127,8 +126,7 @@ JS_PUBLIC_API const char* JS::detail::InitWithFailureDiagnostic(
 
 #ifdef JS_STANDALONE
   // The rust hooks are initialized by Gecko on non-standalone builds.
-  install_rust_panic_hook();
-  install_rust_oom_hook();
+  install_rust_hooks();
 #endif
 
   PRMJ_NowInit();
diff --git a/mozglue/static/rust/build.rs b/mozglue/static/rust/build.rs
index 85a2a2ec96e17..424083d617030 100644
--- a/mozglue/static/rust/build.rs
+++ b/mozglue/static/rust/build.rs
@@ -16,12 +16,15 @@ fn main() {
     println!("cargo:rerun-if-changed=wrappers.cpp");
 
     let ver = version().unwrap();
-    let max_oom_hook_version = Version::parse("1.70.0-alpha").unwrap();
+    let max_oom_hook_version = Version::parse("1.71.0-alpha").unwrap();
+    let max_alloc_error_panic_version = Version::parse("1.72.0-alpha").unwrap();
 
     if ver < max_oom_hook_version {
         println!("cargo:rustc-cfg=feature=\"oom_with_hook\"");
+    } else if ver < max_alloc_error_panic_version {
+        println!("cargo:rustc-cfg=feature=\"oom_with_alloc_error_panic\"");
     } else if std::env::var("MOZ_AUTOMATION").is_ok() {
         panic!("Builds on automation must use a version of rust for which we know how to hook OOM: want < {}, have {}",
-               max_oom_hook_version, ver);
+               max_alloc_error_panic_version, ver);
     }
 }
diff --git a/mozglue/static/rust/lib.rs b/mozglue/static/rust/lib.rs
index 9b73b9d153db3..001b920766323 100644
--- a/mozglue/static/rust/lib.rs
+++ b/mozglue/static/rust/lib.rs
@@ -3,6 +3,7 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 #![cfg_attr(feature = "oom_with_hook", feature(alloc_error_hook))]
+#![cfg_attr(feature = "oom_with_alloc_error_panic", feature(panic_oom_payload))]
 
 use arrayvec::ArrayString;
 use std::cmp;
@@ -68,7 +69,11 @@ impl<const CAP: usize> Deref for ArrayCString<CAP> {
 fn panic_hook(info: &panic::PanicInfo) {
     // Try to handle &str/String payloads, which should handle 99% of cases.
     let payload = info.payload();
-    let message = if let Some(s) = payload.downcast_ref::<&str>() {
+    let message = if let Some(layout) = oom_hook::oom_layout(payload) {
+        unsafe {
+            oom_hook::RustHandleOOM(layout.size());
+        }
+    } else if let Some(s) = payload.downcast_ref::<&str>() {
         s
     } else if let Some(s) = payload.downcast_ref::<String>() {
         s.as_str()
@@ -98,33 +103,40 @@ fn panic_hook(info: &panic::PanicInfo) {
 
 /// Configure a panic hook to redirect rust panics to MFBT's MOZ_Crash.
 #[no_mangle]
-pub extern "C" fn install_rust_panic_hook() {
+pub extern "C" fn install_rust_hooks() {
     panic::set_hook(Box::new(panic_hook));
+    #[cfg(feature = "oom_with_hook")]
+    use std::alloc::set_alloc_error_hook;
+    #[cfg(feature = "oom_with_hook")]
+    set_alloc_error_hook(oom_hook::hook);
 }
 
-#[cfg(feature = "oom_with_hook")]
 mod oom_hook {
-    use std::alloc::{set_alloc_error_hook, Layout};
+    #[cfg(feature = "oom_with_alloc_error_panic")]
+    use std::alloc::AllocErrorPanicPayload;
+    use std::alloc::Layout;
+    use std::any::Any;
+
+    #[inline(always)]
+    pub fn oom_layout(_payload: &dyn Any) -> Option<Layout> {
+        #[cfg(feature = "oom_with_alloc_error_panic")]
+        return _payload
+            .downcast_ref::<AllocErrorPanicPayload>()
+            .map(|p| p.layout());
+        #[cfg(not(feature = "oom_with_alloc_error_panic"))]
+        return None;
+    }
 
     extern "C" {
-        fn RustHandleOOM(size: usize) -> !;
+        pub fn RustHandleOOM(size: usize) -> !;
     }
 
+    #[cfg(feature = "oom_with_hook")]
     pub fn hook(layout: Layout) {
         unsafe {
             RustHandleOOM(layout.size());
         }
     }
-
-    pub fn install() {
-        set_alloc_error_hook(hook);
-    }
-}
-
-#[no_mangle]
-pub extern "C" fn install_rust_oom_hook() {
-    #[cfg(feature = "oom_with_hook")]
-    oom_hook::install();
 }
 
 #[cfg(feature = "moz_memory")]
diff --git a/toolkit/xre/nsAppRunner.cpp b/toolkit/xre/nsAppRunner.cpp
index 093ba5e03b3a9..ef9e068bc9c08 100644
--- a/toolkit/xre/nsAppRunner.cpp
+++ b/toolkit/xre/nsAppRunner.cpp
@@ -6100,14 +6100,10 @@ mozilla::BinPathType XRE_GetChildProcBinPathType(
 }
 
 // From mozglue/static/rust/lib.rs
-extern "C" void install_rust_panic_hook();
-extern "C" void install_rust_oom_hook();
+extern "C" void install_rust_hooks();
 
 struct InstallRustHooks {
-  InstallRustHooks() {
-    install_rust_panic_hook();
-    install_rust_oom_hook();
-  }
+  InstallRustHooks() { install_rust_hooks(); }
 };
 
 InstallRustHooks sInstallRustHooks;