From ee303689706166b51e8020261a6496090e854d69 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 24 Oct 2024 03:13:39 +1100
Subject: [PATCH] refactor(connlib): simplify error handling on crash (#7134)

The `fmt::Display` implementation of `tokio::task::JoinError` already
does exactly what we do here: Extracting the panic message if there is
one. Thus, we can simplify this code why just moving the `JoinError`
into the `DisconnectError` as its source.
---
 rust/connlib/clients/shared/src/callbacks.rs | 20 ++---------
 rust/connlib/clients/shared/src/lib.rs       | 37 +++-----------------
 rust/headless-client/src/lib.rs              |  8 ++++-
 3 files changed, 14 insertions(+), 51 deletions(-)

diff --git a/rust/connlib/clients/shared/src/callbacks.rs b/rust/connlib/clients/shared/src/callbacks.rs
index 93aebe56c..d2c47779b 100644
--- a/rust/connlib/clients/shared/src/callbacks.rs
+++ b/rust/connlib/clients/shared/src/callbacks.rs
@@ -22,28 +22,14 @@ pub trait Callbacks: Clone + Send + Sync {
     fn on_update_resources(&self, _: Vec<ResourceView>) {}
 
     /// Called when the tunnel is disconnected.
-    ///
-    /// If the tunnel disconnected due to a fatal error, `error` is the error
-    /// that caused the disconnect.
-    fn on_disconnect(&self, error: &DisconnectError) {
-        tracing::error!(error = ?error, "tunnel_disconnected");
-        // Note that we can't panic here, since we already hooked the panic to this function.
-        std::process::exit(0);
-    }
+    fn on_disconnect(&self, _: &DisconnectError) {}
 }
 
 /// Unified error type to use across connlib.
 #[derive(thiserror::Error, Debug)]
 pub enum DisconnectError {
-    /// A panic occurred.
-    #[error("Connlib panicked: {0}")]
-    Panic(String),
-    /// The task was cancelled
-    #[error("Connlib task was cancelled")]
-    Cancelled,
-    /// A panic occurred with a non-string payload.
-    #[error("Panicked with a non-string payload")]
-    PanicNonStringPayload,
+    #[error("connlib crashed: {0}")]
+    Crash(#[from] tokio::task::JoinError),
 
     #[error("connection to the portal failed: {0}")]
     PortalConnectionFailed(#[from] phoenix_channel::Error),
diff --git a/rust/connlib/clients/shared/src/lib.rs b/rust/connlib/clients/shared/src/lib.rs
index a9eac86e8..c9a524799 100644
--- a/rust/connlib/clients/shared/src/lib.rs
+++ b/rust/connlib/clients/shared/src/lib.rs
@@ -9,7 +9,6 @@ pub use firezone_tunnel::messages::client::{
 
 use connlib_model::ResourceId;
 use eventloop::Command;
-use firezone_logging::std_dyn_err;
 use firezone_tunnel::ClientTunnel;
 use phoenix_channel::{PhoenixChannel, PublicKeyParam};
 use socket_factory::{SocketFactory, TcpSocket, UdpSocket};
@@ -121,7 +120,7 @@ async fn connect<CB>(
     callbacks: CB,
     portal: PhoenixChannel<(), IngressMessages, ReplyMessages, PublicKeyParam>,
     rx: UnboundedReceiver<Command>,
-) -> Result<(), DisconnectError>
+) -> Result<(), phoenix_channel::Error>
 where
     CB: Callbacks + 'static,
 {
@@ -140,7 +139,7 @@ where
 
 /// A supervisor task that handles, when [`connect`] exits.
 async fn connect_supervisor<CB>(
-    connect_handle: JoinHandle<Result<(), DisconnectError>>,
+    connect_handle: JoinHandle<Result<(), phoenix_channel::Error>>,
     callbacks: CB,
 ) where
     CB: Callbacks,
@@ -149,35 +148,7 @@ async fn connect_supervisor<CB>(
         Ok(Ok(())) => {
             tracing::info!("connlib exited gracefully");
         }
-        Ok(Err(e)) => {
-            if e.is_authentication_error() {
-                tracing::warn!(error = std_dyn_err(&e), "Portal authentication error");
-            } else {
-                tracing::error!(error = std_dyn_err(&e), "connlib failed");
-            }
-
-            callbacks.on_disconnect(&e);
-        }
-        Err(e) => match e.try_into_panic() {
-            Ok(panic) => {
-                if let Some(msg) = panic.downcast_ref::<&str>() {
-                    tracing::error!("connlib panicked: {msg}");
-                    callbacks.on_disconnect(&DisconnectError::Panic(msg.to_string()));
-                    return;
-                }
-                if let Some(msg) = panic.downcast_ref::<String>() {
-                    tracing::error!("connlib panicked: {msg}");
-                    callbacks.on_disconnect(&DisconnectError::Panic(msg.to_string()));
-                    return;
-                }
-
-                tracing::error!("connlib panicked with a non-string payload");
-                callbacks.on_disconnect(&DisconnectError::PanicNonStringPayload);
-            }
-            Err(_) => {
-                tracing::error!("connlib task was cancelled");
-                callbacks.on_disconnect(&DisconnectError::Cancelled);
-            }
-        },
+        Ok(Err(e)) => callbacks.on_disconnect(&DisconnectError::PortalConnectionFailed(e)),
+        Err(e) => callbacks.on_disconnect(&DisconnectError::Crash(e)),
     }
 }
diff --git a/rust/headless-client/src/lib.rs b/rust/headless-client/src/lib.rs
index 69c61a0b4..60c4736f9 100644
--- a/rust/headless-client/src/lib.rs
+++ b/rust/headless-client/src/lib.rs
@@ -12,6 +12,7 @@ use anyhow::{Context as _, Result};
 use connlib_client_shared::Callbacks;
 use connlib_model::ResourceView;
 use firezone_bin_shared::platform::DnsControlMethod;
+use firezone_logging::std_dyn_err;
 use std::{
     net::{IpAddr, Ipv4Addr, Ipv6Addr},
     path::PathBuf,
@@ -96,7 +97,12 @@ pub struct CallbackHandler {
 
 impl Callbacks for CallbackHandler {
     fn on_disconnect(&self, error: &connlib_client_shared::DisconnectError) {
-        tracing::error!(?error, "Got `on_disconnect` from connlib");
+        if error.is_authentication_error() {
+            tracing::warn!(error = std_dyn_err(error));
+        } else {
+            tracing::error!(error = std_dyn_err(error))
+        }
+
         self.cb_tx
             .try_send(ConnlibMsg::OnDisconnect {
                 error_msg: error.to_string(),