fix(neon): Tag Root with an instance id and verify before using

crates/neon-runtime/src/napi/reference.rs

@@ -15,6 +15,8 @@ pub unsafe fn new(env: Env, value: Local) -> napi::Ref {
     result.assume_init()
 }
 
+/// # Safety
+/// Must only be used from the same module context that created the reference
 pub unsafe fn reference(env: Env, value: napi::Ref) -> usize {
     let mut result = MaybeUninit::uninit();
 
@@ -26,6 +28,8 @@ pub unsafe fn reference(env: Env, value: napi::Ref) -> usize {
     result.assume_init() as usize
 }
 
+/// # Safety
+/// Must only be used from the same module context that created the reference
 pub unsafe fn unreference(env: Env, value: napi::Ref) {
     let mut result = MaybeUninit::uninit();
 
@@ -39,6 +43,8 @@ pub unsafe fn unreference(env: Env, value: napi::Ref) {
     }
 }
 
+/// # Safety
+/// Must only be used from the same module context that created the reference
 pub unsafe fn get(env: Env, value: napi::Ref) -> Local {
     let mut result = MaybeUninit::uninit();
 

src/handle/root.rs

@@ -2,6 +2,8 @@ use std::ffi::c_void;
 use std::marker::PhantomData;
 #[cfg(feature = "napi-6")]
 use std::sync::Arc;
+#[cfg(not(feature = "napi-6"))]
+use std::thread::{self, ThreadId};
 
 #[cfg(feature = "napi-6")]
 use neon_runtime::tsfn::ThreadsafeFunction;
@@ -10,15 +12,20 @@ use neon_runtime::{raw, reference};
 use crate::context::Context;
 use crate::handle::Handle;
 #[cfg(feature = "napi-6")]
-use crate::lifecycle::{DropData, InstanceData};
+use crate::lifecycle::{DropData, InstanceData, InstanceId};
 use crate::object::Object;
 use crate::types::boxed::Finalize;
 
+#[cfg(not(feature = "napi-6"))]
+type InstanceId = ThreadId;
+
 #[repr(transparent)]
 #[derive(Clone)]
 pub(crate) struct NapiRef(*mut c_void);
 
 impl NapiRef {
+    /// # Safety
+    /// Must only be used from the same module context that created the reference
     pub(crate) unsafe fn unref(self, env: raw::Env) {
         reference::unreference(env, self.0.cast());
     }
@@ -42,6 +49,7 @@ pub struct Root<T> {
     // `Option` is used to skip `Drop` when `Root::drop` or `Root::into_inner` is used.
     // It will *always* be `Some` when a user is interacting with `Root`.
     internal: Option<NapiRef>,
+    instance_id: InstanceId,
     #[cfg(feature = "napi-6")]
     drop_queue: Arc<ThreadsafeFunction<DropData>>,
     _phantom: PhantomData<T>,
@@ -60,6 +68,16 @@ unsafe impl<T> Send for Root<T> {}
 
 unsafe impl<T> Sync for Root<T> {}
 
+#[cfg(feature = "napi-6")]
+fn instance_id<'a, C: Context<'a>>(cx: &mut C) -> InstanceId {
+    InstanceData::id(cx)
+}
+
+#[cfg(not(feature = "napi-6"))]
+fn instance_id<'a, C: Context<'a>>(_: &mut C) -> InstanceId {
+    thread::current().id()
+}
+
 impl<T: Object> Root<T> {
     /// Create a reference to a JavaScript object. The object will not be
     /// garbage collected until the `Root` is dropped. A `Root<T>` may only
@@ -76,6 +94,7 @@ impl<T: Object> Root<T> {
 
         Self {
             internal: Some(NapiRef(internal as *mut _)),
+            instance_id: instance_id(cx),
             #[cfg(feature = "napi-6")]
             drop_queue: InstanceData::drop_queue(cx),
             _phantom: PhantomData,
@@ -96,14 +115,15 @@ impl<T: Object> Root<T> {
     /// ```
     pub fn clone<'a, C: Context<'a>>(&self, cx: &mut C) -> Self {
         let env = cx.env();
-        let internal = self.as_napi_ref().0 as *mut _;
+        let internal = self.as_napi_ref(cx).0 as *mut _;
 
         unsafe {
             reference::reference(env.to_raw(), internal);
         };
 
         Self {
             internal: self.internal.clone(),
+            instance_id: instance_id(cx),
             #[cfg(feature = "napi-6")]
             drop_queue: Arc::clone(&self.drop_queue),
             _phantom: PhantomData,
@@ -116,14 +136,14 @@ impl<T: Object> Root<T> {
         let env = cx.env().to_raw();
 
         unsafe {
-            self.into_napi_ref().unref(env);
+            self.into_napi_ref(cx).unref(env);
         }
     }
 
     /// Return the referenced JavaScript object and allow it to be garbage collected
     pub fn into_inner<'a, C: Context<'a>>(self, cx: &mut C) -> Handle<'a, T> {
         let env = cx.env();
-        let internal = self.into_napi_ref();
+        let internal = self.into_napi_ref(cx);
         let local = unsafe { reference::get(env.to_raw(), internal.0.cast()) };
 
         unsafe {
@@ -138,25 +158,28 @@ impl<T: Object> Root<T> {
     /// can be used in place of a clone immediately followed by a call to `into_inner`.
     pub fn to_inner<'a, C: Context<'a>>(&self, cx: &mut C) -> Handle<'a, T> {
         let env = cx.env();
-        let local = unsafe { reference::get(env.to_raw(), self.as_napi_ref().0 as *mut _) };
+        let local = unsafe { reference::get(env.to_raw(), self.as_napi_ref(cx).0 as *mut _) };
 
         Handle::new_internal(T::from_raw(env, local))
     }
 
-    fn as_napi_ref(&self) -> &NapiRef {
+    fn as_napi_ref<'a, C: Context<'a>>(&self, cx: &mut C) -> &NapiRef {
+        if self.instance_id != instance_id(cx) {
+            panic!("Attempted to dereference a `neon::handle::Root` from the wrong module ");
+        }
+
         self.internal
             .as_ref()
             // `unwrap` will not `panic` because `internal` will always be `Some`
             // until the `Root` is consumed.
             .unwrap()
     }
 
-    fn into_napi_ref(mut self) -> NapiRef {
-        self.internal
-            .take()
-            // `unwrap` will not `panic` because this is the only method place
-            // `internal` is replaced with `None` and it consumes `self`.
-            .unwrap()
+    fn into_napi_ref<'a, C: Context<'a>>(mut self, cx: &mut C) -> NapiRef {
+        let reference = self.as_napi_ref(cx).clone();
+        // This uses `as_napi_ref` instead of `Option::take` for the instance id safety check
+        self.internal = None;
+        reference
     }
 }
 

src/lifecycle.rs

@@ -8,6 +8,7 @@
 //!
 //! [napi-docs]: https://nodejs.org/api/n-api.html#n_api_environment_life_cycle_apis
 
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use neon_runtime::raw::Env;
@@ -20,10 +21,28 @@ use crate::handle::root::NapiRef;
 #[cfg(feature = "promise-api")]
 use crate::types::promise::NodeApiDeferred;
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(transparent)]
+/// Uniquely identifies an instance of the module
+///
+/// _Note_: Since `InstanceData` is created lazily, the order of `id` may not
+/// reflect the order that instances were created.
+pub(crate) struct InstanceId(u64);
+
+impl InstanceId {
+    fn next() -> Self {
+        static NEXT_ID: AtomicU64 = AtomicU64::new(0);
+
+        Self(NEXT_ID.fetch_add(1, Ordering::SeqCst))
+    }
+}
+
 /// `InstanceData` holds Neon data associated with a particular instance of a
 /// native module. If a module is loaded multiple times (e.g., worker threads), this
 /// data will be unique per instance.
 pub(crate) struct InstanceData {
+    id: InstanceId,
+
     /// Used to free `Root` in the same JavaScript environment that created it
     ///
     /// _Design Note_: An `Arc` ensures the `ThreadsafeFunction` outlives the unloading
@@ -90,6 +109,7 @@ impl InstanceData {
         };
 
         let data = InstanceData {
+            id: InstanceId::next(),
             drop_queue: Arc::new(drop_queue),
             #[cfg(all(feature = "channel-api"))]
             shared_channel,
@@ -111,4 +131,9 @@ impl InstanceData {
         channel.reference(cx);
         channel
     }
+
+    /// Unique identifier for this instance of the module
+    pub(crate) fn id<'a, C: Context<'a>>(cx: &mut C) -> InstanceId {
+        InstanceData::get(cx).id
+    }
 }

test/napi/Cargo.toml

@@ -9,6 +9,9 @@ edition = "2018"
 [lib]
 crate-type = ["cdylib"]
 
+[dependencies]
+once_cell = "1"
+
 [dependencies.neon]
 version = "*"
 path = "../.."

test/napi/lib/workers.js

@@ -0,0 +1,108 @@
+const assert = require("assert");
+const { Worker, isMainThread, parentPort } = require("worker_threads");
+
+const addon = require("..");
+
+// Receive a message, try that method and return the error message
+if (!isMainThread) {
+  parentPort.once("message", (message) => {
+    try {
+      switch (message) {
+        case "get_and_replace":
+          addon.get_and_replace({});
+          break;
+        case "get_or_init":
+          addon.get_or_init(() => ({}));
+          break;
+        case "get_or_init_clone":
+          addon.get_or_init_clone(() => ({}));
+          break;
+        default:
+          throw new Error(`Unexpected message: ${message}`);
+      }
+
+      throw new Error("Did not throw an exception");
+    } catch (err) {
+      parentPort.postMessage(err);
+    }
+  });
+
+  return;
+}
+
+describe("Worker / Root Tagging Tests", () => {
+  describe("Single Threaded", () => {
+    it("should be able to stash a global with `get_and_replace`", () => {
+      const first = {};
+      const second = {};
+
+      assert.strictEqual(addon.get_and_replace(first), undefined);
+      assert.strictEqual(addon.get_and_replace(second), first);
+      assert.strictEqual(addon.get_and_replace({}), second);
+    });
+
+    it("should be able to lazily initialize with `get_or_init`", () => {
+      const o = {};
+
+      assert.strictEqual(
+        addon.get_or_init(() => o),
+        o
+      );
+      assert.strictEqual(
+        addon.get_or_init(() => ({})),
+        o
+      );
+      assert.strictEqual(addon.get_or_init(), o);
+    });
+
+    it("should be able to lazily initialize with `get_or_init_clone`", () => {
+      const o = {};
+
+      assert.strictEqual(
+        addon.get_or_init_clone(() => o),
+        o
+      );
+      assert.strictEqual(
+        addon.get_or_init_clone(() => ({})),
+        o
+      );
+      assert.strictEqual(addon.get_or_init_clone(), o);
+    });
+  });
+
+  // Note: These tests require that the previous set of tests have run or else they will fail
+  describe("Multi-Threaded", () => {
+    it("should fail to use `get_and_replace`", (cb) => {
+      const worker = new Worker(__filename);
+
+      worker.once("message", (message) => {
+        assert.ok(/wrong module/.test(message));
+        cb();
+      });
+
+      worker.postMessage("get_and_replace");
+    });
+
+    it("should fail to use `get_or_init`", (cb) => {
+      const worker = new Worker(__filename);
+
+      worker.once("message", (message) => {
+        assert.ok(/wrong module/.test(message));
+        cb();
+      });
+
+      worker.postMessage("get_or_init");
+    });
+
+    it("should fail to use `get_or_init`", (cb) => {
+      const worker = new Worker(__filename);
+
+      worker.once("message", (message) => {
+        assert.ok(/wrong module/.test(message));
+        cb();
+      });
+
+      worker.postMessage("get_or_init_clone");
+    });
+  });
+});

test/napi/src/js/workers.rs

@@ -0,0 +1,45 @@
+use std::sync::Mutex;
+
+use neon::prelude::*;
+use once_cell::sync::{Lazy, OnceCell};
+
+pub fn get_and_replace(mut cx: FunctionContext) -> JsResult<JsValue> {
+    static OBJECT: Lazy<Mutex<Option<Root<JsObject>>>> = Lazy::new(Default::default);
+
+    let mut global = OBJECT.lock().unwrap_or_else(|err| err.into_inner());
+    let next = cx.argument::<JsObject>(0)?.root(&mut cx);
+    let previous = global.replace(next);
+
+    Ok(match previous {
+        None => cx.undefined().upcast(),
+        Some(previous) => previous.into_inner(&mut cx).upcast(),
+    })
+}
+
+pub fn get_or_init(mut cx: FunctionContext) -> JsResult<JsObject> {
+    static OBJECT: OnceCell<Root<JsObject>> = OnceCell::new();
+
+    let o = OBJECT.get_or_try_init(|| {
+        cx.argument::<JsFunction>(0)?
+            .call_with(&cx)
+            .apply::<JsObject, _>(&mut cx)
+            .map(|v| v.root(&mut cx))
+    })?;
+
+    Ok(o.to_inner(&mut cx))
+}
+
+pub fn get_or_init_clone(mut cx: FunctionContext) -> JsResult<JsObject> {
+    static OBJECT: OnceCell<Root<JsObject>> = OnceCell::new();
+
+    let o = OBJECT.get_or_try_init(|| {
+        cx.argument::<JsFunction>(0)?
+            .call_with(&cx)
+            .apply::<JsObject, _>(&mut cx)
+            .map(|v| v.root(&mut cx))
+    })?;
+
+    // Note: This intentionally uses `clone` instead of `to_inner` in order to
+    // test the `clone` method.
+    Ok(o.clone(&mut cx).into_inner(&mut cx))
+}

test/napi/src/lib.rs

@@ -12,6 +12,7 @@ mod js {
     pub mod strings;
     pub mod threads;
     pub mod types;
+    pub mod workers;
 }
 
 use js::arrays::*;
@@ -337,6 +338,9 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
         "deferred_settle_with_panic_throw",
         deferred_settle_with_panic_throw,
     )?;
+    cx.export_function("get_and_replace", js::workers::get_and_replace)?;
+    cx.export_function("get_or_init", js::workers::get_or_init)?;
+    cx.export_function("get_or_init_clone", js::workers::get_or_init_clone)?;
 
     Ok(())
 }