From 33ab3b8b431ac5630edbf221dec551b7033fbccc Mon Sep 17 00:00:00 2001
From: Kornel <kornel@cloudflare.com>
Date: Thu, 5 Dec 2024 16:37:13 +0000
Subject: [PATCH] Safer wasm references

---
 js-api/src/comment.rs      |  2 +-
 js-api/src/doctype.rs      |  2 +-
 js-api/src/document_end.rs |  2 +-
 js-api/src/element.rs      |  2 +-
 js-api/src/handlers.rs     | 14 ++++----------
 js-api/src/lib.rs          |  8 ++++----
 js-api/src/text_chunk.rs   |  2 +-
 7 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/js-api/src/comment.rs b/js-api/src/comment.rs
index 63227509..71fe8c27 100644
--- a/js-api/src/comment.rs
+++ b/js-api/src/comment.rs
@@ -4,7 +4,7 @@ use lol_html::html_content::Comment as NativeComment;
 #[wasm_bindgen]
 pub struct Comment(NativeRefWrap<NativeComment<'static>>);
 
-impl_from_native!(NativeComment --> Comment);
+impl_from_native!(NativeComment => Comment);
 impl_mutations!(Comment);
 
 #[wasm_bindgen]
diff --git a/js-api/src/doctype.rs b/js-api/src/doctype.rs
index 42510bbf..da498aa9 100644
--- a/js-api/src/doctype.rs
+++ b/js-api/src/doctype.rs
@@ -4,7 +4,7 @@ use lol_html::html_content::Doctype as NativeDoctype;
 #[wasm_bindgen]
 pub struct Doctype(NativeRefWrap<NativeDoctype<'static>>);
 
-impl_from_native!(NativeDoctype --> Doctype);
+impl_from_native!(NativeDoctype => Doctype);
 
 #[wasm_bindgen]
 impl Doctype {
diff --git a/js-api/src/document_end.rs b/js-api/src/document_end.rs
index 6cd0b710..43edb098 100644
--- a/js-api/src/document_end.rs
+++ b/js-api/src/document_end.rs
@@ -4,7 +4,7 @@ use lol_html::html_content::DocumentEnd as NativeDocumentEnd;
 #[wasm_bindgen]
 pub struct DocumentEnd(NativeRefWrap<NativeDocumentEnd<'static>>);
 
-impl_from_native!(NativeDocumentEnd --> DocumentEnd);
+impl_from_native!(NativeDocumentEnd => DocumentEnd);
 
 #[wasm_bindgen]
 impl DocumentEnd {
diff --git a/js-api/src/element.rs b/js-api/src/element.rs
index 5095b439..8b850867 100644
--- a/js-api/src/element.rs
+++ b/js-api/src/element.rs
@@ -21,7 +21,7 @@ impl From<&NativeAttribute<'_>> for Attribute {
 #[wasm_bindgen]
 pub struct Element(NativeRefWrap<NativeElement<'static, 'static>>);
 
-impl_from_native!(NativeElement --> Element);
+impl_from_native!(NativeElement => Element);
 impl_mutations!(Element);
 
 #[wasm_bindgen]
diff --git a/js-api/src/handlers.rs b/js-api/src/handlers.rs
index ee2842de..fcf83a4b 100644
--- a/js-api/src/handlers.rs
+++ b/js-api/src/handlers.rs
@@ -9,7 +9,6 @@ use lol_html::{
     DocumentContentHandlers as NativeDocumentContentHandlers,
     ElementContentHandlers as NativeElementContentHandlers,
 };
-use std::mem;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -26,16 +25,11 @@ macro_rules! make_handler {
             h
         }
         type_hint(Box::new(move |arg: &mut _| {
-            let (js_arg, anchor) = unsafe { $JsArgType::from_native(arg) };
+            $JsArgType::with_native(arg, |js_value| {
+                $handler.call1(&JsValue::NULL, &js_value)
+            }).map_err(|e| HandlerJsErrorWrap(e))?;
 
-            let res = match $handler.call1(&JsValue::NULL, &JsValue::from(js_arg)) {
-                Ok(_) => Ok(()),
-                Err(e) => Err(HandlerJsErrorWrap(e).into()),
-            };
-
-            mem::drop(anchor);
-
-            res
+            Ok(())
         }))
     }};
 }
diff --git a/js-api/src/lib.rs b/js-api/src/lib.rs
index b5f7f98b..a83eb906 100644
--- a/js-api/src/lib.rs
+++ b/js-api/src/lib.rs
@@ -156,12 +156,12 @@ macro_rules! impl_mutations {
 }
 
 macro_rules! impl_from_native {
-    ($Ty:ident --> $JsTy:path) => {
+    ($Ty:ty => $JsTy:path) => {
         impl $JsTy {
-            pub(crate) unsafe fn from_native<'r>(inner: &'r mut $Ty) -> (Self, Anchor<'r>) {
-                let (ref_wrap, anchor) = unsafe { NativeRefWrap::wrap(inner) };
+            pub(crate) fn with_native<'r, R>(inner: &'r mut $Ty, callback: impl FnOnce(&JsValue) -> R) -> R {
+                let (ref_wrap, _anchor) = unsafe { NativeRefWrap::wrap(inner) };
 
-                ($JsTy(ref_wrap), anchor)
+                (callback)(&JsValue::from($JsTy(ref_wrap)))
             }
         }
     };
diff --git a/js-api/src/text_chunk.rs b/js-api/src/text_chunk.rs
index 87ba2f87..1e0adc96 100644
--- a/js-api/src/text_chunk.rs
+++ b/js-api/src/text_chunk.rs
@@ -4,7 +4,7 @@ use lol_html::html_content::TextChunk as NativeTextChunk;
 #[wasm_bindgen]
 pub struct TextChunk(NativeRefWrap<NativeTextChunk<'static>>);
 
-impl_from_native!(NativeTextChunk --> TextChunk);
+impl_from_native!(NativeTextChunk => TextChunk);
 impl_mutations!(TextChunk);
 
 #[wasm_bindgen]