open-standards/joule-code/joule-code-rs/crates/joule-moq-ingest/examples/ingest_energy.rs

//! The two processes an ingest energy measurement reads: a relay, and a gateway that
//! publishes to it, each its own process so that a reader of per-process energy (WAI
//! energy-measurement §9.2, `macos-task`) reads one hop alone.
//!
//! ```text
//! ingest_energy relay                      # prints RELAY <addr> <certificate, DER, hex>
//! ingest_energy gateway <addr> <cert-hex>  # prints PID, then a stream key URL per line read
//! ```
//!
//! The room, its authority, the publisher and the gateway's key are fixed seeds, so both
//! processes hold the same room. The gateway issues each stream key itself (the `/rtmp`
//! resource's admission is not what is measured) and prints `RTMP rtmp://…/live/<key>` for
//! each line on its standard input; it stops at the end of its input. Every listener is
//! loopback.

use ed25519_dalek::SigningKey;
use joule_moq_ingest::room::{build, Room};
use joule_moq_ingest::{GatewayConfig, GatewayServer};
use joule_moq_relay::{RelayConfig, RelayServer};
use std::io::BufRead;
use std::sync::Arc;

const T0: u64 = 1_790_000_000_000_000;

fn keys() -> (SigningKey, SigningKey, SigningKey) {
    (SigningKey::from_bytes(&[0x51; 32]), SigningKey::from_bytes(&[0x52; 32]), SigningKey::from_bytes(&[0x53; 32]))
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().with_writer(std::io::stderr).with_max_level(tracing::Level::INFO).with_target(false).init();
    let args: Vec<String> = std::env::args().skip(1).collect();
    let (authority, publisher, gateway_key) = keys();
    let descriptor = build::room(&authority, T0);
    let room = descriptor.room_id();
    match args.first().map(String::as_str) {
        Some("relay") => {
            let mut rc = RelayConfig::loopback();
            // The door holds the room's media tracks to their receipts under the gateway's key.
            rc.receipts.door = ["audio", "video"]
                .iter()
                .map(|k| joule_moq_relay::config::DoorRule { prefix: vec![hex::encode(room), "*".into(), k.to_string()], trust: vec![hex::encode(gateway_key.verifying_key().to_bytes())], trust_field: None, hold_ms: 10_000 })
                .collect();
            let relay = RelayServer::start(rc).await.expect("the relay starts");
            println!("RELAY {} {}", relay.quic_addr.unwrap(), hex::encode(&relay.identity.chain[0]));
            println!("PID {}", std::process::id());
            let _ = tokio::task::spawn_blocking(|| std::io::stdin().lock().lines().count()).await;
            relay.shutdown().await;
        }
        Some("gateway") => {
            let addr = args.get(1).expect("the relay's address").parse().expect("an address");
            let cert = hex::decode(args.get(2).expect("the relay's certificate")).expect("hex");
            let mut gc = GatewayConfig::loopback(addr);
            gc.relay.trust_der = vec![cert];
            gc.signing.key_seed = Some(gateway_key.to_bytes());
            gc.rooms = vec![joule_moq_ingest::config::RoomConfig { descriptor_path: None, revoked: Vec::new(), descriptor: Some(descriptor.to_bytes()) }];
            gc.rtmp.listen = Some("127.0.0.1:0".parse().unwrap());
            gc.rtmp.stream_key_ttl_s = 3600;
            let server = GatewayServer::start(gc).await.expect("the gateway starts");
            let rtmp = server.rtmp_addr.unwrap();
            println!("PID {}", std::process::id());
            let gw = server.gateway.clone();
            let room_arc: Arc<Room> = gw.rooms.get(&room).unwrap().clone();
            let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
            let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<()>();
            std::thread::spawn(move || {
                for _ in std::io::stdin().lock().lines() {
                    if tx.send(()).is_err() {
                        break;
                    }
                }
            });
            while rx.recv().await.is_some() {
                let g = build::grant(&descriptor, &authority, &publisher.verifying_key().to_bytes(), 0, now - 60, now + 86_000);
                let key = gw.issue_stream_key(room_arc.clone(), g.to_bytes());
                println!("RTMP rtmp://{rtmp}/live/{key}");
            }
            for r in gw.records() {
                eprintln!("RECORD {} {} {}", r.record.ended, r.record.tracks.iter().filter(|t| t.dir == 2).map(|t| t.objects).sum::<u64>(), r.record.detail["counts"]);
            }
            server.shutdown().await;
        }
        _ => {
            eprintln!("usage: ingest_energy relay | ingest_energy gateway <relay addr> <certificate hex>");
            std::process::exit(2);
        }
    }
}