fix(router): gracefully handle client crashes (#1710)

* fix(router): gracefully handle client crashes

* style(comments): remove unused

10 files changed, 90 insertions, 27 deletions

diff --git a/src/sessions.rs b/src/sessions.rs
index 616596ef..420ed9d1 100644
--- a/src/sessions.rs
+++ b/src/sessions.rs
@@ -54,7 +54,7 @@ fn assert_socket(name: &str) -> bool {
     match LocalSocketStream::connect(path) {
         Ok(stream) => {
             let mut sender = IpcSenderWithContext::new(stream);
-            sender.send(ClientToServerMsg::ConnStatus);
+            let _ = sender.send(ClientToServerMsg::ConnStatus);
             let mut receiver: IpcReceiverWithContext<ServerToClientMsg> = sender.get_receiver();
             match receiver.recv() {
                 Some((ServerToClientMsg::Connected, _)) => true,
@@ -115,7 +115,7 @@ pub(crate) fn kill_session(name: &str) {
     let path = &*ZELLIJ_SOCK_DIR.join(name);
     match LocalSocketStream::connect(path) {
         Ok(stream) => {
-            IpcSenderWithContext::new(stream).send(ClientToServerMsg::KillSession);
+            let _ = IpcSenderWithContext::new(stream).send(ClientToServerMsg::KillSession);
         },
         Err(e) => {
             eprintln!("Error occurred: {:?}", e);
diff --git a/zellij-client/src/os_input_output.rs b/zellij-client/src/os_input_output.rs
index f9f3c590..1caf7872 100644
--- a/zellij-client/src/os_input_output.rs
+++ b/zellij-client/src/os_input_output.rs
@@ -145,7 +145,9 @@ impl ClientOsApi for ClientOsInputOutput {
     }
 
     fn send_to_server(&self, msg: ClientToServerMsg) {
-        self.send_instructions_to_server
+        // TODO: handle the error here, right now we silently ignore it
+        let _ = self
+            .send_instructions_to_server
             .lock()
             .unwrap()
             .as_mut()
diff --git a/zellij-client/src/sessions.rs b/zellij-client/src/sessions.rs
index e3e9c4c3..6a3c0799 100644
--- a/zellij-client/src/sessions.rs
+++ b/zellij-client/src/sessions.rs
@@ -8,7 +8,7 @@ pub(crate) fn kill_session(name: &str) {
     let path = &*ZELLIJ_SOCK_DIR.join(name);
     match LocalSocketStream::connect(path) {
         Ok(stream) => {
-            IpcSenderWithContext::new(stream).send(ClientToServerMsg::KillSession);
+            let _ = IpcSenderWithContext::new(stream).send(ClientToServerMsg::KillSession);
         },
         Err(e) => {
             eprintln!("Error occurred: {:?}", e);
diff --git a/zellij-server/src/lib.rs b/zellij-server/src/lib.rs
index 93ea824d..400026b1 100644
--- a/zellij-server/src/lib.rs
+++ b/zellij-server/src/lib.rs
@@ -132,6 +132,16 @@ macro_rules! remove_client {
     };
 }
 
+macro_rules! send_to_client {
+    ($client_id:expr, $os_input:expr, $msg:expr, $session_state:expr) => {
+        let send_to_client_res = $os_input.send_to_client($client_id, $msg);
+        if let Err(_) = send_to_client_res {
+            // failed to send to client, remove it
+            remove_client!($client_id, $os_input, $session_state);
+        }
+    };
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) struct SessionState {
     clients: HashMap<ClientId, Option<Size>>,
@@ -392,15 +402,26 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         Event::ModeUpdate(mode_info),
                     ))
                     .unwrap();
-                os_input.send_to_client(client_id, ServerToClientMsg::SwitchToMode(mode));
+                send_to_client!(
+                    client_id,
+                    os_input,
+                    ServerToClientMsg::SwitchToMode(mode),
+                    session_state
+                );
             },
             ServerInstruction::UnblockInputThread => {
                 for client_id in session_state.read().unwrap().clients.keys() {
-                    os_input.send_to_client(*client_id, ServerToClientMsg::UnblockInputThread);
+                    send_to_client!(
+                        *client_id,
+                        os_input,
+                        ServerToClientMsg::UnblockInputThread,
+                        session_state
+                    );
                 }
             },
             ServerInstruction::ClientExit(client_id) => {
-                os_input.send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
+                let _ =
+                    os_input.send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
                 remove_client!(client_id, os_input, session_state);
                 if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size() {
                     session_data
@@ -465,14 +486,16 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
             ServerInstruction::KillSession => {
                 let client_ids = session_state.read().unwrap().client_ids();
                 for client_id in client_ids {
-                    os_input.send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
+                    let _ = os_input
+                        .send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
                     remove_client!(client_id, os_input, session_state);
                 }
                 break;
             },
             ServerInstruction::DetachSession(client_ids) => {
                 for client_id in client_ids {
-                    os_input.send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
+                    let _ = os_input
+                        .send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
                     remove_client!(client_id, os_input, session_state);
                     if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size()
                     {
@@ -509,14 +532,16 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                 // If `None`- Send an exit instruction. This is the case when a user closes the last Tab/Pane.
                 if let Some(output) = &serialized_output {
                     for (client_id, client_render_instruction) in output.iter() {
-                        os_input.send_to_client(
+                        send_to_client!(
                             *client_id,
+                            os_input,
                             ServerToClientMsg::Render(client_render_instruction.clone()),
+                            session_state
                         );
                     }
                 } else {
                     for client_id in client_ids {
-                        os_input
+                        let _ = os_input
                             .send_to_client(client_id, ServerToClientMsg::Exit(ExitReason::Normal));
                         remove_client!(client_id, os_input, session_state);
                     }
@@ -526,7 +551,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
             ServerInstruction::Error(backtrace) => {
                 let client_ids = session_state.read().unwrap().client_ids();
                 for client_id in client_ids {
-                    os_input.send_to_client(
+                    let _ = os_input.send_to_client(
                         client_id,
                         ServerToClientMsg::Exit(ExitReason::Error(backtrace.clone())),
                     );
@@ -535,7 +560,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                 break;
             },
             ServerInstruction::ConnStatus(client_id) => {
-                os_input.send_to_client(client_id, ServerToClientMsg::Connected);
+                let _ = os_input.send_to_client(client_id, ServerToClientMsg::Connected);
                 remove_client!(client_id, os_input, session_state);
             },
             ServerInstruction::ActiveClients(client_id) => {
@@ -545,7 +570,12 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     client_ids,
                     client_id
                 );
-                os_input.send_to_client(client_id, ServerToClientMsg::ActiveClients(client_ids));
+                send_to_client!(
+                    client_id,
+                    os_input,
+                    ServerToClientMsg::ActiveClients(client_ids),
+                    session_state
+                );
             },
         }
     }
diff --git a/zellij-server/src/os_input_output.rs b/zellij-server/src/os_input_output.rs
index 31f1780b..8bd62a21 100644
--- a/zellij-server/src/os_input_output.rs
+++ b/zellij-server/src/os_input_output.rs
@@ -316,7 +316,11 @@ pub trait ServerOsApi: Send + Sync {
     fn force_kill(&self, pid: Pid) -> Result<(), nix::Error>;
     /// Returns a [`Box`] pointer to this [`ServerOsApi`] struct.
     fn box_clone(&self) -> Box<dyn ServerOsApi>;
-    fn send_to_client(&self, client_id: ClientId, msg: ServerToClientMsg);
+    fn send_to_client(
+        &self,
+        client_id: ClientId,
+        msg: ServerToClientMsg,
+    ) -> Result<(), &'static str>;
     fn new_client(
         &mut self,
         client_id: ClientId,
@@ -373,9 +377,15 @@ impl ServerOsApi for ServerOsInputOutput {
         let _ = kill(pid, Some(Signal::SIGKILL));
         Ok(())
     }
-    fn send_to_client(&self, client_id: ClientId, msg: ServerToClientMsg) {
+    fn send_to_client(
+        &self,
+        client_id: ClientId,
+        msg: ServerToClientMsg,
+    ) -> Result<(), &'static str> {
         if let Some(sender) = self.client_senders.lock().unwrap().get_mut(&client_id) {
-            sender.send(msg);
+            sender.send(msg)
+        } else {
+            Ok(())
         }
     }
     fn new_client(
diff --git a/zellij-server/src/route.rs b/zellij-server/src/route.rs
index e1acbe1b..d82ec42e 100644
--- a/zellij-server/src/route.rs
+++ b/zellij-server/src/route.rs
@@ -514,10 +514,15 @@ pub(crate) fn route_thread_main(
                             let client_id = maybe_client_id.unwrap_or(client_id);
                             if let Some(rlocked_sessions) = rlocked_sessions.as_ref() {
                                 if let Action::SwitchToMode(input_mode) = action {
-                                    os_input.send_to_client(
+                                    let send_res = os_input.send_to_client(
                                         client_id,
                                         ServerToClientMsg::SwitchToMode(input_mode),
                                     );
+                                    if send_res.is_err() {
+                                        let _ = to_server
+                                            .send(ServerInstruction::RemoveClient(client_id));
+                                        return true;
+                                    }
                                 }
                                 if route_action(
                                     action,
@@ -642,7 +647,7 @@ pub(crate) fn route_thread_main(
             },
             None => {
                 log::error!("Received empty message from client");
-                os_input.send_to_client(
+                let _ = os_input.send_to_client(
                     client_id,
                     ServerToClientMsg::Exit(ExitReason::Error(
                         "Received empty message".to_string(),
diff --git a/zellij-server/src/tab/unit/tab_integration_tests.rs b/zellij-server/src/tab/unit/tab_integration_tests.rs
index c19a4f0c..41e2fee4 100644
--- a/zellij-server/src/tab/unit/tab_integration_tests.rs
+++ b/zellij-server/src/tab/unit/tab_integration_tests.rs
@@ -72,7 +72,11 @@ impl ServerOsApi for FakeInputOutput {
     fn box_clone(&self) -> Box<dyn ServerOsApi> {
         Box::new((*self).clone())
     }
-    fn send_to_client(&self, _client_id: ClientId, _msg: ServerToClientMsg) {
+    fn send_to_client(
+        &self,
+        _client_id: ClientId,
+        _msg: ServerToClientMsg,
+    ) -> Result<(), &'static str> {
         unimplemented!()
     }
     fn new_client(
diff --git a/zellij-server/src/tab/unit/tab_tests.rs b/zellij-server/src/tab/unit/tab_tests.rs
index 719f9ce9..2ca6a0cb 100644
--- a/zellij-server/src/tab/unit/tab_tests.rs
+++ b/zellij-server/src/tab/unit/tab_tests.rs
@@ -63,7 +63,11 @@ impl ServerOsApi for FakeInputOutput {
     fn box_clone(&self) -> Box<dyn ServerOsApi> {
         Box::new((*self).clone())
     }
-    fn send_to_client(&self, _client_id: ClientId, _msg: ServerToClientMsg) {
+    fn send_to_client(
+        &self,
+        _client_id: ClientId,
+        _msg: ServerToClientMsg,
+    ) -> Result<(), &'static str> {
         unimplemented!()
     }
     fn new_client(
diff --git a/zellij-server/src/unit/screen_tests.rs b/zellij-server/src/unit/screen_tests.rs
index cbe0c1b5..53cb62d1 100644
--- a/zellij-server/src/unit/screen_tests.rs
+++ b/zellij-server/src/unit/screen_tests.rs
@@ -59,7 +59,11 @@ impl ServerOsApi for FakeInputOutput {
     fn box_clone(&self) -> Box<dyn ServerOsApi> {
         Box::new((*self).clone())
     }
-    fn send_to_client(&self, _client_id: ClientId, _msg: ServerToClientMsg) {
+    fn send_to_client(
+        &self,
+        _client_id: ClientId,
+        _msg: ServerToClientMsg,
+    ) -> Result<(), &'static str> {
         unimplemented!()
     }
     fn new_client(
diff --git a/zellij-utils/src/ipc.rs b/zellij-utils/src/ipc.rs
index c0a7e269..354107f5 100644
--- a/zellij-utils/src/ipc.rs
+++ b/zellij-utils/src/ipc.rs
@@ -156,12 +156,16 @@ impl<T: Serialize> IpcSenderWithContext<T> {
     }
 
     /// Sends an event, along with the current [`ErrorContext`], on this [`IpcSenderWithContext`]'s socket.
-    pub fn send(&mut self, msg: T) {
+    pub fn send(&mut self, msg: T) -> Result<(), &'static str> {
         let err_ctx = get_current_ctx();
-        rmp_serde::encode::write(&mut self.sender, &(msg, err_ctx)).unwrap();
-        // TODO: unwrapping here can cause issues when the server disconnects which we don't mind
-        // do we need to handle errors here in other cases?
-        let _ = self.sender.flush();
+        if rmp_serde::encode::write(&mut self.sender, &(msg, err_ctx)).is_err() {
+            Err("Failed to send message to client")
+        } else {
+            // TODO: unwrapping here can cause issues when the server disconnects which we don't mind
+            // do we need to handle errors here in other cases?
+            let _ = self.sender.flush();
+            Ok(())
+        }
     }
 
     /// Returns an [`IpcReceiverWithContext`] with the same socket as this sender.