From 53a1f19af0c9c7ede84f7193b27ba8923a5302dd Mon Sep 17 00:00:00 2001
From: mfish33 <32677537+mfish33@users.noreply.github.com>
Date: Mon, 20 Dec 2021 17:55:31 -0500
Subject: [PATCH] Fix undefined behavior in the editor wasm bridge (#414)

* removed possible undifined behavior

* changed naming of Editor -> JsEditorHandle

* fix rust formatting

Co-authored-by: Keavon Chambers <keavon@keavon.com>
---
 frontend/src/state/wasm-loader.ts |  4 ++--
 frontend/wasm/src/api.rs          | 37 +++++++++++++++++++++----------
 frontend/wasm/src/lib.rs          |  6 ++++-
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/frontend/src/state/wasm-loader.ts b/frontend/src/state/wasm-loader.ts
index b8aff5a4..d6ac0d48 100644
--- a/frontend/src/state/wasm-loader.ts
+++ b/frontend/src/state/wasm-loader.ts
@@ -4,7 +4,7 @@ import { createJsDispatcher } from "@/dispatcher/js-dispatcher";
 import { JsMessageType } from "@/dispatcher/js-messages";
 
 export type WasmInstance = typeof import("@/../wasm/pkg");
-export type RustEditorInstance = InstanceType<WasmInstance["Editor"]>;
+export type RustEditorInstance = InstanceType<WasmInstance["JsEditorHandle"]>;
 
 let wasmImport: WasmInstance | null = null;
 export async function initWasm() {
@@ -65,7 +65,7 @@ export function createEditorState() {
 		dispatcher.handleJsMessage(messageType, data, rawWasm, instance);
 	};
 
-	const instance = new rawWasm.Editor(rustCallback);
+	const instance = new rawWasm.JsEditorHandle(rustCallback);
 
 	return {
 		dispatcher,
diff --git a/frontend/wasm/src/api.rs b/frontend/wasm/src/api.rs
index 036a2479..d4dd5242 100644
--- a/frontend/wasm/src/api.rs
+++ b/frontend/wasm/src/api.rs
@@ -2,37 +2,40 @@
 // It serves as a thin wrapper over the editor backend API that relies
 // on the dispatcher messaging system and more complex Rust data types.
 
-use std::cell::{Cell, UnsafeCell};
+use std::cell::Cell;
 
 use crate::helpers::Error;
 use crate::type_translators::{translate_blend_mode, translate_key, translate_tool_type};
-use crate::EDITOR_HAS_CRASHED;
+use crate::{EDITOR_HAS_CRASHED, EDITOR_INSTANCES};
 use editor::consts::FILE_SAVE_SUFFIX;
 use editor::input::input_preprocessor::ModifierKeys;
 use editor::input::mouse::{EditorMouseState, ScrollDelta, ViewportBounds};
-use editor::message_prelude::*;
 use editor::misc::EditorError;
 use editor::tool::{tool_options::ToolOptions, tools, ToolType};
 use editor::Color;
 use editor::LayerId;
+use editor::{message_prelude::*, Editor};
 use wasm_bindgen::prelude::*;
 
 // To avoid wasm-bindgen from checking mutable reference issues using WasmRefCell
 // we must make all methods take a non mutable reference to self. Not doing this creates
 // an issue when rust calls into JS which calls back to rust in the same call stack.
 #[wasm_bindgen]
-pub struct Editor {
-	editor: UnsafeCell<editor::Editor>,
+pub struct JsEditorHandle {
+	editor_id: u64,
 	instance_received_crashed: Cell<bool>,
 	handle_response: js_sys::Function,
 }
 
 #[wasm_bindgen]
-impl Editor {
+impl JsEditorHandle {
 	#[wasm_bindgen(constructor)]
-	pub fn new(handle_response: js_sys::Function) -> Editor {
-		Editor {
-			editor: UnsafeCell::new(editor::Editor::new()),
+	pub fn new(handle_response: js_sys::Function) -> JsEditorHandle {
+		let editor_id = generate_uuid();
+		let editor = Editor::new();
+		EDITOR_INSTANCES.with(|instances| instances.borrow_mut().insert(editor_id, editor));
+		JsEditorHandle {
+			editor_id,
 			instance_received_crashed: Cell::new(false),
 			handle_response,
 		}
@@ -50,9 +53,13 @@ impl Editor {
 			return;
 		}
 
-		let editor = unsafe { self.editor.get().as_mut().unwrap() };
-		// Dispatch the message and receive a vector of FrontendMessage responses
-		let responses = editor.handle_message(message.into());
+		let responses = EDITOR_INSTANCES.with(|instances| {
+			instances
+				.borrow_mut()
+				.get_mut(&self.editor_id)
+				.expect("EDITOR_INSTANCES does not contain the current editor_id")
+				.handle_message(message.into())
+		});
 		for response in responses.into_iter() {
 			// Send each FrontendMessage to the JavaScript frontend
 			self.handle_response(response);
@@ -432,6 +439,12 @@ impl Editor {
 	}
 }
 
+impl Drop for JsEditorHandle {
+	fn drop(&mut self) {
+		EDITOR_INSTANCES.with(|instances| instances.borrow_mut().remove(&self.editor_id));
+	}
+}
+
 /// Access a handle to WASM memory
 #[wasm_bindgen]
 pub fn wasm_memory() -> JsValue {
diff --git a/frontend/wasm/src/lib.rs b/frontend/wasm/src/lib.rs
index 572c873c..7faca237 100644
--- a/frontend/wasm/src/lib.rs
+++ b/frontend/wasm/src/lib.rs
@@ -6,12 +6,16 @@ pub mod type_translators;
 use editor::message_prelude::FrontendMessage;
 use logging::WasmLog;
 use std::cell::RefCell;
+use std::collections::HashMap;
 use std::panic;
 use wasm_bindgen::prelude::*;
 
 // Set up the persistent editor backend state
 static LOGGER: WasmLog = WasmLog;
-thread_local! { pub static EDITOR_HAS_CRASHED: RefCell<Option<FrontendMessage>> = RefCell::new(None); }
+thread_local! {
+	pub static EDITOR_HAS_CRASHED: RefCell<Option<FrontendMessage>> = RefCell::new(None);
+	pub static EDITOR_INSTANCES: RefCell<HashMap<u64, editor::Editor>> = RefCell::new(HashMap::new());
+}
 
 // Initialize the backend
 #[wasm_bindgen(start)]