sausage/userland/init/src/lib.rs

//! Unit definitions and the file protocol shared by init and its control tools.
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, fs, io, time::Duration};
/// WASI 0.2's legacy exit call loses numeric error codes. Use the standard extended call.
pub fn exit(code: i32) -> ! {
    wasi::cli::exit::exit_with_code(code as u8);
    unreachable!("exit returned")
}
pub const RUN: &str = "/run/systemd";
pub const UNITS: &str = "/etc/systemd/system";
pub const TRANSIENT: &str = "/run/systemd/transient";
#[derive(Clone, Debug, PartialEq)]
pub struct Unit {
    pub values: BTreeMap<String, String>,
    pub source: String,
}
impl Unit {
    pub fn parse(source: String) -> io::Result<Self> {
        let mut values = BTreeMap::new();
        let mut section = "";
        for line in source.lines().map(str::trim) {
            if line.is_empty() || line.starts_with(['#', ';']) {
                continue;
            }
            if line.starts_with('[') && line.ends_with(']') {
                section = &line[1..line.len() - 1];
                continue;
            }
            let (key, value) = line
                .split_once('=')
                .ok_or_else(|| io::Error::other("expected Key=Value"))?;
            let key = format!("{section}.{}", key.trim());
            if matches!(
                key.as_str(),
                "Unit.After"
                    | "Unit.Before"
                    | "Unit.Requires"
                    | "Unit.Wants"
                    | "Service.Environment"
            ) {
                values
                    .entry(key)
                    .and_modify(|v: &mut String| {
                        v.push(' ');
                        v.push_str(value.trim());
                    })
                    .or_insert(value.trim().into());
            } else {
                values.insert(key, value.trim().into());
            }
        }
        if !values.contains_key("Service.ExecStart") {
            return Err(io::Error::other("missing ExecStart"));
        }
        Ok(Self { values, source })
    }
    pub fn get(&self, key: &str) -> &str {
        self.values.get(key).map(String::as_str).unwrap_or("")
    }
    pub fn words(&self, key: &str) -> Vec<String> {
        self.get(key).split_whitespace().map(name).collect()
    }
    pub fn seconds(&self, key: &str, default: f64) -> Duration {
        Duration::from_secs_f64(
            self.get(key)
                .trim_end_matches('s')
                .parse::<f64>()
                .ok()
                .filter(|n| n.is_finite() && *n >= 0.0)
                .unwrap_or(default),
        )
    }
}
pub fn name(s: &str) -> String {
    if s.ends_with(".service") || s.ends_with(".target") {
        s.into()
    } else {
        format!("{s}.service")
    }
}
pub fn unit_path(name: &str) -> String {
    let persistent = format!("{UNITS}/{name}");
    if std::path::Path::new(&persistent).exists() {
        persistent
    } else {
        format!("{TRANSIENT}/{name}")
    }
}
pub fn enabled(name: &str) -> bool {
    ["default.target", "multi-user.target"]
        .iter()
        .any(|target| std::path::Path::new(&format!("{UNITS}/{target}.wants/{name}")).exists())
}
#[derive(Serialize, Deserialize, Default)]
pub struct Status {
    pub state: String,
    pub pid: Option<u32>,
    pub code: Option<i32>,
}
#[derive(Serialize, Deserialize)]
pub struct Request {
    pub action: String,
    pub unit: String,
}
#[derive(Serialize, Deserialize)]
pub struct Reply {
    pub code: i32,
    pub text: String,
}
pub fn atomic(path: &str, bytes: &[u8]) -> io::Result<()> {
    let tmp = format!("{path}.tmp");
    fs::write(&tmp, bytes)?;
    fs::rename(tmp, path)
}
pub fn request(action: &str, unit: &str) -> io::Result<Reply> {
    let id = format!(
        "{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    );
    let reply = format!("{RUN}/replies/{id}");
    let state = format!("{RUN}/state/{}", name(unit));
    let watcher = sausage_spawn::Watcher::new(&[reply.clone(), state.clone()]);
    atomic(
        &format!("{RUN}/requests/{id}.json"),
        &serde_json::to_vec(&Request {
            action: action.into(),
            unit: name(unit),
        })?,
    )?;
    let deadline = std::time::Instant::now() + Duration::from_secs(60);
    let timeout = wasi::clocks::monotonic_clock::subscribe_duration(
        Duration::from_secs(60).as_nanos() as u64,
    );
    let changed = watcher.subscribe();
    let mut acknowledged = false;
    loop {
        if !acknowledged && let Ok(bytes) = fs::read(&reply) {
            fs::remove_file(&reply)?;
            let reply: Reply = serde_json::from_slice(&bytes)?;
            if reply.code != 0 {
                return Ok(reply);
            }
            acknowledged = true;
        }
        if acknowledged && let Ok(bytes) = fs::read(&state) {
            let status: Status = serde_json::from_slice(&bytes)?;
            let done = if action == "stop" {
                status.state == "inactive"
            } else {
                status.state == "active"
                    || status.state == "failed"
                    || status.state == "inactive" && status.code.is_some()
            };
            if done {
                return Ok(Reply {
                    code: if action != "stop" && status.state == "failed" {
                        1
                    } else {
                        0
                    },
                    text: String::new(),
                });
            }
        }
        if std::time::Instant::now() >= deadline {
            return Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "init did not reply",
            ));
        }
        sausage_spawn::poll(&[&changed, &timeout]);
        let _ = watcher.changes();
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn unit_sections_and_defaults() {
        let unit = Unit::parse("[Unit]\nAfter=prepare\nRequires=prepare\n[Service]\nType=oneshot\nExecStart=/bin/task.wasm\n".into()).unwrap();
        assert_eq!(unit.words("Unit.Requires"), ["prepare.service"]);
        assert_eq!(
            unit.seconds("Service.TimeoutStopSec", 10.),
            Duration::from_secs(10)
        );
        assert!(Unit::parse("[Unit]\nDescription=missing command".into()).is_err());
    }
}