//! The WASM tier (spec §4.1, oxygenflow#75): a job's `exec` names a WASI

//! preview 1 module, or a glob of modules run one after another, which

//! wasmtime runs with the job's workspace at `/src` and nothing else: no

//! network, no other file of the runner. Every runner can offer it, and it

//! is the tier phones and browsers run (docs/handover.md).

//!

//! As in the syscall tier the module runs in a child of the runner

//! (`flow __wasm SPEC`), so the runner streams its output as it comes,

//! stops it at the time limit or when the run is cancelled, and measures

//! its memory the same way; wasmtime itself also holds it to the sandbox's

//! memory limit.

use crate::microvm::{Outcome, Sink};

use crate::syscall::{lines, self_hash, FAILED};

use anyhow::{bail, Context, Result};

use flow_core::api::Record;

use flow_core::{Environment, Job, MountSource, Sandbox};

use serde::{Deserialize, Serialize};

use std::path::{Path, PathBuf};

use std::sync::mpsc;

use std::time::{Duration, Instant};

/// The runtime, named in every receipt this tier writes.

pub const WASMTIME: &str = "49.0.4";

/// A line the helper writes to say something about the job itself (which

/// module starts), shown as the runner's own note rather than the job's.

const NOTE: &str = "\u{1e}N ";

/// The exit code for a module that trapped (crashed) rather than exiting.

const TRAPPED: i32 = 134;

/// What the helper runs, written by the runner for `flow __wasm`.

#[derive(Serialize, Deserialize)]

struct Spec {

    /// The job's files, preopened as `/src` (and as the current directory).

    workspace: PathBuf,

    /// The modules to run in turn, relative to the workspace.

    modules: Vec<String>,

    /// The arguments after the module.

    args: Vec<String>,

    env: Vec<(String, String)>,

    memory_mib: u32,

    /// The job's mounts: a directory of the runner and where the module sees it.

    preopens: Vec<(PathBuf, String)>,

}

/// The receipt environment of a job in this tier: the runtime and its

/// version, this runner's binary, the memory allowed, and every module run

/// with its hash, so the receipt says exactly what code ran.

pub fn environment(sandbox: &Sandbox, job: &Job, modules: &[(String, String)]) -> Environment {

    Environment {

        runtime: "wasm".into(),

        parts: vec![

            ("wasmtime".into(), WASMTIME.into()),

            ("flow".into(), self_hash().to_string()),

            ("memory_mib".into(), sandbox.memory_mib.to_string()),

        ]

        .into_iter()

        .chain(modules.iter().map(|(path, hash)| (format!("module {path}"), hash.clone())))

        .collect(),

        rootfs: String::new(),

        argv: job.exec.clone(),

        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),

    }

}

/// The modules `program` names in `workspace`: one file, or the files a

/// glob matches (`*` in its last part only), sorted. A path may start with

/// `/src/`, where the module sees the workspace.

pub fn modules(workspace: &Path, program: &str) -> Result<Vec<String>> {

    let rel = program.strip_prefix("/src/").unwrap_or(program).trim_start_matches("./");

    if rel.is_empty() || rel.starts_with('/') || rel.split('/').any(|p| p == "..") {

        bail!("module {program:?}: give a path inside the job's files, such as target/wasm32-wasip1/debug/app.wasm");

    }

    let (dir, last) = rel.rsplit_once('/').unwrap_or(("", rel));

    if dir.contains('*') {

        bail!("module {program:?}: only the file name may hold `*`");

    }

    if !last.contains('*') {

        if !workspace.join(rel).is_file() {

            bail!("module {rel} is not among the job's files (its scope or an earlier job's artifact)");

        }

        return Ok(vec![rel.to_string()]);

    }

    let mut found: Vec<String> = std::fs::read_dir(workspace.join(dir))

        .with_context(|| format!("module {program:?}: the directory {dir} is not among the job's files"))?

        .filter_map(|e| e.ok())

        .filter(|e| e.file_type().is_ok_and(|t| t.is_file()))

        .filter_map(|e| e.file_name().into_string().ok())

        .filter(|name| wildcard(last, name))

        .map(|name| if dir.is_empty() { name } else { format!("{dir}/{name}") })

        .collect();

    found.sort();

    if found.is_empty() {

        bail!("no module matches {program:?}");

    }

    Ok(found)

}

/// Whether `name` matches `pattern`, where `*` stands for any run of characters.

fn wildcard(pattern: &str, name: &str) -> bool {

    let parts: Vec<&str> = pattern.split('*').collect();

    let (first, last) = (parts[0], parts[parts.len() - 1]);

    if parts.len() == 1 {

        return pattern == name;

    }

    if !name.starts_with(first) || !name[first.len()..].ends_with(last) || name.len() < first.len() + last.len() {

        return false;

    }

    let mut rest = &name[first.len()..name.len() - last.len()];

    for middle in &parts[1..parts.len() - 1] {

        match rest.find(middle) {

            Some(i) => rest = &rest[i + middle.len()..],

            None => return false,

        }

    }

    true

}

/// Run `job` in this tier with its files in `workspace` (as the syscall

/// tier lays them out); `scratch` holds the helper's spec and the job's

/// mounts.

pub fn run(sandbox: &Sandbox, job: &Job, workspace: &Path, scratch: &Path, timeout: Duration, sink: &mut dyn Sink) -> Result<Outcome> {

    if !job.secrets.is_empty() || !job.identities.is_empty() {

        bail!("job {}: the WASM tier takes no secrets or identities yet; run it in a microvm or syscall sandbox", job.name);

    }

    let Some(program) = job.exec.first() else { bail!("job {}: exec names no module", job.name) };

    let modules = modules(workspace, program)?;

    let hashed: Vec<(String, String)> =

        modules.iter().map(|m| Ok((m.clone(), crate::home::hash_file(&workspace.join(m))?))).collect::<Result<_>>()?;

    let environment = environment(sandbox, job, &hashed);

    let mut preopens = Vec::new();

    let mut disks = 0;

    for (i, m) in job.mounts.iter().enumerate() {

        let dir = match m.source {

            MountSource::Disk => {

                disks += 1;

                crate::syscall::disk_dir(scratch, disks - 1)

            }

            MountSource::Tmpfs => scratch.join(format!("tmpfs-{i}")),

        };

        std::fs::create_dir_all(&dir)?;

        preopens.push((dir, m.target.clone()));

    }

    let spec = Spec {

        workspace: workspace.to_path_buf(),

        modules,

        args: job.exec[1..].to_vec(),

        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),

        memory_mib: sandbox.memory_mib,

        preopens,

    };

    let spec_path = scratch.join("wasm.json");

    std::fs::write(&spec_path, serde_json::to_vec(&spec)?)?;

    let started = Instant::now();

    let mut child = std::process::Command::new("/proc/self/exe")

        .arg("__wasm")

        .arg(&spec_path)

        .env_clear()

        .current_dir(scratch)

        .stdin(std::process::Stdio::null())

        .stdout(std::process::Stdio::piped())

        .stderr(std::process::Stdio::piped())

        .spawn()

        .context("start the WASM helper")?;

    let (tx, rx) = mpsc::channel::<(bool, Vec<u8>)>();

    let out = child.stdout.take().unwrap();

    let err = child.stderr.take().unwrap();

    let tx2 = tx.clone();

    let readers = [std::thread::spawn(move || lines(out, true, tx)), std::thread::spawn(move || lines(err, false, tx2))];

    let t = |s: &Instant| s.elapsed().as_millis() as u64;

    let (mut timed_out, mut cancelled, mut failed_to_start) = (false, false, false);

    let mut peak_memory_mib = None;

    let deadline = started + timeout;

    loop {

        if let Some(m) = crate::facts::peak_rss_mib(child.id()) {

            peak_memory_mib = Some(m);

        }

        match rx.recv_timeout(Duration::from_millis(200)) {

            Ok((stdout, line)) => {

                let text = String::from_utf8_lossy(&line).into_owned();

                let rec = |stream: &str, text: &str| Record { t: t(&started), stream: stream.into(), text: text.into() };

                if let Some(note) = text.strip_prefix(NOTE).filter(|_| !stdout) {

                    sink.record(rec("flow", note));

                } else if let Some(why) = text.strip_prefix(FAILED).filter(|_| !stdout) {

                    failed_to_start = true;

                    sink.record(rec("flow", &format!("the WASM tier could not run the job: {why}")));

                } else {

                    sink.record(rec(if stdout { "out" } else { "err" }, &text));

                }

            }

            Err(mpsc::RecvTimeoutError::Timeout) => {}

            Err(mpsc::RecvTimeoutError::Disconnected) => break,

        }

        if !timed_out && !cancelled && sink.cancelled() {

            cancelled = true;

            let _ = child.kill();

            sink.record(Record { t: t(&started), stream: "flow".into(), text: "stopped: the run was cancelled".into() });

        }

        if Instant::now() >= deadline && !timed_out && !cancelled {

            timed_out = true;

            let _ = child.kill();

            sink.record(Record {

                t: t(&started),

                stream: "flow".into(),

                text: format!("stopped at the time limit of {} s", timeout.as_secs()),

            });

        }

    }

    let status = child.wait()?;

    for r in readers {

        let _ = r.join();

    }

    let _ = std::fs::remove_file(&spec_path);

    let exit_code = status.code();

    if exit_code.is_none() && !timed_out && !cancelled && !failed_to_start {

        sink.record(Record {

            t: t(&started),

            stream: "flow".into(),

            text: format!("the WASM helper ended without an exit code ({status})"),

        });

    }

    Ok(Outcome {

        exit_code: if timed_out || cancelled { None } else { exit_code },

        duration_ms: t(&started),

        environment,

        timed_out,

        cancelled,

        peak_memory_mib,

        pruned: Vec::new(),

    })

}

/// `flow __wasm SPEC`: run the spec's modules in turn, each in a fresh

/// instance; the exit code is the first non-zero one, after all have run,

/// so one failing module does not hide another's results.

pub fn wasm_main(spec_path: &Path) -> i32 {

    match helper(spec_path) {

        Ok(code) => code,

        Err(e) => {

            eprintln!("{FAILED}{e:#}");

            125

        }

    }

}

/// wasmtime's own error, as this crate's: with its whole chain of causes.

fn ww<T>(r: wasmtime::Result<T>, what: impl FnOnce() -> String) -> Result<T> {

    r.map_err(|e| anyhow::anyhow!("{}: {e:#}", what()))

}

struct State {

    wasi: wasmtime_wasi::p1::WasiP1Ctx,

    limits: wasmtime::StoreLimits,

}

fn helper(spec_path: &Path) -> Result<i32> {

    let spec: Spec = serde_json::from_slice(&std::fs::read(spec_path)?).context("read the spec")?;

    let engine = ww(wasmtime::Engine::new(&wasmtime::Config::new()), || "start wasmtime".into())?;

    let mut linker: wasmtime::Linker<State> = wasmtime::Linker::new(&engine);

    ww(wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |s: &mut State| &mut s.wasi), || "link WASI".into())?;

    let mut first_failure = 0;

    for path in &spec.modules {

        let code = run_module(&engine, &linker, &spec, path)?;

        if spec.modules.len() > 1 {

            eprintln!("{NOTE}{path}: exit {code}");

        }

        if first_failure == 0 && code != 0 {

            first_failure = code;

        }

    }

    Ok(first_failure)

}

fn run_module(engine: &wasmtime::Engine, linker: &wasmtime::Linker<State>, spec: &Spec, path: &str) -> Result<i32> {

    use wasmtime_wasi::{FsPerms, WasiCtxBuilder};

    if spec.modules.len() > 1 {

        eprintln!("{NOTE}running {path}");

    }

    let module = ww(wasmtime::Module::from_file(engine, spec.workspace.join(path)), || format!("module {path} does not load"))?;

    let name = path.rsplit('/').next().unwrap_or(path);

    let mut argv = vec![name.to_string()];

    argv.extend(spec.args.iter().cloned());

    let mut wasi = WasiCtxBuilder::new();

    wasi.inherit_stdout().inherit_stderr().args(&argv).envs(&spec.env).env("PWD", "/src").allow_blocking_current_thread(true);

    let open = |wasi: &mut WasiCtxBuilder, dir: &Path, at: &str| {

        ww(wasi.preopened_dir(dir, at, FsPerms::ReadWrite).map(|_| ()), || format!("open {} as {at}", dir.display()))

    };

    open(&mut wasi, &spec.workspace, "/src")?;

    // Relative paths resolve against the workspace too.

    open(&mut wasi, &spec.workspace, ".")?;

    for (dir, target) in &spec.preopens {

        open(&mut wasi, dir, target)?;

    }

    let limits = wasmtime::StoreLimitsBuilder::new().memory_size(spec.memory_mib as usize * 1024 * 1024).build();

    let mut store = wasmtime::Store::new(engine, State { wasi: wasi.build_p1(), limits });

    store.limiter(|s| &mut s.limits);

    let instance = ww(linker.instantiate(&mut store, &module), || format!("module {path} cannot start here"))?;

    let start =

        ww(instance.get_typed_func::<(), ()>(&mut store, "_start"), || format!("module {path} has no _start: is it a WASI command?"))?;

    Ok(match start.call(&mut store, ()) {

        Ok(()) => 0,

        Err(e) => match e.downcast_ref::<wasmtime_wasi::I32Exit>() {

            Some(exit) => exit.0,

            None => {

                eprintln!("{NOTE}{path} trapped: {e:#}");

                TRAPPED

            }

        },

    })

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn globs_match_names_with_stars() {

        assert!(wildcard("*.wasm", "oxy_core-1a2b.wasm"));

        assert!(wildcard("oxy_core-*.wasm", "oxy_core-1a2b.wasm"));

        assert!(!wildcard("oxy_core-*.wasm", "oxy_core-1a2b.d"));

        assert!(wildcard("a*b*c", "aXbYc") && !wildcard("a*b*c", "aXcYb"));

        assert!(wildcard("plain.wasm", "plain.wasm") && !wildcard("plain.wasm", "other.wasm"));

    }

    #[test]

    fn modules_are_files_of_the_job() {

        let dir = tempfile::tempdir().unwrap();

        let deps = dir.path().join("target/wasm32-wasip1/debug/deps");

        std::fs::create_dir_all(&deps).unwrap();

        for f in ["b-2.wasm", "a-1.wasm", "a-1.d"] {

            std::fs::write(deps.join(f), b"").unwrap();

        }

        let ws = dir.path();

        assert_eq!(

            modules(ws, "/src/target/wasm32-wasip1/debug/deps/*.wasm").unwrap(),

            ["target/wasm32-wasip1/debug/deps/a-1.wasm", "target/wasm32-wasip1/debug/deps/b-2.wasm"]

        );

        assert_eq!(modules(ws, "./target/wasm32-wasip1/debug/deps/b-2.wasm").unwrap(), ["target/wasm32-wasip1/debug/deps/b-2.wasm"]);

        assert!(modules(ws, "target/wasm32-wasip1/debug/deps/*.so").unwrap_err().to_string().contains("no module matches"));

        assert!(modules(ws, "../etc/passwd").is_err());

        assert!(modules(ws, "/etc/x.wasm").is_err());

        assert!(modules(ws, "target/*/debug/a.wasm").unwrap_err().to_string().contains("only the file name"));

    }

}