open-standards/joule-code/joule-code-rs/crates/joule-moq-room/tests/co_performance.rs

//! A room with co-performance (WAI co-performance §8) kept by the service, against a relay on
//! loopback that holds the floor and the audience's inputs:
//!
//! - a lead's schedule that marks, asked for over HTTP with the lead's grant and its proof of
//!   possession, signed by the clock; a request without a proof, with a proof used before, and
//!   from a grant without the direct power, refused;
//! - two performers trading turns: an object in turn, one out of turn and one later than its
//!   budget, each of the last two marked by the clock as it takes them;
//! - an AI participant's intents: one played as announced, one not;
//! - two audience members, recorded joining by the clock, voting in a window a lead opens; the
//!   clock commits to the window once it closes and signs its selection, credited;
//! - a fidelity statement asked for by the lead;
//!
//! and the room as the relay serves it, with the clock's log of objects, verifies as
//! co-performance §9 verifies a room: accept, every object's marks founded, no object out of
//! turn or late without its mark.

use ed25519_dalek::SigningKey;
use joule_moq::draft22::data::{FetchEntry, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::kvp::{prop, KeyValuePair, Properties};
use joule_moq::draft22::Parameters;
use joule_moq::grant_proof::{self as gp, GrantProof};
use joule_moq::{TrackName, TrackNamespace};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::config::{Authority, FloorConfig, Grant};
use joule_moq_relay::transport::mux::MuxConfig;
use joule_moq_relay::transport::websocket;
use joule_moq_relay::{RelayConfig, RelayServer};
use joule_moq_room::keys::Keys;
use joule_moq_room::{http, Config, Service};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::time::timeout;
use wai::co_performance::{self as co, Statement};
use wai::live_events::{kind, Caption, Event, Payload};

const STEP: Duration = Duration::from_secs(20);
const S: u64 = 1_000_000;

async fn relay(keys: &Keys) -> RelayServer {
    let mut c = RelayConfig::loopback();
    c.quic_listen = None;
    c.auth.required = true;
    c.auth.authorities = vec![Authority { key: hex::encode(keys.authority.verifying_key().to_bytes()), rooms: vec![] }];
    let t = hex::encode(Sha256::digest(keys.clock_token));
    let clock = hex::encode(keys.clock.verifying_key().to_bytes());
    c.auth.grants = vec![
        Grant { prefix: vec!["*".into(), clock.clone(), "events".into()], token_sha256: t.clone(), publish: true, subscribe: false, proof: false, unlimited: true, seek: false },
        Grant { prefix: vec![], token_sha256: t, publish: false, subscribe: true, proof: false, unlimited: false, seek: true },
    ];
    c.floor = Some(FloorConfig { room_keys: vec![clock], ..FloorConfig::default() });
    RelayServer::start(c).await.expect("the relay starts")
}

/// One HTTP/1.1 request: its status, its headers (lowercase names) and its JSON.
async fn call(addr: std::net::SocketAddr, method: &str, path: &str, body: Option<&Value>, headers: &[(&str, String)]) -> (u16, BTreeMap<String, String>, Value) {
    let mut s = tokio::net::TcpStream::connect(addr).await.unwrap();
    let b = body.map(|v| v.to_string()).unwrap_or_default();
    let mut req = format!("{method} {path} HTTP/1.1\r\nHost: room.test\r\nCF-Connecting-IP: 10.7.0.1\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n", b.len());
    for (k, v) in headers {
        req.push_str(&format!("{k}: {v}\r\n"));
    }
    req.push_str("\r\n");
    req.push_str(&b);
    s.write_all(req.as_bytes()).await.unwrap();
    let mut out = Vec::new();
    timeout(STEP, s.read_to_end(&mut out)).await.expect("an answer in time").unwrap();
    let text = String::from_utf8(out).unwrap();
    let (head, body) = text.split_once("\r\n\r\n").expect("a response");
    let status: u16 = head.split(' ').nth(1).unwrap().parse().unwrap();
    let hs = head.lines().skip(1).filter_map(|l| l.split_once(':')).map(|(k, v)| (k.trim().to_ascii_lowercase(), v.trim().to_string())).collect();
    (status, hs, serde_json::from_str(body).unwrap_or(Value::Null))
}

fn key(i: u8) -> SigningKey {
    SigningKey::from_bytes(&[i; 32])
}

fn pk(k: &SigningKey) -> [u8; 32] {
    k.verifying_key().to_bytes()
}

fn pubhex(k: &SigningKey) -> String {
    hex::encode(pk(k))
}

/// A page's identity in the room: its key and its grant.
struct Who {
    key: SigningKey,
    grant: Vec<u8>,
    challenge: Option<[u8; 32]>,
}

impl Who {
    /// A co-performance request with this grant and a proof over the last challenge the
    /// service gave it, asking for one first when it holds none.
    async fn ask(&mut self, addr: std::net::SocketAddr, path: &str, body: &Value) -> (u16, Value) {
        for _ in 0..2 {
            let mut headers = vec![("Authorization", format!("Bearer {}", gp::base64url(&self.grant)))];
            if let Some(c) = self.challenge.take() {
                let (digest, _) = gp::grant_digest_and_principal(&self.grant).unwrap();
                let p = GrantProof::sign(digest, &self.key, c, &gp::context("POST", "room.test", path)).unwrap();
                headers.push((gp::HEADER_PROOF, gp::base64url(&p.to_bytes())));
            }
            let (st, hs, v) = call(addr, "POST", path, Some(body), &headers).await;
            self.challenge = hs.get(&gp::HEADER_CHALLENGE.to_ascii_lowercase()).and_then(|c| gp::from_base64url(c)).and_then(|b| b.try_into().ok());
            if st != 401 {
                return (st, v);
            }
        }
        (401, Value::Null)
    }
}

async fn session(server: &RelayServer, who: &Who) -> Client {
    let url = format!("ws://{}{}", server.websocket_addr.unwrap(), server.relay().config().path);
    let dial = || async { Ok::<joule_moq_relay::transport::SharedConnection, String>(Arc::new(websocket::connect(&url, MuxConfig::default()).await.map_err(|e| e.to_string())?)) };
    let opts = ClientOptions { grants: vec![(who.grant.clone(), who.key.clone())], target: format!("//{}{}", server.websocket_addr.unwrap(), server.relay().config().path), ..ClientOptions::default() };
    // The relay reads this SETUP before it sends its own, so the grant is proved over a
    // challenge its SETUP gave: on a second session, where the first held none.
    Client::connect_proving(dial, opts).await.expect("SETUP")
}

/// Publish `fields` of `who`'s own namespace, each track `0` at an alias of its own.
async fn publish(c: &Client, room: &str, who: &SigningKey, fields: &[(&str, u64)]) -> Vec<(joule_moq_relay::wire::ReqTx, joule_moq_relay::wire::ReqRx)> {
    let mut held = Vec::new();
    for (field, alias) in fields {
        let ns = TrackNamespace::from_strs([room, pubhex(who).as_str(), field]);
        let mut p = Properties::new();
        p.push(KeyValuePair::varint(prop::DEFAULT_PUBLISHER_PRIORITY, 64));
        held.push(c.publish(ns, TrackName::from_str("0"), *alias, Parameters::new(), p).await.expect("PUBLISH"));
    }
    held
}

/// One object of group `g` at `alias`, on a stream of its own.
async fn send(c: &Client, alias: u64, g: u64, payload: Vec<u8>) {
    let h = SubgroupHeader { track_alias: alias, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(64), properties: false, end_of_group: true, first_object: true };
    c.send_subgroup(h, &[SubgroupObject::normal(0, payload)]).await.expect("an object");
}

/// A `WAI1` envelope whose `target` states `due` on the room clock, as a page in a room with
/// co-performance stamps its media.
fn envelope(due: u64, body: &[u8]) -> Vec<u8> {
    let m = json!({"wai": "1.0", "media": "audio", "intent": "replicate", "model_requirement": {"capability": "wai.audio.opus", "fallback": null}, "conditioning": {"kind": "opus"}, "target": {"timestamp_us": due}}).to_string();
    let mut b = b"WAI1".to_vec();
    b.extend_from_slice(&(m.len() as u32).to_le_bytes());
    b.extend_from_slice(m.as_bytes());
    b.extend_from_slice(&(body.len() as u32).to_le_bytes());
    b.extend_from_slice(body);
    b
}

/// A principal's events track as a page writes it: each record an object of its own.
struct Events {
    key: SigningKey,
    room: [u8; 32],
    t0: u64,
    seq: u64,
    prev: [u8; 32],
    alias: u64,
}

impl Events {
    async fn emit(&mut self, c: &Client, t: u64, k: u8, payload: Vec<u8>) -> [u8; 32] {
        let ev = Event::sign(&self.key, self.room, self.seq, self.prev, t, None, k, payload);
        let rec = ev.to_bytes();
        let entry = wai::live_events::to_msf_record(&rec, self.t0).unwrap();
        send(c, self.alias, self.seq, format!("[{entry}]").into_bytes()).await;
        self.prev = ev.digest();
        self.seq += 1;
        ev.digest()
    }
}

/// Every record of a principal's events track as the relay serves it, by `seq`.
async fn chain(c: &Client, room: &str, who: &str, t0: u64) -> Vec<Vec<u8>> {
    let ns = TrackNamespace::from_strs([room, who, "events"]);
    let Ok((_, entries)) = c.fetch(ns, TrackName::from_str("0"), Parameters::new()).await else { return Vec::new() };
    let mut by_seq = BTreeMap::new();
    for e in entries {
        let FetchEntry::Object(o) = e else { continue };
        let arr: Vec<Value> = serde_json::from_slice(&o.payload).unwrap();
        for r in arr {
            let rec = wai::live_events::from_msf_record(&r.to_string(), t0).unwrap();
            by_seq.insert(Event::parse(&rec).unwrap().seq, rec);
        }
    }
    by_seq.into_values().collect()
}

fn statements(chain: &[Vec<u8>]) -> Vec<Statement> {
    chain.iter().filter_map(|r| Event::parse(r).ok()).filter(|e| e.kind == kind::META).filter_map(|e| Statement::parse(&e.payload).ok()).collect()
}

fn now_us() -> u64 {
    joule_moq_room::now_us()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_room_with_co_performance_marks_commits_selects_and_verifies() {
    timeout(Duration::from_secs(240), run()).await.expect("the test finishes in time");
}

async fn run() {
    let keys = Keys::generate();
    let server = relay(&keys).await;
    let cfg = Config::from_toml(&format!(
        "listen = \"127.0.0.1:0\"\nkeys = \"/nonexistent\"\nrelay_url = \"ws://{}{}\"\n[room]\nmax_publishers = 3\njoin_wait_s = 30\n[rate]\nclient_burst = 20\n[co]\naudience_input = 2\nlatency_budget_ms = 1500\nlead = \"first\"\nai = true\nfidelity_every_s = 0\ncommit_after_ms = 800\nlead_ms = 200\n",
        server.websocket_addr.unwrap(),
        server.relay().config().path
    ))
    .unwrap();
    let svc = Service::new(Keys::from_toml(&keys.to_toml()).unwrap(), &cfg);
    let listener = tokio::net::TcpListener::bind(cfg.listen).await.unwrap();
    let addr = listener.local_addr().unwrap();
    svc.start();
    tokio::spawn(http::serve(svc.clone(), listener));
    timeout(STEP, async {
        while !svc.connected.load(Ordering::Acquire) {
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    })
    .await
    .expect("the clock's session opens");

    // Alice (the lead: the room's first publisher), Bob (an AI participant), two audience members.
    let model = [0x9d; 32];
    let grant = |k: SigningKey, body: Value| {
        async move {
            let (st, _, g) = call(addr, "POST", "/grant", Some(&body), &[]).await;
            assert_eq!(st, 200, "{g}");
            (Who { key: k, grant: hex::decode(g["grant"].as_str().unwrap()).unwrap(), challenge: None }, g)
        }
    };
    let (alice_k, bob_k, vera_k, val_k) = (key(0x31), key(0x32), key(0x33), key(0x34));
    let (mut alice, ga) = grant(alice_k.clone(), json!({"principal": pubhex(&alice_k), "role": "publisher"})).await;
    let (mut bob, _) = grant(bob_k.clone(), json!({"principal": pubhex(&bob_k), "role": "publisher", "ai": hex::encode(model)})).await;
    let (mut vera, gv) = grant(vera_k.clone(), json!({"principal": pubhex(&vera_k), "role": "viewer"})).await;
    let (val, _) = grant(val_k.clone(), json!({"principal": pubhex(&val_k), "role": "viewer"})).await;
    let room = ga["room"]["room"].as_str().unwrap().to_string();
    let room_id: [u8; 32] = hex::decode(&room).unwrap().try_into().unwrap();
    let t0 = ga["room"]["t0_us"].as_u64().unwrap();
    assert!(gv["room"]["co"]["policy"].is_string(), "a room with co-performance states its policy");
    let rc = || now_us() - t0;

    // Everyone enters; the clock records the publishers joining when their video appears, and
    // the audience members when their events track does.
    let ca = session(&server, &alice).await;
    let cb = session(&server, &bob).await;
    let cv = session(&server, &vera).await;
    let cl = session(&server, &val).await;
    let _ha = publish(&ca, &room, &alice_k, &[("audio", 10), ("video", 11), ("events", 12)]).await;
    let _hb = publish(&cb, &room, &bob_k, &[("audio", 10), ("video", 11), ("events", 12)]).await;
    let _hv = publish(&cv, &room, &vera_k, &[("events", 12)]).await;
    let _hl = publish(&cl, &room, &val_k, &[("events", 12)]).await;
    let clock = hex::encode(keys.clock.verifying_key().to_bytes());
    let watcher = session(&server, &Who { key: val_k.clone(), grant: val.grant.clone(), challenge: None }).await;
    let mut joins = 0;
    for _ in 0..200 {
        joins = chain(&watcher, &room, &clock, t0).await.iter().filter(|r| Event::parse(r).unwrap().kind == kind::JOIN).count();
        if joins == 4 {
            break;
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    assert_eq!(joins, 4, "two publishers and two audience members joined");

    // Requests the service refuses: no proof (answered with a challenge), a schedule from a
    // grant without the direct power, a proof presented twice.
    let sched = |from: u64| json!({"mode": "mark", "from_us": from, "segments": [{"start_us": from, "length": 2_000_000, "turns": 4, "pairs": false, "holders": [pubhex(&alice_k), pubhex(&bob_k)]}]});
    let (st, hs, v) = call(addr, "POST", "/co/schedule", Some(&sched(rc() + S)), &[("Authorization", format!("Bearer {}", gp::base64url(&alice.grant)))]).await;
    assert!(st == 401 && v["reason"] == "pop_missing" && hs.contains_key("wai-room-challenge"), "{st} {v}");
    let (st, v) = bob.ask(addr, "/co/schedule", &sched(rc() + S)).await;
    assert_eq!(st, 403, "{v}");
    let c = gp::from_base64url(&hs["wai-room-challenge"]).unwrap().try_into().unwrap();
    let (digest, _) = gp::grant_digest_and_principal(&alice.grant).unwrap();
    let proof = gp::base64url(&GrantProof::sign(digest, &alice_k, c, &gp::context("POST", "room.test", "/co/fidelity")).unwrap().to_bytes());
    let headers = [("Authorization", format!("Bearer {}", gp::base64url(&alice.grant))), (gp::HEADER_PROOF, proof)];
    let (st, _, _) = call(addr, "POST", "/co/fidelity", Some(&json!({})), &headers).await;
    assert_eq!(st, 200, "the proof passes once");
    let (st, _, v) = call(addr, "POST", "/co/fidelity", Some(&json!({})), &headers).await;
    assert!(st == 401 && v["reason"] == "pop_replayed", "{st} {v}");

    // Alice states a schedule that marks: Alice and Bob trade two-second turns, Alice first.
    let from = rc() + S;
    let (st, v) = alice.ask(addr, "/co/schedule", &sched(from)).await;
    assert_eq!(st, 200, "{v}");
    let from = Statement::parse(&hex::decode(v["statement"].as_str().unwrap()).unwrap()).map(|s| match s {
        Statement::Schedule(x) => x.from,
        _ => panic!("a schedule"),
    });
    let from = from.unwrap();
    let until = |x: u64| async move {
        let now = now_us() - t0;
        if x > now {
            tokio::time::sleep(Duration::from_micros(x - now)).await;
        }
    };

    // Bob, an AI participant, announces two lines; it says the first as announced, the second
    // not. Its captions state that a machine made them, with its model.
    let mut eb = Events { key: bob_k.clone(), room: room_id, t0, seq: 0, prev: [0; 32], alias: 12 };
    let line = |at: u64, text: &str| Statement::Intent(co::Intent { track: [0; 32], at, span: S as u32, plan: co::Plan::Text(text.into()) }).encode();
    let caption = |text: &str| Payload::Caption(Caption { duration_ms: 1000, origin: 2, model: Some(model), lang: "en".into(), text: text.into() }).encode();
    let (at1, at2) = (from + 2_500_000, from + 3_500_000);
    eb.emit(&cb, rc(), kind::META, line(at1, "here comes the bridge")).await;
    eb.emit(&cb, rc(), kind::META, line(at2, "and back to the top")).await;

    // Alice's turn: her object in turn, Bob's out of turn.
    until(from + 500_000).await;
    send(&ca, 10, 1, envelope(rc(), b"alice in her turn")).await;
    until(from + 1_000_000).await;
    send(&cb, 10, 1, envelope(rc(), b"bob out of turn")).await;
    // Bob's turn: its first line as announced, a late object, its second line otherwise.
    until(at1).await;
    eb.emit(&cb, at1, kind::CAPTION, caption("here comes the bridge")).await;
    send(&cb, 10, 2, envelope(rc().saturating_sub(3 * S), b"bob, three seconds late")).await;
    until(at2).await;
    eb.emit(&cb, at2, kind::CAPTION, caption("and somewhere else")).await;

    // The audience: a vote of three options, closing in three seconds; both vote for option 1,
    // and Vera cannot open a window or hand a turn over.
    let (st, v) = vera.ask(addr, "/co/window", &json!({"rule": "plurality", "options": 3, "k": 1, "seconds": 3, "question": "which way?"})).await;
    assert_eq!(st, 403, "{v}");
    let (st, v) = alice.ask(addr, "/co/window", &json!({"rule": "plurality", "options": 3, "k": 1, "seconds": 3, "question": "which way?"})).await;
    assert_eq!(st, 200, "{v}");
    let w: [u8; 32] = hex::decode(v["event"].as_str().unwrap()).unwrap().try_into().unwrap();
    let vote = |o| Statement::Input(co::AudienceInput { window: w, body: co::InputBody::Vote(o) }).encode();
    tokio::time::sleep(Duration::from_millis(300)).await;
    let mut ev = Events { key: vera_k.clone(), room: room_id, t0, seq: 0, prev: [0; 32], alias: 12 };
    let mut el = Events { key: val_k.clone(), room: room_id, t0, seq: 0, prev: [0; 32], alias: 12 };
    ev.emit(&cv, rc(), kind::META, vote(1)).await;
    el.emit(&cl, rc(), kind::META, vote(1)).await;
    let mut sel = None;
    for _ in 0..100 {
        sel = statements(&chain(&watcher, &room, &clock, t0).await).into_iter().find_map(|s| match s {
            Statement::Selection(x) if x.window == w => Some(x),
            _ => None,
        });
        if sel.is_some() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(200)).await;
    }
    let sel = sel.expect("the clock selects once the window closes");
    assert_eq!((sel.picks, sel.credits[0].count), (co::Picks::Options(vec![1]), 2));

    // Once the lines are some seconds past, Alice asks for the fidelity statement.
    until(at2 + co_lag() + S).await;
    let (st, v) = alice.ask(addr, "/co/fidelity", &json!({})).await;
    assert!(st == 200 && v["ok"] == true, "{v}");

    // The room as the relay serves it, and the clock's log of objects.
    let (_, _, r) = call(addr, "GET", "/room", None, &[]).await;
    let descriptor = hex::decode(r["descriptor"].as_str().unwrap()).unwrap();
    let grants: Vec<Vec<u8>> = r["grants"].as_array().unwrap().iter().chain(r["co"]["audience"].as_array().unwrap()).map(|g| hex::decode(g.as_str().unwrap()).unwrap()).collect();
    let policies = vec![hex::decode(r["co"]["policy"].as_str().unwrap()).unwrap()];
    let mut chains = vec![chain(&watcher, &room, &clock, t0).await];
    for k in [&alice_k, &bob_k, &vera_k, &val_k] {
        chains.push(chain(&watcher, &room, &pubhex(k), t0).await);
    }
    let (_, _, log) = call(addr, "GET", "/co/objects?from=0", None, &[]).await;
    let objects: Vec<co::Object> = log["objects"]
        .as_array()
        .unwrap()
        .iter()
        .map(|o| co::Object {
            principal: hex::decode(o["principal"].as_str().unwrap()).unwrap().try_into().unwrap(),
            track: hex::decode(o["track"].as_str().unwrap()).unwrap().try_into().unwrap(),
            group: o["group"].as_u64().unwrap(),
            object: o["object"].as_u64().unwrap(),
            digest: hex::decode(o["digest"].as_str().unwrap()).unwrap().try_into().unwrap(),
            arrival: o["arrival"].as_u64().unwrap(),
            due: o["due"].as_u64().unwrap(),
        })
        .collect();
    assert_eq!(objects.len(), 3, "{log}");
    let run = co::verify(&co::Input { descriptor: &descriptor, grants: &grants, head: None, chains: &chains, policies: &policies, objects: &objects, payloads: &[] });
    assert_eq!(run.outcome, co::Outcome::Accept, "the room verifies");
    let floor: Vec<(co::Floor, bool, u8)> = run.objects.iter().map(|o| (o.floor, o.late, o.marks)).collect();
    assert_eq!(floor, vec![(co::Floor::InTurn, false, 0), (co::Floor::OutOfTurn, false, co::flag::OUT_OF_TURN), (co::Floor::InTurn, true, co::flag::LATE)]);
    let kept: Vec<&str> = run.intents.iter().map(|i| i.outcome.name()).collect();
    assert_eq!(kept, ["kept", "diverged"]);
    assert_eq!(run.fidelity, vec![(pk(&bob_k), co::Fidelity { announced: 2, kept: 1, diverged: 1, missed: 0 })]);
    assert!(statements(&chains[0]).iter().any(|s| matches!(s, Statement::Fidelity(f) if f.entries == run.fidelity)), "the clock's fidelity statement, recomputed");
    assert_eq!(run.windows[0].picks, Some(co::Picks::Options(vec![1])));

    // An object the clock did not take, stated out of turn with no mark, is found out.
    let mut stray = objects.clone();
    stray.push(co::Object { group: 9, arrival: objects[1].arrival, due: objects[1].arrival, ..objects[1] });
    let run = co::verify(&co::Input { descriptor: &descriptor, grants: &grants, head: None, chains: &chains, policies: &policies, objects: &stray, payloads: &[] });
    assert_eq!(run.outcome, co::Outcome::Refuse { at: co::At::Object(3), reason: co::Reason::TurnUnmarked });
    for c in [&ca, &cb, &cv, &cl, &watcher] {
        c.close(0, "done");
    }
    server.shutdown().await;
}

fn co_lag() -> u64 {
    joule_moq_room::co::FIDELITY_LAG_US
}