tiletopia/src-tauri/src/pty.rs

//! PTY backend. Spawns `wsl.exe` (or any command) 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;
use std::io::{Read, Write};
use std::sync::atomic::{AtomicU64, Ordering};

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

pub type PaneId = u64;

/// 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: Box<dyn Write + Send>,
    #[allow(dead_code)]
    child: Box<dyn portable_pty::Child + Send + Sync>,
}

pub struct PtyManager {
    panes: Mutex<HashMap<PaneId, PaneHandle>>,
    next_id: AtomicU64,
}

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

    /// Spawn `wsl.exe` (optionally `-d <distro>`, optionally `--cd <cwd>`).
    /// Returns the new pane id. A background thread starts reading the PTY
    /// immediately and emits `pane://{id}/data` events.
    pub fn spawn_wsl(
        &self,
        app: AppHandle,
        distro: Option<String>,
        cwd: Option<String>,
        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")?;

        let mut cmd = CommandBuilder::new("wsl.exe");
        if let Some(d) = distro.as_deref() {
            cmd.arg("-d");
            cmd.arg(d);
        }
        if let Some(c) = cwd.as_deref() {
            cmd.arg("--cd");
            cmd.arg(c);
        }
        // Force a login shell so .bashrc etc. run and PATH is populated.
        // wsl.exe without an explicit command launches the default shell
        // interactively, which is exactly what we want.

        let child = pair
            .slave
            .spawn_command(cmd)
            .context("failed to spawn wsl.exe; is WSL installed?")?;

        // 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 = pair
            .master
            .take_writer()
            .context("take_writer failed")?;

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

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

        // Reader thread: pump bytes -> base64 -> emit.
        let app_for_reader = app.clone();
        let event_name = format!("pane://{id}/data");
        std::thread::spawn(move || {
            let mut buf = [0u8; 8192];
            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) => {
                        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 mut panes = self.panes.lock();
        let pane = panes
            .get_mut(&id)
            .ok_or_else(|| anyhow!("no pane with id {id}"))?;
        pane.writer.write_all(bytes).context("pty write failed")?;
        pane.writer.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<()> {
        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(())
    }
}

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

// ---- 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)
}