From 216c42716570b774e1bf45da6bcae64f5680d308 Mon Sep 17 00:00:00 2001
From: jspspike <jspspike@gmail.com>
Date: Mon, 6 Jun 2022 16:52:23 -0500
Subject: [PATCH] Added TryFuture to catch unwind async panic

---
 worker-build/src/main.rs     | 14 +++++++++++++
 worker-macros/src/event.rs   | 24 +++++++++++++---------
 worker-sys/Cargo.toml        |  2 ++
 worker-sys/src/lib.rs        |  1 +
 worker-sys/src/try_future.rs | 40 ++++++++++++++++++++++++++++++++++++
 5 files changed, 71 insertions(+), 10 deletions(-)
 create mode 100644 worker-sys/src/try_future.rs

diff --git a/worker-build/src/main.rs b/worker-build/src/main.rs
index fd642dfc..ae3f62bd 100644
--- a/worker-build/src/main.rs
+++ b/worker-build/src/main.rs
@@ -184,6 +184,20 @@ wasm.__wbindgen_start?.();
     // Add the code that initializes the `wasm` variable we declared at the top of the file.
     fixed_bindgen_glue.push_str(&initialization_glue);
 
+    let t = format!(
+        r#"
+export function try_(a) {{
+    try {{
+        a()
+    }} catch (e) {{
+        // handle error here
+    }}
+}}
+        "#
+    );
+
+    fixed_bindgen_glue.push_str(&t);
+
     write_string_to_file(bindgen_glue_path, fixed_bindgen_glue)?;
 
     Ok(())
diff --git a/worker-macros/src/event.rs b/worker-macros/src/event.rs
index a59455a2..9fe02bbb 100644
--- a/worker-macros/src/event.rs
+++ b/worker-macros/src/event.rs
@@ -46,16 +46,9 @@ pub fn expand_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
             // rename the original attributed fn
             input_fn.sig.ident = input_fn_ident.clone();
 
-            let error_handling = match respond_with_errors {
-                true => {
-                    quote! {
+            let error_handling = quote! {
                         ::worker::Response::error(e.to_string(), 500).unwrap().into()
-                    }
-                }
-                false => {
-                    quote! { panic!("{}", e) }
-                }
-            };
+                    };
 
             // create a new "main" function that takes the worker_sys::Request, and calls the
             // original attributed function, passing in a converted worker::Request
@@ -66,9 +59,20 @@ pub fn expand_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
                     ctx: ::worker::worker_sys::Context
                 ) -> ::worker::worker_sys::Response {
                     let ctx = worker::Context::new(ctx);
+                    let f = #input_fn_ident(::worker::Request::from(req), env, ctx);
+                    let resp = ::worker::worker_sys::try_future::TryFuture{ f };
                     // get the worker::Result<worker::Response> by calling the original fn
-                    match #input_fn_ident(::worker::Request::from(req), env, ctx).await.map(::worker::worker_sys::Response::from) {
+                    //
+
+                    let r = match resp.await {
                         Ok(res) => res,
+                        Err(e) => {
+                            todo!("Reset wasm instance here?");
+                            Err(::worker::Error::PanicError) 
+                        },
+                    };
+                    match r {
+                        Ok(res) => ::worker::worker_sys::Response::from(res),
                         Err(e) => {
                             ::worker::console_log!("{}", &e);
                             #error_handling
diff --git a/worker-sys/Cargo.toml b/worker-sys/Cargo.toml
index 45539123..10c90d4b 100644
--- a/worker-sys/Cargo.toml
+++ b/worker-sys/Cargo.toml
@@ -9,7 +9,9 @@ description = "Low-level extern definitions / FFI bindings to the Cloudflare Wor
 
 [dependencies]
 cfg-if = "1.0.0"
+futures = "0.3.21"
 js-sys = "0.3.57"
+pin-project = "1.0.10"
 wasm-bindgen = "0.2.80"
 
 [dependencies.web-sys]
diff --git a/worker-sys/src/lib.rs b/worker-sys/src/lib.rs
index b0f6764e..be9140c0 100644
--- a/worker-sys/src/lib.rs
+++ b/worker-sys/src/lib.rs
@@ -14,6 +14,7 @@ pub mod request_init;
 pub mod response;
 pub mod response_init;
 pub mod schedule;
+pub mod try_future;
 pub mod websocket;
 
 /// When debugging your Worker via `wrangler dev`, `wrangler tail`, or from the Workers Dashboard,
diff --git a/worker-sys/src/try_future.rs b/worker-sys/src/try_future.rs
new file mode 100644
index 00000000..4544304b
--- /dev/null
+++ b/worker-sys/src/try_future.rs
@@ -0,0 +1,40 @@
+use futures::task::Poll;
+use pin_project::pin_project;
+use wasm_bindgen::prelude::*;
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+
+#[pin_project]
+pub struct TryFuture<F: ?Sized> {
+    #[pin]
+    pub f: F,
+}
+impl<F: ?Sized> Future for TryFuture<F>
+where
+    F: Future,
+{
+    type Output = Result<F::Output, ()>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match try__(|| self.project().f.poll(cx)) {
+            Ok(Poll::Ready(t)) => Poll::Ready(Ok(t)),
+            Ok(Poll::Pending) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(e)),
+        }
+    }
+}
+
+fn try__<O>(f: impl FnOnce() -> O) -> Result<O, ()> {
+    let mut f = Some(f);
+    let mut o = None;
+    try_(&mut || o = Some(f.take().unwrap()()));
+    o.ok_or(())
+}
+
+#[wasm_bindgen]
+extern "C" {
+    fn try_(a: &mut dyn FnMut());
+    pub fn new_wasm_instance();
+}