sausage/cli/tests/actors/toolbox/src/main.rs

use std::collections::HashMap;
use std::io::{self, BufRead, Write};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use serde_json::{Value, json};

fn send(out: &mut impl Write, message: Value) {
    writeln!(out, "{message}").unwrap();
    out.flush().unwrap();
}

fn reply(out: &mut impl Write, message: &Value, body: Value) {
    send(
        out,
        json!({"to": message["from"], "reply_to": message["id"], "body": body}),
    );
}

/// Connect to `address`, send one line, and return the line that comes back.
fn tcp_line(address: &str, line: &str) -> io::Result<String> {
    let mut stream = std::net::TcpStream::connect(address)?;
    writeln!(stream, "{line}")?;
    let mut answer = String::new();
    io::BufReader::new(stream).read_line(&mut answer)?;
    Ok(answer.trim_end().to_string())
}

fn io_error(error: io::Error) -> Value {
    json!({"type": "error", "error": error.to_string()})
}

fn main() {
    let actor = std::env::var("SAUSAGE_ACTOR").unwrap_or_default();
    let version = std::env::var("SAUSAGE_VERSION").unwrap_or_else(|_| "candidate".to_string());
    let candidate = std::env::var_os("SAUSAGE_CANDIDATE").is_some();
    if actor == "crashloop" && version == "2" && !candidate {
        std::fs::write("/data/state", "bad-v2").unwrap();
        panic!("version 2 crashes on startup");
    }

    let mut open = HashMap::<String, std::fs::File>::new();
    // Files that this actor serves at /actor/<id>/fs/ (path -> text).
    let mut served = std::collections::BTreeMap::<String, String>::new();
    let stdin = io::stdin();
    let mut out = io::stdout().lock();
    let mut emit_shutdown = false;
    if actor == "spawn-race" && !candidate {
        let stamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        std::fs::write(format!("/data/start-{stamp}"), "started").unwrap();
    }
    for line in stdin.lock().lines() {
        let Ok(line) = line else { break };
        let Ok(message) = serde_json::from_str::<Value>(&line) else {
            continue;
        };
        let body = &message["body"];
        let path = body["path"].as_str().unwrap_or_default().to_string();
        let text = body["text"].as_str().unwrap_or_default();
        let response = match body["type"].as_str() {
            Some("shutdown") => {
                if actor == "queue" {
                    let _ = std::fs::write("/data/shutdown", "stopping");
                    std::thread::sleep(Duration::from_millis(800));
                }
                if emit_shutdown {
                    send(
                        &mut out,
                        json!({"to": "console", "reply_to": message["id"], "body": {"type": "shutdown-reply"}}),
                    );
                    send(&mut out, json!({"body": {"type": "shutdown-publish"}}));
                }
                break;
            }
            Some("echo") => json!({"type": "echo", "value": body["value"]}),
            Some("fs_read") => {
                let dir = if path == "/" { "/".to_string() } else { format!("{path}/") };
                let mut names: Vec<String> = served
                    .keys()
                    .filter_map(|k| k.strip_prefix(&dir))
                    .map(|rest| match rest.split_once('/') {
                        Some((first, _)) => format!("{first}/"),
                        None => rest.to_string(),
                    })
                    .collect();
                names.dedup();
                if let Some(text) = served.get(&path) {
                    json!({"type": "fs_file", "text": text})
                } else if path == "/" || !names.is_empty() {
                    let entries: Vec<Value> = names
                        .iter()
                        .map(|n| match n.strip_suffix('/') {
                            Some(d) => json!({"name": d, "type": "dir"}),
                            None => json!({"name": n, "type": "file"}),
                        })
                        .collect();
                    json!({"type": "fs_dir", "entries": entries})
                } else {
                    json!({"type": "fs_error", "errno": "no-entry"})
                }
            }
            Some("fs_write") => {
                served.insert(path.clone(), text.to_string());
                json!({"type": "fs_done"})
            }
            Some("fs_remove") => match served.remove(&path) {
                Some(_) => json!({"type": "fs_done"}),
                None => json!({"type": "fs_error", "errno": "no-entry"}),
            },
            Some("publish") => {
                send(
                    &mut out,
                    json!({"body": {"type": "published", "value": body["value"]}}),
                );
                continue;
            }
            Some("version") => json!({"type": "version", "version": version}),
            Some("ordered") => json!({"type": "ordered", "n": body["n"], "version": version}),
            Some("write") => match std::fs::write(&path, text) {
                Ok(()) => json!({"type": "ok"}),
                Err(error) => io_error(error),
            },
            Some("mkdir") => match std::fs::create_dir_all(&path) {
                Ok(()) => json!({"type": "ok"}),
                Err(error) => io_error(error),
            },
            Some("rename") => match std::fs::rename(&path, body["to"].as_str().unwrap_or_default())
            {
                Ok(()) => json!({"type": "ok"}),
                Err(error) => io_error(error),
            },
            Some("open") => {
                let result = (|| -> io::Result<()> {
                    let mut file = std::fs::File::create(&path)?;
                    file.write_all(text.as_bytes())?;
                    open.insert(path.clone(), file);
                    Ok(())
                })();
                match result {
                    Ok(()) => json!({"type": "opened"}),
                    Err(error) => io_error(error),
                }
            }
            Some("close") => json!({"type": "closed", "closed": open.remove(&path).is_some()}),
            Some("arm-shutdown") => {
                emit_shutdown = true;
                json!({"type": "armed"})
            }
            Some("ignore-stdin") => json!({"type": "ignoring"}),
            Some("tcp") => match tcp_line(body["address"].as_str().unwrap_or_default(), text) {
                Ok(line) => json!({"type": "tcp", "line": line}),
                Err(error) => io_error(error),
            },
            Some("read") => match std::fs::read_to_string(&path) {
                Ok(text) => json!({"type": "content", "text": text}),
                Err(error) => io_error(error),
            },
            Some("list") => match std::fs::read_dir(&path) {
                Ok(entries) => {
                    let mut names: Vec<String> = entries
                        .filter_map(|entry| entry.ok())
                        .map(|entry| entry.file_name().to_string_lossy().to_string())
                        .collect();
                    names.sort();
                    json!({"type": "entries", "names": names})
                }
                Err(error) => io_error(error),
            },
            _ => json!({"type": "error", "error": "unknown message"}),
        };
        reply(&mut out, &message, response);
        if body["type"] == "ignore-stdin" {
            loop {
                std::thread::sleep(Duration::from_secs(60));
            }
        }
    }
}