From 747dcf561d710d324b02249806fa0b113178c3d2 Mon Sep 17 00:00:00 2001
From: Behdad Esfahbod <behdad@behdad.org>
Date: Thu, 23 Feb 2023 11:53:08 -0700
Subject: [PATCH] [wasm] Strong typing for object references

Such that wasm cannot crash us by passing wrong object refs.

https://github.com/bytecodealliance/wasm-micro-runtime/discussions/1987

It still is unsafe if some code in the process other than HarfBuzz
registers refs with wasm-micro-runtime, since wasm_externref_ref2obj()
takes no context variable and looks up refs globally :(.

Maybe I fix that later by keeping a hash table of ref->obj-type instead.
---
 src/hb-wasm-api.cc   | 11 ++---------
 src/hb-wasm-api.hh   | 32 ++++++++++++++++++++++++++++++++
 src/hb-wasm-shape.cc | 41 +++++++++++++++++++++--------------------
 3 files changed, 55 insertions(+), 29 deletions(-)

diff --git a/src/hb-wasm-api.cc b/src/hb-wasm-api.cc
index 37264900f..74b18435d 100644
--- a/src/hb-wasm-api.cc
+++ b/src/hb-wasm-api.cc
@@ -24,22 +24,15 @@
 
 #include "hb-wasm-api.hh"
 
-
-#define nullref 0
 #define module_inst wasm_runtime_get_module_inst (exec_env)
-#define HB_REF2OBJ(obj) \
-  hb_##obj##_t *obj = nullptr; \
-  (void) wasm_externref_ref2obj (obj##ref, (void **) &obj)
-#define HB_OBJ2REF(obj) \
-  uint32_t obj##ref = nullref; \
-  (void) wasm_externref_obj2ref (module_inst, obj, &obj##ref)
 
 
 #include "hb-wasm-api-font.hh"
 
 
-#undef nullref
 #undef module_inst
 #undef HB_WASM_EXEC_ENV
 #undef HB_REF2OBJ
 #undef HB_OBJ2REF
+
+hb_user_data_key_t _hb_wasm_ref_type_key = {};
diff --git a/src/hb-wasm-api.hh b/src/hb-wasm-api.hh
index c696a9471..a53f3f5b0 100644
--- a/src/hb-wasm-api.hh
+++ b/src/hb-wasm-api.hh
@@ -42,4 +42,36 @@
 #undef HB_WASM_END_DECLS
 
 
+enum {
+  hb_wasm_ref_type_none,
+  hb_wasm_ref_type_face,
+  hb_wasm_ref_type_font,
+  hb_wasm_ref_type_buffer,
+};
+
+HB_INTERNAL extern hb_user_data_key_t _hb_wasm_ref_type_key;
+
+#define nullref 0
+
+#define HB_REF2OBJ(obj) \
+  hb_##obj##_t *obj = nullptr; \
+  HB_STMT_START { \
+    (void) wasm_externref_ref2obj (obj##ref, (void **) &obj); \
+    /* Check object type. */ \
+    /* This works because all our objects have the same hb_object_t layout. */ \
+    if (unlikely (hb_##obj##_get_user_data (obj, &_hb_wasm_ref_type_key) != \
+		  (void *) hb_wasm_ref_type_##obj)) \
+      obj = nullptr; \
+  } HB_STMT_END
+
+#define HB_OBJ2REF(obj) \
+  uint32_t obj##ref = nullref; \
+  HB_STMT_START { \
+    hb_##obj##_set_user_data (obj, &_hb_wasm_ref_type_key, \
+			      (void *) hb_wasm_ref_type_##obj, \
+			      nullptr, false); \
+    (void) wasm_externref_obj2ref (module_inst, obj, &obj##ref); \
+  } HB_STMT_END
+
+
 #endif /* HB_WASM_API_HH */
diff --git a/src/hb-wasm-shape.cc b/src/hb-wasm-shape.cc
index ecb42dbe1..ea31cce70 100644
--- a/src/hb-wasm-shape.cc
+++ b/src/hb-wasm-shape.cc
@@ -67,7 +67,7 @@ init_wasm ()
   init_args.native_module_name = "env";
   init_args.native_symbols = _hb_wasm_native_symbols;
 
-  if (!wasm_runtime_full_init (&init_args))
+  if (unlikely (!wasm_runtime_full_init (&init_args)))
   {
     DEBUG_MSG (WASM, nullptr, "Init runtime environment failed.");
     return false;
@@ -95,7 +95,7 @@ _hb_wasm_shaper_face_data_create (hb_face_t *face)
 
   wasm_module = wasm_runtime_load ((uint8_t *) hb_blob_get_data_writable (wasm_blob, nullptr),
 				   length, nullptr, 0);
-  if (!wasm_module)
+  if (unlikely (!wasm_module))
   {
     DEBUG_MSG (WASM, nullptr, "Load wasm module failed.");
     goto fail;
@@ -168,46 +168,47 @@ _hb_wasm_shape (hb_shape_plan_t    *shape_plan,
   module_inst = wasm_runtime_instantiate(face_data->wasm_module,
 					 stack_size, heap_size,
 					 nullptr, 0);
-  if (!module_inst)
+  if (unlikely (!module_inst))
   {
     DEBUG_MSG (WASM, face_data->wasm_module, "Instantiate wasm module failed.");
+    return false;
+  }
+
+  // cmake -DWAMR_BUILD_REF_TYPES=1 for these to work
+  HB_OBJ2REF (font);
+  HB_OBJ2REF (buffer);
+  if (unlikely (!fontref || !bufferref))
+  {
+    DEBUG_MSG (WASM, module_inst, "Failed to register objects.");
     goto fail;
   }
 
   exec_env = wasm_runtime_create_exec_env (module_inst, stack_size);
-  if (!exec_env) {
+  if (unlikely (!exec_env)) {
     DEBUG_MSG (WASM, module_inst, "Create wasm execution environment failed.");
     goto fail;
   }
 
-  if (!(shape_func = wasm_runtime_lookup_function (module_inst, "shape", nullptr)))
+  shape_func = wasm_runtime_lookup_function (module_inst, "shape", nullptr);
+  if (unlikely (!shape_func))
   {
     DEBUG_MSG (WASM, module_inst, "Shape function not found.");
     goto fail;
   }
 
-  uint32_t font_ref;
-  uint32_t buffer_ref;
-  // cmake -DWAMR_BUILD_REF_TYPES=1 for these to work
-  if (!wasm_externref_obj2ref (module_inst, font, &font_ref) ||
-      !wasm_externref_obj2ref (module_inst, buffer, &buffer_ref))
-  {
-    DEBUG_MSG (WASM, module_inst, "Failed to register objects.");
-    goto fail;
-  }
-
   wasm_val_t results[1];
   wasm_val_t arguments[2];
 
   results[0].kind = WASM_I32;
   arguments[0].kind = WASM_I32;
-  arguments[0].of.i32 = font_ref;
+  arguments[0].of.i32 = fontref;
   arguments[1].kind = WASM_I32;
-  arguments[1].of.i32 = buffer_ref;
+  arguments[1].of.i32 = bufferref;
 
-  if (!wasm_runtime_call_wasm_a (exec_env, shape_func,
-				 ARRAY_LENGTH (results), results,
-				 ARRAY_LENGTH (arguments), arguments))
+   ret = wasm_runtime_call_wasm_a (exec_env, shape_func,
+				   ARRAY_LENGTH (results), results,
+				   ARRAY_LENGTH (arguments), arguments);
+  if (unlikely (!ret))
   {
     DEBUG_MSG (WASM, module_inst, "Calling shape function failed: %s",
 	       wasm_runtime_get_exception(module_inst));