From dba5dcbd8301913155d1b2df3b6c875a735668f2 Mon Sep 17 00:00:00 2001
From: har7an <99636919+har7an@users.noreply.github.com>
Date: Thu, 11 Aug 2022 09:10:12 +0000
Subject: [PATCH] fix (screen): don't crash when first tab doesn't exist
 (#1648)

* screen: Don't crash when first tab doesn't exist

while trying to attach a new client. Instead, check whether the first
tab does exist and if not, take the first tab index from the tabs
present in the session. If no tabs exist, panic with a better error
message.

* changelog: Add PR #1648

* add test

* fix(tabs): send actual default mode info to new tab

Co-authored-by: Thomas Linford <linford.t@gmail.com>
Co-authored-by: Aram Drevekenin <aram@poor.dev>
---
 CHANGELOG.md                           |  1 +
 zellij-server/src/screen.rs            | 27 ++++++++++++++------------
 zellij-server/src/unit/screen_tests.rs | 22 +++++++++++++++++++++
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 08ba5acbb..f6f3f3301 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased]
+* fix: crash when attaching to a session without the first tab (https://github.com/zellij-org/zellij/pull/1648)
 
 ## [0.31.1] - 2022-08-02
 * add: `solarized-light` theme to the example theme directory (https://github.com/zellij-org/zellij/pull/1608)
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 79cbdb570..c184e29fb 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -641,11 +641,6 @@ impl Screen {
     pub fn new_tab(&mut self, layout: Layout, new_pids: Vec<RawFd>, client_id: ClientId) {
         let tab_index = self.get_new_tab_index();
         let position = self.tabs.len();
-        let client_mode_info = self
-            .mode_info
-            .get(&client_id)
-            .unwrap_or(&self.default_mode_info)
-            .clone();
         let mut tab = Tab::new(
             tab_index,
             position,
@@ -657,7 +652,7 @@ impl Screen {
             self.bus.senders.clone(),
             self.max_panes,
             self.style,
-            client_mode_info,
+            self.default_mode_info.clone(),
             self.draw_pane_frames,
             self.connected_clients.clone(),
             self.session_is_mirrored,
@@ -702,21 +697,29 @@ impl Screen {
     }
 
     pub fn add_client(&mut self, client_id: ClientId) {
-        let mut tab_index = 0;
         let mut tab_history = vec![];
-        if let Some((_first_client, first_active_tab_index)) = self.active_tab_indices.iter().next()
-        {
-            tab_index = *first_active_tab_index;
-        }
         if let Some((_first_client, first_tab_history)) = self.tab_history.iter().next() {
             tab_history = first_tab_history.clone();
         }
+
+        let tab_index = if let Some((_first_client, first_active_tab_index)) =
+            self.active_tab_indices.iter().next()
+        {
+            *first_active_tab_index
+        } else if self.tabs.contains_key(&0) {
+            0
+        } else if let Some(tab_index) = self.tabs.keys().next() {
+            tab_index.to_owned()
+        } else {
+            panic!("Can't find a valid tab to attach client to!");
+        };
+
         self.active_tab_indices.insert(client_id, tab_index);
         self.connected_clients.borrow_mut().insert(client_id);
         self.tab_history.insert(client_id, tab_history);
         self.tabs
             .get_mut(&tab_index)
-            .unwrap()
+            .unwrap_or_else(|| panic!("Failed to attach client to tab with index {tab_index}"))
             .add_client(client_id, None);
     }
     pub fn remove_client(&mut self, client_id: ClientId) {
diff --git a/zellij-server/src/unit/screen_tests.rs b/zellij-server/src/unit/screen_tests.rs
index 0cf8d4535..cbe0c1b57 100644
--- a/zellij-server/src/unit/screen_tests.rs
+++ b/zellij-server/src/unit/screen_tests.rs
@@ -544,3 +544,25 @@ fn update_screen_pixel_dimensions() {
         "empty update does not delete existing data",
     );
 }
+
+#[test]
+fn attach_after_first_tab_closed() {
+    // ensure https://github.com/zellij-org/zellij/issues/1645 is fixed
+    let size = Size {
+        cols: 121,
+        rows: 20,
+    };
+    let mut screen = create_new_screen(size);
+
+    new_tab(&mut screen, 1);
+    {
+        let active_tab = screen.get_active_tab_mut(1).unwrap();
+        active_tab.new_pane(PaneId::Terminal(2), Some(1));
+        active_tab.toggle_active_pane_fullscreen(1);
+    }
+    new_tab(&mut screen, 2);
+
+    screen.close_tab_at_index(0);
+    screen.remove_client(1);
+    screen.add_client(1);
+}