xy/crates/xy-protocol/src/kdl_parse.rs

use crate::{
    ConfigError, RestartConfig, RestartPolicy, ServerConfig, StopConfig, WaitCondition,
    WaitForConfig, default_wait_interval, default_wait_timeout,
};
use kdl::{KdlDocument, KdlNode};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::Duration;

pub fn parse_server_config(
    name: &str,
    text: &str,
    source_path: &Path,
) -> Result<ServerConfig, ConfigError> {
    validate_name(name).map_err(|_| ConfigError::InvalidName {
        name: name.to_string(),
    })?;

    let doc: KdlDocument = text
        .parse()
        .map_err(|err: kdl::KdlError| ConfigError::Parse {
            path: source_path.to_path_buf(),
            message: err.to_string(),
        })?;

    let command = require_string_arg(&doc, "command", source_path)?;
    let args = optional_string_args(&doc, "args");
    let port = require_u16_arg(&doc, "port", source_path)?;
    let env = optional_string_map(&doc, "env");
    let working_dir = optional_string_arg(&doc, "working-dir").map(PathBuf::from);
    let restart = parse_restart(&doc, source_path)?;
    let stop = parse_stop(&doc, source_path)?;
    let wait_for = parse_wait_for(&doc, source_path)?;

    Ok(ServerConfig {
        name: name.to_string(),
        command: PathBuf::from(command),
        args,
        port,
        env,
        working_dir,
        restart,
        stop,
        wait_for,
    })
}

fn validate_name(name: &str) -> Result<(), ()> {
    if name.is_empty() {
        return Err(());
    }

    if name
        .chars()
        .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_')
    {
        Ok(())
    } else {
        Err(())
    }
}

fn require_string_arg(
    doc: &KdlDocument,
    name: &'static str,
    path: &Path,
) -> Result<String, ConfigError> {
    let node = find_node(doc, name).ok_or(ConfigError::MissingField {
        path: path.to_path_buf(),
        field: name,
    })?;

    node.entries()
        .first()
        .and_then(|e| e.value().as_string().map(str::to_string))
        .ok_or(ConfigError::InvalidValue {
            path: path.to_path_buf(),
            field: name,
            message: "expected string argument".into(),
        })
}

fn require_u16_arg(doc: &KdlDocument, name: &'static str, path: &Path) -> Result<u16, ConfigError> {
    let node = find_node(doc, name).ok_or(ConfigError::MissingField {
        path: path.to_path_buf(),
        field: name,
    })?;

    let v = node
        .entries()
        .first()
        .and_then(|e| e.value().as_integer())
        .ok_or(ConfigError::InvalidValue {
            path: path.to_path_buf(),
            field: name,
            message: "expected integer".into(),
        })?;

    u16::try_from(v).map_err(|_| ConfigError::InvalidValue {
        path: path.to_path_buf(),
        field: name,
        message: format!("port {v} out of u16 range"),
    })
}

fn optional_string_arg(doc: &KdlDocument, name: &str) -> Option<String> {
    find_node(doc, name)
        .and_then(|n| n.entries().first())
        .and_then(|e| e.value().as_string().map(str::to_string))
}

fn optional_string_args(doc: &KdlDocument, name: &str) -> Vec<String> {
    find_node(doc, name)
        .map(|n| {
            n.entries()
                .iter()
                .filter_map(|e| e.value().as_string().map(str::to_string))
                .collect()
        })
        .unwrap_or_default()
}

fn optional_string_map(doc: &KdlDocument, name: &str) -> BTreeMap<String, String> {
    let Some(node) = find_node(doc, name) else {
        return BTreeMap::new();
    };

    let mut out = BTreeMap::new();

    if let Some(children) = node.children() {
        for child in children.nodes() {
            let key = child.name().value().to_string();

            if let Some(val) = child.entries().first().and_then(|e| e.value().as_string()) {
                out.insert(key, val.to_string());
            }
        }
    }

    out
}

fn parse_restart(doc: &KdlDocument, path: &Path) -> Result<RestartConfig, ConfigError> {
    let Some(node) = find_node(doc, "restart") else {
        return Ok(RestartConfig::default());
    };

    let Some(children) = node.children() else {
        return Ok(RestartConfig::default());
    };

    let mut out = RestartConfig::default();

    for child in children.nodes() {
        match child.name().value() {
            "policy" => {
                let s = string_arg(child, "policy", path)?;

                out.policy = match s.as_str() {
                    "always" => RestartPolicy::Always,
                    "on-failure" => RestartPolicy::OnFailure,
                    "never" => RestartPolicy::Never,
                    other => {
                        return Err(ConfigError::InvalidValue {
                            path: path.to_path_buf(),
                            field: "restart.policy",
                            message: format!("unknown policy `{other}`"),
                        });
                    }
                };
            }
            "backoff-initial" => {
                out.backoff_initial = parse_duration_arg(child, "restart.backoff-initial", path)?;
            }
            "backoff-max" => {
                out.backoff_max = parse_duration_arg(child, "restart.backoff-max", path)?;
            }
            "max-retries-per-minute" => {
                let v = child
                    .entries()
                    .first()
                    .and_then(|e| e.value().as_integer())
                    .ok_or(ConfigError::InvalidValue {
                        path: path.to_path_buf(),
                        field: "restart.max-retries-per-minute",
                        message: "expected integer".into(),
                    })?;

                out.max_retries_per_minute =
                    u32::try_from(v).map_err(|_| ConfigError::InvalidValue {
                        path: path.to_path_buf(),
                        field: "restart.max-retries-per-minute",
                        message: format!("out of u32 range: {v}"),
                    })?;
            }
            other => {
                return Err(ConfigError::InvalidValue {
                    path: path.to_path_buf(),
                    field: "restart",
                    message: format!("unknown key `{other}`"),
                });
            }
        }
    }

    Ok(out)
}

fn parse_stop(doc: &KdlDocument, path: &Path) -> Result<StopConfig, ConfigError> {
    let Some(node) = find_node(doc, "stop") else {
        return Ok(StopConfig::default());
    };

    let Some(children) = node.children() else {
        return Ok(StopConfig::default());
    };

    let mut out = StopConfig::default();

    for child in children.nodes() {
        match child.name().value() {
            "grace" => {
                out.grace = parse_duration_arg(child, "stop.grace", path)?;
            }
            other => {
                return Err(ConfigError::InvalidValue {
                    path: path.to_path_buf(),
                    field: "stop",
                    message: format!("unknown key `{other}`"),
                });
            }
        }
    }

    Ok(out)
}

fn expand_tilde(raw: &str, home: Option<&Path>) -> PathBuf {
    let Some(rest) = raw.strip_prefix("~/") else {
        return PathBuf::from(raw);
    };

    match home {
        Some(home) => home.join(rest),
        None => PathBuf::from(raw),
    }
}

fn parse_wait_for(doc: &KdlDocument, path: &Path) -> Result<Option<WaitForConfig>, ConfigError> {
    let Some(node) = find_node(doc, "wait-for") else {
        return Ok(None);
    };

    let invalid = |message: String| ConfigError::InvalidValue {
        path: path.to_path_buf(),
        field: "wait-for",
        message,
    };

    let Some(children) = node.children() else {
        return Err(invalid(
            "expected a block with exactly one condition".into(),
        ));
    };

    let home = std::env::var_os("HOME").map(PathBuf::from);

    let mut conditions: Vec<WaitCondition> = Vec::new();
    let mut command: Option<PathBuf> = None;
    let mut args: Option<Vec<String>> = None;
    let mut timeout = default_wait_timeout();
    let mut interval = default_wait_interval();

    for child in children.nodes() {
        match child.name().value() {
            "path" => {
                let raw = string_arg(child, "wait-for", path)?;

                conditions.push(WaitCondition::Path(expand_tilde(&raw, home.as_deref())));
            }
            "tcp" => {
                let addr = string_arg(child, "wait-for", path)?;

                conditions.push(WaitCondition::Tcp(addr));
            }
            "command" => {
                if command.is_some() {
                    return Err(invalid("duplicate key `command`".into()));
                }

                command = Some(PathBuf::from(string_arg(child, "wait-for", path)?));
            }
            "args" => {
                if args.is_some() {
                    return Err(invalid("duplicate key `args`".into()));
                }

                args = Some(
                    child
                        .entries()
                        .iter()
                        .filter_map(|e| e.value().as_string().map(str::to_string))
                        .collect(),
                );
            }
            "timeout" => {
                timeout = parse_duration_arg(child, "wait-for", path)?;
            }
            "interval" => {
                interval = parse_duration_arg(child, "wait-for", path)?;
            }
            other => {
                return Err(invalid(format!("unknown key `{other}`")));
            }
        }
    }

    if let Some(command) = command {
        conditions.push(WaitCondition::Command {
            command,
            args: args.unwrap_or_default(),
        });
    } else if args.is_some() {
        return Err(invalid("`args` requires `command`".into()));
    }

    match conditions.len() {
        1 => Ok(Some(WaitForConfig {
            condition: conditions.remove(0),
            timeout,
            interval,
        })),
        0 => Err(invalid(
            "expected exactly one of `path`, `tcp` or `command`".into(),
        )),
        n => Err(invalid(format!(
            "expected exactly one condition, found {n}"
        ))),
    }
}

fn string_arg(node: &KdlNode, field: &'static str, path: &Path) -> Result<String, ConfigError> {
    node.entries()
        .first()
        .and_then(|e| e.value().as_string().map(str::to_string))
        .ok_or(ConfigError::InvalidValue {
            path: path.to_path_buf(),
            field,
            message: "expected string".into(),
        })
}

fn parse_duration_arg(
    node: &KdlNode,
    field: &'static str,
    path: &Path,
) -> Result<Duration, ConfigError> {
    let s = string_arg(node, field, path)?;

    humantime::parse_duration(&s).map_err(|err| ConfigError::InvalidValue {
        path: path.to_path_buf(),
        field,
        message: format!("invalid duration `{s}`: {err}"),
    })
}

fn find_node<'a>(doc: &'a KdlDocument, name: &str) -> Option<&'a KdlNode> {
    doc.nodes().iter().find(|n| n.name().value() == name)
}

pub fn load_all_configs(dir: &Path) -> Result<Vec<ServerConfig>, ConfigError> {
    if !dir.exists() {
        return Ok(Vec::new());
    }

    let entries = std::fs::read_dir(dir).map_err(|e| ConfigError::Io {
        path: dir.to_path_buf(),
        source: e,
    })?;

    let mut configs = Vec::new();
    for entry in entries {
        let entry = entry.map_err(|e| ConfigError::Io {
            path: dir.to_path_buf(),
            source: e,
        })?;
        let path = entry.path();
        if path.extension().and_then(|s| s.to_str()) != Some("kdl") {
            continue;
        }
        let name = path
            .file_stem()
            .and_then(|s| s.to_str())
            .ok_or(ConfigError::InvalidName {
                name: path.display().to_string(),
            })?
            .to_string();
        let text = std::fs::read_to_string(&path).map_err(|e| ConfigError::Io {
            path: path.clone(),
            source: e,
        })?;
        configs.push(parse_server_config(&name, &text, &path)?);
    }

    check_duplicate_ports(&configs)?;
    Ok(configs)
}

fn check_duplicate_ports(configs: &[ServerConfig]) -> Result<(), ConfigError> {
    let mut seen: std::collections::HashMap<u16, String> = std::collections::HashMap::new();
    for c in configs {
        if let Some(other) = seen.insert(c.port, c.name.clone()) {
            return Err(ConfigError::DuplicatePort {
                name_a: other,
                name_b: c.name.clone(),
                port: c.port,
            });
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;

    fn p() -> &'static Path {
        Path::new("/tmp/test.kdl")
    }

    #[test]
    fn parses_minimal_config() {
        let text = "command \"/usr/local/bin/foo\"\nport 8421\n";
        let cfg = parse_server_config("foo", text, p()).unwrap();

        assert_eq!(cfg.name, "foo");
        assert_eq!(cfg.command, PathBuf::from("/usr/local/bin/foo"));
        assert_eq!(cfg.port, 8421);
        assert!(cfg.args.is_empty());
        assert!(cfg.env.is_empty());
    }

    #[test]
    fn parses_full_config() {
        let text = r#"
command "/usr/local/bin/foo"
args "--http" "--port" "8421"
port 8421
env {
    RUST_LOG "info"
    FOO_BAR "baz"
}
working-dir "/tmp/work"
restart {
    policy "always"
    backoff-initial "2s"
    backoff-max "1m"
    max-retries-per-minute 10
}
stop {
    grace "30s"
}
"#;
        let cfg = parse_server_config("foo", text, p()).unwrap();

        assert_eq!(cfg.args, vec!["--http", "--port", "8421"]);
        assert_eq!(cfg.env.get("RUST_LOG").map(String::as_str), Some("info"));
        assert_eq!(cfg.working_dir, Some(PathBuf::from("/tmp/work")));
        assert_eq!(cfg.restart.policy, RestartPolicy::Always);
        assert_eq!(cfg.restart.backoff_initial, Duration::from_secs(2));
        assert_eq!(cfg.restart.backoff_max, Duration::from_secs(60));
        assert_eq!(cfg.restart.max_retries_per_minute, 10);
        assert_eq!(cfg.stop.grace, Duration::from_secs(30));
    }

    #[test]
    fn missing_command_fails() {
        let err = parse_server_config("foo", "port 8421", p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::MissingField {
                field: "command",
                ..
            }
        ));
    }

    #[test]
    fn missing_port_fails() {
        let err = parse_server_config("foo", "command \"/bin/x\"", p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::MissingField { field: "port", .. }
        ));
    }

    #[test]
    fn unknown_restart_policy_fails() {
        let text = "command \"/bin/x\"\nport 1\nrestart { policy \"maybe\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "restart.policy",
                ..
            }
        ));
    }

    #[test]
    fn invalid_name_rejected() {
        let text = "command \"/bin/x\"\nport 1";
        let err = parse_server_config("Foo Bar", text, p()).unwrap_err();

        assert!(matches!(err, ConfigError::InvalidName { .. }));
    }

    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn load_all_finds_and_parses_files() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("a.kdl"), "command \"/bin/a\"\nport 8001").unwrap();
        fs::write(dir.path().join("b.kdl"), "command \"/bin/b\"\nport 8002").unwrap();
        fs::write(dir.path().join("ignored.txt"), "not a config").unwrap();
        let mut configs = load_all_configs(dir.path()).unwrap();
        configs.sort_by(|x, y| x.name.cmp(&y.name));
        assert_eq!(configs.len(), 2);
        assert_eq!(configs[0].name, "a");
        assert_eq!(configs[1].port, 8002);
    }

    #[test]
    fn load_all_returns_empty_for_missing_dir() {
        let dir = tempdir().unwrap();
        let configs = load_all_configs(&dir.path().join("does-not-exist")).unwrap();
        assert!(configs.is_empty());
    }

    #[test]
    fn duplicate_ports_detected() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("a.kdl"), "command \"/bin/a\"\nport 8001").unwrap();
        fs::write(dir.path().join("b.kdl"), "command \"/bin/b\"\nport 8001").unwrap();
        let err = load_all_configs(dir.path()).unwrap_err();
        match err {
            ConfigError::DuplicatePort { port, .. } => assert_eq!(port, 8001),
            other => panic!("unexpected error: {other:?}"),
        }
    }

    #[test]
    fn parses_a_path_condition_with_defaults() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { path \"/var/run/d.sock\" }";
        let cfg = parse_server_config("foo", text, p()).unwrap();

        let wait = cfg.wait_for.unwrap();

        assert_eq!(
            wait.condition,
            WaitCondition::Path(PathBuf::from("/var/run/d.sock"))
        );
        assert_eq!(wait.timeout, Duration::from_secs(120));
        assert_eq!(wait.interval, Duration::from_secs(1));
    }

    #[test]
    fn expands_a_leading_tilde_against_the_given_home() {
        let home = PathBuf::from("/home/someone");

        assert_eq!(
            expand_tilde("~/.orbstack/run/docker.sock", Some(&home)),
            PathBuf::from("/home/someone/.orbstack/run/docker.sock")
        );
    }

    #[test]
    fn leaves_non_tilde_paths_and_unknown_home_alone() {
        let home = PathBuf::from("/home/someone");

        assert_eq!(
            expand_tilde("/var/run/d.sock", Some(&home)),
            PathBuf::from("/var/run/d.sock")
        );
        assert_eq!(expand_tilde("~/a", None), PathBuf::from("~/a"));
        assert_eq!(
            expand_tilde("~weird/a", Some(&home)),
            PathBuf::from("~weird/a")
        );
    }

    #[test]
    fn parses_a_tcp_condition_with_overrides() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { tcp \"127.0.0.1:5432\"\ntimeout \"30s\"\ninterval \"250ms\" }";
        let cfg = parse_server_config("foo", text, p()).unwrap();

        let wait = cfg.wait_for.unwrap();

        assert_eq!(wait.condition, WaitCondition::Tcp("127.0.0.1:5432".into()));
        assert_eq!(wait.timeout, Duration::from_secs(30));
        assert_eq!(wait.interval, Duration::from_millis(250));
    }

    #[test]
    fn parses_a_command_condition_from_sibling_nodes() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { command \"docker\"\nargs \"info\" }";
        let cfg = parse_server_config("foo", text, p()).unwrap();

        assert_eq!(
            cfg.wait_for.unwrap().condition,
            WaitCondition::Command {
                command: PathBuf::from("docker"),
                args: vec!["info".to_string()],
            }
        );
    }

    #[test]
    fn absent_wait_for_leaves_no_gate() {
        let text = "command \"/bin/x\"\nport 1";
        let cfg = parse_server_config("foo", text, p()).unwrap();

        assert!(cfg.wait_for.is_none());
    }

    #[test]
    fn wait_for_without_a_condition_fails() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { timeout \"5s\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn wait_for_with_two_conditions_fails() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { path \"/a\"\ntcp \"127.0.0.1:1\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn wait_for_args_without_command_fails() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { args \"info\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn unknown_wait_for_key_fails() {
        let text = "command \"/bin/x\"\nport 1\nwait-for { path \"/a\"\nnope \"x\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn duplicate_command_key_fails() {
        let text =
            "command \"/bin/x\"\nport 1\nwait-for { command \"/bin/a\"\ncommand \"/bin/b\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn duplicate_args_key_fails() {
        let text =
            "command \"/bin/x\"\nport 1\nwait-for { command \"docker\"\nargs \"a\"\nargs \"b\" }";
        let err = parse_server_config("foo", text, p()).unwrap_err();

        assert!(matches!(
            err,
            ConfigError::InvalidValue {
                field: "wait-for",
                ..
            }
        ));
    }

    #[test]
    fn expand_tilde_alone_is_left_unchanged() {
        let home = PathBuf::from("/home/someone");

        assert_eq!(expand_tilde("~", Some(&home)), PathBuf::from("~"));
    }

    #[test]
    fn expand_tilde_of_empty_string_is_left_unchanged() {
        let home = PathBuf::from("/home/someone");

        assert_eq!(expand_tilde("", Some(&home)), PathBuf::from(""));
    }
}