Fix memory overflow error and add panic hook for server

6 files changed, 96 insertions, 50 deletions

diff --git a/src/cli.rs b/src/cli.rs
index 72666dd66..9ea84374c 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -13,22 +13,23 @@ pub struct CliArgs {
     pub max_panes: Option<usize>,
 
     /// Change where zellij looks for layouts and plugins
-    #[structopt(long)]
+    #[structopt(long, parse(from_os_str))]
     pub data_dir: Option<PathBuf>,
 
-    #[structopt(long)]
+    /// Run server listening at the specified socket path
+    #[structopt(long, parse(from_os_str))]
     pub server: Option<PathBuf>,
 
     /// Path to a layout yaml file
-    #[structopt(short, long)]
+    #[structopt(short, long, parse(from_os_str))]
     pub layout: Option<PathBuf>,
 
     /// Change where zellij looks for the configuration
-    #[structopt(short, long, env=ZELLIJ_CONFIG_FILE_ENV)]
+    #[structopt(short, long, env=ZELLIJ_CONFIG_FILE_ENV, parse(from_os_str))]
     pub config: Option<PathBuf>,
 
     /// Change where zellij looks for the configuration
-    #[structopt(long, env=ZELLIJ_CONFIG_DIR_ENV)]
+    #[structopt(long, env=ZELLIJ_CONFIG_DIR_ENV, parse(from_os_str))]
     pub config_dir: Option<PathBuf>,
 
     #[structopt(subcommand)]
diff --git a/src/client/mod.rs b/src/client/mod.rs
index d24b66d3a..594960eb6 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -32,6 +32,7 @@ pub enum ClientInstruction {
     Render(Option<String>),
     UnblockInputThread,
     Exit,
+    ServerError(String),
 }
 
 fn spawn_server(socket_path: &Path) -> io::Result<()> {
@@ -52,7 +53,6 @@ fn spawn_server(socket_path: &Path) -> io::Result<()> {
 }
 
 pub fn start_client(mut os_input: Box<dyn ClientOsApi>, opts: CliArgs, config: Config) {
-    spawn_server(&*ZELLIJ_IPC_PIPE).unwrap();
     let take_snapshot = "\u{1b}[?1049h";
     let bracketed_paste = "\u{1b}[?2004h";
     os_input.unset_raw_mode(0);
@@ -60,12 +60,11 @@ pub fn start_client(mut os_input: Box<dyn ClientOsApi>, opts: CliArgs, config: C
         .get_stdout_writer()
         .write(take_snapshot.as_bytes())
         .unwrap();
-    let _ = os_input
-        .get_stdout_writer()
-        .write(clear_client_terminal_attributes.as_bytes())
-        .unwrap();
+
     std::env::set_var(&"ZELLIJ", "0");
 
+    spawn_server(&*ZELLIJ_IPC_PIPE).unwrap();
+
     let mut command_is_executing = CommandIsExecuting::new();
 
     let config_options = Options::from_cli(&config.options, opts.option.clone());
@@ -131,22 +130,39 @@ pub fn start_client(mut os_input: Box<dyn ClientOsApi>, opts: CliArgs, config: C
         .name("router".to_string())
         .spawn({
             let os_input = os_input.clone();
-            move || {
-                loop {
-                    let (instruction, mut err_ctx) = os_input.recv_from_server();
-                    err_ctx.add_call(ContextType::Client(ClientContext::from(&instruction)));
-                    if let ClientInstruction::Exit = instruction {
-                        break;
+            let mut should_break = false;
+            move || loop {
+                let (instruction, mut err_ctx) = os_input.recv_from_server();
+                err_ctx.add_call(ContextType::Client(ClientContext::from(&instruction)));
+                match instruction {
+                    ClientInstruction::Exit | ClientInstruction::ServerError(_) => {
+                        should_break = true;
                     }
-                    send_client_instructions.send(instruction).unwrap();
+                    _ => {}
+                }
+                send_client_instructions.send(instruction).unwrap();
+                if should_break {
+                    break;
                 }
-                send_client_instructions
-                    .send(ClientInstruction::Exit)
-                    .unwrap();
             }
         })
         .unwrap();
 
+    let handle_error = |backtrace: String| {
+        os_input.unset_raw_mode(0);
+        let goto_start_of_last_line = format!("\u{1b}[{};{}H", full_screen_ws.rows, 1);
+        let restore_snapshot = "\u{1b}[?1049l";
+        let error = format!(
+            "{}\n{}{}",
+            goto_start_of_last_line, restore_snapshot, backtrace
+        );
+        let _ = os_input
+            .get_stdout_writer()
+            .write(error.as_bytes())
+            .unwrap();
+        std::process::exit(1);
+    };
+
     loop {
         let (client_instruction, mut err_ctx) = receive_client_instructions
             .recv()
@@ -159,18 +175,10 @@ pub fn start_client(mut os_input: Box<dyn ClientOsApi>, opts: CliArgs, config: C
             ClientInstruction::Exit => break,
             ClientInstruction::Error(backtrace) => {
                 let _ = os_input.send_to_server(ServerInstruction::ClientExit);
-                os_input.unset_raw_mode(0);
-                let goto_start_of_last_line = format!("\u{1b}[{};{}H", full_screen_ws.rows, 1);
-                let restore_snapshot = "\u{1b}[?1049l";
-                let error = format!(
-                    "{}\n{}{}",
-                    goto_start_of_last_line, restore_snapshot, backtrace
-                );
-                let _ = os_input
-                    .get_stdout_writer()
-                    .write(error.as_bytes())
-                    .unwrap();
-                std::process::exit(1);
+                handle_error(backtrace);
+            }
+            ClientInstruction::ServerError(backtrace) => {
+                handle_error(backtrace);
             }
             ClientInstruction::Render(output) => {
                 if output.is_none() {
diff --git a/src/common/errors.rs b/src/common/errors.rs
index c2caf19d6..c92b9b7f7 100644
--- a/src/common/errors.rs
+++ b/src/common/errors.rs
@@ -19,12 +19,29 @@ const MAX_THREAD_CALL_STACK: usize = 6;
 use super::thread_bus::SenderWithContext;
 #[cfg(not(test))]
 use std::panic::PanicInfo;
+
+pub trait ErrorInstruction {
+    fn error(err: String) -> Self;
+}
+
+impl ErrorInstruction for ClientInstruction {
+    fn error(err: String) -> Self {
+        ClientInstruction::Error(err)
+    }
+}
+
+impl ErrorInstruction for ServerInstruction {
+    fn error(err: String) -> Self {
+        ServerInstruction::Error(err)
+    }
+}
+
 /// Custom panic handler/hook. Prints the [`ErrorContext`].
 #[cfg(not(test))]
-pub fn handle_panic(
-    info: &PanicInfo<'_>,
-    send_app_instructions: &SenderWithContext<ClientInstruction>,
-) {
+pub fn handle_panic<T>(info: &PanicInfo<'_>, sender: &SenderWithContext<T>)
+where
+    T: ErrorInstruction + Clone,
+{
     use backtrace::Backtrace;
     use std::{process, thread};
     let backtrace = Backtrace::new();
@@ -70,7 +87,7 @@ pub fn handle_panic(
         println!("{}", backtrace);
         process::exit(1);
     } else {
-        let _ = send_app_instructions.send(ClientInstruction::Error(backtrace));
+        let _ = sender.send(T::error(backtrace));
     }
 }
 
@@ -333,6 +350,7 @@ pub enum ClientContext {
     Error,
     UnblockInputThread,
     Render,
+    ServerError,
 }
 
 impl From<&ClientInstruction> for ClientContext {
@@ -340,6 +358,7 @@ impl From<&ClientInstruction> for ClientContext {
         match *client_instruction {
             ClientInstruction::Exit => ClientContext::Exit,
             ClientInstruction::Error(_) => ClientContext::Error,
+            ClientInstruction::ServerError(_) => ClientContext::ServerError,
             ClientInstruction::Render(_) => ClientContext::Render,
             ClientInstruction::UnblockInputThread => ClientContext::UnblockInputThread,
         }
@@ -355,6 +374,7 @@ pub enum ServerContext {
     TerminalResize,
     UnblockInputThread,
     ClientExit,
+    Error,
 }
 
 impl From<&ServerInstruction> for ServerContext {
@@ -366,6 +386,7 @@ impl From<&ServerInstruction> for ServerContext {
             ServerInstruction::Render(_) => ServerContext::Render,
             ServerInstruction::UnblockInputThread => ServerContext::UnblockInputThread,
             ServerInstruction::ClientExit => ServerContext::ClientExit,
+            ServerInstruction::Error(_) => ServerContext::Error,
         }
     }
 }
diff --git a/src/common/os_input_output.rs b/src/common/os_input_output.rs
index 2a9a14e07..a8d674785 100644
--- a/src/common/os_input_output.rs
+++ b/src/common/os_input_output.rs
@@ -302,7 +302,7 @@ pub trait ClientOsApi: Send + Sync {
     fn set_raw_mode(&mut self, fd: RawFd);
     /// Set the terminal associated to file descriptor `fd` to
     /// [cooked mode](https://en.wikipedia.org/wiki/Terminal_mode).
-    fn unset_raw_mode(&mut self, fd: RawFd);
+    fn unset_raw_mode(&self, fd: RawFd);
     /// Returns the writer that allows writing to standard output.
     fn get_stdout_writer(&self) -> Box<dyn io::Write>;
     /// Returns the raw contents of standard input.
@@ -326,7 +326,7 @@ impl ClientOsApi for ClientOsInputOutput {
     fn set_raw_mode(&mut self, fd: RawFd) {
         into_raw_mode(fd);
     }
-    fn unset_raw_mode(&mut self, fd: RawFd) {
+    fn unset_raw_mode(&self, fd: RawFd) {
         let orig_termios = self.orig_termios.lock().unwrap();
         unset_raw_mode(fd, orig_termios.clone());
     }
diff --git a/src/server/mod.rs b/src/server/mod.rs
index f85a4c7f2..884dbee14 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -5,7 +5,7 @@ use interprocess::local_socket::LocalSocketListener;
 use serde::{Deserialize, Serialize};
 use std::sync::{Arc, RwLock};
 use std::thread;
-use std::{path::PathBuf, sync::mpsc::channel};
+use std::{path::PathBuf, sync::mpsc};
 use wasmer::Store;
 use zellij_tile::data::PluginCapabilities;
 
@@ -18,8 +18,8 @@ use crate::common::{
     os_input_output::{set_permissions, ServerOsApi},
     pty::{pty_thread_main, Pty, PtyInstruction},
     screen::{screen_thread_main, ScreenInstruction},
-    setup::{get_default_data_dir, install::populate_data_dir},
-    thread_bus::{ChannelWithContext, SenderType, SenderWithContext},
+    setup::install::populate_data_dir,
+    thread_bus::{ChannelWithContext, SenderType, SenderWithContext, SyncChannelWithContext},
     utils::consts::ZELLIJ_PROJ_DIR,
     wasm_vm::{wasm_thread_main, PluginInstruction},
 };
@@ -37,6 +37,7 @@ pub enum ServerInstruction {
     Render(Option<String>),
     UnblockInputThread,
     ClientExit,
+    Error(String),
 }
 
 pub struct SessionMetaData {
@@ -63,10 +64,21 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
         .umask(0o077)
         .start()
         .expect("could not daemonize the server process");
-    let (to_server, server_receiver): ChannelWithContext<ServerInstruction> = channel();
-    let to_server = SenderWithContext::new(SenderType::Sender(to_server));
+
+    let (to_server, server_receiver): SyncChannelWithContext<ServerInstruction> =
+        mpsc::sync_channel(50);
+    let to_server = SenderWithContext::new(SenderType::SyncSender(to_server));
     let sessions: Arc<RwLock<Option<SessionMetaData>>> = Arc::new(RwLock::new(None));
 
+    #[cfg(not(test))]
+    std::panic::set_hook({
+        use crate::errors::handle_panic;
+        let to_server = to_server.clone();
+        Box::new(move |info| {
+            handle_panic(info, &to_server);
+        })
+    });
+
     #[cfg(test)]
     thread::Builder::new()
         .name("server_router".to_string())
@@ -149,15 +161,19 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
             ServerInstruction::ClientExit => {
                 *sessions.write().unwrap() = None;
                 os_input.send_to_client(ClientInstruction::Exit);
-                drop(std::fs::remove_file(&socket_path));
                 break;
             }
             ServerInstruction::Render(output) => {
                 os_input.send_to_client(ClientInstruction::Render(output))
             }
+            ServerInstruction::Error(backtrace) => {
+                os_input.send_to_client(ClientInstruction::ServerError(backtrace));
+                break;
+            }
             _ => panic!("Received unexpected instruction."),
         }
     }
+    drop(std::fs::remove_file(&socket_path));
 }
 
 fn init_session(
@@ -167,12 +183,12 @@ fn init_session(
     to_server: SenderWithContext<ServerInstruction>,
     full_screen_ws: PositionAndSize,
 ) -> SessionMetaData {
-    let (to_screen, screen_receiver): ChannelWithContext<ScreenInstruction> = channel();
+    let (to_screen, screen_receiver): ChannelWithContext<ScreenInstruction> = mpsc::channel();
     let to_screen = SenderWithContext::new(SenderType::Sender(to_screen));
 
-    let (to_plugin, plugin_receiver): ChannelWithContext<PluginInstruction> = channel();
+    let (to_plugin, plugin_receiver): ChannelWithContext<PluginInstruction> = mpsc::channel();
     let to_plugin = SenderWithContext::new(SenderType::Sender(to_plugin));
-    let (to_pty, pty_receiver): ChannelWithContext<PtyInstruction> = channel();
+    let (to_pty, pty_receiver): ChannelWithContext<PtyInstruction> = mpsc::channel();
     let to_pty = SenderWithContext::new(SenderType::Sender(to_pty));
 
     // Determine and initialize the data directory
diff --git a/src/tests/fakes.rs b/src/tests/fakes.rs
index 1697e65f7..2612848dd 100644
--- a/src/tests/fakes.rs
+++ b/src/tests/fakes.rs
@@ -153,7 +153,7 @@ impl ClientOsApi for FakeInputOutput {
             .unwrap()
             .push(IoEvent::IntoRawMode(pid));
     }
-    fn unset_raw_mode(&mut self, pid: RawFd) {
+    fn unset_raw_mode(&self, pid: RawFd) {
         self.io_events
             .lock()
             .unwrap()
@@ -217,7 +217,7 @@ impl ClientOsApi for FakeInputOutput {
             cb();
         }
     }
-    fn connect_to_server(&self, path: &std::path::Path) {}
+    fn connect_to_server(&self, _path: &std::path::Path) {}
 }
 
 impl ServerOsApi for FakeInputOutput {