From be558164f9a1d546d6ac271e289b530c0323bb01 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Mon, 7 Mar 2022 10:15:14 -0800
Subject: [PATCH] fix(console): exit crossterm before printing panic messages

If the `tokio-console` CLI panics, we will eventually exit `crossterm`'s
terminal capturing and return to the user's shell. However, this
currently happens *after* the panic message is printed. This is because
we `crossterm` in a drop guard that's held in `main`, but panic messages
are printed out _before_ unwinding.

This means that for most panics, the panic message is never actually
displayed --- instead, the console just crashes to a blank screen, and
then returns to the user's shell with no output. Some manual testing
(e.g. putting `panic!("fake panic");` in a few different places)
indicates that panics will only be displayed nicely if they occur prior
to the call to `term::init_crossterm()` in the main function.

This branch fixes this issue by changing the panic hook to explicitly
exit `crossterm` _before_ printing the panic. This way, panics should
always be printed to stdout, regardless of where they occur in the
console. I factored out the code for exiting crossterm's terminal
capturing into a function that's now called in the drop guard _and_ in
the panic hook.

I also added some code for logging the panic using `tracing`. That way,
even if we fail to exit crossterm's terminal capturing, but the
console's debug logs are being redirected to a file, we'll still get the
panic in the logs.
---
 tokio-console/src/term.rs        | 21 ++++++++++-------
 tokio-console/src/view/styles.rs | 40 +++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/tokio-console/src/term.rs b/tokio-console/src/term.rs
index eb8adece6..605119315 100644
--- a/tokio-console/src/term.rs
+++ b/tokio-console/src/term.rs
@@ -4,7 +4,7 @@ pub use tui::{backend::CrosstermBackend, Terminal};
 
 pub fn init_crossterm() -> color_eyre::Result<(Terminal<CrosstermBackend<io::Stdout>>, OnShutdown)>
 {
-    use crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen};
+    use crossterm::terminal::{self, EnterAlternateScreen};
     terminal::enable_raw_mode().wrap_err("Failed to enable crossterm raw mode")?;
 
     let mut stdout = std::io::stdout();
@@ -13,18 +13,21 @@ pub fn init_crossterm() -> color_eyre::Result<(Terminal<CrosstermBackend<io::Std
     let backend = CrosstermBackend::new(io::stdout());
     let term = Terminal::new(backend).wrap_err("Failed to create crossterm terminal")?;
 
-    let cleanup = OnShutdown::new(|| {
-        // Be a good terminal citizen...
-        let mut stdout = std::io::stdout();
-        crossterm::execute!(stdout, LeaveAlternateScreen)
-            .wrap_err("Failed to disable crossterm alternate screen")?;
-        terminal::disable_raw_mode().wrap_err("Failed to enable crossterm raw mode")?;
-        Ok(())
-    });
+    let cleanup = OnShutdown::new(exit_crossterm);
 
     Ok((term, cleanup))
 }
 
+pub(crate) fn exit_crossterm() -> color_eyre::Result<()> {
+    use crossterm::terminal::{self, LeaveAlternateScreen};
+    // Be a good terminal citizen...
+    let mut stdout = std::io::stdout();
+    crossterm::execute!(stdout, LeaveAlternateScreen)
+        .wrap_err("Failed to disable crossterm alternate screen")?;
+    terminal::disable_raw_mode().wrap_err("Failed to enable crossterm raw mode")?;
+    Ok(())
+}
+
 pub struct OnShutdown {
     action: fn() -> color_eyre::Result<()>,
 }
diff --git a/tokio-console/src/view/styles.rs b/tokio-console/src/view/styles.rs
index 9d861eb6c..6ace49e64 100644
--- a/tokio-console/src/view/styles.rs
+++ b/tokio-console/src/view/styles.rs
@@ -51,7 +51,45 @@ impl Styles {
             // disable colors in error reports
             builder = builder.theme(Theme::new());
         }
-        builder.install()
+
+        // We're going to wrap the panic hook in some extra code of our own, so
+        // we can't use `HookBuilder::install` --- instead, split it into an
+        // error hook and a panic hook.
+        let (panic_hook, error_hook) = builder.into_hooks();
+
+        // Set the panic hook.
+        std::panic::set_hook(Box::new(move |panic_info| {
+            // First, try to log the panic. This way, even if the more
+            // user-friendly panic message isn't displayed correctly, the panic
+            // will still be recorded somewhere.
+            if let Some(location) = panic_info.location() {
+                // If the panic has a source location, record it as structured fields.
+                tracing::error!(
+                    message = %panic_info,
+                    panic.file = location.file(),
+                    panic.line = location.line(),
+                    panic.column = location.column(),
+                );
+            } else {
+                // Otherwise, just log the whole thing.
+                tracing::error!(message = %panic_info);
+            }
+
+            // After logging the panic, use the `color_eyre` panic hook to print
+            // a nice, user-friendly panic message.
+
+            // Leave crossterm before printing panic messages; otherwise, they
+            // may not be displayed and the app will just crash to a blank
+            // screen, which isn't great.
+            let _ = crate::term::exit_crossterm();
+            // Print the panic message.
+            eprintln!("{}", panic_hook.panic_report(panic_info));
+        }));
+
+        // Set the error hook.
+        error_hook.install()?;
+
+        Ok(())
     }
 
     pub fn if_utf8<'a>(&self, utf8: &'a str, ascii: &'a str) -> &'a str {