From 3294997f94cf8f5c0b59159bd6156ddbbf0cfdce Mon Sep 17 00:00:00 2001
From: Benedikt Reinartz <filmor@gmail.com>
Date: Fri, 11 Oct 2024 22:14:47 +0200
Subject: [PATCH] Catch panics in encoding NIF results (#656)

Also adds tests for panicking in parameters, return values and the NIF
function itself.

Fixes #655
---
 CHANGELOG.md                                  |  1 +
 rustler/src/codegen_runtime.rs                | 13 ++++++--
 rustler_tests/lib/rustler_test.ex             |  4 +++
 rustler_tests/native/rustler_test/src/lib.rs  |  1 +
 .../native/rustler_test/src/test_panic.rs     | 30 +++++++++++++++++++
 rustler_tests/test/panic_test.exs             | 12 ++++++++
 6 files changed, 58 insertions(+), 3 deletions(-)
 create mode 100644 rustler_tests/native/rustler_test/src/test_panic.rs
 create mode 100644 rustler_tests/test/panic_test.exs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69553a7f..59cbeee6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,7 @@ versions.
   now (#638)
 - API functions for Windows are correctly assigned now for NIF version 2.15 and
   above (#635)
+- Panics in encoding the result of NIF function are caught (#656)
 
 ### Changed
 
diff --git a/rustler/src/codegen_runtime.rs b/rustler/src/codegen_runtime.rs
index 6c86cb23..c6fbb0dc 100644
--- a/rustler/src/codegen_runtime.rs
+++ b/rustler/src/codegen_runtime.rs
@@ -3,6 +3,7 @@
 use std::ffi::CString;
 use std::fmt;
 
+use crate::types::atom;
 use crate::{Encoder, Env, OwnedBinary, Term};
 
 // Re-export of inventory
@@ -26,10 +27,16 @@ pub unsafe trait NifReturnable {
 
 unsafe impl<T> NifReturnable for T
 where
-    T: crate::Encoder,
+    T: crate::Encoder + std::panic::RefUnwindSafe,
 {
     unsafe fn into_returned(self, env: Env) -> NifReturned {
-        NifReturned::Term(self.encode(env).as_c_arg())
+        if let Ok(res) = std::panic::catch_unwind(|| NifReturned::Term(self.encode(env).as_c_arg()))
+        {
+            res
+        } else {
+            let term = atom::nif_panicked().as_c_arg();
+            NifReturned::Raise(term)
+        }
     }
 }
 
@@ -126,7 +133,7 @@ where
             Err(err) => match err.downcast::<NifReturned>() {
                 Ok(ty) => NifReturned::Term(ty.apply(env)),
                 Err(_) => {
-                    let term = crate::types::atom::nif_panicked().as_c_arg();
+                    let term = atom::nif_panicked().as_c_arg();
                     NifReturned::Raise(term)
                 }
             },
diff --git a/rustler_tests/lib/rustler_test.ex b/rustler_tests/lib/rustler_test.ex
index f2c7b58c..b2c2232f 100644
--- a/rustler_tests/lib/rustler_test.ex
+++ b/rustler_tests/lib/rustler_test.ex
@@ -148,6 +148,10 @@ defmodule RustlerTest do
 
   def append_to_path(_path, _to_append), do: err()
 
+  def panic_in_nif(), do: err()
+  def panic_in_encode(), do: err()
+  def panic_in_decode(_), do: err()
+
   if Helper.has_nif_version("2.16") do
     def perform_dyncall(_res, _a, _b, _c), do: err()
   end
diff --git a/rustler_tests/native/rustler_test/src/lib.rs b/rustler_tests/native/rustler_test/src/lib.rs
index bb44c685..e37e35be 100644
--- a/rustler_tests/native/rustler_test/src/lib.rs
+++ b/rustler_tests/native/rustler_test/src/lib.rs
@@ -10,6 +10,7 @@ mod test_list;
 mod test_local_pid;
 mod test_map;
 mod test_nif_attrs;
+mod test_panic;
 mod test_path;
 mod test_primitives;
 mod test_range;
diff --git a/rustler_tests/native/rustler_test/src/test_panic.rs b/rustler_tests/native/rustler_test/src/test_panic.rs
new file mode 100644
index 00000000..6c89b3a4
--- /dev/null
+++ b/rustler_tests/native/rustler_test/src/test_panic.rs
@@ -0,0 +1,30 @@
+use rustler::{Decoder, Encoder, Env, Term};
+
+#[rustler::nif]
+pub fn panic_in_nif() -> i32 {
+    panic!("panic!")
+}
+
+struct Panicking;
+
+impl Encoder for Panicking {
+    fn encode<'a>(&self, _env: Env<'a>) -> Term<'a> {
+        panic!("panic in encode!")
+    }
+}
+
+impl<'a> Decoder<'a> for Panicking {
+    fn decode(_term: Term<'a>) -> rustler::NifResult<Self> {
+        panic!("panic in decode!")
+    }
+}
+
+#[rustler::nif]
+pub fn panic_in_encode() -> Panicking {
+    Panicking
+}
+
+#[rustler::nif]
+pub fn panic_in_decode(_p: Panicking) -> i32 {
+    0
+}
diff --git a/rustler_tests/test/panic_test.exs b/rustler_tests/test/panic_test.exs
new file mode 100644
index 00000000..fda43635
--- /dev/null
+++ b/rustler_tests/test/panic_test.exs
@@ -0,0 +1,12 @@
+defmodule PanicTest do
+  use ExUnit.Case
+
+  test "panics in NIFs are caught" do
+    assert_raise ErlangError, &RustlerTest.panic_in_nif/0
+    assert_raise ErlangError, &RustlerTest.panic_in_encode/0
+
+    assert_raise ErlangError, fn ->
+      RustlerTest.panic_in_decode(:anything)
+    end
+  end
+end