bssh 3.0.1

// Copyright 2025 Lablup Inc. and Jeongkyu Shin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Core PTY session management implementation

use super::constants::*;
use super::escape_filter::EscapeSequenceFilter;
use super::output_delivery::deliver_session_output;
use super::raw_input_task;
use super::terminal_modes::configure_terminal_modes;
use crate::pty::{
    PtyConfig, PtyMessage, PtyState,
    terminal::{TerminalOps, TerminalStateGuard},
};
use anyhow::{Context, Result};
use russh::{Channel, ChannelMsg, client::Msg};
use std::io::{self, IsTerminal, Write};
use tokio::sync::{mpsc, watch};
use tokio::time::Duration;

#[derive(Debug, Clone, PartialEq, Eq)]
enum SessionSetupStep {
    Environment(usize),
    TerminalEnvironment,
    RequestPty,
    RequestShell,
}

fn session_setup_plan(config: &PtyConfig) -> Vec<SessionSetupStep> {
    let mut steps = (0..config.environment.len())
        .map(SessionSetupStep::Environment)
        .collect::<Vec<_>>();
    if !config.disable_pty {
        steps.push(SessionSetupStep::TerminalEnvironment);
        steps.push(SessionSetupStep::RequestPty);
    }
    steps.push(SessionSetupStep::RequestShell);
    steps
}

/// A PTY session managing the bidirectional communication between
/// local terminal and remote SSH session.
pub struct PtySession {
    /// Unique session identifier
    pub session_id: usize,
    /// SSH channel for communication
    channel: Channel<Msg>,
    /// PTY configuration
    config: PtyConfig,
    /// Current session state
    state: PtyState,
    /// Terminal state guard for proper cleanup
    terminal_guard: Option<TerminalStateGuard>,
    /// Cancellation signal for graceful shutdown
    cancel_tx: watch::Sender<bool>,
    cancel_rx: watch::Receiver<bool>,
    /// Message channels for internal communication (bounded to prevent memory exhaustion)
    msg_tx: Option<mpsc::Sender<PtyMessage>>,
    msg_rx: Option<mpsc::Receiver<PtyMessage>>,
    /// Filter for terminal escape sequence responses
    escape_filter: EscapeSequenceFilter,
}

impl PtySession {
    /// Create a new PTY session
    pub async fn new(session_id: usize, channel: Channel<Msg>, config: PtyConfig) -> Result<Self> {
        // Use bounded channel with reasonable buffer size to prevent memory exhaustion
        let (msg_tx, msg_rx) = mpsc::channel(PTY_MESSAGE_CHANNEL_SIZE);

        // Create cancellation channel
        let (cancel_tx, cancel_rx) = watch::channel(false);

        Ok(Self {
            session_id,
            channel,
            config,
            state: PtyState::Inactive,
            terminal_guard: None,
            cancel_tx,
            cancel_rx,
            msg_tx: Some(msg_tx),
            msg_rx: Some(msg_rx),
            escape_filter: EscapeSequenceFilter::new(),
        })
    }

    /// Get the current session state
    pub fn state(&self) -> PtyState {
        self.state
    }

    pub(crate) fn requests_remote_pty(&self) -> bool {
        !self.config.disable_pty
    }

    /// Initialize the raw session, optionally requesting a remote PTY.
    pub async fn initialize(&mut self) -> Result<()> {
        self.state = PtyState::Initializing;

        for step in session_setup_plan(&self.config) {
            match step {
                SessionSetupStep::Environment(index) => {
                    let (name, value) = &self.config.environment[index];
                    self.channel
                        .set_env(false, name, value)
                        .await
                        .with_context(|| {
                            format!("Failed to send SSH environment variable {name}")
                        })?;
                }
                SessionSetupStep::TerminalEnvironment => {
                    if let Err(error) = self
                        .channel
                        .set_env(false, "TERM", &self.config.term_type)
                        .await
                    {
                        tracing::debug!("Server did not accept TERM environment variable: {error}");
                    }
                    if let Err(error) = self.channel.set_env(false, "COLORTERM", "truecolor").await
                    {
                        tracing::trace!(
                            "Server did not accept COLORTERM environment variable: {error}"
                        );
                    }
                }
                SessionSetupStep::RequestPty => {
                    let (width, height) = crate::pty::utils::get_terminal_size()?;
                    let terminal_modes = configure_terminal_modes();
                    self.channel
                        .request_pty(
                            false,
                            &self.config.term_type,
                            width,
                            height,
                            0,
                            0,
                            &terminal_modes,
                        )
                        .await
                        .with_context(|| "Failed to request PTY on SSH channel")?;
                }
                SessionSetupStep::RequestShell => {
                    self.channel
                        .request_shell(false)
                        .await
                        .with_context(|| "Failed to request shell on SSH channel")?;
                }
            }
        }

        self.state = PtyState::Active;
        tracing::debug!(
            "Raw SSH session {} initialized (remote PTY: {})",
            self.session_id,
            !self.config.disable_pty
        );
        Ok(())
    }

    /// Run the main PTY session loop
    pub async fn run(&mut self) -> Result<()> {
        if self.state == PtyState::Inactive {
            self.initialize().await?;
        }

        if self.state != PtyState::Active {
            anyhow::bail!("PTY session is not in active state");
        }

        // Match OpenSSH client_loop semantics: only a session with a remote PTY
        // changes the local terminal to raw mode. No-PTY shells retain the OS
        // line discipline and forward the bytes that stdin delivers unchanged.
        let mut terminal_guard = if io::stdin().is_terminal() && io::stdout().is_terminal() {
            if self.config.disable_pty {
                TerminalStateGuard::new_without_raw_mode()?
            } else {
                TerminalStateGuard::new()?
            }
        } else {
            TerminalStateGuard::new_without_raw_mode()?
        };
        let pending_input = terminal_guard.take_pending_input();
        self.terminal_guard = Some(terminal_guard);

        // Enable mouse support if requested
        if self.config.enable_mouse && !self.config.disable_pty {
            TerminalOps::enable_mouse()?;
        }

        // Get message receiver
        let mut msg_rx = self
            .msg_rx
            .take()
            .ok_or_else(|| anyhow::anyhow!("Message receiver already taken"))?;

        // A no-PTY shell keeps the raw byte stream but must not send remote
        // window-change requests.
        let resize_task = if self.config.disable_pty {
            None
        } else {
            let mut resize_signals = crate::pty::utils::setup_resize_handler()?;
            let mut cancel_for_resize = self.cancel_rx.clone();
            let resize_tx = self
                .msg_tx
                .as_ref()
                .ok_or_else(|| anyhow::anyhow!("Message sender not available"))?
                .clone();

            Some(tokio::spawn(async move {
                loop {
                    tokio::select! {
                        signal = async {
                            for signal in resize_signals.forever() {
                                if signal == signal_hook::consts::SIGWINCH {
                                    return signal;
                                }
                            }
                            signal_hook::consts::SIGWINCH
                        } => {
                            if signal == signal_hook::consts::SIGWINCH
                                && let Ok((width, height)) = crate::pty::utils::get_terminal_size()
                                && resize_tx
                                    .try_send(PtyMessage::Resize { width, height })
                                    .is_err()
                            {
                                break;
                            }
                        }
                        _ = cancel_for_resize.changed() => {
                            if *cancel_for_resize.borrow() {
                                break;
                            }
                        }
                    }
                }
            }))
        };

        // Spawn input reader task
        let input_tx = self
            .msg_tx
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Message sender not available"))?
            .clone();
        let cancel_for_input = self.cancel_rx.clone();

        // Spawn input reader in the blocking thread pool. Local escape handling
        // is an SSH PTY feature; a no-PTY shell forwards stdin without consuming
        // sequences such as "~.".
        let escape_enabled = !self.config.disable_pty;
        let input_task = tokio::task::spawn_blocking(move || {
            raw_input_task::run(input_tx, cancel_for_input, pending_input, escape_enabled);
        });

        // We'll integrate channel reading into the main loop since russh Channel doesn't clone

        // Main message handling loop using tokio::select! for efficient event multiplexing
        let mut should_terminate = false;
        let mut cancel_rx = self.cancel_rx.clone();

        // Track last activity time for connection health monitoring
        let mut last_activity = std::time::Instant::now();
        let health_check_interval = Duration::from_secs(CONNECTION_HEALTH_CHECK_INTERVAL_SECS);
        let max_idle_time = Duration::from_secs(MAX_IDLE_TIME_BEFORE_WARNING_SECS);
        let mut idle_warning_shown = false;

        while !should_terminate {
            tokio::select! {
                // Handle SSH channel messages
                msg = self.channel.wait() => {
                    // Reset activity timer on any channel activity
                    last_activity = std::time::Instant::now();
                    idle_warning_shown = false;

                    match msg {
                        Some(ChannelMsg::Data { ref data }) => {
                            let terminal_guard = &mut self.terminal_guard;
                            if let Err(e) = deliver_session_output(
                                self.config.disable_pty,
                                &mut self.escape_filter,
                                &mut io::stdout(),
                                data,
                                |delivered| {
                                    if let Some(guard) = terminal_guard.as_mut() {
                                        guard.observe_remote_output(delivered);
                                    }
                                },
                            ) {
                                tracing::error!("Failed to write to stdout: {e}");
                                should_terminate = true;
                            }
                        }
                        Some(ChannelMsg::ExtendedData { ref data, ext }) => {
                            if ext == 1 {
                                let result = if self.config.disable_pty {
                                    deliver_session_output(
                                        true,
                                        &mut self.escape_filter,
                                        &mut io::stderr(),
                                        data,
                                        |_| {},
                                    )
                                } else {
                                    let terminal_guard = &mut self.terminal_guard;
                                    deliver_session_output(
                                        false,
                                        &mut self.escape_filter,
                                        &mut io::stdout(),
                                        data,
                                        |delivered| {
                                            if let Some(guard) = terminal_guard.as_mut() {
                                                guard.observe_remote_output(delivered);
                                            }
                                        },
                                    )
                                };
                                if let Err(e) = result {
                                    tracing::error!("Failed to write stderr: {e}");
                                    should_terminate = true;
                                }
                            }
                        }
                        Some(ChannelMsg::Eof) | Some(ChannelMsg::Close) => {
                            tracing::debug!("SSH channel closed");
                            // Signal cancellation to all child tasks before terminating
                            let _ = self.cancel_tx.send(true);
                            should_terminate = true;
                        }
                        Some(_) => {
                            // Handle other channel messages if needed
                        }
                        None => {
                            // Channel ended - connection is dead
                            tracing::warn!(
                                "SSH channel returned None - connection may have dropped"
                            );
                            should_terminate = true;
                        }
                    }
                }

                // Handle local messages (input, resize, etc.)
                message = msg_rx.recv() => {
                    // Reset activity timer for local input (user is active)
                    if matches!(message, Some(PtyMessage::LocalInput(_))) {
                        last_activity = std::time::Instant::now();
                        idle_warning_shown = false;
                    }

                    match message {
                        Some(PtyMessage::LocalInput(data)) => {
                            if let Err(e) = self.channel.data(data.as_slice()).await {
                                tracing::error!("Failed to send data to SSH channel: {e}");
                                // Connection likely dead - terminate gracefully
                                crate::diagnosticln!(
                                    "\r\n[bssh] Connection lost: failed to send data to remote host\r"
                                );
                                should_terminate = true;
                            }
                        }
                        Some(PtyMessage::RemoteOutput(data)) => {
                            let terminal_guard = &mut self.terminal_guard;
                            if let Err(e) = deliver_session_output(
                                self.config.disable_pty,
                                &mut self.escape_filter,
                                &mut io::stdout(),
                                &data,
                                |delivered| {
                                    if let Some(guard) = terminal_guard.as_mut() {
                                        guard.observe_remote_output(delivered);
                                    }
                                },
                            ) {
                                tracing::error!("Failed to write to stdout: {e}");
                                should_terminate = true;
                            }
                        }
                        Some(PtyMessage::Resize { width, height }) => {
                            if !self.config.disable_pty {
                                if let Err(e) = self.channel.window_change(width, height, 0, 0).await {
                                    tracing::warn!("Failed to send window resize to remote: {e}");
                                } else {
                                    tracing::debug!("Terminal resized to {width}x{height}");
                                }
                            }
                        }
                        Some(PtyMessage::Terminate) => {
                            tracing::debug!("PTY session {} terminating", self.session_id);
                            should_terminate = true;
                        }
                        Some(PtyMessage::Error(error)) => {
                            tracing::error!("PTY error: {error}");
                            should_terminate = true;
                        }
                        None => {
                            // Message channel closed
                            should_terminate = true;
                        }
                    }
                }

                // Handle cancellation signal
                _ = cancel_rx.changed() => {
                    if *cancel_rx.borrow() {
                        tracing::debug!("PTY session {} received cancellation signal", self.session_id);
                        should_terminate = true;
                    }
                }

                // Periodic health check to detect dead connections
                _ = tokio::time::sleep(health_check_interval) => {
                    let idle_duration = last_activity.elapsed();

                    // Check if the session has been idle for too long
                    if idle_duration > max_idle_time && !idle_warning_shown {
                        tracing::debug!(
                            "PTY session {} idle for {:?}, connection may be stale",
                            self.session_id,
                            idle_duration
                        );
                        // Don't terminate, but log for debugging
                        // SSH keepalive should handle actual connection detection
                        idle_warning_shown = true;
                    }

                    // Periodic trace logging for debugging long sessions
                    tracing::trace!(
                        "PTY session {} health check: idle for {:?}",
                        self.session_id,
                        idle_duration
                    );
                }
            }
        }

        // Signal cancellation to all tasks
        let _ = self.cancel_tx.send(true);

        // Tasks will exit gracefully on cancellation
        // No need to abort since they check cancellation signal

        // Wait for tasks to complete gracefully with select!
        let resize_cleanup = async {
            if let Some(task) = resize_task {
                let _ = task.await;
            } else {
                std::future::pending::<()>().await;
            }
        };
        let _ = tokio::time::timeout(Duration::from_millis(TASK_CLEANUP_TIMEOUT_MS), async {
            tokio::select! {
                _ = resize_cleanup => {},
                _ = input_task => {},
                _ = tokio::time::sleep(Duration::from_millis(TASK_CLEANUP_TIMEOUT_MS)) => {}
            }
        })
        .await;

        // Disable mouse support if we enabled it
        if self.config.enable_mouse && !self.config.disable_pty {
            let _ = TerminalOps::disable_mouse();
        }

        // IMPORTANT: Explicitly restore terminal state by dropping the guard
        // The guard's drop implementation handles synchronized cleanup
        self.terminal_guard = None;

        // Flush stdout to ensure all output is written
        let _ = io::stdout().flush();

        self.state = PtyState::Closed;
        Ok(())
    }

    /// Shutdown the PTY session
    pub async fn shutdown(&mut self) -> Result<()> {
        self.state = PtyState::ShuttingDown;

        // Signal cancellation to all tasks
        let _ = self.cancel_tx.send(true);

        // Send EOF to close the channel gracefully
        if let Err(e) = self.channel.eof().await {
            tracing::warn!("Failed to send EOF to SSH channel: {e}");
        }

        // Drop terminal guard to restore terminal state
        self.terminal_guard = None;

        self.state = PtyState::Closed;
        Ok(())
    }
}

impl Drop for PtySession {
    fn drop(&mut self) {
        // Signal cancellation to all tasks when session is dropped
        let _ = self.cancel_tx.send(true);
        // Terminal guard will be dropped automatically, restoring terminal state
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn no_pty_setup_sends_policy_env_once_before_one_shell_request() {
        let config = PtyConfig {
            disable_pty: true,
            environment: vec![
                ("FIRST".into(), "one".into()),
                ("SECOND".into(), "two".into()),
            ],
            ..Default::default()
        };
        let plan = session_setup_plan(&config);
        assert_eq!(
            plan,
            [
                SessionSetupStep::Environment(0),
                SessionSetupStep::Environment(1),
                SessionSetupStep::RequestShell,
            ]
        );
        assert_eq!(
            plan.iter()
                .filter(|step| matches!(step, SessionSetupStep::RequestPty))
                .count(),
            0
        );
        assert_eq!(
            plan.iter()
                .filter(|step| matches!(step, SessionSetupStep::RequestShell))
                .count(),
            1
        );
        let environment_positions = plan
            .iter()
            .enumerate()
            .filter_map(|(index, step)| {
                matches!(step, SessionSetupStep::Environment(_)).then_some(index)
            })
            .collect::<Vec<_>>();
        let shell_position = plan
            .iter()
            .position(|step| matches!(step, SessionSetupStep::RequestShell))
            .unwrap();
        assert!(
            environment_positions
                .iter()
                .all(|position| *position < shell_position)
        );
    }

    #[test]
    fn pty_setup_orders_environment_terminal_pty_and_shell() {
        let config = PtyConfig {
            environment: vec![("POLICY".into(), "value".into())],
            ..Default::default()
        };
        assert_eq!(
            session_setup_plan(&config),
            [
                SessionSetupStep::Environment(0),
                SessionSetupStep::TerminalEnvironment,
                SessionSetupStep::RequestPty,
                SessionSetupStep::RequestShell,
            ]
        );
    }
}