path: root/src/tests/e2e/remote_runner.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use vte;
use zellij_server::panes::sixel::SixelImageStore;
use zellij_server::panes::{LinkHandler, TerminalPane};
use zellij_utils::data::{Palette, Style};
use zellij_utils::pane_size::{Dimension, PaneGeom, Size, SizeInPixels};

use ssh2::Session;
use std::io::prelude::*;
use std::net::TcpStream;

use std::cell::RefCell;
use std::rc::Rc;

const ZELLIJ_EXECUTABLE_LOCATION: &str = "/usr/src/zellij/zellij";
const SET_ENV_VARIABLES: &str = "EDITOR=/usr/bin/vi";
const ZELLIJ_DATA_DIR: &str = "/usr/src/zellij/e2e-data";
const ZELLIJ_FIXTURE_PATH: &str = "/usr/src/zellij/fixtures";
const CONNECTION_STRING: &str = "127.0.0.1:2222";
const CONNECTION_USERNAME: &str = "test";
const CONNECTION_PASSWORD: &str = "test";
const SESSION_NAME: &str = "e2e-test";
const RETRIES: usize = 10;
const CLEANUP_DONE_MARKER: &str = "__ZELLIJ_CLEANUP_DONE__";

fn ssh_connect() -> ssh2::Session {
    let tcp = TcpStream::connect(CONNECTION_STRING).unwrap();
    let mut sess = Session::new().unwrap();
    sess.set_tcp_stream(tcp);
    sess.handshake().unwrap();
    sess.userauth_password(CONNECTION_USERNAME, CONNECTION_PASSWORD)
        .unwrap();
    sess
}

fn setup_remote_environment(channel: &mut ssh2::Channel, win_size: Size) {
    let (columns, rows) = (win_size.cols as u32, win_size.rows as u32);
    channel
        .request_pty("xterm", None, Some((columns, rows, 0, 0)))
        .unwrap();
    channel.shell().unwrap();
    channel.write_all(b"export PS1=\"$ \"\n").unwrap();
    channel.flush().unwrap();
}

fn stop_zellij(channel: &mut ssh2::Channel) {
    // here we remove the status-bar-tips cache to make sure only the quicknav tip is loaded
    channel
        .write_all(b"find /tmp | grep status-bar-tips | xargs rm\n")
        .unwrap();
    channel.write_all(b"killall -KILL zellij\n").unwrap();
    channel.write_all(b"rm -rf /tmp/*\n").unwrap(); // remove temporary artifacts from previous
                                                    // tests
    channel.write_all(b"rm -rf /tmp/*\n").unwrap(); // remove temporary artifacts from previous
    channel.write_all(b"rm -rf /tmp/*\n").unwrap(); // remove temporary artifacts from previous
    channel
        .write_all(b"rm -rf ~/.cache/zellij/*/session_info\n")
        .unwrap();
    channel
        .write_all(b"rm -rf ~/.cache/zellij/permissions.kdl\n")
        .unwrap();
    // create an arch-independent symlink so the binary path in snapshots is stable across
    // x86_64 (CI) and aarch64 (Apple Silicon local dev)
    channel
        .write_all(
            b"ln -sf /usr/src/zellij/$(uname -m)-unknown-linux-musl/release/zellij /usr/src/zellij/zellij\n",
        )
        .unwrap();
}

fn start_zellij(channel: &mut ssh2::Channel) {
    stop_zellij(channel);
    channel
        .write_all(
            format!(
                "{} {} --session {} --data-dir {} options --show-release-notes false --show-startup-tips false\n",
                SET_ENV_VARIABLES, ZELLIJ_EXECUTABLE_LOCATION, SESSION_NAME, ZELLIJ_DATA_DIR
            )
            .as_bytes(),
        )
        .unwrap();
    channel.flush().unwrap();
}

fn start_zellij_mirrored_session_with_layout(channel: &mut ssh2::Channel, layout_file_name: &str) {
    stop_zellij(channel);
    channel
        .write_all(
            format!(
                "{} {} --session {} --data-dir {} --new-session-with-layout {} options --show-release-notes false --show-startup-tips false --mirror-session true --serialization-interval 1\n",
                SET_ENV_VARIABLES,
                ZELLIJ_EXECUTABLE_LOCATION,
                SESSION_NAME,
                ZELLIJ_DATA_DIR,
                format!("{}/{}", ZELLIJ_FIXTURE_PATH, layout_file_name)
            )
            .as_bytes(),
        )
        .unwrap();
    channel.flush().unwrap();
}

fn wait_for_startup(last_snapshot: &Arc<Mutex<String>>) {
    let start = std::time::Instant::now();
    let timeout = std::time::Duration::from_secs(10);
    loop {
        if last_snapshot.lock().unwrap().contains("Ctrl +") {
            break;
        }
        if start.elapsed() > timeout {
            break; // timed out — let the test proceed and fail naturally
        }
        std::thread::sleep(std::time::Duration::from_millis(50));
    }
}

fn wait_for_shell_output(channel: &mut ssh2::Channel, needle: &str) {
    let start = std::time::Instant::now();
    let timeout = std::time::Duration::from_secs(10);
    let mut output = String::new();
    let mut buf = [0u8; 4096];
    loop {
        if output.contains(needle) || start.elapsed() > timeout {
            break;
        }
        match channel.read(&mut buf) {
            Ok(0) => std::thread::sleep(std::time::Duration::from_millis(50)),
            Ok(count) => output.push_str(&String::from_utf8_lossy(&buf[..count])),
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                std::thread::sleep(std::time::Duration::from_millis(50));
            },
            Err(_) => break,
        }
    }
}

fn read_from_channel(
    channel: &Arc<Mutex<ssh2::Channel>>,
    last_snapshot: &Arc<Mutex<String>>,
    cursor_coordinates: &Arc<Mutex<(usize, usize)>>,
    pane_geom: &PaneGeom,
) -> (Arc<AtomicBool>, std::thread::JoinHandle<()>) {
    let should_keep_running = Arc::new(AtomicBool::new(true));
    let thread = std::thread::Builder::new()
        .name("read_thread".into())
        .spawn({
            let pane_geom = *pane_geom;
            let should_keep_running = should_keep_running.clone();
            let channel = channel.clone();
            let last_snapshot = last_snapshot.clone();
            let cursor_coordinates = cursor_coordinates.clone();
            move || {
                let mut retries_left = 3;
                let mut should_sleep = false;
                let mut vte_parser = vte::Parser::new();
                let character_cell_size = Rc::new(RefCell::new(Some(SizeInPixels {
                    height: 21,
                    width: 8,
                })));
                let sixel_image_store = Rc::new(RefCell::new(SixelImageStore::default()));
                let debug = false;
                let arrow_fonts = true;
                let styled_underlines = true;
                let explicitly_disable_kitty_keyboard_protocol = false;
                let mut terminal_output = TerminalPane::new(
                    0,
                    pane_geom,
                    Style::default(),
                    0,
                    String::new(),
                    Rc::new(RefCell::new(LinkHandler::new())),
                    character_cell_size,
                    sixel_image_store,
                    Rc::new(RefCell::new(Palette::default())),
                    Rc::new(RefCell::new(HashMap::new())),
                    None,
                    None,
                    debug,
                    arrow_fonts,
                    styled_underlines,
                    true, // osc8_hyperlinks
                    explicitly_disable_kitty_keyboard_protocol,
                    None,
                ); // 0 is the pane index
                loop {
                    if !should_keep_running.load(Ordering::SeqCst) {
                        break;
                    }
                    if should_sleep {
                        std::thread::sleep(std::time::Duration::from_millis(100));
                        should_sleep = false;
                    }
                    let mut buf = [0u8; 1280000];
                    match channel.lock().unwrap().read(&mut buf) {
                        Ok(0) => {
                            let current_snapshot = take_snapshot(&mut terminal_output);
                            let mut last_snapshot = last_snapshot.lock().unwrap();
                            *cursor_coordinates.lock().unwrap() = terminal_output
                                .cursor_coordinates()
                                .map(|(x, y, _)| (x, y))
                                .unwrap_or((0, 0));
                            *last_snapshot = current_snapshot;
                            should_sleep = true;
                        },
                        Ok(count) => {
                            vte_parser.advance(&mut terminal_output.grid, &buf[..count]);
                            let current_snapshot = take_snapshot(&mut terminal_output);
                            let mut last_snapshot = last_snapshot.lock().unwrap();
                            *cursor_coordinates.lock().unwrap() = terminal_output
                                .grid
                                .cursor_coordinates()
                                .map(|(x, y, _)| (x, y))
                                .unwrap_or((0, 0));
                            *last_snapshot = current_snapshot;
                            should_sleep = true;
                        },
                        Err(e) => {
                            if e.kind() == std::io::ErrorKind::WouldBlock {
                                let current_snapshot = take_snapshot(&mut terminal_output);
                                let mut last_snapshot = last_snapshot.lock().unwrap();
                                *cursor_coordinates.lock().unwrap() = terminal_output
                                    .cursor_coordinates()
                                    .map(|(x, y, _)| (x, y))
                                    .unwrap_or((0, 0));
                                *last_snapshot = current_snapshot;
                                should_sleep = true;
                            } else if retries_left > 0 {
                                retries_left -= 1;
                            } else {
                                break;
                            }
                        },
                    }
                }
            }
        })
        .unwrap();
    (should_keep_running, thread)
}

pub fn take_snapshot(terminal_output: &mut TerminalPane) -> String {
    let output_lines = terminal_output.read_buffer_as_lines();
    let cursor_coordinates = terminal_output
        .cursor_coordinates()
        .and_then(|(x, y, visible)| if visible { Some((x, y)) } else { None });
    let mut snapshot = String::new();
    for (line_index, line) in output_lines.iter().enumerate() {
        for (character_index, terminal_character) in line.iter().enumerate() {
            if let Some((cursor_x, cursor_y)) = cursor_coordinates {
                if line_index == cursor_y && character_index == cursor_x {
                    snapshot.push('█');
                    continue;
                }
            }
            snapshot.push(terminal_character.character);
        }
        if line_index != output_lines.len() - 1 {
            snapshot.push('\n');
        }
    }
    snapshot
}

pub struct RemoteTerminal {
    channel: Arc<Mutex<ssh2::Channel>>,
    cursor_x: usize,
    cursor_y: usize,
    last_snapshot: Arc<Mutex<String>>,
}

impl std::fmt::Debug for RemoteTerminal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "cursor x: {}\ncursor_y: {}\ncurrent_snapshot:\n{}",
            self.cursor_x,
            self.cursor_y,
            *self.last_snapshot.lock().unwrap()
        )
    }
}

impl RemoteTerminal {
    pub fn cursor_position_is(&self, x: usize, y: usize) -> bool {
        x == self.cursor_x && y == self.cursor_y
    }
    pub fn status_bar_appears(&self) -> bool {
        self.last_snapshot.lock().unwrap().contains("Ctrl +")
            && self.last_snapshot.lock().unwrap().contains("LOCK")
    }
    pub fn ctrl_plus_appears(&self) -> bool {
        self.last_snapshot.lock().unwrap().contains("Ctrl +")
    }
    pub fn tab_bar_appears(&self) -> bool {
        self.last_snapshot.lock().unwrap().contains("Tab #1")
    }
    pub fn snapshot_contains(&self, text: &str) -> bool {
        self.last_snapshot.lock().unwrap().contains(text)
    }
    #[allow(unused)]
    pub fn current_snapshot(&self) -> String {
        // convenience method for writing tests,
        // this should only be used when developing,
        // please prefer "snapsht_contains" instead
        self.last_snapshot.lock().unwrap().clone()
    }
    #[allow(unused)]
    pub fn current_cursor_position(&self) -> String {
        // convenience method for writing tests,
        // this should only be used when developing,
        // please prefer "cursor_position_is" instead
        format!("x: {}, y: {}", self.cursor_x, self.cursor_y)
    }
    pub fn send_key(&mut self, key: &[u8]) {
        let mut channel = self.channel.lock().unwrap();
        channel.write_all(key).unwrap();
        channel.flush().unwrap();
    }
    pub fn change_size(&mut self, cols: u32, rows: u32) {
        self.channel
            .lock()
            .unwrap()
            .request_pty_size(cols, rows, Some(cols), Some(rows))
            .unwrap();
    }
    pub fn attach_to_original_session(&mut self) {
        {
            let mut channel = self.channel.lock().unwrap();
            channel
                .write_all(
                    format!("{} attach {}\n", ZELLIJ_EXECUTABLE_LOCATION, SESSION_NAME).as_bytes(),
                )
                .unwrap();
            channel.flush().unwrap();
        } // release mutex before sleeping so the reader thread can process Zellij's startup output
        std::thread::sleep(std::time::Duration::from_secs(1)); // wait until Zellij stops parsing startup ANSI codes from the terminal STDIN
    }
    pub fn send_blocking_command_through_the_cli(&mut self, command: &str) {
        let mut channel = self.channel.lock().unwrap();
        channel
            .write_all(
                format!(
                    "{} run --blocking --floating --close-on-exit -- {}",
                    ZELLIJ_EXECUTABLE_LOCATION, command
                )
                .as_bytes(),
            )
            .unwrap();
        channel.flush().unwrap();
    }
}

#[derive(Clone)]
pub struct Step {
    pub instruction: fn(RemoteTerminal) -> bool,
    pub name: &'static str,
}

pub struct RemoteRunner {
    steps: Vec<Step>,
    current_step_index: usize,
    channel: Arc<Mutex<ssh2::Channel>>,
    currently_running_step: Option<String>,
    retries_left: usize,
    retry_pause_ms: usize,
    last_snapshot: Arc<Mutex<String>>,
    cursor_coordinates: Arc<Mutex<(usize, usize)>>, // x, y
    reader_thread: (Arc<AtomicBool>, std::thread::JoinHandle<()>),
    pub test_timed_out: bool,
}

impl RemoteRunner {
    pub fn new(win_size: Size) -> Self {
        let sess = ssh_connect();
        let mut channel = sess.channel_session().unwrap();
        let mut rows = Dimension::fixed(win_size.rows);
        let mut cols = Dimension::fixed(win_size.cols);
        rows.set_inner(win_size.rows);
        cols.set_inner(win_size.cols);
        let pane_geom = PaneGeom {
            x: 0,
            y: 0,
            rows,
            cols,
            stacked: None,
            is_pinned: false,
            logical_position: None,
        };
        setup_remote_environment(&mut channel, win_size);
        start_zellij(&mut channel);
        let channel = Arc::new(Mutex::new(channel));
        let last_snapshot = Arc::new(Mutex::new(String::new()));
        let cursor_coordinates = Arc::new(Mutex::new((0, 0)));
        sess.set_blocking(false);
        let reader_thread =
            read_from_channel(&channel, &last_snapshot, &cursor_coordinates, &pane_geom);
        wait_for_startup(&last_snapshot);
        RemoteRunner {
            steps: vec![],
            channel,
            currently_running_step: None,
            current_step_index: 0,
            retries_left: RETRIES,
            retry_pause_ms: 100,
            test_timed_out: false,
            last_snapshot,
            cursor_coordinates,
            reader_thread,
        }
    }
    pub fn new_mirrored_session_with_layout(win_size: Size, layout_file_name: &str) -> Self {
        let sess = ssh_connect();
        let mut channel = sess.channel_session().unwrap();
        let mut rows = Dimension::fixed(win_size.rows);
        let mut cols = Dimension::fixed(win_size.cols);
        rows.set_inner(win_size.rows);
        cols.set_inner(win_size.cols);
        let pane_geom = PaneGeom {
            x: 0,
            y: 0,
            rows,
            cols,
            stacked: None,
            is_pinned: false,
            logical_position: None,
        };
        setup_remote_environment(&mut channel, win_size);
        start_zellij_mirrored_session_with_layout(&mut channel, layout_file_name);
        let channel = Arc::new(Mutex::new(channel));
        let last_snapshot = Arc::new(Mutex::new(String::new()));
        let cursor_coordinates = Arc::new(Mutex::new((0, 0)));
        sess.set_blocking(false);
        let reader_thread =
            read_from_channel(&channel, &last_snapshot, &cursor_coordinates, &pane_geom);
        wait_for_startup(&last_snapshot);
        RemoteRunner {
            steps: vec![],
            channel,
            currently_running_step: None,
            current_step_index: 0,
            retries_left: RETRIES,
            retry_pause_ms: 100,
            test_timed_out: false,
            last_snapshot,
            cursor_coordinates,
            reader_thread,
        }
    }
    pub fn kill_running_sessions(win_size: Size) {
        let sess = ssh_connect();
        let mut channel = sess.channel_session().unwrap();
        setup_remote_environment(&mut channel, win_size);
        stop_zellij(&mut channel);
        channel
            .write_all(format!("printf '{}\\n'\n", CLEANUP_DONE_MARKER).as_bytes())
            .unwrap();
        channel.flush().unwrap();
        sess.set_blocking(false);
        wait_for_shell_output(&mut channel, CLEANUP_DONE_MARKER);
    }
    #[allow(unused)]
    pub fn retry_pause_ms(mut self, retry_pause_ms: usize) -> Self {
        self.retry_pause_ms = retry_pause_ms;
        self
    }
    pub fn add_step(mut self, step: Step) -> Self {
        self.steps.push(step);
        self
    }
    pub fn run_next_step(&mut self) {
        if let Some(next_step) = self.steps.get(self.current_step_index) {
            println!(
                "running step: {}, retries left: {}",
                next_step.name, self.retries_left
            );
            let (cursor_x, cursor_y) = *self.cursor_coordinates.lock().unwrap();
            let remote_terminal = RemoteTerminal {
                cursor_x,
                cursor_y,
                last_snapshot: self.last_snapshot.clone(),
                channel: self.channel.clone(),
            };
            let instruction = next_step.instruction;
            self.currently_running_step = Some(String::from(next_step.name));
            if instruction(remote_terminal) {
                self.retries_left = RETRIES;
                self.current_step_index += 1;
            } else {
                self.retries_left = self.retries_left.saturating_sub(1);
                std::thread::sleep(std::time::Duration::from_millis(self.retry_pause_ms as u64));
            }
        }
    }
    pub fn steps_left(&self) -> bool {
        self.steps.get(self.current_step_index).is_some()
    }
    pub fn take_snapshot_after(&mut self, step: Step) -> String {
        let mut retries_left = RETRIES;
        let instruction = step.instruction;
        loop {
            println!(
                "taking snapshot: {}, retries left: {}",
                step.name, retries_left
            );
            if retries_left == 0 {
                self.test_timed_out = true;
                return self.last_snapshot.lock().unwrap().clone();
            }
            let (cursor_x, cursor_y) = *self.cursor_coordinates.lock().unwrap();
            let remote_terminal = RemoteTerminal {
                cursor_x,
                cursor_y,
                last_snapshot: self.last_snapshot.clone(),
                channel: self.channel.clone(),
            };
            if instruction(remote_terminal) {
                return self.last_snapshot.lock().unwrap().clone();
            } else {
                retries_left -= 1;
                std::thread::sleep(std::time::Duration::from_millis(100));
                continue;
            }
        }
    }
    pub fn run_all_steps(&mut self) {
        println!();
        loop {
            self.run_next_step();
            if !self.steps_left() {
                break;
            } else if self.retries_left == 0 {
                self.test_timed_out = true;
                break;
            }
        }
    }
}

impl Drop for RemoteRunner {
    fn drop(&mut self) {
        let _ = self.channel.lock().unwrap().close();
        let reader_thread_running = &mut self.reader_thread.0;
        reader_thread_running.store(false, Ordering::SeqCst);
    }
}