sausage/userland/ssh/tests/server/src/main.rs

//! Native localhost sshd fixture and host-only signing oracle for WASI tests.
use base64::{Engine, engine::general_purpose::STANDARD};
use serde_json::{Value, json};
use signature::Signer;
use ssh_key::{Algorithm, LineEnding, PrivateKey};
use std::{
    fs,
    io::{BufRead, BufReader, Read, Write},
    net::{TcpListener, TcpStream},
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
    time::{Duration, Instant},
};
fn key(path: &Path) -> PrivateKey {
    let key = PrivateKey::random(&mut rand_core::OsRng, Algorithm::Ed25519).unwrap();
    key.write_openssh_file(path, LineEnding::LF).unwrap();
    key
}
struct Daemon {
    child: Child,
    config: PathBuf,
    binary: PathBuf,
    host_key: PathBuf,
    log: PathBuf,
    port: u16,
}
impl Daemon {
    fn start(binary: &Path, config: &Path, log: &Path, port: u16) -> Child {
        let child = Command::new(binary)
            .args(["-D", "-e", "-f"])
            .arg(config)
            .stdout(Stdio::null())
            .stderr(
                fs::OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(log)
                    .unwrap(),
            )
            .spawn()
            .unwrap();
        let deadline = Instant::now() + Duration::from_secs(5);
        while TcpStream::connect(("127.0.0.1", port)).is_err() {
            assert!(
                Instant::now() < deadline,
                "sshd failed to start; inspect {}",
                log.display()
            );
            std::thread::yield_now();
        }
        child
    }
    fn rotate(&mut self) {
        self.child.kill().unwrap();
        self.child.wait().unwrap();
        key(&self.host_key);
        self.child = Self::start(&self.binary, &self.config, &self.log, self.port);
    }
}
impl Drop for Daemon {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}
fn main() {
    let root = PathBuf::from(std::env::args().nth(1).expect("fixture root"));
    let user = std::env::args().nth(2).expect("SSH user");
    let host = root.join("host");
    fs::create_dir_all(&host).unwrap();
    fs::create_dir_all(root.join("guest/home")).unwrap();
    let client = key(&host.join("client_key"));
    fs::write(
        host.join("authorized_keys"),
        client.public_key().to_openssh().unwrap(),
    )
    .unwrap();
    let host_key = host.join("server_key");
    key(&host_key);
    let socket = TcpListener::bind("127.0.0.1:0").unwrap();
    let port = socket.local_addr().unwrap().port();
    drop(socket);
    let binary = std::env::split_paths(&std::env::var_os("PATH").unwrap())
        .map(|p| p.join("sshd"))
        .find(|p| p.is_file())
        .unwrap()
        .canonicalize()
        .unwrap();
    let config = host.join("sshd_config");
    fs::write(&config, format!("Port {port}\nListenAddress 127.0.0.1\nHostKey {}\nPidFile {}\nAuthorizedKeysFile {}\nStrictModes no\nUsePAM no\nPasswordAuthentication no\nKbdInteractiveAuthentication no\nPermitRootLogin no\nAllowUsers {user}\nLogLevel VERBOSE\n", host_key.display(), host.join("sshd.pid").display(), host.join("authorized_keys").display())).unwrap();
    let log = host.join("sshd.log");
    let child = Daemon::start(&binary, &config, &log, port);
    let mut daemon = Daemon {
        child,
        config,
        binary,
        host_key,
        log,
        port,
    };
    let api = TcpListener::bind("127.0.0.1:0").unwrap();
    fs::write(
        root.join("ready.json"),
        json!({"ssh":port,"api":api.local_addr().unwrap().port(),"user":user}).to_string(),
    )
    .unwrap();
    for connection in api.incoming() {
        let mut stream = connection.unwrap();
        let mut reader = BufReader::new(stream.try_clone().unwrap());
        let mut first = String::new();
        reader.read_line(&mut first).unwrap();
        let route = first.split_whitespace().nth(1).unwrap_or("");
        let mut length = 0;
        loop {
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();
            if line == "\r\n" || line.is_empty() {
                break;
            }
            if let Some((name, value)) = line.split_once(':') {
                if name.eq_ignore_ascii_case("content-length") {
                    length = value.trim().parse().unwrap();
                }
            }
        }
        let mut bytes = vec![0; length];
        reader.read_exact(&mut bytes).unwrap();
        let request: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
        let (status, body) = match route {
            "/fixture/rotate" => {
                daemon.rotate();
                (200, json!({"ok":true}))
            }
            "/ssh/identity" | "/ssh/sign" if host.join("deny").exists() => (
                403,
                json!({"item":"test-key","request":"test-request","message":"approval required"}),
            ),
            "/ssh/identity" => (
                200,
                json!({"item":"test-key","public_key":STANDARD.encode(client.public_key().to_bytes().unwrap()),"user":user}),
            ),
            "/ssh/sign" if request["item"] == "test-key" && request["host"] == "127.0.0.1" => {
                let bytes = STANDARD
                    .decode(request["payload"].as_str().unwrap())
                    .unwrap();
                let signature: ssh_key::Signature = client.try_sign(&bytes).unwrap();
                let signature: Vec<u8> = signature.try_into().unwrap();
                (200, json!({"signature":STANDARD.encode(signature)}))
            }
            _ => (400, json!({"error":"invalid request"})),
        };
        let body = body.to_string();
        write!(stream, "HTTP/1.1 {status} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).unwrap();
    }
}