From 13a8b0afc1df48dce217c938b1f3ec692c392309 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sun, 13 Aug 2023 04:30:49 -0400
Subject: [PATCH] Add panic-handlers which exit on any panic

By default, tokio-spawned worker panics will only kill the task, not the
program. Due to our extensive use of panicking on invariants, we should ensure
the program exits.
---
 coordinator/src/main.rs   | 14 ++++++++++++++
 message-queue/src/main.rs | 12 ++++++++++++
 processor/src/main.rs     | 12 ++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 675051aa..b0bfd767 100644
--- a/coordinator/src/main.rs
+++ b/coordinator/src/main.rs
@@ -458,6 +458,8 @@ pub async fn handle_processors<D: Db, Pro: Processors, P: P2p>(
               Ok(hash) => {
                 log::info!("voted on key pair for {:?} in TX {}", id.set, hex::encode(hash))
               }
+              // This is assumed to be some ephemeral error due to the assumed fault-free creation
+              // TODO: Differentiate connection errors from already published to an invariant
               Err(e) => {
                 log::error!("couldn't connect to Serai node to publish vote TX: {:?}", e);
                 tokio::time::sleep(Duration::from_secs(10)).await;
@@ -718,6 +720,18 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
 
 #[tokio::main]
 async fn main() {
+  // Override the panic handler with one which will panic if any tokio task panics
+  {
+    let existing = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic| {
+      existing(panic);
+      const MSG: &str = "exiting the process due to a task panicking";
+      println!("{MSG}");
+      log::error!("{MSG}");
+      std::process::exit(1);
+    }));
+  }
+
   if std::env::var("RUST_LOG").is_err() {
     std::env::set_var("RUST_LOG", serai_env::var("RUST_LOG").unwrap_or_else(|| "info".to_string()));
   }
diff --git a/message-queue/src/main.rs b/message-queue/src/main.rs
index e5c24e62..d807ddc8 100644
--- a/message-queue/src/main.rs
+++ b/message-queue/src/main.rs
@@ -118,6 +118,18 @@ mod binaries {
 async fn main() {
   use binaries::*;
 
+  // Override the panic handler with one which will panic if any tokio task panics
+  {
+    let existing = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic| {
+      existing(panic);
+      const MSG: &str = "exiting the process due to a task panicking";
+      println!("{MSG}");
+      log::error!("{MSG}");
+      std::process::exit(1);
+    }));
+  }
+
   if std::env::var("RUST_LOG").is_err() {
     std::env::set_var("RUST_LOG", serai_env::var("RUST_LOG").unwrap_or_else(|| "info".to_string()));
   }
diff --git a/processor/src/main.rs b/processor/src/main.rs
index 37cc7ce0..b03eeef8 100644
--- a/processor/src/main.rs
+++ b/processor/src/main.rs
@@ -726,6 +726,18 @@ async fn run<N: Network, D: Db, Co: Coordinator>(mut raw_db: D, network: N, mut
 
 #[tokio::main]
 async fn main() {
+  // Override the panic handler with one which will panic if any tokio task panics
+  {
+    let existing = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic| {
+      existing(panic);
+      const MSG: &str = "exiting the process due to a task panicking";
+      println!("{MSG}");
+      log::error!("{MSG}");
+      std::process::exit(1);
+    }));
+  }
+
   if std::env::var("RUST_LOG").is_err() {
     std::env::set_var("RUST_LOG", serai_env::var("RUST_LOG").unwrap_or_else(|| "info".to_string()));
   }