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

//! The C ABI's envelope surface (feature `ffi`, no codec library): unpack returns
//! each manifest exactly as the envelope carries it, the `WAI2` rendition table
//! carries every component's identity and whether SPEC §7.1 may select it, packing
//! keeps members this library does not model, `wai_envelope_select` picks by §7.1
//! and reports each skipped entry's reason, and the packers refuse a candidate
//! capability (SPEC §5).
#![cfg(feature = "ffi")]

use std::ffi::{CStr, CString};
use std::ptr;

use wai::container::{Prior, Rendition, WaiMulti};
use wai::ffi::*;

fn empty() -> WaiBuffer {
    WaiBuffer { data: ptr::null_mut(), len: 0, cap: 0 }
}

fn take(b: WaiBuffer) -> Vec<u8> {
    let v = if b.data.is_null() { Vec::new() } else { unsafe { std::slice::from_raw_parts(b.data, b.len) }.to_vec() };
    unsafe { wai_buffer_free(b) };
    v
}

fn last_error() -> String {
    let p = wai_last_error();
    if p.is_null() { String::new() } else { unsafe { CStr::from_ptr(p) }.to_string_lossy().into_owned() }
}

fn stream_wai() -> Vec<u8> {
    std::fs::read(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_mlicpp/stream.wai")).unwrap()
}

#[test]
fn unpack_returns_the_original_manifest_bytes() {
    let env = stream_wai();
    let (mut m, mut p) = (empty(), empty());
    assert_eq!(unsafe { wai_envelope_unpack(env.as_ptr(), env.len(), &mut m, &mut p) }, WAI_OK);
    let m = take(m);
    assert_eq!(m.len(), 339);
    assert_eq!(m, env[8..8 + 339]);
    assert_eq!(take(p).len(), env.len() - 8 - 339 - 4);

    // A manifest with `security`, whitespace and a member this library does not
    // model comes back unchanged.
    let man = br#"{"wai":"1.0", "media":"text","intent":"replicate","model_requirement":{"capability":"wai.text.zstd"},"conditioning":{"kind":"zstd"},"security":{"encryption":"mls-sframe","kid":2},"target":null,"later":1}"#;
    let mut e = b"WAI1".to_vec();
    e.extend_from_slice(&(man.len() as u32).to_le_bytes());
    e.extend_from_slice(man);
    e.extend_from_slice(&1u32.to_le_bytes());
    e.push(7);
    let (mut m, mut p) = (empty(), empty());
    assert_eq!(unsafe { wai_envelope_unpack(e.as_ptr(), e.len(), &mut m, &mut p) }, WAI_OK);
    assert_eq!(take(m), man);
    take(p);
}

fn multi() -> Vec<u8> {
    WaiMulti::new("video", "replicate", serde_json::Value::Null, vec![
        Rendition::media("wai.video.lcevc", "lcevc", b"MISLABELLED".to_vec()),
        Rendition::media("wai.neural.int_mlicpp", "int_mlicpp", b"PINNED".to_vec())
            .with_id("pinned")
            .with_prior(Prior::new("m", "ab".repeat(32), "int")),
        Rendition::media("wai.video.hevc", "hevc", b"HEVC".to_vec()).with_id("base").with_payload_digest(),
        Rendition::companion("wai.meta.hdr10plus", b"META".to_vec(), "wai.video.hevc").with_companion_of_id("base"),
    ]).to_bytes().unwrap()
}

#[test]
fn unpack_multi_table_carries_identity_and_selectability() {
    let env = multi();
    let (mut m, mut t, mut b) = (empty(), empty(), empty());
    assert_eq!(unsafe { wai_envelope_unpack_multi(env.as_ptr(), env.len(), &mut m, &mut t, &mut b) }, WAI_OK);
    let ml = u32::from_le_bytes(env[4..8].try_into().unwrap()) as usize;
    assert_eq!(take(m), env[8..8 + ml]);
    let table: serde_json::Value = serde_json::from_slice(&take(t)).unwrap();
    take(b);
    let sel: Vec<bool> = table.as_array().unwrap().iter().map(|e| e["selectable"].as_bool().unwrap()).collect();
    assert_eq!(sel, [false, true, true, false], "a companion-only entry and a companion are never selectable");
    assert_eq!(table[1]["prior"]["sha256"], "ab".repeat(32));
    assert_eq!(table[1]["id"], "pinned");
    assert_eq!(table[2]["sha256"], wai::pin::sha256_hex(b"HEVC"));
    assert_eq!(table[3]["role"], "companion");
    assert_eq!(table[3]["companion_of"], "wai.video.hevc");
    assert_eq!(table[3]["companion_of_id"], "base");
    assert!(table[0].get("prior").is_none() && table[0].get("id").is_none());
}

#[test]
fn envelope_select_follows_spec_7_1() {
    let env = multi();
    let all = CString::new(r#"["wai.video.lcevc","wai.neural.int_mlicpp","wai.video.hevc","wai.meta.hdr10plus"]"#).unwrap();
    let mut idx = u32::MAX;
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), all.as_ptr(), &mut idx) }, WAI_OK);
    // Entry 0 is companion-only, entry 1's pin applies and this call holds no set.
    assert_eq!(idx, 2);
    let none = CString::new(r#"["wai.neural.int_mlicpp"]"#).unwrap();
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), none.as_ptr(), &mut idx) }, WAI_ERR_NO_RENDITION);
    let why = last_error();
    assert!(why.contains("0: not_selectable") && why.contains("1: prior_not_held: ") && why.contains("2: missing_capability"), "{why}");
    // A bad argument is reported, not a panic.
    let bad = CString::new("not json").unwrap();
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), bad.as_ptr(), &mut idx) }, WAI_ERR_INVALID_ARG);
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), all.as_ptr(), ptr::null_mut()) }, WAI_ERR_INVALID_ARG);
}

#[test]
fn pack_multi_keeps_every_member() {
    let manifest = CString::new(r#"{"wai":"1.1","media":"image","intent":"replicate","renditions":[{"capability":"wai.image.png","kind":"png","id":"p","sha256":"0000000000000000000000000000000000000000000000000000000000000000","later":{"x":1}}],"target":null}"#).unwrap();
    let payload = b"PNG";
    let ptrs = [payload.as_ptr()];
    let lens = [payload.len()];
    let mut out = empty();
    assert_eq!(unsafe { wai_envelope_pack_multi(manifest.as_ptr(), ptrs.as_ptr(), lens.as_ptr(), 1, &mut out) }, WAI_OK);
    let env = take(out);
    let ml = u32::from_le_bytes(env[4..8].try_into().unwrap()) as usize;
    assert_eq!(&env[8..8 + ml], manifest.as_bytes());
    // A mismatched payload count is an argument error.
    let mut out = empty();
    assert_eq!(unsafe { wai_envelope_pack_multi(manifest.as_ptr(), ptrs.as_ptr(), lens.as_ptr(), 0, &mut out) }, WAI_ERR_INVALID_ARG);
}

/// SPEC §4, §5 and §7.1: a candidate is never selected, whatever the sink
/// advertises; the table marks it not selectable, and selection reports why.
#[test]
fn envelope_select_never_picks_a_candidate() {
    let env = WaiMulti::new("binary", "improve", serde_json::Value::Null, vec![
        Rendition::media("wai.tensor.nnc", "bytes", b"TENSOR".to_vec()),
        Rendition::media("wai.image.png", "png", b"PNG".to_vec()),
    ]).to_bytes().unwrap();
    let (mut m, mut t, mut b) = (empty(), empty(), empty());
    assert_eq!(unsafe { wai_envelope_unpack_multi(env.as_ptr(), env.len(), &mut m, &mut t, &mut b) }, WAI_OK);
    take(m);
    take(b);
    let table: serde_json::Value = serde_json::from_slice(&take(t)).unwrap();
    let sel: Vec<bool> = table.as_array().unwrap().iter().map(|e| e["selectable"].as_bool().unwrap()).collect();
    assert_eq!(sel, [false, true], "a candidate is never selectable");
    let all = CString::new(r#"["wai.tensor.nnc","wai.image.png"]"#).unwrap();
    let mut idx = u32::MAX;
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), all.as_ptr(), &mut idx) }, WAI_OK);
    assert_eq!(idx, 1);
    let only = CString::new(r#"["wai.tensor.nnc"]"#).unwrap();
    assert_eq!(unsafe { wai_envelope_select(env.as_ptr(), env.len(), only.as_ptr(), &mut idx) }, WAI_ERR_NO_RENDITION);
    let why = last_error();
    assert!(why.contains("0: not_selectable") && why.contains("1: missing_capability"), "{why}");
}

/// SPEC §5, *Candidate capabilities*: an encoder MUST NOT emit one, as the
/// capability, the fallback, or a `WAI2` component, and the C packers refuse it.
#[test]
fn envelope_pack_refuses_a_candidate() {
    let man = |cap: &str, fb: &str| CString::new(format!(
        r#"{{"wai":"1.0","media":"binary","intent":"improve","model_requirement":{{"capability":"{cap}","fallback":{fb}}},"conditioning":{{"kind":"bytes"}},"target":null}}"#)).unwrap();
    let payload = b"bytes";
    for (cap, fb) in [("wai.splat.sog", "null"), ("wai.image.png", "\"wai.tensor.nnc\"")] {
        let m = man(cap, fb);
        let mut out = empty();
        let rc = unsafe { wai_envelope_pack(m.as_ptr(), payload.as_ptr(), payload.len(), &mut out) };
        assert_eq!(rc, WAI_ERR_INVALID_ARG, "{cap} / {fb}");
        assert!(out.data.is_null());
        assert!(last_error().contains("MUST NOT emit"), "{}", last_error());
    }
    let ok = man("wai.image.png", "null");
    let mut out = empty();
    assert_eq!(unsafe { wai_envelope_pack(ok.as_ptr(), payload.as_ptr(), payload.len(), &mut out) }, WAI_OK);
    take(out);

    let ptrs = [payload.as_ptr(), payload.as_ptr()];
    let lens = [payload.len(), payload.len()];
    for comps in [
        r#"[{"capability":"wai.splat.sog","kind":"bytes"}]"#,
        r#"[{"capability":"wai.video.av1","kind":"av1"},{"capability":"wai.tensor.nnc","kind":"bytes","role":"companion","companion_of":"wai.video.av1"}]"#,
    ] {
        let multi = CString::new(format!(
            r#"{{"wai":"1.1","media":"binary","intent":"improve","renditions":{comps},"target":null}}"#)).unwrap();
        let n = serde_json::from_str::<serde_json::Value>(comps).unwrap().as_array().unwrap().len();
        let mut out = empty();
        let rc = unsafe { wai_envelope_pack_multi(multi.as_ptr(), ptrs.as_ptr(), lens.as_ptr(), n, &mut out) };
        assert_eq!(rc, WAI_ERR_INVALID_ARG, "{comps}");
        assert!(out.data.is_null());
    }
}

/// SPEC §5, *Payloads of the WAI-native integer codecs*: the C packers refuse a payload
/// that is not the format SPEC §5 pins its integer capability (or fallback, or `WAI2`
/// component) to, parsed as a sink parses it: a tool's JSON form, a `WIV1` cut inside
/// its header, and a keyframe-0 `WIV1` under `wai.video.int_hyper`, whose magic and
/// version are right. The committed payloads are written.
#[test]
fn envelope_pack_refuses_an_integer_payload_that_is_not_its_format() {
    let vec = |rel: &str| std::fs::read(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors").join(rel)).unwrap();
    let (k0, k1, json) = (vec("int_video/int_video_bitstream.wiv1"), vec("int_video/video_nk_bitstream.wiv1"), vec("int_video/int_video_bitstream.json"));
    let man = |cap: &str, fb: &str| CString::new(format!(
        r#"{{"wai":"1.0","media":"video","intent":"improve","model_requirement":{{"capability":"{cap}","fallback":{fb}}},"conditioning":{{"kind":"bytes"}},"target":null}}"#)).unwrap();
    let pack = |m: &CString, p: &[u8]| {
        let mut out = empty();
        let rc = unsafe { wai_envelope_pack(m.as_ptr(), p.as_ptr(), p.len(), &mut out) };
        if rc == WAI_OK { take(out); } else { assert!(out.data.is_null()); }
        (rc, if rc == WAI_OK { String::new() } else { last_error() })
    };
    let (mot, hyp) = ("wai.video.int_motion", "wai.video.int_hyper");
    for (m, p, code) in [
        (man(mot, "null"), &json[..], "bad_magic"),
        (man(mot, "null"), &k0[..40], "stream_length"),
        (man(hyp, "null"), &k0[..], "keyframe_capability_mismatch"),
        (man("wai.video.av1", &format!("\"{mot}\"")), &k1[..], "keyframe_capability_mismatch"),
    ] {
        let (rc, why) = pack(&m, p);
        assert_eq!(rc, WAI_ERR_INVALID_ARG, "{code}");
        assert!(why.contains(&format!("({code}:")), "{why}");
    }
    assert_eq!(pack(&man(mot, "null"), &k0).0, WAI_OK);
    assert_eq!(pack(&man(hyp, "null"), &k1).0, WAI_OK);

    let multi = CString::new(format!(
        r#"{{"wai":"1.1","media":"video","intent":"improve","renditions":[{{"capability":"{mot}","kind":"bytes"}},{{"capability":"{hyp}","kind":"bytes"}}],"target":null}}"#)).unwrap();
    let pack2 = |a: &[u8], b: &[u8]| {
        let (ptrs, lens) = ([a.as_ptr(), b.as_ptr()], [a.len(), b.len()]);
        let mut out = empty();
        let rc = unsafe { wai_envelope_pack_multi(multi.as_ptr(), ptrs.as_ptr(), lens.as_ptr(), 2, &mut out) };
        if rc == WAI_OK { take(out); } else { assert!(out.data.is_null()); }
        (rc, if rc == WAI_OK { String::new() } else { last_error() })
    };
    assert_eq!(pack2(&k0, &k1).0, WAI_OK);
    let (rc, why) = pack2(&k0, &k0);
    assert_eq!(rc, WAI_ERR_INVALID_ARG);
    assert!(why.contains("component 1:") && why.contains("(keyframe_capability_mismatch:"), "{why}");
}

/// `wai_int_payload_declared` reads what an integer payload's header declares, before
/// any decode: a `WIV1`'s clip bytes and symbols, and, for keyframe 0, which runs no
/// network, its working bytes and no multiply-accumulates; a `WIH1`'s symbols, its
/// image size and its cost being its networks'; and refuses a payload that is not the
/// capability's format, or a capability with none.
#[test]
fn int_payload_declared_reads_the_header() {
    let vec = |rel: &str| std::fs::read(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors").join(rel)).unwrap();
    let declared = |cap: &str, p: &[u8]| {
        let c = CString::new(cap).unwrap();
        let mut out = WaiIntDeclared::default();
        let rc = unsafe { wai_int_payload_declared(c.as_ptr(), p.as_ptr(), p.len(), std::ptr::null(), &mut out) };
        (rc, out, if rc == WAI_OK { String::new() } else { last_error() })
    };
    let k0 = vec("int_video/int_video_bitstream.wiv1");
    let v = wai::int_wire::parse_wiv1(&k0, Some(0)).unwrap();
    let clip = v.n_frames as u64 * v.h as u64 * v.w as u64 * 3;
    let (rc, out, _) = declared("wai.video.int_motion", &k0);
    assert_eq!((rc, out.output_bytes, out.output_known, out.symbols), (WAI_OK, clip, 1, clip));
    assert_eq!((out.cost_known, out.working_bytes, out.macs_lo, out.macs_hi), (1, v.keyframe0_working_bytes(), 0, 0));
    assert!(out.working_bytes > clip, "the clip, its table, its motion vectors and its frame list");
    let wih1 = vec("int_hyper/tiny_bitstream.wih1");
    let k = wai::int_wire::parse_wih1(&wih1).unwrap();
    let (rc, out, _) = declared("wai.neural.int_hyper", &wih1);
    assert_eq!((rc, out.output_known, out.symbols), (WAI_OK, 0, (k.n_z * k.zh * k.zw + k.m * k.sh * k.sw) as u64));
    assert_eq!(out.cost_known, 0, "a WIH1's cost is its networks'");
    let (rc, _, why) = declared("wai.video.int_motion", &vec("int_video/int_video_bitstream.json"));
    assert!(rc == WAI_ERR_INVALID_ARG && why.contains("(bad_magic:"), "{why}");
    let (rc, _, _) = declared("wai.image.png", &k0);
    assert_eq!(rc, WAI_ERR_INVALID_ARG);
}