open-standards/joule-code/joule-code-rs/crates/joule-moq-web/src/e2ee.rs

//! The member of an end-to-end encrypted room (WAI `extensions/confidentiality.md`), for a
//! page: feature `e2ee`, built as a wasm of its own (`joule_moq_web_e2ee`) that a page loads
//! for such a room alone.
//!
//! - [`check_record`]: the room's confidentiality record (`WRE1`, §3) checked against the
//!   room id, the descriptor's authority and its recording mode.
//! - [`E2eeMember`]: `joule-moq-mls`'s [`RoomMember`], whose MLS state and every key it
//!   derives live in the page's memory alone and are never written anywhere. It reads the
//!   room's lifecycle and its MLS messages in the order the room's carrier gives them (§5.1),
//!   follows the group from outside until a Welcome makes it a member, creates the group or
//!   commits when the rule names it (§5.4, §5.5), staging each until the carrier returns it,
//!   seals each encoded frame before it is published (§4.9), and opens each object it
//!   receives through the receiving checks in their order (§6), so that a payload that does
//!   not open never reaches a decoder (§10).

use crate::{js, namespace_of};
use joule_moq::carriage::prop as wai_prop;
use joule_moq::draft22::kvp::{prop, KeyValuePair, Properties};
use joule_moq::secure_objects::room::{split_key_id, Incoming, NotOpened, Outgoing};
use joule_moq::secure_objects::{Plaintext, Suite};
use joule_moq::TrackName;
use joule_moq_mls::group::Error as MlsError;
use joule_moq_mls::{Event, Received, Record, RoomMember};
use serde_json::{json, Value};
use wasm_bindgen::prelude::*;

fn key32(b: &[u8], what: &str) -> Result<[u8; 32], JsError> {
    b.try_into().map_err(|_| JsError::new(&format!("{what} is 32 bytes")))
}

/// The confidentiality record `record` checked as §3.3 orders it, for the room `room`
/// whose descriptor names `authority` and the recording `mode`: `{"ok": true, "suite"}` or
/// `{"ok": false, "reason"}` with the extension's reason.
pub fn record_outcome(record: &[u8], room: &[u8; 32], authority: &[u8; 32], mode: u8) -> Value {
    match Record::check(record, room, authority, mode) {
        Ok((r, _)) => json!({"ok": true, "suite": r.suite, "mls_suite": r.mls_suite}),
        Err(r) => json!({"ok": false, "reason": r.code()}),
    }
}

#[wasm_bindgen(js_name = e2eeCheckRecord)]
pub fn check_record(record: &[u8], room: &[u8], authority: &[u8], mode: u8) -> Result<String, JsError> {
    Ok(record_outcome(record, &key32(room, "a room id")?, &key32(authority, "an authority's key")?, mode).to_string())
}

fn error_json(e: &MlsError) -> Value {
    match e {
        MlsError::Refused(r) => json!({"ok": false, "reason": r.code()}),
        MlsError::Mls(m) => json!({"ok": false, "error": "mls", "detail": m}),
        MlsError::NotNow(m) => json!({"ok": false, "error": "not_now", "detail": m}),
    }
}

/// A sealed object, for a page to sign and send: its properties encoded (MOQT §8.3), its
/// payload, and its KID.
#[wasm_bindgen]
pub struct SealedOut {
    properties: Vec<u8>,
    payload: Vec<u8>,
    kid: u64,
}

#[wasm_bindgen]
impl SealedOut {
    /// The object's property list: the Immutable Properties with the Key ID and the
    /// envelope flag, and the delivery timeout beside them.
    pub fn properties(&mut self) -> Vec<u8> {
        std::mem::take(&mut self.properties)
    }
    /// The sealed payload, ciphertext and tag, moved out.
    pub fn payload(&mut self) -> Vec<u8> {
        std::mem::take(&mut self.payload)
    }
    #[wasm_bindgen(getter)]
    pub fn kid(&self) -> u64 {
        self.kid
    }
}

/// What opening an object gave: [`OpenOut::json`] says how it went, and when it opened,
/// [`OpenOut::payload`] is the original payload.
#[wasm_bindgen]
pub struct OpenOut {
    json: String,
    payload: Vec<u8>,
}

#[wasm_bindgen]
impl OpenOut {
    /// `{"ok": true, "epoch", "leaf", "sender", "digest"}` (`digest`: whether the plaintext
    /// carried `WAI_INNER_DIGEST`, checked); `{"ok": false, "reason"}`, refused with the
    /// extension's reason; or `{"ok": false, "pending": <epoch>}`, sealed in an epoch this
    /// member has not reached, to be held until it applies the commit that opens it (§5.8).
    #[wasm_bindgen(getter)]
    pub fn json(&self) -> String {
        self.json.clone()
    }
    /// The original payload, moved out; empty unless it opened.
    pub fn payload(&mut self) -> Vec<u8> {
        std::mem::take(&mut self.payload)
    }
}

/// A commit staged for the carrier: the commit and, when it adds members, their Welcome.
#[wasm_bindgen]
pub struct StagedCommit {
    commit: Vec<u8>,
    welcome: Option<Vec<u8>>,
    epoch: u64,
}

#[wasm_bindgen]
impl StagedCommit {
    pub fn commit(&self) -> Vec<u8> {
        self.commit.clone()
    }
    pub fn welcome(&self) -> Option<Vec<u8>> {
        self.welcome.clone()
    }
    /// The epoch the commit opens.
    #[wasm_bindgen(getter)]
    pub fn epoch(&self) -> u64 {
        self.epoch
    }
}

/// A principal of an end-to-end encrypted room, as a member of its group.
#[wasm_bindgen]
pub struct E2eeMember {
    m: RoomMember,
}

impl E2eeMember {
    pub fn member(&self) -> &RoomMember {
        &self.m
    }

    pub fn member_mut(&mut self) -> &mut RoomMember {
        &mut self.m
    }

    /// Seal `payload` as object (`group`, `object`) of the track `name` in the namespace
    /// `ns` (§4.9), under this member's latest epoch: with `WAI_INNER_DIGEST` when `digest`,
    /// the envelope flag in the Immutable Properties when `envelope`, and an
    /// OBJECT_DELIVERY_TIMEOUT beside them when `timeout_ms` is not 0.
    #[allow(clippy::too_many_arguments)]
    pub fn seal_object(&self, ns: &joule_moq::TrackNamespace, name: &TrackName, group: u64, object: u64, priority: u8, payload: &[u8], digest: bool, envelope: bool, timeout_ms: u64) -> Result<(Properties, Vec<u8>, u64), String> {
        let pt = if digest { Plaintext::with_digest(payload.to_vec()) } else { Plaintext { payload: payload.to_vec(), ..Default::default() } };
        let immutable: Vec<KeyValuePair> = if envelope { vec![KeyValuePair::varint(wai_prop::WAI_ENVELOPE, 1)] } else { vec![] };
        let mutable: Vec<KeyValuePair> = if timeout_ms > 0 { vec![KeyValuePair::varint(prop::OBJECT_DELIVERY_TIMEOUT, timeout_ms)] } else { vec![] };
        let o = Outgoing { ns, name, group, object, priority, plaintext: &pt, immutable: &immutable, mutable: &mutable };
        let s = self.m.seal(&o).map_err(|e| e.to_string())?;
        Ok((s.properties, s.payload, s.kid))
    }

    /// Open an object through the checks of §6, in their order.
    #[allow(clippy::too_many_arguments)]
    pub fn open_object(&mut self, ns: &joule_moq::TrackNamespace, name: &TrackName, group: u64, object: u64, priority: u8, properties: &Properties, payload: &[u8]) -> (Value, Vec<u8>) {
        let inc = Incoming { ns, name, group, object, priority, properties, payload };
        match self.m.open(&inc) {
            Ok(o) => (json!({"ok": true, "epoch": o.epoch, "leaf": o.leaf, "sender": hex::encode(o.sender), "digest": o.plaintext.inner_digest.is_some()}), o.plaintext.payload),
            Err(NotOpened::Refused(r)) => (json!({"ok": false, "reason": r.code()}), Vec::new()),
            Err(NotOpened::Pending(e)) => (json!({"ok": false, "pending": e}), Vec::new()),
        }
    }

    /// The member as a page shows it and a test reads it.
    pub fn state_json(&self) -> Value {
        let k = self.m.keyring();
        let latest = k.latest();
        let roster = latest.and_then(|e| k.roster(e)).map(|r| r.iter().map(|(l, p)| (l.to_string(), json!(hex::encode(p)))).collect::<serde_json::Map<_, _>>());
        let g = self.m.model().group();
        json!({
            "epoch": self.m.epoch(),
            "joined": k.joined_at(),
            "removed": k.removed(),
            "held": k.held(),
            "leaf": k.my_leaf().filter(|_| self.m.epoch().is_some()),
            "roster": roster,
            "group": g.map(|g| json!({
                "generation": g.generation,
                "epoch": g.epoch,
                "roster": g.roster.iter().map(|(l, p)| (l.to_string(), json!(hex::encode(p)))).collect::<serde_json::Map<_, _>>(),
            })),
            "pending": self.m.model().pending().len(),
            "staged": self.m.staged(),
            "following": self.m.following(),
            "committer": self.m.is_committer(),
        })
    }
}

#[wasm_bindgen]
impl E2eeMember {
    /// The principal of the 32-byte Ed25519 `seed` in the room `room`, whose clock's key is
    /// `clock`, sealing under the record's `suite` (`0x0004` or `0x0005`).
    #[wasm_bindgen(constructor)]
    pub fn new(room: &[u8], clock: &[u8], seed: &[u8], suite: u16) -> Result<E2eeMember, JsError> {
        let suite = match Suite::from_id(suite) {
            Some(s @ (Suite::AesGcm128Sha256 | Suite::AesGcm256Sha512)) => s,
            _ => return Err(JsError::new("a room's suite is 0x0004 or 0x0005")),
        };
        let key = ed25519_dalek::SigningKey::from_bytes(&key32(seed, "a seed")?);
        Ok(E2eeMember { m: RoomMember::new(key32(room, "a room id")?, key32(clock, "a clock's key")?, key, suite) })
    }

    /// A KeyPackage for the clock (the keying draft §4.7), framed; its private keys stay
    /// here.
    #[wasm_bindgen(js_name = keyPackage)]
    pub fn key_package(&mut self) -> Result<Vec<u8>, JsError> {
        self.m.key_package().map_err(js)
    }

    /// A lifecycle event of the room, in `joule-moq-mls`'s form:
    /// `{"event": "join", "subject": <hex>, "keyed": bool}`,
    /// `{"event": "leave", "subject", "reason"}` or
    /// `{"event": "grant_exchange", "subject", "keyed_before", "keyed_after"}`.
    pub fn event(&mut self, ev: &str) -> Result<(), JsError> {
        let ev: Event = serde_json::from_str(ev).map_err(js)?;
        self.m.event(&ev);
        Ok(())
    }

    /// A framed MLS message from the room's carrier, in its order: `{"ok": true,
    /// "received": "proposal" | "epoch" | "removed" | "joined" | "created" | "followed" |
    /// "ignored", "epoch"}`, or why it was refused.
    pub fn receive(&mut self, frame: &[u8]) -> String {
        let v = match self.m.receive(frame) {
            Ok(r) => {
                let (what, epoch) = match r {
                    Received::Proposal => ("proposal", None),
                    Received::Epoch(e) => ("epoch", Some(e)),
                    Received::Removed(e) => ("removed", Some(e)),
                    Received::Joined(e) => ("joined", Some(e)),
                    Received::Created => ("created", Some(0)),
                    Received::Followed => ("followed", None),
                    Received::Ignored => ("ignored", None),
                };
                json!({"ok": true, "received": what, "epoch": epoch})
            }
            Err(e) => error_json(&e),
        };
        v.to_string()
    }

    /// Whether the room owes this principal the group's creation now (§5.4).
    #[wasm_bindgen(js_name = owesCreate)]
    pub fn owes_create(&self) -> bool {
        self.m.owes_create()
    }

    /// The group made and held until the carrier returns its GroupInfo: the GroupInfo.
    #[wasm_bindgen(js_name = stageCreate)]
    pub fn stage_create(&mut self) -> Result<Vec<u8>, JsError> {
        self.m.stage_create().map_err(js)
    }

    /// A member of a conversation (WAI live-messaging §7.2): any member a commit does not
    /// remove may commit, the room service's log ordering the commits.
    #[wasm_bindgen(js_name = anyMemberCommits)]
    pub fn any_member_commits(&mut self) {
        self.m.set_commit_rule(joule_moq_mls::CommitRule::AnyMember);
    }

    /// Whether this member should commit now: it is the committer (§5.5), the clock has
    /// proposals pending, and nothing of its own is staged.
    #[wasm_bindgen(js_name = shouldCommit)]
    pub fn should_commit(&self) -> bool {
        !self.m.staged() && self.m.is_committer() && self.m.clock_pending()
    }

    /// The commit of the clock's pending proposals, held until the carrier returns it.
    #[wasm_bindgen(js_name = stageCommit)]
    pub fn stage_commit(&mut self) -> Result<StagedCommit, JsError> {
        let c = self.m.stage_commit(false).map_err(js)?;
        Ok(StagedCommit { commit: c.commit, welcome: c.welcome, epoch: c.epoch })
    }

    /// Drop what is staged: the carrier refused it.
    pub fn withdraw(&mut self) {
        self.m.withdraw();
    }

    pub fn staged(&self) -> bool {
        self.m.staged()
    }

    /// The member's current epoch, or nothing while it is not a member.
    pub fn epoch(&self) -> Option<u64> {
        self.m.epoch()
    }

    /// Delete `epoch`'s keys (§5.8): its objects are refused `epoch_not_joined` from now on.
    pub fn forget(&mut self, epoch: u64) {
        self.m.keyring_mut().forget(epoch);
    }

    /// Seal for a page: `ns` a JSON array of the namespace's fields.
    #[allow(clippy::too_many_arguments)]
    pub fn seal(&self, ns: &str, name: &str, group: u64, object: u64, priority: u8, payload: &[u8], digest: bool, envelope: bool, timeout_ms: u64) -> Result<SealedOut, JsError> {
        let ns = namespace_of(&serde_json::from_str(ns).map_err(js)?).map_err(js)?;
        let (p, payload, kid) = self.seal_object(&ns, &TrackName::from_str(name), group, object, priority, payload, digest, envelope, timeout_ms).map_err(js)?;
        Ok(SealedOut { properties: p.to_bytes().map_err(|e| js(e.to_string()))?, payload, kid })
    }

    /// Open for a page: `properties` as [`crate::UniItem::properties`] gives them.
    #[allow(clippy::too_many_arguments)]
    pub fn open(&mut self, ns: &str, name: &str, group: u64, object: u64, priority: u8, properties: &[u8], payload: &[u8]) -> Result<OpenOut, JsError> {
        let ns = namespace_of(&serde_json::from_str(ns).map_err(js)?).map_err(js)?;
        // A property list that does not read is §6 step 1's.
        let Ok(props) = Properties::decode(properties, "properties") else {
            return Ok(OpenOut { json: json!({"ok": false, "reason": "properties_malformed"}).to_string(), payload: Vec::new() });
        };
        let (v, payload) = self.open_object(&ns, &TrackName::from_str(name), group, object, priority, &props, payload);
        Ok(OpenOut { json: v.to_string(), payload })
    }

    pub fn state(&self) -> String {
        self.state_json().to_string()
    }
}

/// An object's property list changed as a relay could change it, for a test that plays
/// that relay: `kind` is `drop_key_id` (the Immutable Properties removed, so no Key ID),
/// `key_id_outside` (a Key ID added outside them as well), or `key_id` (the Key ID inside
/// them replaced by `kid`, every other property kept).
pub fn altered_properties(properties: &[u8], kind: &str, kid: u64) -> Result<Vec<u8>, String> {
    let mut p = Properties::decode(properties, "properties").map_err(|e| e.to_string())?;
    match kind {
        "drop_key_id" => p.0.retain(|k| k.ty != prop::IMMUTABLE_PROPERTIES),
        "key_id_outside" => {
            p.push(KeyValuePair::varint(joule_moq::secure_objects::KEY_ID_PROPERTY, kid));
        }
        "key_id" => {
            let inner = p.immutable().map_err(|e| e.to_string())?.unwrap_or_default();
            let others: Vec<KeyValuePair> = inner.0.into_iter().filter(|k| k.ty != joule_moq::secure_objects::KEY_ID_PROPERTY).collect();
            let sip = joule_moq::secure_objects::immutable_properties(kid, &others);
            let mutable: Vec<KeyValuePair> = p.0.into_iter().filter(|k| k.ty != prop::IMMUTABLE_PROPERTIES).collect();
            p = joule_moq::secure_objects::object_properties(&sip, &mutable);
        }
        other => return Err(format!("no alteration `{other}`")),
    }
    p.to_bytes().map_err(|e| e.to_string())
}

#[wasm_bindgen(js_name = e2eeAlteredProperties)]
pub fn altered_properties_js(properties: &[u8], kind: &str, kid: u64) -> Result<Vec<u8>, JsError> {
    altered_properties(properties, kind, kid).map_err(js)
}

/// The KID of a member at leaf `leaf` in `epoch` (§5.7).
#[wasm_bindgen(js_name = e2eeKid)]
pub fn kid(epoch: u64, leaf: u32) -> Option<u64> {
    joule_moq::secure_objects::room::key_id(epoch, leaf)
}

/// The epoch and leaf index a KID names (§5.7), as `[epoch, leaf]`.
#[wasm_bindgen(js_name = e2eeSplitKid)]
pub fn split_kid(kid: u64) -> Vec<u64> {
    let (e, l) = split_key_id(kid);
    vec![e, l as u64]
}

#[cfg(test)]
#[path = "e2ee_tests.rs"]
mod tests;