From 34efda72d128b336f3afd72d70098deccd79c6c2 Mon Sep 17 00:00:00 2001
From: Dmitry Novikov <novikov.dm.al@gmail.com>
Date: Mon, 24 Apr 2023 15:36:40 +0300
Subject: [PATCH] fix(`sp-io`): Proper alloc panc handling (#31)

---
 primitives/application-crypto/Cargo.toml |  1 -
 primitives/io/Cargo.toml                 |  8 +++---
 primitives/io/src/lib.rs                 | 32 +++++++++++++++++++-----
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/primitives/application-crypto/Cargo.toml b/primitives/application-crypto/Cargo.toml
index e7aa118dfd6ac..02b85a0d96913 100644
--- a/primitives/application-crypto/Cargo.toml
+++ b/primitives/application-crypto/Cargo.toml
@@ -42,5 +42,4 @@ full_crypto = [
 	# Don't add `panic_handler` and `alloc_error_handler` since they are expected to be provided
 	# by the user anyway.
 	"sp-io/disable_panic_handler",
-	"sp-io/disable_oom",
 ]
diff --git a/primitives/io/Cargo.toml b/primitives/io/Cargo.toml
index 4accb0bf5c3f8..3a54dbb833003 100644
--- a/primitives/io/Cargo.toml
+++ b/primitives/io/Cargo.toml
@@ -66,12 +66,11 @@ with-tracing = [
 	"sp-tracing/with-tracing"
 ]
 
-# These two features are used for `no_std` builds for the environments which already provides
-# `#[panic_handler]`, `#[alloc_error_handler]` and `#[global_allocator]`.
+# These two features are used for `no_std` builds for the environments which
+# already provides `#[panic_handler]` and `#[global_allocator]`.
 #
 # For the regular wasm runtime builds those are not used.
 disable_panic_handler = []
-disable_oom = []
 disable_allocator = []
 
 # This feature flag controls the runtime's behavior when encountering
@@ -85,8 +84,7 @@ disable_allocator = []
 # logs, with the caller receving a generic "wasm `unreachable` instruction executed"
 # error message.
 #
-# This has no effect if both `disable_panic_handler` and `disable_oom`
-# are enabled.
+# This has no effect if `disable_panic_handler` is enabled.
 #
 # WARNING: Enabling this feature flag requires the `PanicHandler::abort_on_panic`
 #          host function to be supported by the host. Do *not* enable it for your
diff --git a/primitives/io/src/lib.rs b/primitives/io/src/lib.rs
index b40d773ef016c..bfc8e06e036cd 100644
--- a/primitives/io/src/lib.rs
+++ b/primitives/io/src/lib.rs
@@ -19,6 +19,10 @@
 
 #![warn(missing_docs)]
 #![cfg_attr(not(feature = "std"), no_std)]
+#![cfg_attr(
+	all(not(feature = "disable_panic_handler"), not(feature = "std")),
+	feature(panic_oom_payload)
+)]
 #![cfg_attr(
 	feature = "std",
 	doc = "Substrate runtime standard library as compiled when linked with Rust's standard library."
@@ -1771,13 +1775,29 @@ mod allocator_impl {
 #[panic_handler]
 #[no_mangle]
 pub fn panic(info: &core::panic::PanicInfo) -> ! {
-	let message = sp_std::alloc::format!("{}", info);
-	#[cfg(feature = "improved_panic_error_reporting")]
-	{
+	use sp_std::alloc::{alloc::AllocErrorPanicPayload, format, string::String};
+
+	let improved_panic_error_reporting = match () {
+		#[cfg(feature = "improved_panic_error_reporting")]
+		() => true,
+		#[cfg(not(feature = "improved_panic_error_reporting"))]
+		() => false,
+	};
+
+	let message = info
+		.payload()
+		.downcast_ref::<AllocErrorPanicPayload>()
+		.map(|_| {
+			let msg = improved_panic_error_reporting
+				.then_some("Runtime memory exhausted.")
+				.unwrap_or("Runtime memory exhausted. Aborting");
+			String::from(msg)
+		})
+		.unwrap_or_else(|| format!("{info}"));
+
+	if improved_panic_error_reporting {
 		panic_handler::abort_on_panic(&message);
-	}
-	#[cfg(not(feature = "improved_panic_error_reporting"))]
-	{
+	} else {
 		logging::log(LogLevel::Error, "runtime", message.as_bytes());
 		core::arch::wasm32::unreachable();
 	}