open-standards/wai/wai-key-proxy/tests/end_to_end.rs

//! The signing proxy and the key server end to end, over HTTP on loopback, against a
//! licence server written here (a fixture: it answers every request with fixed bytes).
//!
//! It holds: a granted request is forwarded and its licence returned unchanged, with a
//! receipt that commits to the listed keys; a refused one is answered 403 and never reaches
//! the licence server; a replay is refused, also when the first request is still being
//! forwarded; an offline licence and its renewal; a licence server's refusal releases
//! nothing; a clear key is released only by a ticket the key server redeems once, so no key
//! byte passes through the proxy; the log is held by one process, survives a restart and a
//! torn last line; and the whole log verifies under the reference verifier.

use std::collections::BTreeMap;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU16, AtomicU64, Ordering};
use std::time::Duration;

use axum::Router;
use axum::body::Bytes;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use ed25519_dalek::SigningKey;
use tokio::sync::Mutex;
use wai::key_release::corpus::{content_key, kid};
use wai::key_release::{
    CommitmentList, Decision, Delivery, DeliveryKind, Entitlement, LogCase, ProxyReceipt, Request, RequestKind, b64url, cpix,
    hex, key_commitment, prove, sha256, unb64url, verify_log,
};
use wai_key_proxy::config::{KeyServerFile, Proxy, ProxyFile, RouteFile};
use wai_key_proxy::http::{Answer, Target, send};
use wai_key_proxy::{keyserver, proxy};

const SYSTEM: &str = "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed";

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

struct Upstream {
    hits: Mutex<Vec<(HeaderMap, Vec<u8>)>>,
    status: AtomicU16,
    delay_ms: AtomicU64,
}

async fn licence_server(State(u): State<Arc<Upstream>>, headers: HeaderMap, body: Bytes) -> (StatusCode, Vec<u8>) {
    tokio::time::sleep(Duration::from_millis(u.delay_ms.load(Ordering::SeqCst))).await;
    u.hits.lock().await.push((headers, body.to_vec()));
    let mut licence = b"LICENCE:".to_vec();
    licence.extend_from_slice(&body);
    (StatusCode::from_u16(u.status.load(Ordering::SeqCst)).unwrap(), licence)
}

async fn spawn(router: Router) -> (SocketAddr, tokio::task::JoinHandle<()>) {
    let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let a = l.local_addr().unwrap();
    let h = tokio::spawn(async move {
        let _ = axum::serve(l, router).await;
    });
    (a, h)
}

struct World {
    dir: PathBuf,
    upstream: Arc<Upstream>,
    upstream_addr: SocketAddr,
    key_addr: SocketAddr,
    film_list: CommitmentList,
}

impl World {
    async fn new(name: &str) -> World {
        let dir = std::env::temp_dir().join(format!("wai-key-proxy-{}-{name}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let docs = wai::key_release::corpus::documents();
        let mut lists = Vec::new();
        for (doc, required) in [("film.cpix.xml", "svod:premium"), ("episode.cpix.xml", "svod:basic")] {
            std::fs::write(dir.join(doc), &docs[doc]).unwrap();
            let d = cpix::derive(docs[doc].as_bytes()).unwrap();
            let l = CommitmentList::seal(&key(1), "holder", &d.content_id, required, Some(d.sha256), 0, d.list_keys());
            std::fs::write(dir.join(format!("{}.wkl1", d.content_id)), l.to_bytes()).unwrap();
            lists.push(l);
        }
        std::fs::write(dir.join("proxy.key"), hex(&[3; 32])).unwrap();
        let upstream = Arc::new(Upstream { hits: Mutex::new(Vec::new()), status: AtomicU16::new(200), delay_ms: AtomicU64::new(0) });
        let (upstream_addr, _) =
            spawn(Router::new().route("/licence", axum::routing::post(licence_server)).with_state(upstream.clone())).await;
        let ks = KeyServerFile {
            listen: "127.0.0.1:0".parse().unwrap(),
            cpix: vec![dir.join("film.cpix.xml"), dir.join("episode.cpix.xml")],
            proxies: vec![hex(&key(3).verifying_key().to_bytes())],
            redeemed: dir.join("redeemed.txt"),
        };
        let (key_addr, _) = spawn(keyserver::router(keyserver::KeyApp::load(&ks).unwrap())).await;
        World { dir, upstream, upstream_addr, key_addr, film_list: lists.remove(0) }
    }

    fn config(&self) -> Proxy {
        let route = |name: &str, content: &str, delivery: &str, sys: Option<&str>, up: bool| RouteFile {
            name: name.into(),
            content_id: content.into(),
            delivery: delivery.into(),
            system_id: sys.map(String::from),
            upstream: up.then(|| format!("http://{}/licence", self.upstream_addr)),
            upstream_headers: BTreeMap::from([("x-licence-server-token".to_string(), "fixture".to_string())]),
            key_url: (!up).then(|| format!("http://{}/v1/key", self.key_addr)),
        };
        let mut routes = vec![
            route("film.a", "film-0001", "drm", Some(SYSTEM), true),
            route("film.ck", "film-0001", "clear-key", None, false),
            route("episode.hls", "episode-0002", "hls-key-uri", None, false),
        ];
        for r in &mut routes[1..] {
            r.upstream_headers.clear();
        }
        Proxy::from_file(ProxyFile {
            listen: "127.0.0.1:0".parse().unwrap(),
            admin_listen: None,
            signer_key: self.dir.join("proxy.key"),
            signer_id: "https://keys.proxy.example".into(),
            authorities: vec![hex(&key(2).verifying_key().to_bytes())],
            lists: vec![self.dir.join("film-0001.wkl1"), self.dir.join("episode-0002.wkl1")],
            routes,
            log: self.dir.join("receipts.jsonl"),
            ticket_ttl_s: 60,
            upstream_timeout_ms: 5_000,
            max_body: 1 << 16,
        })
        .unwrap()
    }

    async fn start(&self) -> (Arc<proxy::App>, SocketAddr, SocketAddr, tokio::task::JoinHandle<()>, tokio::task::JoinHandle<()>) {
        // A restarted proxy waits for the old one's log to be let go, as a supervisor would.
        let mut tries = 0;
        let app = loop {
            match proxy::App::open(self.config()) {
                Ok(a) => break a,
                Err(e) if tries < 100 => {
                    tries += 1;
                    let _ = e;
                    tokio::time::sleep(Duration::from_millis(30)).await;
                }
                Err(e) => panic!("{e}"),
            }
        };
        let (a, h) = spawn(proxy::router(app.clone())).await;
        let (b, h2) = spawn(proxy::admin_router(app.clone())).await;
        (app, a, b, h, h2)
    }
}

fn entitlement(class: &str, offline_s: u64, renewals: u32) -> Entitlement {
    let now = wai_key_proxy::now();
    Entitlement::seal(&key(2), "authority", [9; 16], key(4).verifying_key().to_bytes(), class, now - 60, now + 3_600, offline_s, renewals)
}

struct Ask {
    route: &'static str,
    content: &'static str,
    delivery: Delivery,
    kids: Vec<[u8; 16]>,
    kind: RequestKind,
    renews: Option<[u8; 32]>,
    body: Vec<u8>,
}

/// The header fields of a request, signed by the device (key 4).
fn signed(a: &Ask, e: &Entitlement) -> Vec<(String, String)> {
    let mut kids = a.kids.clone();
    kids.sort_unstable();
    let t = wai_key_proxy::now();
    let req = Request {
        content_id: a.content.into(),
        delivery: a.delivery,
        kids: kids.clone(),
        kind: a.kind,
        renews: a.renews,
        request_sha256: sha256(&a.body),
        request_t: t,
    };
    let eb = e.to_bytes();
    let mut h = vec![
        ("wai-key-ids".to_string(), kids.iter().map(|k| hex(k)).collect::<Vec<_>>().join(",")),
        ("wai-key-kind".into(), a.kind.label().into()),
        ("wai-key-time".into(), t.to_string()),
        ("wai-key-entitlement".into(), b64url(&eb)),
        ("wai-key-proof".into(), b64url(&prove(&key(4), &eb, &req))),
        ("content-type".into(), "application/octet-stream".into()),
    ];
    if let Some(r) = a.renews {
        h.push(("wai-key-renews".into(), hex(&r)));
    }
    h
}

async fn call(addr: SocketAddr, method: &str, path: &str, headers: &[(String, String)], body: Vec<u8>) -> Answer {
    let t = Target { authority: addr.to_string(), path: path.into() };
    send(method, &t, headers, body, Duration::from_secs(10)).await.expect("an answer")
}

async fn receipt_of(addr: SocketAddr, a: &Answer) -> ProxyReceipt {
    let id = a.header("wai-key-receipt").expect("every decided request names its receipt");
    let got = call(addr, "GET", &format!("/v1/receipts/{id}"), &[], vec![]).await;
    assert_eq!(got.status, 200);
    let r = ProxyReceipt::from_json(&serde_json::from_slice(&got.body).unwrap()).unwrap();
    assert_eq!(hex(&r.identity()), id);
    assert_eq!(r.check(), Ok(()));
    r
}

fn film_drm(kind: RequestKind, renews: Option<[u8; 32]>, body: &str) -> Ask {
    Ask {
        route: "film.a",
        content: "film-0001",
        delivery: Delivery { kind: DeliveryKind::Drm, system_id: cpix::uuid(SYSTEM) },
        kids: vec![kid("film/video"), kid("film/audio")],
        kind,
        renews,
        body: body.as_bytes().to_vec(),
    }
}

fn path(a: &Ask) -> String {
    format!("/v1/licence/{}", a.route)
}

/// The whole log, from the listener that is not public, under the reference verifier.
async fn verify_whole_log(admin: SocketAddr, w: &World) -> usize {
    let log = call(admin, "GET", "/v1/log", &[], vec![]).await;
    assert_eq!(log.status, 200);
    let receipts: Vec<serde_json::Value> =
        String::from_utf8(log.body).unwrap().lines().map(|l| serde_json::from_str(l).unwrap()).collect();
    let lists = vec![std::fs::read(w.dir.join("film-0001.wkl1")).unwrap(), std::fs::read(w.dir.join("episode-0002.wkl1")).unwrap()];
    let cpix = vec![(0, std::fs::read(w.dir.join("film.cpix.xml")).unwrap()), (1, std::fs::read(w.dir.join("episode.cpix.xml")).unwrap())];
    let case = LogCase {
        authorities: &[key(2).verifying_key().to_bytes()],
        proxies: &[key(3).verifying_key().to_bytes()],
        lists: &lists,
        cpix: &cpix,
        receipts: &receipts,
        tickets: &[],
    };
    verify_log(&case).unwrap_or_else(|r| panic!("the proxy's own log is refused: {} at {} {}", r.reason.code(), r.place, r.index)).len()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn releases_refusals_tickets_and_the_log() {
    let w = World::new("main").await;
    let (_app, addr, admin, _, _) = w.start().await;
    let premium = entitlement("svod:premium", 600, 1);

    // A granted request is forwarded, and its licence returned unchanged.
    let ask = film_drm(RequestKind::Streaming, None, "challenge one");
    let headers = signed(&ask, &premium);
    let a = call(addr, "POST", &path(&ask), &headers, ask.body.clone()).await;
    assert_eq!(a.status, 200);
    assert_eq!(a.body, b"LICENCE:challenge one");
    let r = receipt_of(addr, &a).await;
    assert_eq!((r.decision, r.released, r.upstream_status), (Decision::Granted, true, Some(200)));
    assert_eq!(r.response_sha256, Some(sha256(b"LICENCE:challenge one")));
    for k in &r.keys {
        let label = if k.kid == kid("film/video") { "film/video" } else { "film/audio" };
        assert_eq!(k.key_commit, Some(key_commitment(&content_key(label, 16))), "the commitment is the key's");
    }
    {
        let hits = w.upstream.hits.lock().await;
        assert_eq!(hits.len(), 1);
        let (h, body) = &hits[0];
        assert_eq!(body, b"challenge one", "the body is forwarded unchanged");
        assert!(h.get("wai-key-ids").is_some() && h.get("x-licence-server-token").is_some());
        assert!(h.get("wai-key-entitlement").is_none() && h.get("wai-key-proof").is_none(), "the entitlement stays at the proxy");
    }

    // The same request again: refused, and not forwarded.
    let again = call(addr, "POST", &path(&ask), &headers, ask.body.clone()).await;
    assert_eq!(again.status, 403);
    assert_eq!(receipt_of(addr, &again).await.decision, Decision::DeniedReplayed);
    // The wrong class: refused, and not forwarded.
    let basic = entitlement("svod:basic", 0, 0);
    let ask2 = film_drm(RequestKind::Streaming, None, "challenge two");
    let refused = call(addr, "POST", &path(&ask2), &signed(&ask2, &basic), ask2.body.clone()).await;
    assert_eq!(refused.status, 403);
    let rr: ProxyReceipt = ProxyReceipt::from_json(&serde_json::from_slice(&refused.body).unwrap()).unwrap();
    assert_eq!((rr.decision, rr.released, rr.upstream_status), (Decision::DeniedNoEntitlement, false, None));
    assert_eq!(w.upstream.hits.lock().await.len(), 1, "no refused request reaches the licence server");

    // A replay sent while the first request is still at the licence server waits, and is
    // refused after it.
    w.upstream.delay_ms.store(300, Ordering::SeqCst);
    let ask3 = film_drm(RequestKind::Streaming, None, "challenge three");
    let (h3, p3) = (signed(&ask3, &premium), path(&ask3));
    let (x, y) = tokio::join!(call(addr, "POST", &p3, &h3, ask3.body.clone()), async {
        tokio::time::sleep(Duration::from_millis(50)).await;
        call(addr, "POST", &p3, &h3, ask3.body.clone()).await
    });
    w.upstream.delay_ms.store(0, Ordering::SeqCst);
    assert_eq!((x.status, y.status), (200, 403));
    assert_eq!(receipt_of(addr, &y).await.decision, Decision::DeniedReplayed);

    // An offline licence and its renewal; the licence server is told the expiry.
    let off = film_drm(RequestKind::Offline, None, "offline");
    let o = call(addr, "POST", &path(&off), &signed(&off, &premium), off.body.clone()).await;
    assert_eq!(o.status, 200);
    let or = receipt_of(addr, &o).await;
    let expiry = or.licence_expiry.expect("an offline licence has an expiry");
    assert_eq!(expiry, (or.issued + 600).min(premium.not_after));
    let told = w.upstream.hits.lock().await.last().unwrap().0.get("wai-key-licence-expiry").map(|v| v.to_str().unwrap().to_string());
    assert_eq!(told, Some(expiry.to_string()));
    let ren = film_drm(RequestKind::Renewal, Some(or.identity()), "renewal");
    let n = call(addr, "POST", &path(&ren), &signed(&ren, &premium), ren.body.clone()).await;
    assert_eq!(n.status, 200);
    assert_eq!(receipt_of(addr, &n).await.renewal, Some(1));

    // The licence server refuses: nothing is released.
    w.upstream.status.store(403, Ordering::SeqCst);
    let ask4 = film_drm(RequestKind::Streaming, None, "challenge four");
    let up = call(addr, "POST", &path(&ask4), &signed(&ask4, &premium), ask4.body.clone()).await;
    w.upstream.status.store(200, Ordering::SeqCst);
    assert_eq!(up.status, 403);
    let ur = receipt_of(addr, &up).await;
    assert_eq!((ur.decision, ur.released, ur.upstream_status), (Decision::Granted, false, Some(403)));
    assert!(ur.keys.iter().all(|k| k.key_commit.is_none()));

    // An HLS key file: a ticket, redeemed once at the key server.
    let basic_hls = entitlement("svod:basic", 0, 0);
    let hls = Ask {
        route: "episode.hls",
        content: "episode-0002",
        delivery: Delivery { kind: DeliveryKind::HlsKeyUri, system_id: None },
        kids: vec![kid("episode/main")],
        kind: RequestKind::Streaming,
        renews: None,
        body: vec![],
    };
    let k = call(addr, "GET", &path(&hls), &signed(&hls, &basic_hls), vec![]).await;
    assert_eq!(k.status, 307);
    let kr = receipt_of(addr, &k).await;
    assert!(kr.released && kr.upstream_status.is_none());
    let loc = k.header("location").unwrap().to_string();
    let key_path = loc.split_once(&w.key_addr.to_string()).unwrap().1.to_string();
    let got = call(w.key_addr, "GET", &key_path, &[], vec![]).await;
    assert_eq!(got.status, 200);
    assert_eq!(got.body, content_key("episode/main", 16), "the key server returns the key file");
    assert_eq!(call(w.key_addr, "GET", &key_path, &[], vec![]).await.status, 409, "a ticket is redeemed once");
    let mut forged = loc.split_once("wai-ticket=").unwrap().1.to_string();
    let last = forged.pop().unwrap();
    forged.push(if last == 'A' { 'B' } else { 'A' });
    assert_eq!(call(w.key_addr, "GET", &format!("/v1/key?wai-ticket={forged}"), &[], vec![]).await.status, 403);

    // Clear Key: the licence request reaches the key server by the redirect; a JWK set back.
    let ck = Ask {
        route: "film.ck",
        content: "film-0001",
        delivery: Delivery { kind: DeliveryKind::ClearKey, system_id: None },
        kids: vec![kid("film/video")],
        kind: RequestKind::Offline,
        renews: None,
        body: format!("{{\"kids\":[\"{}\"],\"type\":\"persistent-license\"}}", b64url(&kid("film/video"))).into_bytes(),
    };
    let c = call(addr, "POST", &path(&ck), &signed(&ck, &premium), ck.body.clone()).await;
    assert_eq!(c.status, 307);
    let loc = c.header("location").unwrap().to_string();
    let jwk = call(w.key_addr, "POST", loc.split_once(&w.key_addr.to_string()).unwrap().1, &[], ck.body.clone()).await;
    assert_eq!(jwk.status, 200);
    let j: serde_json::Value = serde_json::from_slice(&jwk.body).unwrap();
    assert_eq!(j["type"], "persistent-license");
    assert_eq!(unb64url(j["keys"][0]["k"].as_str().unwrap()), Some(content_key("film/video", 16)));

    // Requests that name nothing to decide have no receipt.
    assert_eq!(call(addr, "POST", "/v1/licence/nowhere", &headers, vec![]).await.status, 404);
    assert_eq!(call(addr, "POST", &path(&ask), &[], vec![]).await.status, 400);

    // No key byte, in any spelling, is in the proxy's log.
    let log = std::fs::read(w.dir.join("receipts.jsonl")).unwrap();
    for label in ["film/video", "film/audio", "episode/main"] {
        let k = content_key(label, 16);
        let text = String::from_utf8_lossy(&log);
        assert!(!text.contains(&hex(&k)) && !text.contains(&b64url(&k)) && !text.contains(&wai::key_release::b64(&k)));
        assert!(!log.windows(16).any(|w| w == k.as_slice()));
    }
    let n = verify_whole_log(admin, &w).await;
    assert_eq!(n, 10, "every decided request is in the log");
    assert!(w.film_list.keys.len() == 2);
    let _ = std::fs::remove_dir_all(&w.dir);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn the_log_is_held_once_and_survives_a_restart() {
    let w = World::new("restart").await;
    let premium = entitlement("svod:premium", 600, 1);
    let first = film_drm(RequestKind::Streaming, None, "before the restart");
    let first_headers = signed(&first, &premium);
    let off = film_drm(RequestKind::Offline, None, "offline before the restart");
    let off_id = {
        let (app, addr, _, h, h2) = w.start().await;
        assert_eq!(call(addr, "POST", &path(&first), &first_headers, first.body.clone()).await.status, 200);
        let o = call(addr, "POST", &path(&off), &signed(&off, &premium), off.body.clone()).await;
        let id = receipt_of(addr, &o).await.identity();
        assert!(proxy::App::open(w.config()).is_err(), "a second process does not take the log");
        h.abort();
        h2.abort();
        let _ = h.await;
        let _ = h2.await;
        drop(app);
        id
    };
    // A write that never finished, as a crash would leave it.
    use std::io::Write;
    std::fs::OpenOptions::new().append(true).open(w.dir.join("receipts.jsonl")).unwrap().write_all(b"{\"form\":2,\"cont").unwrap();
    let (app, addr, admin, _, _) = w.start().await;
    assert_eq!(app.log_len().await, 2, "the log is read back, its unfinished line cut");
    let again = call(addr, "POST", &path(&first), &first_headers, first.body.clone()).await;
    assert_eq!((again.status, receipt_of(addr, &again).await.decision), (403, Decision::DeniedReplayed), "a replay across a restart");
    let ren = film_drm(RequestKind::Renewal, Some(off_id), "renewal after the restart");
    let n = call(addr, "POST", &path(&ren), &signed(&ren, &premium), ren.body.clone()).await;
    assert_eq!((n.status, receipt_of(addr, &n).await.renewal), (200, Some(1)), "a renewal of a licence released before the restart");
    assert_eq!(verify_whole_log(admin, &w).await, 4);
    let _ = std::fs::remove_dir_all(&w.dir);
}