Fixes napi-rs#1338 use after free with async

Xaeroxe · Oct 7, 2022 · 6c6149a · 6c6149a
1 parent a12bdc4
commit 6c6149a
Showing 1 changed file with 19 additions and 8 deletions.
diff --git a/crates/backend/src/codegen/fn.rs b/crates/backend/src/codegen/fn.rs
@@ -35,19 +35,30 @@ impl TryToTokens for NapiFn {
         quote! { Ok(#receiver(#(#arg_names),*).await) }
       };
       quote! {
-        napi::bindgen_prelude::execute_tokio_future(env, async move { #call }, |env, #receiver_ret_name| {
+        struct NapiRefContainer([napi::bindgen_prelude::sys::napi_ref; #args_len + 1]);
+        impl NapiRefContainer {
+          fn drop(self, env: napi::bindgen_prelude::sys::napi_env) {
+            for r in self.0 {
+              assert_eq!(unsafe { napi::bindgen_prelude::sys::napi_delete_reference(env, r) }, napi::bindgen_prelude::sys::Status::napi_ok);
+            }
+          }
+        }
+        unsafe impl Send for NapiRefContainer {}
+        unsafe impl Sync for NapiRefContainer {}
+        let mut arg_iter = cb.args.iter().chain(Some(&cb.this)).map(|a| {
+          let mut node_ref = ::std::mem::MaybeUninit::uninit();
+          assert_eq!(unsafe { napi::bindgen_prelude::sys::napi_create_reference(env, *a, 1, node_ref.as_mut_ptr()) }, napi::bindgen_prelude::sys::Status::napi_ok);
+          unsafe { node_ref.assume_init() }
+        });
+        let args_ref = NapiRefContainer(::std::array::from_fn(|_| arg_iter.next().expect("infallible")));
+        napi::bindgen_prelude::execute_tokio_future(env, async move { #call }, move |env, #receiver_ret_name| {
+          args_ref.drop(env);
           #ret
         })
       }
     };
 
-    let function_call = if args_len == 0
-      && self.fn_self.is_none()
-      && self.kind != FnKind::Constructor
-      && self.kind != FnKind::Factory
-    {
-      quote! { #native_call }
-    } else if self.kind == FnKind::Constructor {
+    let function_call = if self.kind == FnKind::Constructor {
       quote! {
         // constructor function is called from class `factory`
         // so we should skip the original `constructor` logic