rustlavel_http/
panic.rs

1//! Catching panics inside a handler so one bad request cannot take down the
2//! connection — and so the dev error page can show where it happened.
3
4use std::any::Any;
5use std::cell::RefCell;
6use std::future::Future;
7use std::panic::{AssertUnwindSafe, catch_unwind};
8use std::pin::Pin;
9use std::sync::Once;
10use std::task::{Context, Poll};
11
12/// Where the last panic on this thread happened, captured by our hook because
13/// the payload alone does not carry a location.
14#[derive(Debug, Clone)]
15pub struct PanicLocation {
16    pub file: String,
17    pub line: u32,
18    pub column: u32,
19}
20
21thread_local! {
22    static LAST_LOCATION: RefCell<Option<PanicLocation>> = const { RefCell::new(None) };
23}
24
25/// Install a panic hook that records the location and stays quiet.
26///
27/// Without this the default hook prints a backtrace to stderr for every caught
28/// panic, which is noise when the error page is about to show the same thing.
29pub fn install_hook() {
30    static INSTALLED: Once = Once::new();
31    INSTALLED.call_once(|| {
32        let previous = std::panic::take_hook();
33        std::panic::set_hook(Box::new(move |info| {
34            if let Some(location) = info.location() {
35                LAST_LOCATION.with(|slot| {
36                    *slot.borrow_mut() = Some(PanicLocation {
37                        file: location.file().to_string(),
38                        line: location.line(),
39                        column: location.column(),
40                    });
41                });
42            }
43            // Keep the default reporting for panics outside a request.
44            if !in_request() {
45                previous(info);
46            }
47        }));
48    });
49}
50
51thread_local! {
52    static IN_REQUEST: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
53}
54
55fn in_request() -> bool {
56    IN_REQUEST.with(std::cell::Cell::get)
57}
58
59pub fn take_location() -> Option<PanicLocation> {
60    LAST_LOCATION.with(|slot| slot.borrow_mut().take())
61}
62
63/// Turn a panic payload into a readable message.
64pub fn message_of(payload: &(dyn Any + Send)) -> String {
65    if let Some(text) = payload.downcast_ref::<&str>() {
66        (*text).to_string()
67    } else if let Some(text) = payload.downcast_ref::<String>() {
68        text.clone()
69    } else {
70        "the handler panicked".to_string()
71    }
72}
73
74/// A future that catches a panic raised while polling the inner future.
75pub struct CatchUnwind<F> {
76    inner: F,
77}
78
79impl<F> CatchUnwind<F> {
80    pub fn new(inner: F) -> Self {
81        CatchUnwind { inner }
82    }
83}
84
85impl<F: Future> Future for CatchUnwind<F> {
86    type Output = Result<F::Output, Box<dyn Any + Send>>;
87
88    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
89        // SAFETY: `inner` is never moved out of the pinned projection, and the
90        // future is dropped in place along with this wrapper.
91        let inner = unsafe { self.map_unchecked_mut(|this| &mut this.inner) };
92
93        let was_in_request = IN_REQUEST.with(|flag| flag.replace(true));
94        let result = catch_unwind(AssertUnwindSafe(|| inner.poll(cx)));
95        IN_REQUEST.with(|flag| flag.set(was_in_request));
96
97        match result {
98            Ok(Poll::Pending) => Poll::Pending,
99            Ok(Poll::Ready(value)) => Poll::Ready(Ok(value)),
100            Err(payload) => Poll::Ready(Err(payload)),
101        }
102    }
103}
104
105/// Run a future, converting a panic into an `Err`.
106pub async fn catch<F: Future>(future: F) -> Result<F::Output, String> {
107    CatchUnwind::new(future).await.map_err(|payload| message_of(payload.as_ref()))
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[tokio::test]
115    async fn a_panicking_future_is_caught() {
116        install_hook();
117        let caught = catch(async { panic!("boom") }).await;
118
119        assert_eq!(caught.unwrap_err(), "boom");
120        assert!(take_location().is_some());
121    }
122
123    #[tokio::test]
124    async fn a_healthy_future_passes_through() {
125        install_hook();
126        let value = catch(async { 42 }).await;
127
128        assert_eq!(value.unwrap(), 42);
129    }
130
131    #[tokio::test]
132    async fn panics_across_await_points_are_caught() {
133        install_hook();
134        let caught = catch(async {
135            tokio::task::yield_now().await;
136            panic!("after yielding");
137        })
138        .await;
139
140        assert_eq!(caught.unwrap_err(), "after yielding");
141    }
142}
rustlavel_http/panic.rs

rustlavel_http/
panic.rs
