open-standards/wai/wai-rs/tests/cli_derive.rs

//! `wai upscale` and `wai tonemap` end to end: what each runs and writes is what its extension's verifier
//! proves under `require-proof`, from the files it wrote alone.
#![cfg(all(feature = "exact_upscale", feature = "exact_tonemap"))]

use std::path::{Path, PathBuf};
use std::process::{Command, Output};

use wai::exact_tonemap as tm;
use wai::exact_upscale as ex;

const SEED: &str = "0707070707070707070707070707070707070707070707070707070707070707";

fn wai_in(d: &Path, args: &[&str]) -> Output {
    Command::new(env!("CARGO_BIN_EXE_wai")).current_dir(d).args(args).output().expect("wai runs")
}

fn scratch(name: &str) -> PathBuf {
    let d = std::env::temp_dir().join(format!("wai-cli-derive-{name}-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&d);
    std::fs::create_dir_all(&d).unwrap();
    d
}

fn ok(o: &Output) -> &Output {
    assert!(o.status.success(), "{}{}", String::from_utf8_lossy(&o.stdout), String::from_utf8_lossy(&o.stderr));
    o
}

fn prior_path() -> String {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("../live-conformance/upscale/priors/generic.wsr").to_str().unwrap().to_string()
}

#[test]
fn an_upscale_writes_what_exact_upscale_proves() {
    let d = scratch("upscale");
    let rgb: Vec<u8> = (0..8u32 * 6).flat_map(|i| [(i * 7 % 256) as u8, (i * 3 % 256) as u8, (i % 256) as u8]).collect();
    std::fs::write(d.join("in.png"), wai::png_still::encode(&rgb, 8, 6)).unwrap();
    let prior = prior_path();
    let o = wai_in(&d, &["upscale", "run", "in.png", "--prior", &prior, "-o", "out.png", "--record", "r.json", "--envelope", "in.wai", "--key", SEED, "--json"]);
    let v: serde_json::Value = serde_json::from_slice(&ok(&o).stdout).unwrap();
    let v = &v["data"];
    assert_eq!((&v["width"], &v["height"], &v["subpixels"]), (&serde_json::json!(16), &serde_json::json!(12), &serde_json::json!(16 * 12 * 3)));

    let text = std::fs::read_to_string(d.join("r.json")).unwrap();
    let envelope = std::fs::read(d.join("in.wai")).unwrap();
    let (pw, ph, samples) = wai::png_still::decode(&std::fs::read(d.join("out.png")).unwrap()).unwrap();
    assert_eq!((pw, ph), (16, 12));
    let u = ex::read_upscale(&text).unwrap();
    use sha2::{Digest, Sha256};
    assert_eq!(Sha256::digest(&samples).as_slice(), u.output, "the picture is the record's output");
    let case = |priors: Vec<Vec<u8>>| ex::Case {
        upscales: vec![text.clone()],
        objects: vec![envelope.clone()],
        decoded: vec![(wai::jwp_hash::content_hash(&envelope), rgb.clone())],
        priors,
        policy: ex::Policy { require_proof: true, ..Default::default() },
        ..Default::default()
    };
    assert!(matches!(ex::verify(&case(vec![std::fs::read(&prior).unwrap()])), ex::Outcome::Accept(_)));
    assert!(!matches!(ex::verify(&case(vec![])), ex::Outcome::Accept(_)), "with no prior, nothing is proved");

    // `wai upscale verify` holds what it wrote, and refuses what it did not.
    let verify = |d: &Path, args: &[&str]| -> (bool, serde_json::Value) {
        let o = wai_in(d, &[&["upscale", "verify", "r.json", "--envelope", "in.wai"], args].concat());
        (o.status.success(), serde_json::from_slice(&o.stdout).unwrap_or(serde_json::Value::Null))
    };
    let (good, v) = verify(&d, &["--prior", &prior, "--picture", "out.png"]);
    assert!(good, "{v}");
    assert_eq!((&v["outcome"], &v["subpixels"]), (&serde_json::json!("accept"), &serde_json::json!(16 * 12 * 3)));
    let mut other = samples.clone();
    other[0] ^= 1;
    std::fs::write(d.join("other.png"), wai::png_still::encode(&other, 16, 12)).unwrap();
    let (good, v) = verify(&d, &["--prior", &prior, "--picture", "other.png"]);
    assert!(!good && v["reason"] == "picture_mismatch", "{v}");
    let title = Path::new(env!("CARGO_MANIFEST_DIR")).join("../live-conformance/upscale/priors/title.wsr");
    let (good, v) = verify(&d, &["--prior", title.to_str().unwrap()]);
    assert!(!good && v["reason"] == "upscale_not_proved", "another prior proves nothing: {v}");

    // The picture alone needs no key; a record does, and names the envelope.
    ok(&wai_in(&d, &["upscale", "run", "in.png", "--prior", &prior, "-o", "bare.png"]));
    assert_eq!(std::fs::read(d.join("bare.png")).unwrap(), std::fs::read(d.join("out.png")).unwrap(), "the derivation is exact");
    assert!(!wai_in(&d, &["upscale", "run", "in.png", "--prior", &prior, "-o", "x.png", "--record", "x.json", "--envelope", "x.wai"]).status.success(), "a record needs a key");
    assert!(!wai_in(&d, &["upscale", "run", "in.png", "--prior", &prior, "-o", "x.png", "--record", "x.json", "--key", SEED]).status.success(), "a record needs its envelope");
    assert!(!wai_in(&d, &["upscale", "run", "in.png", "--prior", &prior, "-o", "x.png", "--derivation", "sr3"]).status.success());
    std::fs::write(d.join("grey.png"), wai::png_still::encode_grey(&[9u8; 48], 8, 6)).unwrap();
    assert!(!wai_in(&d, &["upscale", "run", "grey.png", "--prior", &prior, "-o", "x.png"]).status.success(), "the reader holds 8-bit RGB");
}

fn hdr_envelope(w: u32, h: u32) -> Vec<u8> {
    use wai::container::{Rendition, WaiMulti};
    let e = tm::corpus::encode(&tm::corpus::scene(0, w, h), w, h, 4, false, (1, 1), false);
    let companion = tm::wgm1(&e.metadata.to_bytes(), &wai::png_still::encode_grey(&e.map, e.map_w, e.map_h));
    WaiMulti::new(
        "image",
        "replicate",
        serde_json::Value::Null,
        vec![
            Rendition::media(wai::codecs::CAP_IMAGE_PNG, "png", wai::png_still::encode(&e.base, w, h)).with_id("base").with_payload_digest(),
            Rendition::companion(tm::CAP_GAINMAP, companion, wai::codecs::CAP_IMAGE_PNG).with_id("gainmap").with_companion_of_id("base").with_payload_digest(),
        ],
    )
    .to_bytes()
    .unwrap()
}

#[test]
fn a_tone_mapping_writes_what_exact_tonemap_proves() {
    let d = scratch("tonemap");
    let envelope = hdr_envelope(16, 12);
    std::fs::write(d.join("in.wai"), &envelope).unwrap();
    let o = wai_in(&d, &["tonemap", "run", "in.wai", "--headrooms", "20000,10000", "--out", "r", "--record", "r.json", "--key", SEED, "--json"]);
    let v: serde_json::Value = serde_json::from_slice(&ok(&o).stdout).unwrap();
    assert_eq!(v["data"]["subpixels"], 2 * 16 * 12 * 3);

    let text = std::fs::read_to_string(d.join("r.json")).unwrap();
    let r = tm::read_record(&text).unwrap();
    // Ascending by headroom, each file the bytes its picture states.
    assert_eq!(r.pictures.iter().map(|p| p.headroom).collect::<Vec<_>>(), vec![10_000, 20_000]);
    for p in &r.pictures {
        let b = std::fs::read(d.join(format!("r/rendition-h{}.rgb16", p.headroom))).unwrap();
        assert_eq!(tm::digest(&b), p.output, "headroom {}", p.headroom);
    }
    let mut held = tm::Held::default();
    held.hold_envelope(envelope);
    assert!(matches!(tm::verify(std::slice::from_ref(&text), &held, tm::Policy::RequireProof), tm::Outcome::Accept(_)));
    assert!(!matches!(tm::verify(std::slice::from_ref(&text), &tm::Held::default(), tm::Policy::RequireProof), tm::Outcome::Accept(_)), "with no envelope, nothing is proved");

    // `wai tonemap verify` holds what it wrote, and refuses what it did not.
    let o = wai_in(&d, &["tonemap", "verify", "r.json", "--envelope", "in.wai", "--renditions", "r"]);
    let v: serde_json::Value = serde_json::from_slice(&ok(&o).stdout).unwrap();
    assert_eq!((&v["outcome"], &v["subpixels"]), (&serde_json::json!("accept"), &serde_json::json!(2 * 16 * 12 * 3)));
    let path = d.join("r/rendition-h20000.rgb16");
    let mut b = std::fs::read(&path).unwrap();
    b[0] ^= 1;
    std::fs::write(&path, b).unwrap();
    let o = wai_in(&d, &["tonemap", "verify", "r.json", "--envelope", "in.wai", "--renditions", "r"]);
    let v: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap();
    assert!(!o.status.success() && v["reason"] == "rendition_mismatch" && v["at"] == serde_json::json!(["renditions", 1]), "{v}");
    std::fs::write(d.join("other.wai"), hdr_envelope(8, 6)).unwrap();
    let o = wai_in(&d, &["tonemap", "verify", "r.json", "--envelope", "other.wai"]);
    let v: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap();
    assert!(!o.status.success() && v["reason"] == "rendition_not_proved", "another envelope proves nothing: {v}");

    assert!(!wai_in(&d, &["tonemap", "run", "in.wai", "--headrooms", "10,10", "--out", "x"]).status.success(), "headrooms are distinct");
    assert!(!wai_in(&d, &["tonemap", "run", "in.wai", "--headrooms", "60001", "--out", "x"]).status.success(), "past the largest headroom");
    assert!(!wai_in(&d, &["tonemap", "run", "in.wai", "--headrooms", "ten", "--out", "x"]).status.success());
    assert!(!wai_in(&d, &["tonemap", "run", "in.wai", "--headrooms", "10", "--out", "x", "--input-class", "invented"]).status.success());
    std::fs::write(d.join("one.wai"), b"WAI1 not this").unwrap();
    assert!(!wai_in(&d, &["tonemap", "run", "one.wai", "--headrooms", "10", "--out", "x"]).status.success(), "a gain map needs a WAI2");
}