open-standards/joule-code/joule-code-rs/crates/joule-moq-remote/tests/loopback.rs

//! A remote desktop session over loopback: the host agent, a relay and three viewers, every
//! record of the session then verified whole by `wai-rs` (WAI remote-desktop §11).
//!
//! The relay requires grants with their proofs and knows the room's authority and the clock's
//! token alone. The host keeps a floor and gives it on request. Alice (`remote_control` 3:
//! pointer and keys) requests it, and her click is applied; bob (no control power) is refused
//! `input_not_permitted`; carol (pointer) is refused `input_no_floor`; alice types a string and
//! composes a text, which go sealed to the host's sealing key (remote-desktop §5.5) and are
//! injected as typed, while bob, who follows her events, receives no trace of them, and her key
//! sent in the clear is refused `input_unsealed`; the floor is handed to carol, and alice is
//! refused while carol's move is applied; the host pauses, and carol is refused
//! `input_screen_paused`; it resumes. Each viewer verifies every frame against its receipt, and
//! alice and carol sign latency records of the frames their outcomes name. Then the session's
//! chains, as the host holds them, verify at the commitment level, as a viewer verifies them, and
//! at the opened level under the host's opening key: every outcome the host signed is what the
//! rules give, every seal opens to what alice typed, the floor changes as it did, and each
//! latency's legs are stated.
//!
//! One test runs with the pure-function encoder anywhere; on macOS a second runs with the
//! platform's hardware encoder and checks the units are the Annex B the capability carries.
//! They take a lock, so they run one at a time on a loaded host.

use joule_moq_relay::config::{Authority, Grant as RelayGrant};
use joule_moq_relay::{RelayConfig, RelayServer};
use joule_moq_remote::agent::{self, Command, Config, Decided, Room};
use joule_moq_remote::capture::Pattern;
use joule_moq_remote::encode::{Encoder, Ipcm};
use joule_moq_remote::inject::Recorder;
use joule_moq_remote::room::{new_room, unhex32, Spec};
use joule_moq_remote::viewer::Viewer;
use sha2::{Digest, Sha256};
use std::time::Duration;
use wai::remote_desktop::{self as rd, class, floor_action, screen_state, Action, Fate, Refusal, SealFact, TimedAction};

static ONE_AT_A_TIME: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

const STEP: Duration = Duration::from_secs(20);

fn at(dt_us: u32, action: Action) -> TimedAction {
    TimedAction { dt_us, action }
}

async fn session(make_encoder: impl FnOnce(u32, u32) -> Box<dyn Encoder>, check_units: bool) {
    let _one = ONE_AT_A_TIME.lock().await;
    let (w, h) = (320u32, 192u32);
    let t0 = joule_moq_remote::capture::now_unix_us() - 1_000_000;
    let spec = Spec { width: w as u16, height: h as u16, fps: 20, participants: vec![("alice".into(), class::POINTER | class::KEYS), ("bob".into(), 0), ("carol".into(), class::POINTER)], grant_s: 3600 };
    let (rf, keys, _authority, files) = new_room(&[0x5a; 32], t0, &spec).expect("a room");
    let mut c = RelayConfig::loopback();
    c.auth.required = true;
    c.auth.authorities = vec![Authority { key: rf.authority.clone(), rooms: vec![rf.room.clone()] }];
    let token = hex::decode(&keys.clock_token).unwrap();
    c.auth.grants = vec![RelayGrant { prefix: vec![rf.room.clone()], token_sha256: hex::encode(Sha256::digest(&token)), publish: true, subscribe: true, proof: false, unlimited: true, seek: false }];
    let server = RelayServer::start(c).await.expect("the relay starts");
    let url = format!("ws://{}{}", server.websocket_addr.expect("a WebSocket listener"), server.relay().config().path);

    let pattern = Pattern::new(w, h);
    let recorder = Recorder::new(pattern.drawn.clone());
    let injected = recorder.log.clone();
    let encoder = make_encoder(w, h);
    let identity = encoder.identity();
    let cfg = Config {
        relay_url: url.clone(),
        trust: joule_moq_remote::tls::Trust::system(),
        opening_key: None,
        room: Room::Desk { room: rf.clone(), keys },
        fps: 20,
        max_age_ms: 3_000,
        floor: rd::floor_mode::FLOOR,
        give_on_request: true,
        redactions: vec![rd::Rect { x: 0, y: 0, w: 64, h: 16 }],
        clip_max: 1000,
        pointer: true,
        key_every: Duration::from_secs(2),
        run_for: Some(Duration::from_secs(90)),
    };
    let (cmd, cmds) = tokio::sync::mpsc::unbounded_channel();
    let host = tokio::spawn(agent::run(cfg, Box::new(pattern), encoder, Box::new(recorder), cmds));

    let file = |n: &str| files.iter().find(|f| f.me == n).cloned().expect("a participant");
    let mut alice = Viewer::join(&file("alice"), &url).await.expect("alice joins");
    let mut bob = Viewer::join(&file("bob"), &url).await.expect("bob joins");
    let mut carol = Viewer::join(&file("carol"), &url).await.expect("carol joins");
    for (n, v) in [("alice", &mut alice), ("bob", &mut bob), ("carol", &mut carol)] {
        assert!(v.pump(STEP, |v| v.joined && v.screen.is_some() && v.frames.len() >= 3).await, "{n}: joined, the screen statement and frames ({} frames, joined {})", v.frames.len(), v.joined);
    }
    let screen = alice.screen.clone().unwrap().1;
    assert_eq!(screen.encoder, identity.encoder, "the screen statement names the encoder");
    assert_eq!((screen.width as u32, screen.height as u32), (w, h));

    // Alice asks for the floor; the host gives it on request.
    alice.floor(floor_action::REQUEST, [0; 32]).await.unwrap();
    let me_a = alice.me;
    assert!(alice.pump(STEP, |v| v.holder == Some(me_a)).await, "the floor given to alice");
    let a1 = alice.input(vec![at(6_000, Action::Move { x: 100, y: 50 }), at(3_000, Action::Button { button: 0, down: true, x: 100, y: 50 }), at(0, Action::Button { button: 0, down: false, x: 100, y: 50 })]).await.unwrap();
    assert!(alice.pump(STEP, |v| v.outcomes.get(&a1).is_some_and(|o| o.2.is_some_and(|an| v.frames.contains_key(&an.group)))).await, "alice's click applied and its frame held");
    let la = alice.latency(a1, 0).await.expect("alice's latency record");
    assert!(la.latency_us > 0, "a latency is measured");
    assert!(screen.seal.is_some(), "the host states a sealing key");

    // Alice types a string and composes a text: sealed to the host, injected as typed. Bob
    // follows her events, as any viewer of the room may.
    bob.follow(alice.me).await.expect("bob follows alice's events");
    let mut typed = Vec::new();
    for (k, u) in [0x0Bu16, 0x18, 0x11, 0x17, 0x08, 0x15, 0x1F].iter().enumerate() {
        typed.push(at(20_000 - 2_000 * k as u32, Action::Key { usage: *u, down: true, modifiers: 0 }));
        typed.push(at(19_000 - 2_000 * k as u32, Action::Key { usage: *u, down: false, modifiers: 0 }));
    }
    typed.push(at(0, Action::Text("pässwörd-hunter2".into())));
    let a3 = alice.input(typed.clone()).await.unwrap();
    assert!(alice.pump(STEP, |v| v.outcomes.contains_key(&a3)).await, "alice's typing answered");
    assert_eq!(alice.outcomes[&a3].1.reason, 0, "applied");
    // Her key sent in the clear to a host that states a sealing key: refused, never injected.
    let a4 = alice.input_on(alice.screen.as_ref().unwrap().0, vec![at(0, Action::Text("in the clear".into()))]).await.unwrap();
    assert!(alice.pump(STEP, |v| v.outcomes.contains_key(&a4)).await);
    assert_eq!(alice.outcomes[&a4].1.reason, Refusal::Unsealed.number(), "a text in the clear is refused");
    assert!(bob.pump(STEP, |v| v.received.iter().any(|r| wai::live_events::Event::parse(r).is_ok_and(|e| e.principal == me_a && rd::read_statement(&e.payload).is_some_and(|s| matches!(s, Ok(rd::Statement::Input(i)) if i.sealed_count() > 0))))).await, "bob receives alice's sealed input");

    // Bob carries no control power; carol has it and not the floor.
    let b1 = bob.input(vec![at(0, Action::Move { x: 10, y: 10 })]).await.unwrap();
    assert!(bob.pump(STEP, |v| v.outcomes.contains_key(&b1)).await, "bob's input answered");
    assert_eq!(bob.outcomes[&b1].1.reason, Refusal::NotPermitted.number());
    let c1 = carol.input(vec![at(0, Action::Move { x: 20, y: 20 })]).await.unwrap();
    assert!(carol.pump(STEP, |v| v.outcomes.contains_key(&c1)).await, "carol's input answered");
    assert_eq!(carol.outcomes[&c1].1.reason, Refusal::NoFloor.number());

    // The floor handed to carol.
    carol.floor(floor_action::REQUEST, [0; 32]).await.unwrap();
    cmd.send(Command::GiveFloor(carol.me)).unwrap();
    let me_c = carol.me;
    assert!(carol.pump(STEP, |v| v.holder == Some(me_c)).await, "the floor handed to carol");
    assert!(alice.pump(STEP, |v| v.holder == Some(me_c)).await);
    let a2 = alice.input(vec![at(0, Action::Key { usage: 0x04, down: true, modifiers: 0 })]).await.unwrap();
    assert!(alice.pump(STEP, |v| v.outcomes.contains_key(&a2)).await);
    assert_eq!(alice.outcomes[&a2].1.reason, Refusal::NoFloor.number(), "alice, off the floor");
    let c2 = carol.input(vec![at(0, Action::Move { x: 200, y: 120 })]).await.unwrap();
    assert!(carol.pump(STEP, |v| v.outcomes.get(&c2).is_some_and(|o| o.2.is_some_and(|an| v.frames.contains_key(&an.group)))).await, "carol's move applied");
    carol.latency(c2, 0).await.expect("carol's latency record");

    // Paused, then live again.
    cmd.send(Command::Pause).unwrap();
    assert!(carol.pump(STEP, |v| v.screen.as_ref().is_some_and(|s| s.1.state == screen_state::PAUSED)).await, "the host pauses");
    let c3 = carol.input(vec![at(0, Action::Move { x: 5, y: 5 })]).await.unwrap();
    assert!(carol.pump(STEP, |v| v.outcomes.contains_key(&c3)).await);
    assert_eq!(carol.outcomes[&c3].1.reason, Refusal::ScreenPaused.number());
    cmd.send(Command::Resume).unwrap();
    assert!(carol.pump(STEP, |v| v.screen.as_ref().is_some_and(|s| s.1.state == screen_state::LIVE)).await, "the host resumes");
    // Let every viewer's last records reach the host before it stops.
    for v in [&mut alice, &mut bob, &mut carol] {
        v.pump(Duration::from_millis(700), |_| false).await;
    }
    cmd.send(Command::Stop).unwrap();
    let report = tokio::time::timeout(Duration::from_secs(60), host).await.expect("the host stops").expect("the task").expect("the session");
    for v in [&alice, &bob, &carol] {
        assert_eq!(v.refused_frames, 0, "every frame and receipt verified");
        v.close();
    }
    server.shutdown().await;

    // What the host decided, and what it injected.
    let decided = |d: &[u8; 32]| report.decided.iter().find(|(x, _)| x == d).map(|(_, o)| o.clone());
    assert!(matches!(decided(&a1), Some(Decided::Applied { .. })), "{:?}", decided(&a1));
    assert_eq!(decided(&b1), Some(Decided::Refused(Refusal::NotPermitted)));
    assert_eq!(decided(&c1), Some(Decided::Refused(Refusal::NoFloor)));
    assert_eq!(decided(&a2), Some(Decided::Refused(Refusal::NoFloor)));
    assert!(matches!(decided(&c2), Some(Decided::Applied { .. })));
    assert_eq!(decided(&c3), Some(Decided::Refused(Refusal::ScreenPaused)));
    let log = injected.lock().unwrap().clone();
    assert_eq!(log.len(), 4 + typed.len(), "alice's three actions, what she typed and carol's move, nothing else: {log:?}");
    let opened: Vec<Action> = typed.iter().map(|a| a.action.clone()).collect();
    assert_eq!(&log[3..3 + typed.len()], &opened[..], "the host injected what alice typed, opened");
    // Bob's received records hold no trace of what alice typed: not the text, not a key's
    // fields as an action in the clear carries them.
    for needle in [&b"hunter2"[..], "pässwörd".as_bytes(), &[action_type_key(), 0x0B, 0x00, 0x01, 0x00]] {
        assert!(!bob.received.iter().any(|r| r.windows(needle.len()).any(|w| w == needle)), "bob's records hold {needle:?}");
    }
    assert!(bob.received.iter().any(|r| r.windows(b"in the clear".len()).any(|w| w == b"in the clear")), "the search finds what was sent in the clear");
    assert!(report.frames >= 20 && report.key_frames >= 1, "{} frames", report.frames);

    // The session, verified whole from the records the host holds: at the commitment level, and
    // at the opened level under the host's opening key, every seal opened to what was typed.
    let vi = rd::VerifyInput { descriptor: &report.descriptor, grants: &report.grants, head: None, chains: &report.chains };
    let opened = rd::verify_opened(&vi, &[report.opening_key.0]);
    assert_eq!(opened.outcome, rd::Outcome::Accept, "the session verifies at the opened level");
    let fact = opened.inputs.iter().find(|f| f.digest == a3).expect("alice's typing");
    let Some(SealFact::Opened(st)) = &fact.seal else { panic!("its seal opened: {:?}", fact.seal) };
    assert_eq!(rd::read_statement(st), Some(Ok(rd::Statement::Input(rd::Input::clear(rf_host(&rf), alice.sent[&a3].1.screen, typed.clone())))), "it opens to what alice typed");
    let run = rd::verify(&vi);
    assert_eq!(run.outcome, rd::Outcome::Accept, "the session verifies");
    let fate = |d: &[u8; 32]| run.inputs.iter().find(|f| f.digest == *d).map(|f| f.fate);
    assert!(matches!(fate(&a1), Some(Fate::Applied { .. })));
    assert_eq!(fate(&b1), Some(Fate::Refused(Refusal::NotPermitted)));
    assert_eq!(fate(&c1), Some(Fate::Refused(Refusal::NoFloor)));
    assert_eq!(fate(&c3), Some(Fate::Refused(Refusal::ScreenPaused)));
    let holders: Vec<Option<[u8; 32]>> = run.floor.iter().map(|c| c.holder).collect();
    assert!(holders.starts_with(&[Some(alice.me), Some(carol.me)]), "the floor to alice, then carol: {holders:?}");
    assert_eq!(run.screens.len(), 3, "live, paused, live");
    assert_eq!(run.latencies.len(), 2);
    for l in &run.latencies {
        let legs = l.legs.expect("checked");
        assert_eq!(legs.up + legs.host + legs.down, l.latency_us as i64);
        assert_eq!(legs.uncertainty, Some(0), "the host is the room's clock and the viewers read it on one machine");
        eprintln!("latency {} µs: up {} µs, host {} µs, down {} µs ({})", l.latency_us, legs.up, legs.host, legs.down, identity.encoder);
    }
    let host_pk = unhex32(&rf.host.principal).unwrap();
    assert_eq!(run.totals[&host_pk].applied, 3);
    assert_eq!(run.totals[&host_pk].refused[Refusal::Unsealed.number() as usize - 1], 1);
    if check_units {
        assert!(report.bytes > 0);
    }
}

fn action_type_key() -> u8 {
    rd::action_type::KEY
}

fn rf_host(rf: &joule_moq_remote::room::RoomFile) -> [u8; 32] {
    unhex32(&rf.host.principal).unwrap()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_session_with_the_pure_function_encoder() {
    tokio::time::timeout(Duration::from_secs(240), session(|w, h| Box::new(Ipcm::new(w, h).unwrap()), false)).await.expect("the session finishes in time");
}

#[cfg(target_os = "macos")]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_session_with_the_platform_hardware_encoder() {
    use joule_moq_remote::adapter_macos_encode::{Codec, PlatformEncoder};
    let e = PlatformEncoder::new(320, 192, 20, 1_000, Codec::Avc).expect("the platform encoder");
    eprintln!("platform encoder: {:?}", e.identity());
    drop(e);
    tokio::time::timeout(Duration::from_secs(240), session(|w, h| Box::new(PlatformEncoder::new(w, h, 20, 1_000, Codec::Avc).unwrap()), true)).await.expect("the session finishes in time");
}

/// The platform encoder's units are what `wai.video.avc` carries: one picture each, Annex B, the
/// parameter sets before every IDR picture, an IDR picture when asked.
#[cfg(target_os = "macos")]
#[test]
fn the_platform_encoder_writes_annex_b_access_units() {
    use joule_moq_remote::adapter_macos_encode::{Codec, PlatformEncoder};
    use joule_moq_remote::capture::Capture;
    use joule_moq_remote::encode::{avc_types, avc_unit_ok};
    let mut e = PlatformEncoder::new(320, 192, 30, 1_000, Codec::Avc).expect("the platform encoder");
    let mut p = Pattern::new(320, 192);
    let mut keys = 0;
    for i in 0..30 {
        p.drawn.lock().unwrap().pointer = (i * 7, i * 3);
        let u = e.encode(&p.grab().unwrap(), i % 10 == 0).expect("a unit");
        assert!(avc_unit_ok(&u), "picture {i}: {:?}", avc_types(&u.data));
        if i % 10 == 0 {
            assert!(u.key, "picture {i}: an IDR picture when asked");
        }
        keys += u.key as u32;
    }
    assert!(keys >= 3);
    let id = e.identity();
    assert!(rd::is_encoder_name(id.encoder.as_bytes()), "{}", id.encoder);
    eprintln!("{} on the media engine: {}", id.encoder, id.hardware == rd::hardware::MEDIA_ENGINE);
}