pikasTech-unidesk/src/components/backend-core/src/cli.rs

use std::env;
use std::time::Duration;

use anyhow::bail;
use base64::Engine;
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::watch;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;

use crate::http::internal_fetch_json;

async fn write_stdout(bytes: &[u8]) {
    let mut stdout = tokio::io::stdout();
    let _ = stdout.write_all(bytes).await;
    let _ = stdout.flush().await;
}

async fn write_stderr(bytes: &[u8]) {
    let mut stderr = tokio::io::stderr();
    let _ = stderr.write_all(bytes).await;
    let _ = stderr.flush().await;
}

pub async fn handle_cli() -> anyhow::Result<bool> {
    let mut args = env::args().skip(1).collect::<Vec<_>>();
    if args.is_empty() {
        return Ok(false);
    }
    match args.remove(0).as_str() {
        "--fetch-json" => {
            fetch_json_cli(args).await?;
            Ok(true)
        }
        "--ssh-broker" => {
            ssh_broker_cli(args).await?;
            Ok(true)
        }
        _ => Ok(false),
    }
}

async fn fetch_json_cli(args: Vec<String>) -> anyhow::Result<()> {
    let mut url = None;
    let mut method = "GET".to_string();
    let mut body = None;
    let mut max_response_bytes = 5_000_000_usize;
    let mut require_ok = false;
    let mut index = 0;
    while index < args.len() {
        match args[index].as_str() {
            "--require-ok" => {
                require_ok = true;
                index += 1;
            }
            "--method" => {
                method = args
                    .get(index + 1)
                    .cloned()
                    .unwrap_or_else(|| "GET".to_string());
                index += 2;
            }
            "--body-json" => {
                let raw = args
                    .get(index + 1)
                    .cloned()
                    .unwrap_or_else(|| "{}".to_string());
                body = Some(serde_json::from_str::<Value>(&raw)?);
                index += 2;
            }
            "--max-response-bytes" => {
                max_response_bytes = args
                    .get(index + 1)
                    .and_then(|value| value.parse::<usize>().ok())
                    .unwrap_or(max_response_bytes);
                index += 2;
            }
            value if url.is_none() => {
                url = Some(value.to_string());
                index += 1;
            }
            other => bail!("unexpected --fetch-json argument: {other}"),
        }
    }
    let Some(url) = url else {
        bail!("--fetch-json requires URL");
    };
    let result = internal_fetch_json(&url, &method, body, max_response_bytes).await?;
    println!("{result}");
    if require_ok && !result.get("ok").and_then(Value::as_bool).unwrap_or(false) {
        std::process::exit(1);
    }
    Ok(())
}

async fn ssh_broker_cli(args: Vec<String>) -> anyhow::Result<()> {
    let raw_payload = args
        .first()
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("--ssh-broker requires payload JSON"))?;
    let payload: Value = serde_json::from_str(&raw_payload)?;
    let token = env::var("PROVIDER_TOKEN")
        .or_else(|_| env::var("UNIDESK_PROVIDER_TOKEN"))
        .unwrap_or_default();
    let url = format!(
        "ws://127.0.0.1:8080/ws/ssh?token={}",
        percent_encoding::utf8_percent_encode(&token, percent_encoding::NON_ALPHANUMERIC)
    );
    let (ws, _) = connect_async(url).await?;
    let (mut ws_write, mut ws_read) = ws.split();
    let provider_id = payload
        .get("providerId")
        .and_then(Value::as_str)
        .unwrap_or("");
    let open = json!({
        "type": "ssh.open",
        "providerId": provider_id,
        "cwd": payload.get("cwd").cloned().unwrap_or(Value::Null),
        "command": payload.get("command").cloned().unwrap_or(Value::Null),
        "tty": payload.get("tty").cloned().unwrap_or(Value::Null),
        "cols": payload.get("cols").cloned().unwrap_or(json!(100)),
        "rows": payload.get("rows").cloned().unwrap_or(json!(30)),
    });
    ws_write.send(Message::Text(open.to_string())).await?;

    let (opened_tx, mut opened_rx) = watch::channel(false);
    let stdin_eot_on_end = payload
        .get("stdinEotOnEnd")
        .and_then(Value::as_bool)
        .unwrap_or(false);
    let write_task = tokio::spawn(async move {
        // tcp-pool input frames sent before provider-gateway starts the host
        // SSH process are dropped because the provider has no session yet.
        if opened_rx
            .wait_for(|opened| *opened)
            .await
            .is_err()
        {
            return;
        }
        let mut stdin = tokio::io::stdin();
        let mut buffer = vec![0_u8; 8192];
        loop {
            match stdin.read(&mut buffer).await {
                Ok(0) => break,
                Ok(size) => {
                    let data = base64::engine::general_purpose::STANDARD.encode(&buffer[..size]);
                    if ws_write
                        .send(Message::Text(
                            json!({ "type": "ssh.input", "data": data, "encoding": "base64" })
                                .to_string(),
                        ))
                        .await
                        .is_err()
                    {
                        return;
                    }
                }
                Err(_) => break,
            }
        }
        if stdin_eot_on_end {
            let data = base64::engine::general_purpose::STANDARD.encode([4_u8]);
            let _ = ws_write
                .send(Message::Text(
                    json!({ "type": "ssh.input", "data": data, "encoding": "base64" }).to_string(),
                ))
                .await;
        }
        let _ = ws_write
            .send(Message::Text(json!({ "type": "ssh.eof" }).to_string()))
            .await;
    });

    let open_timeout_ms = payload
        .get("openTimeoutMs")
        .and_then(Value::as_u64)
        .unwrap_or(15_000)
        .clamp(1_000, 300_000);
    let runtime_timeout_ms = payload
        .get("runtimeTimeoutMs")
        .and_then(Value::as_u64)
        .unwrap_or(60_000)
        .clamp(1_000, 3_600_000);
    let open_timer = tokio::time::sleep(Duration::from_millis(open_timeout_ms));
    let runtime_timer = tokio::time::sleep(Duration::from_millis(runtime_timeout_ms));
    tokio::pin!(open_timer);
    tokio::pin!(runtime_timer);

    let mut opened = false;
    let mut exit_code = 255_i32;
    loop {
        let message = tokio::select! {
            _ = &mut runtime_timer => {
                write_stderr(format!("unidesk ssh bridge runtime timeout after {runtime_timeout_ms}ms\n").as_bytes()).await;
                exit_code = 124;
                break;
            }
            _ = &mut open_timer, if !opened => {
                write_stderr(format!("unidesk ssh bridge timed out waiting for provider session after {open_timeout_ms}ms\n").as_bytes()).await;
                exit_code = 255;
                break;
            }
            message = ws_read.next() => message,
        };
        let Some(message) = message else {
            break;
        };
        let message = message?;
        let text = match message {
            Message::Text(text) => text,
            Message::Close(_) => break,
            _ => continue,
        };
        let parsed: Value = serde_json::from_str(&text).unwrap_or_else(|_| json!({}));
        match parsed.get("type").and_then(Value::as_str).unwrap_or("") {
            "ssh.opened" | "ssh.dispatched" => {
                opened = true;
                if parsed.get("type").and_then(Value::as_str) == Some("ssh.opened") {
                    let _ = opened_tx.send(true);
                }
            }
            "ssh.data" => {
                opened = true;
                let data = parsed.get("data").and_then(Value::as_str).unwrap_or("");
                let bytes = base64::engine::general_purpose::STANDARD
                    .decode(data)
                    .unwrap_or_default();
                if parsed.get("stream").and_then(Value::as_str) == Some("stderr") {
                    write_stderr(&bytes).await;
                } else {
                    write_stdout(&bytes).await;
                }
            }
            "ssh.exit" => {
                exit_code = parsed
                    .get("exitCode")
                    .and_then(Value::as_i64)
                    .unwrap_or(255) as i32;
                break;
            }
            "ssh.error" => {
                write_stderr(
                    format!(
                        "{}\n",
                        parsed
                            .get("message")
                            .and_then(Value::as_str)
                            .unwrap_or("ssh error")
                    )
                    .as_bytes(),
                )
                .await;
                exit_code = 255;
                break;
            }
            _ => {}
        }
    }
    write_task.abort();
    std::process::exit(exit_code);
}