tiletopia/src-tauri/src/pty.rs

//! PTY backend. Spawns a shell (`wsl.exe`, `powershell.exe`, or `ssh.exe`)
//! through portable-pty, reads its output on a background thread, and
//! forwards chunks to the frontend as `pane://{id}/data` events.

use std::collections::{HashMap, VecDeque};
use std::io::{Read, Write};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{anyhow, Context, Result};
use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use parking_lot::Mutex;
use portable_pty::{native_pty_system, CommandBuilder, MasterPty, PtySize};
use serde::{Deserialize, Serialize};
use tauri::{AppHandle, Emitter};

use crate::creds;

pub type PaneId = u64;

/// Discriminated union describing what to spawn into a fresh PTY. Serialized
/// as `{ kind: "wsl" | "powershell" | "ssh", ... }` from the frontend.
/// Also reused as the schema for the MCP `spawn_pane` tool — `JsonSchema`
/// lets rmcp render it for Claude; `Serialize` lets the backend bounce it
/// back into the `mcp://request` event payload for the frontend handler.
#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub enum SpawnSpec {
    Wsl {
        distro: Option<String>,
        cwd: Option<String>,
    },
    Powershell,
    Ssh {
        host: String,
        user: Option<String>,
        port: Option<u16>,
        #[serde(rename = "identityFile")]
        identity_file: Option<String>,
        #[serde(rename = "jumpHost")]
        jump_host: Option<String>,
        #[serde(rename = "extraArgs")]
        extra_args: Option<Vec<String>>,
        /// SshHost.id (if any) — backend uses this to fetch a saved
        /// password from keyring at spawn time. Never sent back to the
        /// frontend.
        #[serde(rename = "hostId")]
        host_id: Option<String>,
    },
}

/// Type alias for the shared writer handle. Wrapped in Arc<Mutex<>> so the
/// reader thread can also take it briefly to autotype a saved password at
/// the SSH prompt.
type SharedWriter = Arc<Mutex<Box<dyn Write + Send>>>;

/// Per-pane scrollback ring exposed to the MCP server. Capped — we drop the
/// oldest bytes when full. `seq` is a monotonic byte counter that wraps at
/// u64; the MCP `read_pane` tool uses it for incremental polling and the
/// `wait_for_idle` tool uses it to detect silence.
pub const PANE_RING_CAPACITY: usize = 256 * 1024;

pub struct PaneRing {
    buf: VecDeque<u8>,
    seq: u64,
}

impl PaneRing {
    fn new() -> Self {
        Self {
            buf: VecDeque::with_capacity(PANE_RING_CAPACITY),
            seq: 0,
        }
    }

    fn push(&mut self, bytes: &[u8]) {
        for &b in bytes {
            if self.buf.len() == PANE_RING_CAPACITY {
                self.buf.pop_front();
            }
            self.buf.push_back(b);
        }
        self.seq = self.seq.wrapping_add(bytes.len() as u64);
    }

    /// Snapshot: current contents (oldest-first) + the seq counter.
    pub fn snapshot(&self) -> (Vec<u8>, u64) {
        (self.buf.iter().copied().collect(), self.seq)
    }
}

/// What we keep alive for each spawned PTY.
///
/// `master` stays in scope to keep the PTY alive; we never write through it
/// directly (we use `writer` instead) and we never read through it directly
/// (the reader thread holds its own clone via `try_clone_reader`).
struct PaneHandle {
    #[allow(dead_code)]
    master: Box<dyn MasterPty + Send>,
    writer: SharedWriter,
    #[allow(dead_code)]
    child: Box<dyn portable_pty::Child + Send + Sync>,
    /// Same Arc the reader thread appends into; the MCP server reads via
    /// {@link PtyManager::ring}.
    ring: Arc<Mutex<PaneRing>>,
}

pub struct PtyManager {
    panes: Mutex<HashMap<PaneId, PaneHandle>>,
    next_id: AtomicU64,
    /// Per-pane "this PTY is mid-transfer between windows; do not kill it
    /// even if some window's XtermPane unmounts" refcount. Incremented by
    /// {@link mark_transferring} when a transfer begins; decremented by
    /// {@link claim} when the target window finishes mounting. While >0,
    /// {@link kill} is a no-op for that id.
    ///
    /// Refcount (vs. plain flag) so concurrent transfers — or the rare
    /// case where a transfer is retried before the previous one fully
    /// releases — don't drop the suppression early.
    transferring: Mutex<HashMap<PaneId, u32>>,
}

impl PtyManager {
    pub fn new() -> Self {
        Self {
            panes: Mutex::new(HashMap::new()),
            next_id: AtomicU64::new(1),
            transferring: Mutex::new(HashMap::new()),
        }
    }

    /// Bump the transferring refcount for a pane. While >0, {@link kill} is
    /// a no-op so the source window's React unmount-cleanup can't tear
    /// down the PTY mid-transfer.
    pub fn mark_transferring(&self, id: PaneId) {
        *self.transferring.lock().entry(id).or_insert(0) += 1;
    }

    /// Decrement the transferring refcount. When it reaches zero the entry
    /// is removed and {@link kill} can act on this pane again.
    pub fn claim(&self, id: PaneId) {
        let mut map = self.transferring.lock();
        if let Some(rc) = map.get_mut(&id) {
            if *rc > 1 {
                *rc -= 1;
            } else {
                map.remove(&id);
            }
        }
    }

    /// Spawn the shell described by `spec` into a fresh PTY. Returns the
    /// new pane id; a background thread immediately starts reading and
    /// emits `pane://{id}/data` events.
    pub fn spawn(
        &self,
        app: AppHandle,
        spec: SpawnSpec,
        cols: u16,
        rows: u16,
    ) -> Result<PaneId> {
        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows,
                cols,
                pixel_width: 0,
                pixel_height: 0,
            })
            .context("openpty failed")?;

        // Look up any saved password BEFORE building the command (cheap, no
        // bytes-on-the-wire involved). If this is an SSH spawn with a host
        // id and the user has stored a credential, the reader thread will
        // autotype it when ssh prompts.
        let saved_password = match &spec {
            SpawnSpec::Ssh { host_id: Some(id), .. } => match creds::get(id) {
                Ok(p) => p,
                Err(e) => {
                    tracing::warn!("keyring lookup for {id} failed: {e}");
                    None
                }
            },
            _ => None,
        };

        let (cmd, spawn_err) = build_command(&spec)?;
        let child = pair.slave.spawn_command(cmd).context(spawn_err)?;

        // We need to keep the master alive (drop = close the PTY), but we
        // also need the reader and writer split from it.
        let mut reader = pair
            .master
            .try_clone_reader()
            .context("try_clone_reader failed")?;
        let writer_raw = pair
            .master
            .take_writer()
            .context("take_writer failed")?;
        let writer: SharedWriter = Arc::new(Mutex::new(writer_raw));
        let ring: Arc<Mutex<PaneRing>> = Arc::new(Mutex::new(PaneRing::new()));

        let id = self.next_id.fetch_add(1, Ordering::Relaxed);

        self.panes.lock().insert(
            id,
            PaneHandle {
                master: pair.master,
                writer: writer.clone(),
                child,
                ring: ring.clone(),
            },
        );

        // Reader thread: pump bytes -> base64 -> emit. Also handles the
        // password-prompt autotype state machine if `saved_password` is set,
        // and pushes raw bytes into the per-pane scrollback ring for the
        // MCP server to read.
        let app_for_reader = app.clone();
        let event_name = format!("pane://{id}/data");
        let writer_for_reader = writer.clone();
        let ring_for_reader = ring.clone();
        std::thread::spawn(move || {
            let mut buf = [0u8; 8192];
            let mut pw_state = PasswordState::from(saved_password);
            loop {
                match reader.read(&mut buf) {
                    Ok(0) => {
                        tracing::info!("pane {id}: EOF");
                        let _ = app_for_reader.emit(&format!("pane://{id}/exit"), ());
                        break;
                    }
                    Ok(n) => {
                        // Try to autotype before emitting so we don't wait
                        // on the renderer; pw_state mutates here.
                        pw_state.observe(&buf[..n], &writer_for_reader, id);

                        // Mirror bytes into the scrollback ring (MCP source).
                        ring_for_reader.lock().push(&buf[..n]);

                        let chunk_b64 = B64.encode(&buf[..n]);
                        if let Err(e) =
                            app_for_reader.emit(&event_name, DataChunk { b64: chunk_b64 })
                        {
                            tracing::warn!("emit failed for pane {id}: {e}");
                        }
                    }
                    Err(e) => {
                        tracing::warn!("pane {id} read error: {e}");
                        let _ = app_for_reader.emit(&format!("pane://{id}/exit"), ());
                        break;
                    }
                }
            }
        });

        Ok(id)
    }

    pub fn write(&self, id: PaneId, bytes: &[u8]) -> Result<()> {
        let writer = {
            let panes = self.panes.lock();
            let pane = panes
                .get(&id)
                .ok_or_else(|| anyhow!("no pane with id {id}"))?;
            pane.writer.clone()
        };
        let mut w = writer.lock();
        w.write_all(bytes).context("pty write failed")?;
        w.flush().ok();
        Ok(())
    }

    pub fn resize(&self, id: PaneId, cols: u16, rows: u16) -> Result<()> {
        let panes = self.panes.lock();
        let pane = panes
            .get(&id)
            .ok_or_else(|| anyhow!("no pane with id {id}"))?;
        pane.master
            .resize(PtySize {
                rows,
                cols,
                pixel_width: 0,
                pixel_height: 0,
            })
            .context("pty resize failed")?;
        Ok(())
    }

    pub fn kill(&self, id: PaneId) -> Result<()> {
        // If a transfer is in flight for this pane, suppress the kill so
        // the source window's unmount-cleanup can't race the target
        // window's mount-claim. The target's claim() will decrement the
        // refcount; the next caller of kill() (if any) will actually kill.
        if self.transferring.lock().contains_key(&id) {
            tracing::debug!("pty kill suppressed during transfer for pane {id}");
            return Ok(());
        }
        let mut panes = self.panes.lock();
        if let Some(mut pane) = panes.remove(&id) {
            // Best-effort: ask the child to die. Dropping `master` after this
            // closes the PTY which will unblock the reader thread.
            let _ = pane.child.kill();
        }
        Ok(())
    }

    /// Borrow the per-pane scrollback ring. Returns None if the pane has
    /// been killed. The Arc lets callers hold the ring even after the
    /// PaneHandle is dropped (reader thread will stop pushing into it).
    pub fn ring(&self, id: PaneId) -> Option<Arc<Mutex<PaneRing>>> {
        self.panes.lock().get(&id).map(|p| p.ring.clone())
    }
}

#[derive(Serialize, Clone)]
struct DataChunk {
    b64: String,
}

// ---- command construction ---------------------------------------------------

/// Reject hostnames / usernames that would let an attacker smuggle in a
/// flag (`-oProxyCommand=...`) or a shell metacharacter via OpenSSH's token
/// expansion. We additionally pass `--` before the host on the command line,
/// but rejecting up front gives a clearer error and avoids ever handing the
/// bad value to ssh.exe.
pub fn validate_ssh_token(label: &str, value: &str) -> Result<()> {
    if value.is_empty() {
        return Err(anyhow!("ssh: {label} must not be empty"));
    }
    if value.starts_with('-') {
        return Err(anyhow!("ssh: {label} must not start with '-' (got {value:?})"));
    }
    if value.chars().any(|c| c.is_control() || c == '\n' || c == '\r') {
        return Err(anyhow!("ssh: {label} must not contain control characters"));
    }
    Ok(())
}

fn build_command(spec: &SpawnSpec) -> Result<(CommandBuilder, &'static str)> {
    match spec {
        SpawnSpec::Wsl { distro, cwd } => {
            let mut c = CommandBuilder::new("wsl.exe");
            if let Some(d) = distro.as_deref() {
                c.arg("-d");
                c.arg(d);
            }
            // Default new panes to the WSL user's home (~) rather than the
            // Windows-side cwd we inherit from the launcher (typically
            // C:\Users\<you>, which shows up as /mnt/c/Users/<you> inside WSL).
            // wsl.exe resolves `~` against the distro's default shell.
            let resolved_cwd = cwd.as_deref().unwrap_or("~");
            c.arg("--cd");
            c.arg(resolved_cwd);
            Ok((c, "failed to spawn wsl.exe; is WSL installed?"))
        }
        SpawnSpec::Powershell => {
            // cwd intentionally ignored — see commit history.
            let mut c = CommandBuilder::new("powershell.exe");
            c.arg("-NoLogo");
            Ok((c, "failed to spawn powershell.exe"))
        }
        SpawnSpec::Ssh {
            host,
            user,
            port,
            identity_file,
            jump_host,
            extra_args,
            // Read in `spawn()` to look up the saved password; not needed
            // when building the command line.
            host_id: _,
        } => {
            validate_ssh_token("host", host)?;
            if let Some(u) = user.as_deref() {
                validate_ssh_token("user", u)?;
            }
            if let Some(jh) = jump_host.as_deref() {
                validate_ssh_token("jump host", jh)?;
            }

            let mut c = CommandBuilder::new("ssh.exe");
            // ssh would auto-detect a tty here, but force it explicitly so
            // remote-side TUI apps don't accidentally see a non-tty stdin.
            c.arg("-t");
            if let Some(u) = user.as_deref() {
                c.arg("-l");
                c.arg(u);
            }
            if let Some(p) = port {
                c.arg("-p");
                c.arg(p.to_string());
            }
            if let Some(idf) = identity_file.as_deref() {
                c.arg("-i");
                c.arg(idf);
            }
            if let Some(jh) = jump_host.as_deref() {
                c.arg("-J");
                c.arg(jh);
            }
            if let Some(extra) = extra_args.as_deref() {
                for a in extra {
                    c.arg(a);
                }
            }
            // `--` ends option parsing — a hostname starting with `-` can't
            // smuggle in flags via OpenSSH's option parser.
            c.arg("--");
            c.arg(host);
            // Some Windows OpenSSH builds otherwise advertise a TERM the
            // remote side doesn't recognise; xterm.js speaks xterm-256color.
            c.env("TERM", "xterm-256color");
            Ok((c, "failed to spawn ssh.exe; is OpenSSH installed?"))
        }
    }
}

// ---- password-prompt autotype ----------------------------------------------

/// How long after spawn we keep watching for a password prompt. If nothing
/// matches in this window, we disarm and never autotype — so a remote shell
/// that prints "password" hours later can't get our credential injected.
const PASSWORD_AUTOTYPE_WINDOW: Duration = Duration::from_secs(30);
/// Sliding window of recent PTY output we scan for the prompt. Keeps the
/// scan bounded; matches don't need much context.
const PROMPT_SCAN_TAIL: usize = 256;

enum PasswordState {
    Disabled,
    Armed {
        password: String,
        deadline: Instant,
        tail: Vec<u8>,
    },
}

impl PasswordState {
    fn from(password: Option<String>) -> Self {
        match password {
            None => Self::Disabled,
            Some(p) => Self::Armed {
                password: p,
                deadline: Instant::now() + PASSWORD_AUTOTYPE_WINDOW,
                tail: Vec::with_capacity(PROMPT_SCAN_TAIL * 2),
            },
        }
    }

    /// Called for each chunk of PTY output. Mutates state — once we write
    /// the password (or time out) the state collapses to Disabled and this
    /// becomes a no-op for the rest of the connection.
    fn observe(&mut self, chunk: &[u8], writer: &SharedWriter, pane_id: PaneId) {
        let (password, tail, deadline) = match self {
            PasswordState::Disabled => return,
            PasswordState::Armed { password, tail, deadline } => (password, tail, deadline),
        };

        if Instant::now() > *deadline {
            *self = PasswordState::Disabled;
            return;
        }

        tail.extend_from_slice(chunk);
        if tail.len() > PROMPT_SCAN_TAIL {
            let drop = tail.len() - PROMPT_SCAN_TAIL;
            tail.drain(..drop);
        }

        if !looks_like_password_prompt(tail) {
            return;
        }

        // Match — write the password + Enter, then collapse to Disabled.
        let mut w = writer.lock();
        if let Err(e) = w.write_all(password.as_bytes()) {
            tracing::warn!("pane {pane_id}: password autotype write failed: {e}");
        }
        let _ = w.write_all(b"\n");
        let _ = w.flush();
        *self = PasswordState::Disabled;
    }
}

fn looks_like_password_prompt(buf: &[u8]) -> bool {
    // OpenSSH prompts: `<user>@<host>'s password:`, `Permission denied,
    // please try again. password:`, `Enter passphrase for key '...':`.
    // Lowercase the recent tail and substring-match — cheap and good enough.
    let s = String::from_utf8_lossy(buf).to_ascii_lowercase();
    s.contains("password:") || s.contains("passphrase")
}

// ---- distro enumeration -----------------------------------------------------

/// Run a process without flashing a console window on Windows.
fn quiet_command(program: &str) -> std::process::Command {
    let mut c = std::process::Command::new(program);
    #[cfg(windows)]
    {
        use std::os::windows::process::CommandExt;
        const CREATE_NO_WINDOW: u32 = 0x08000000;
        c.creation_flags(CREATE_NO_WINDOW);
    }
    c
}

/// `wsl.exe -l -q` lists installed distros (one per line, UTF-16LE).
/// Returns Ok(empty) on non-Windows or when wsl.exe isn't on PATH.
pub fn list_wsl_distros() -> Result<Vec<String>> {
    if !cfg!(windows) {
        return Ok(Vec::new());
    }

    let out = match quiet_command("wsl.exe").args(["-l", "-q"]).output() {
        Ok(o) => o,
        Err(e) => {
            tracing::debug!("wsl.exe not available: {e}");
            return Ok(Vec::new());
        }
    };

    if !out.status.success() {
        return Ok(Vec::new());
    }

    let raw_u16: Vec<u16> = out
        .stdout
        .chunks_exact(2)
        .map(|b| u16::from_le_bytes([b[0], b[1]]))
        .collect();
    let decoded = String::from_utf16_lossy(&raw_u16);

    let distros: Vec<String> = decoded
        .lines()
        .map(|l| {
            l.trim_matches(|c: char| c == '\u{FEFF}' || c.is_whitespace())
                .to_string()
        })
        .filter(|l| !l.is_empty())
        .collect();

    Ok(distros)
}