open-standards/wai/wai-rs/tests/cli_studio.rs

//! `wai dev` serves WAI Studio at `/_wai/`, on the origin it already serves on loopback, and
//! nothing new listens beyond loopback: run as a person runs it, against a stand-in relay and
//! room service (tests/fixtures/dev-stack/fake-stack.py), with Studio's files from
//! `$WAI_STUDIO`.
//!
//! - `/_wai/` answers with Studio's page, and `/_wai/api/state` with the stack;
//! - every TCP socket the stack listens on, `wai dev`'s and each service's, is bound to a
//!   loopback address, and `wai dev` itself listens on one socket, the origin;
//! - the local wai runs an operation for the page: `envelope.inspect` on an envelope the
//!   page sends, and, in a build with the floor, `media.encode` and `media.decode`.
//!
//! It runs where a process's listening sockets can be read: Linux (/proc) and macOS (lsof).
#![cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code, unused_imports))]

use std::io::{BufRead, BufReader, Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};

fn tmp(name: &str) -> PathBuf {
    let d = std::env::temp_dir().join(format!("wai-studio-{name}-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&d);
    std::fs::create_dir_all(&d).unwrap();
    d
}

fn fake() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/dev-stack/fake-stack.py")
}

/// `wai dev` started in `dir`, and the origin it printed.
fn start(dir: &Path, studio: &Path) -> (Child, SocketAddr) {
    let web = dir.join("web");
    std::fs::create_dir_all(web.join("room")).unwrap();
    std::fs::write(web.join("room/index.html"), "<script type=module src=app.js></script>").unwrap();
    let mut child = Command::new(env!("CARGO_BIN_EXE_wai"))
        .current_dir(dir)
        .args(["dev", "--port", "0", "--for", "120", "--dir", "dev"])
        .arg("--web").arg(&web)
        .arg("--relay").arg(fake())
        .arg("--room").arg(fake())
        .env("WAI_STUDIO", studio)
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
        .expect("wai dev runs");
    let mut lines = BufReader::new(child.stdout.take().unwrap()).lines();
    let end = Instant::now() + Duration::from_secs(60);
    let origin = loop {
        assert!(Instant::now() < end, "wai dev did not come up");
        let l = lines.next().expect("wai dev printed its origin").unwrap();
        if let Some(rest) = l.strip_prefix("wai dev: ready on http://") {
            break rest.trim().parse::<SocketAddr>().unwrap();
        }
    };
    // The rest of what it prints is drained, so it never blocks on a full pipe.
    std::thread::spawn(move || for _ in lines {});
    (child, origin)
}

/// One HTTP/1.1 request, and the answer's status and body.
fn http(addr: SocketAddr, method: &str, path: &str, headers: &[(&str, &str)], body: &[u8]) -> (u16, Vec<u8>) {
    let mut s = TcpStream::connect(addr).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(150))).unwrap();
    let mut req = format!("{method} {path} HTTP/1.1\r\nHost: {addr}\r\nConnection: close\r\nContent-Length: {}\r\n", body.len());
    for (k, v) in headers {
        req += &format!("{k}: {v}\r\n");
    }
    s.write_all(format!("{req}\r\n").as_bytes()).unwrap();
    s.write_all(body).unwrap();
    let mut all = Vec::new();
    s.read_to_end(&mut all).unwrap();
    let end = all.windows(4).position(|w| w == b"\r\n\r\n").expect("a response head");
    let status = String::from_utf8_lossy(&all[..end]).split_whitespace().nth(1).unwrap().parse().unwrap();
    (status, all[end + 4..].to_vec())
}

const API: [(&str, &str); 1] = [("x-wai-studio", "1")];

fn json(addr: SocketAddr, method: &str, path: &str, body: &[u8]) -> serde_json::Value {
    let (s, b) = http(addr, method, path, &API, body);
    assert_eq!(s, 200, "{method} {path}: {}", String::from_utf8_lossy(&b));
    serde_json::from_slice(&b).unwrap()
}

/// The local addresses of the TCP sockets process `pid` listens on.
#[cfg(target_os = "linux")]
fn listening(pid: u32) -> Vec<String> {
    let mut inodes = std::collections::HashSet::new();
    for e in std::fs::read_dir(format!("/proc/{pid}/fd")).unwrap().flatten() {
        if let Ok(l) = std::fs::read_link(e.path())
            && let Some(i) = l.to_string_lossy().strip_prefix("socket:[").and_then(|r| r.strip_suffix(']'))
        {
            inodes.insert(i.to_string());
        }
    }
    let mut out = Vec::new();
    for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
        for line in std::fs::read_to_string(table).unwrap_or_default().lines().skip(1) {
            let f: Vec<&str> = line.split_whitespace().collect();
            // local_address, rem_address, st (0A: LISTEN), …, inode.
            if f.len() > 9 && f[3] == "0A" && inodes.contains(f[9]) {
                let (ip, port) = f[1].split_once(':').unwrap();
                let port = u16::from_str_radix(port, 16).unwrap();
                let ip = if ip.len() == 8 {
                    let n = u32::from_str_radix(ip, 16).unwrap();
                    std::net::Ipv4Addr::from(n.to_le_bytes()).to_string()
                } else {
                    let b: Vec<u8> = (0..16).map(|i| u8::from_str_radix(&ip[2 * i..2 * i + 2], 16).unwrap()).collect();
                    let mut w = [0u8; 16];
                    for g in 0..4 {
                        for k in 0..4 {
                            w[4 * g + k] = b[4 * g + 3 - k];
                        }
                    }
                    format!("[{}]", std::net::Ipv6Addr::from(w))
                };
                out.push(format!("{ip}:{port}"));
            }
        }
    }
    out
}

#[cfg(target_os = "macos")]
fn listening(pid: u32) -> Vec<String> {
    let o = Command::new("lsof").args(["-nP", "-a", "-p", &pid.to_string(), "-iTCP", "-sTCP:LISTEN", "-Fn"]).output().expect("lsof");
    String::from_utf8_lossy(&o.stdout).lines().filter_map(|l| l.strip_prefix('n')).map(str::to_string).collect()
}

fn loopback(addr: &str) -> bool {
    addr.starts_with("127.") || addr.starts_with("[::1]:")
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
#[test]
fn wai_dev_serves_studio_at_wai_and_nothing_listens_beyond_loopback() {
    let d = tmp("dev");
    let studio = d.join("studio");
    std::fs::create_dir_all(&studio).unwrap();
    std::fs::write(studio.join("index.html"), "<main id=wai-studio data-host=local></main><script type=module src=wai-studio.js></script>").unwrap();
    let (mut wai, origin) = start(&d, &studio);
    assert!(origin.ip().is_loopback(), "the origin: {origin}");

    let (s, page) = http(origin, "GET", "/_wai/", &[], b"");
    assert_eq!((s, String::from_utf8_lossy(&page).contains("wai-studio.js")), (200, true), "/_wai/ answers with Studio's page");
    let state = json(origin, "GET", "/_wai/api/state", b"");
    assert_eq!(state["ready"], true, "{state}");
    let services = state["services"].as_array().unwrap();
    assert_eq!(services.len(), 3, "the origin, the relay and the room service: {state}");

    // Every socket each process listens on is on loopback; wai dev's own is the origin alone.
    for svc in services {
        let pid = svc["pid"].as_u64().unwrap() as u32;
        let addrs = listening(pid);
        assert!(!addrs.is_empty(), "{} (pid {pid}) listens", svc["name"]);
        for a in &addrs {
            assert!(loopback(a), "{} (pid {pid}) listens on {a}, beyond loopback", svc["name"]);
        }
        assert!(loopback(svc["listen"].as_str().unwrap()), "{svc}");
        if pid == wai.id() {
            assert_eq!(addrs, vec![origin.to_string()], "wai dev listens on its origin alone");
        }
    }

    // The local wai runs an operation for the page: an envelope it wraps itself, inspected.
    let env = d.join("x.wai");
    std::fs::write(d.join("x.txt"), "WAI Studio").unwrap();
    let o = Command::new(env!("CARGO_BIN_EXE_wai")).current_dir(&d).args(["wrap", "wai.text.zstd", "x.txt", "-o"]).arg(&env).output().unwrap();
    assert!(o.status.success(), "{}", String::from_utf8_lossy(&o.stderr));
    let r = json(origin, "POST", "/_wai/api/op/envelope.inspect", &std::fs::read(&env).unwrap());
    assert_eq!((r["op"].as_str(), r["ok"].as_bool()), (Some("envelope.inspect"), Some(true)), "{r}");
    assert!(r["stdout"].as_str().unwrap().contains("wai.text.zstd"), "{r}");
    assert_eq!(r["argv"], serde_json::json!(["wai", "inspect", "in.wai"]));

    if cfg!(feature = "floor") {
        let text = "the floor, through the local wai ".repeat(40);
        let r = json(origin, "POST", "/_wai/api/op/media.encode?codec=zstd&quality=5", text.as_bytes());
        assert_eq!(r["ok"], true, "{r}");
        let wai_b64 = r["outputs"][0]["data"].as_str().unwrap();
        let envelope = b64(wai_b64);
        let r = json(origin, "POST", "/_wai/api/op/media.decode", &envelope);
        assert_eq!(r["ok"], true, "{r}");
        let out = r["outputs"].as_array().unwrap().iter().find(|o| o["name"] == "out").unwrap();
        assert_eq!(b64(out["data"].as_str().unwrap()), text.as_bytes(), "decoded byte for byte");
    }
    // Nothing of an operation is left behind.
    let left: Vec<_> = std::fs::read_dir(d.join("dev/studio")).map(|r| r.flatten().map(|e| e.file_name()).collect()).unwrap_or_default();
    assert!(left.is_empty(), "{left:?}");

    let _ = wai.kill();
    let _ = wai.wait();
    for svc in services.iter().skip(1) {
        let _ = Command::new("kill").arg(svc["pid"].as_u64().unwrap().to_string()).status();
    }
}

fn b64(s: &str) -> Vec<u8> {
    let val = |c: u8| match c {
        b'A'..=b'Z' => c - b'A',
        b'a'..=b'z' => c - b'a' + 26,
        b'0'..=b'9' => c - b'0' + 52,
        b'+' => 62,
        b'/' => 63,
        _ => panic!("not base64: {c}"),
    };
    let b: Vec<u8> = s.bytes().filter(|c| *c != b'=').collect();
    let mut out = Vec::new();
    for ch in b.chunks(4) {
        let n = ch.iter().enumerate().fold(0u32, |n, (i, c)| n | (u32::from(val(*c)) << (18 - 6 * i)));
        out.extend_from_slice(&n.to_be_bytes()[1..ch.len()]);
    }
    out
}