panic unwinding

Cargo.toml

@@ -22,6 +22,7 @@ members = [
     "macros",
     "method-override",
     "native-tls",
+    "panic-boundary",
     "proxy",
     "redirect",
     "router",

panic-boundary/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "trillium-panic-boundary"
+version = "0.1.0"
+authors = ["Jacob Rothstein <[email protected]>"]
+edition = "2021"
+description = "experimental panic boundary for trillium.rs"
+license = "MIT OR Apache-2.0"
+repository = "https://github.com/trillium-rs/trillium"
+readme = "../README.md"
+keywords = ["trillium", "framework", "async"]
+categories = ["web-programming::http-server", "web-programming"]
+
+[dependencies]
+async-channel = "1.9.0"
+futures-lite = "1.13.0"
+trillium = { path = "../trillium", version = "^0.2.0"}
+
+[dev-dependencies]
+trillium-smol = { path = "../smol" }
+trillium-testing = { path = "../testing" }

panic-boundary/examples/panic.rs

@@ -0,0 +1,21 @@
+use trillium::Conn;
+use trillium_panic_boundary::{Unwind, UnwindConnExt};
+
+fn main() {
+    trillium_smol::run(Unwind::new(
+        |conn: Conn| async move {
+            if conn.path().starts_with("/panic") {
+                panic!("PANIC: {} {}", conn.method(), conn.path());
+            } else {
+                conn.ok("no panic")
+            }
+        },
+        |mut conn: Conn| async move {
+            if let Some(s) = conn.take_panic_message() {
+                conn.with_status(500).with_body(s.to_string())
+            } else {
+                conn.with_status(500).with_body("panic")
+            }
+        },
+    ));
+}

panic-boundary/src/lib.rs

@@ -0,0 +1,105 @@
+use futures_lite::FutureExt;
+use std::{
+    borrow::Cow,
+    ops::Deref,
+    panic::{resume_unwind, AssertUnwindSafe},
+};
+use trillium::{Conn, Handler};
+
+pub struct Unwind<H, PH> {
+    handler: H,
+    panic_handler: PH,
+}
+
+struct PanicMessage(Cow<'static, str>);
+impl Deref for PanicMessage {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        &*self.0
+    }
+}
+
+impl<H, PH> Unwind<H, PH>
+where
+    H: Handler,
+    PH: Handler,
+{
+    pub fn new(handler: H, panic_handler: PH) -> Self {
+        Self {
+            handler,
+            panic_handler,
+        }
+    }
+}
+
+pub trait UnwindConnExt {
+    fn panic_message(&self) -> Option<&str>;
+    fn take_panic_message(&mut self) -> Option<Cow<'static, str>>;
+}
+impl UnwindConnExt for Conn {
+    fn panic_message(&self) -> Option<&str> {
+        self.state().map(|PanicMessage(ref cow)| &**cow)
+    }
+
+    fn take_panic_message(&mut self) -> Option<Cow<'static, str>> {
+        self.take_state().map(|PanicMessage(cow)| cow)
+    }
+}
+
+#[trillium::async_trait]
+impl<H: Handler, PH: Handler> Handler for Unwind<H, PH> {
+    async fn run(&self, mut conn: Conn) -> Conn {
+        let (tx, rx) = async_channel::bounded(1);
+        conn.on_drop(move |conn| {
+            let _ = tx.try_send(conn);
+        });
+
+        match AssertUnwindSafe(self.handler.run(conn))
+            .catch_unwind()
+            .await
+        {
+            Ok(conn) => conn,
+            Err(e) => match rx.recv().await {
+                Ok(mut conn) => {
+                    if let Some(s) = e.downcast_ref::<&str>() {
+                        conn.set_state(PanicMessage(Cow::from(*s)));
+                    } else if let Some(s) = e.downcast_ref::<String>() {
+                        conn.set_state(PanicMessage(Cow::from(s.clone())));
+                    }
+
+                    self.panic_handler.run(conn).await
+                }
+
+                Err(_) => resume_unwind(e),
+            },
+        }
+    }
+
+    async fn before_send(&self, mut conn: Conn) -> Conn {
+        let (tx, rx) = async_channel::bounded(1);
+        conn.on_drop(move |conn| {
+            let _ = tx.try_send(conn);
+        });
+
+        match AssertUnwindSafe(self.handler.before_send(conn))
+            .catch_unwind()
+            .await
+        {
+            Ok(conn) => conn,
+            Err(e) => match rx.recv().await {
+                Ok(mut conn) => {
+                    if let Some(s) = e.downcast_ref::<&str>() {
+                        conn.set_state(PanicMessage(Cow::from(*s)));
+                    } else if let Some(s) = e.downcast_ref::<String>() {
+                        conn.set_state(PanicMessage(Cow::from(s.clone())));
+                    }
+
+                    self.panic_handler.before_send(conn).await
+                }
+
+                Err(_) => resume_unwind(e),
+            },
+        }
+    }
+}