johnandersen777.bsky.social/atlibs

Switch repository

Loading…

//! XRPC subscriptions over WebSocket (javascript/src/run-module/subscribe.js): atproto event
//! streams send binary frames, a DAG-CBOR header { op, t } then a DAG-CBOR body; jetstream-style
//! streams send JSON text. Messages come out in the atproto JSON data model ($bytes, $link).

use std::time::Duration;

use base64::Engine;
use base64::engine::general_purpose::STANDARD_NO_PAD;
use futures_util::{SinkExt, StreamExt};
use serde_json::{Map, Value, json};
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::protocol::CloseFrame;
use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;

use crate::crypto::{base32, unbase32};
use crate::run::{Sub, SubCmd};
use crate::store::unbase64;
use crate::{Signal, js_string};

/// Arrays, maps and tags nest at most this deep: frames come from the network, and each level recurses.
pub const MAX_DEPTH: usize = 128;

/// Consecutive DAG-CBOR items as JSON: bytes -> { $bytes }, tag 42 -> { $link }.
pub fn decode_cbor(bytes: &[u8]) -> Result<Vec<Value>, String> {
    let mut at = 0;
    let mut items = vec![];
    while at < bytes.len() {
        items.push(item(bytes, &mut at, 0)?);
    }
    Ok(items)
}

fn take<'a>(bytes: &'a [u8], at: &mut usize, n: usize) -> Result<&'a [u8], String> {
    let out = at
        .checked_add(n)
        .and_then(|end| bytes.get(*at..end))
        .ok_or("CBOR: truncated")?;
    *at += n;
    Ok(out)
}

fn item(bytes: &[u8], at: &mut usize, depth: usize) -> Result<Value, String> {
    if depth > MAX_DEPTH {
        return Err(format!("CBOR: nested deeper than {MAX_DEPTH}"));
    }
    let head = take(bytes, at, 1)?[0];
    let (major, info) = (head >> 5, head & 31);
    if major == 7 {
        return Ok(match info {
            20 => json!(false),
            21 => json!(true),
            22 | 23 => Value::Null,
            25 => {
                let b = take(bytes, at, 2)?;
                json!(f16(u16::from_be_bytes([b[0], b[1]])))
            }
            26 => json!(f32::from_be_bytes(take(bytes, at, 4)?.try_into().unwrap()) as f64),
            27 => json!(f64::from_be_bytes(take(bytes, at, 8)?.try_into().unwrap())),
            _ => return Err(format!("CBOR: simple value {info}")),
        });
    }
    let n = match info {
        0..24 => info as u64,
        24 => take(bytes, at, 1)?[0] as u64,
        25 => u16::from_be_bytes(take(bytes, at, 2)?.try_into().unwrap()) as u64,
        26 => u32::from_be_bytes(take(bytes, at, 4)?.try_into().unwrap()) as u64,
        27 => u64::from_be_bytes(take(bytes, at, 8)?.try_into().unwrap()),
        _ => return Err("CBOR: indefinite lengths are not DAG-CBOR".into()),
    };
    Ok(match major {
        0 => json!(n),
        1 => json!(-1 - n as i128),
        2 => json!({ "$bytes": STANDARD_NO_PAD.encode(take(bytes, at, n as usize)?) }),
        3 => json!(String::from_utf8_lossy(take(bytes, at, n as usize)?)),
        4 => Value::Array(
            (0..n)
                .map(|_| item(bytes, at, depth + 1))
                .collect::<Result<_, _>>()?,
        ),
        5 => {
            let mut map = Map::new();
            for _ in 0..n {
                let key = item(bytes, at, depth + 1)?;
                map.insert(js_string(Some(&key)), item(bytes, at, depth + 1)?);
            }
            Value::Object(map)
        }
        _ => {
            // tag 42 is a CID: bytes with a 0x00 multibase-identity prefix
            let tagged = item(bytes, at, depth + 1)?;
            match (n, tagged.get("$bytes").and_then(Value::as_str)) {
                (42, Some(b)) => {
                    json!({ "$link": format!("b{}", base32(unbase64(b)?.get(1..).unwrap_or_default())) })
                }
                _ => return Err(format!("CBOR: tag {n}")),
            }
        }
    })
}

fn f16(h: u16) -> f64 {
    let (sign, exp, frac) = (
        if h >> 15 == 1 { -1.0 } else { 1.0 },
        (h >> 10) & 0x1f,
        (h & 0x3ff) as f64,
    );
    sign * match exp {
        0 => frac * 2f64.powi(-24),
        31 => {
            if frac == 0.0 {
                f64::INFINITY
            } else {
                f64::NAN
            }
        }
        e => (1.0 + frac / 1024.0) * 2f64.powi(e as i32 - 15),
    }
}

/// Values (the atproto JSON data model) as consecutive DAG-CBOR items.
pub fn encode_cbor(values: &[Value]) -> Result<Vec<u8>, String> {
    let mut out = vec![];
    for v in values {
        encode(v, &mut out)?;
    }
    Ok(out)
}

fn head(out: &mut Vec<u8>, major: u8, n: u64) -> Result<(), String> {
    let m = major << 5;
    match n {
        0..24 => out.push(m | n as u8),
        24..0x100 => out.extend([m | 24, n as u8]),
        0x100..0x10000 => out.extend([m | 25].into_iter().chain((n as u16).to_be_bytes())),
        0x10000..0x1_0000_0000 => out.extend([m | 26].into_iter().chain((n as u32).to_be_bytes())),
        _ => out.extend([m | 27].into_iter().chain(n.to_be_bytes())),
    }
    Ok(())
}

fn encode(v: &Value, out: &mut Vec<u8>) -> Result<(), String> {
    match v {
        Value::Null => out.push(0xf6),
        Value::Bool(b) => out.push(if *b { 0xf5 } else { 0xf4 }),
        Value::Number(n) => {
            let int = n.as_i64().or_else(|| {
                n.as_f64()
                    .filter(|f| f.fract() == 0.0 && f.abs() <= 9_007_199_254_740_991.0)
                    .map(|f| f as i64)
            });
            match int.filter(|i| i.unsigned_abs() <= 9_007_199_254_740_991) {
                Some(i) if i >= 0 => head(out, 0, i as u64)?,
                Some(i) => head(out, 1, (-1 - i) as u64)?,
                None => return Err("CBOR: DAG-CBOR in atproto has no floats".into()),
            }
        }
        Value::String(s) => {
            head(out, 3, s.len() as u64)?;
            out.extend(s.as_bytes());
        }
        Value::Array(items) => {
            head(out, 4, items.len() as u64)?;
            for item in items {
                encode(item, out)?;
            }
        }
        Value::Object(map) => {
            if map.len() == 1 {
                if let Some(Value::String(b)) = map.get("$bytes") {
                    let bytes = unbase64(b)?;
                    head(out, 2, bytes.len() as u64)?;
                    out.extend(bytes);
                    return Ok(());
                }
                if let Some(Value::String(link)) = map.get("$link") {
                    let bytes = [&[0][..], &unbase32(link.get(1..).unwrap_or_default())].concat();
                    head(out, 6, 42)?;
                    head(out, 2, bytes.len() as u64)?;
                    out.extend(bytes);
                    return Ok(());
                }
            }
            // DAG-CBOR map keys sort by encoded length, then bytewise.
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort_by(|a, b| a.len().cmp(&b.len()).then(a.cmp(b)));
            head(out, 5, keys.len() as u64)?;
            for key in keys {
                encode(&json!(key), out)?;
                encode(&map[key], out)?;
            }
        }
    }
    Ok(())
}

/// One event-stream frame as a message: { $type: nsid + t, ...body } for op 1; an error frame
/// (op -1) is Err. Text is JSON.
pub fn decode_frame(data: &Message, nsid: &str) -> Result<Option<Value>, String> {
    match data {
        Message::Text(text) => serde_json::from_str(text)
            .map(Some)
            .map_err(|e| e.to_string()),
        Message::Binary(bytes) => {
            let items = decode_cbor(bytes)?;
            let (header, body) = (
                items.first().cloned().unwrap_or(Value::Null),
                items.get(1).cloned().unwrap_or(Value::Null),
            );
            if header.get("op").and_then(Value::as_i64) == Some(-1) {
                let error = body
                    .get("error")
                    .and_then(Value::as_str)
                    .unwrap_or("error frame");
                let message = body
                    .get("message")
                    .and_then(Value::as_str)
                    .map(|m| format!(": {m}"))
                    .unwrap_or_default();
                return Err(format!("{error}{message}"));
            }
            let t = header
                .get("t")
                .map(|t| js_string(Some(t)))
                .unwrap_or_default();
            let mut out = Map::new();
            out.insert("$type".into(), json!(format!("{nsid}{t}")));
            if let Value::Object(body) = body {
                out.extend(body);
            }
            Ok(Some(Value::Object(out)))
        }
        _ => Ok(None),
    }
}

/// What send(data) writes: a string as a text frame; { $bytes } as a binary frame; with framing
/// "cbor" { $type: "...#t", ...body } as an event-stream frame (header { op: 1, t }); else JSON text.
pub fn encode_frame(data: &Value, framing: Option<&str>) -> Result<Message, String> {
    if let Value::String(s) = data {
        return Ok(Message::text(s.clone()));
    }
    if let Some(Value::String(b)) = data
        .get("$bytes")
        .filter(|_| data.as_object().is_some_and(|m| m.len() == 1))
    {
        return Ok(Message::binary(unbase64(b)?));
    }
    if framing == Some("cbor") {
        let mut body = data.as_object().cloned().unwrap_or_default();
        let kind = body
            .shift_remove("$type")
            .map(|t| js_string(Some(&t)))
            .unwrap_or_default();
        // t is the "#fragment" of $type; with no "#" the header has no t ($type is the nsid itself)
        let header = match kind.find('#') {
            Some(i) => json!({ "op": 1, "t": &kind[i..] }),
            None => json!({ "op": 1 }),
        };
        return Ok(Message::binary(encode_cbor(&[header, Value::Object(body)])?));
    }
    Ok(Message::text(data.to_string()))
}

/// wss://host/xrpc/nsid?params from an https (or http) base.
pub fn ws_url(
    base: &str,
    nsid: &str,
    params: Option<&Map<String, Value>>,
) -> Result<String, String> {
    let mut url = reqwest::Url::parse(base)
        .and_then(|b| b.join(&format!("/xrpc/{nsid}")))
        .map_err(|e| format!("{base}: {e}"))?;
    let scheme = if url.scheme() == "http" { "ws" } else { "wss" };
    url.set_scheme(scheme)
        .map_err(|_| format!("{base}: cannot be a WebSocket URL"))?;
    if let Some(params) = params {
        url.query_pairs_mut()
            .clear()
            .extend_pairs(params.iter().map(|(k, v)| (k, js_string(Some(v)))));
    }
    Ok(url.to_string())
}

/// ATLIBS_EXTRA_CA, test-only: a PEM file of roots trusted for wss beside webpki's, so a test can subscribe
/// to a local TLS server (rust/host/tests/wss.rs, cli.test.mjs). Unset (always, outside tests): tokio-tungstenite's
/// own connector. ClientConfig::builder() asks the process default provider, as that connector does: crate::tls().
fn extra_ca() -> Result<Option<tokio_tungstenite::Connector>, String> {
    use rustls::pki_types::{CertificateDer, pem::PemObject};
    let Some(path) = std::env::var_os("ATLIBS_EXTRA_CA") else { return Ok(None) };
    let mut roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() };
    for cert in CertificateDer::pem_file_iter(&path).map_err(|e| format!("ATLIBS_EXTRA_CA: {e}"))? {
        roots.add(cert.map_err(|e| format!("ATLIBS_EXTRA_CA: {e}"))?).map_err(|e| format!("ATLIBS_EXTRA_CA: {e}"))?;
    }
    let config = rustls::ClientConfig::builder().with_root_certificates(roots).with_no_client_auth();
    Ok(Some(tokio_tungstenite::Connector::Rustls(std::sync::Arc::new(config))))
}

/// subscribeXrpc: open a WebSocket; messages end when it closes cleanly and fail when it fails or
/// an error frame arrives. keepalive { ms, data } sends data every ms.
pub fn subscribe(
    url: String,
    nsid: String,
    framing: Option<String>,
    keepalive: Option<Value>,
    headers: Vec<(String, String)>,
    signal: Signal,
) -> Sub {
    crate::tls();
    let (msg_tx, messages) = mpsc::unbounded_channel();
    let (commands, mut cmd_rx) = mpsc::unbounded_channel::<SubCmd>();
    tokio::spawn(async move {
        use tokio_tungstenite::tungstenite::client::IntoClientRequest;
        let request = url.as_str().into_client_request().map(|mut r| {
            for (name, value) in &headers {
                if let (Ok(n), Ok(v)) = (
                    tokio_tungstenite::tungstenite::http::HeaderName::from_bytes(name.as_bytes()),
                    tokio_tungstenite::tungstenite::http::HeaderValue::from_str(value),
                ) {
                    r.headers_mut().insert(n, v);
                }
            }
            r
        });
        let Ok(request) = request else {
            let _ = msg_tx.send(Err(format!("{nsid}: bad WebSocket URL")));
            return;
        };
        let connector = match extra_ca() {
            Ok(connector) => connector,
            Err(error) => {
                let _ = msg_tx.send(Err(error));
                return;
            }
        };
        let socket = tokio::select! {
            s = tokio_tungstenite::connect_async_tls_with_config(request, None, false, connector) => s,
            _ = signal.aborted() => return,
        };
        let Ok((socket, _)) = socket else {
            let _ = msg_tx.send(Err(format!("{nsid}: WebSocket error")));
            return;
        };
        let (mut sink, mut stream) = socket.split();
        let ping_ms = keepalive
            .as_ref()
            .and_then(|k| k.get("ms"))
            .and_then(Value::as_f64)
            .filter(|ms| *ms > 0.0);
        let ping_data = keepalive
            .as_ref()
            .and_then(|k| k.get("data"))
            .cloned()
            .unwrap_or(Value::Null);
        // clamped (crate::run::ms_duration): a huge keepalive.ms must not panic this task
        let period = crate::run::ms_duration(ping_ms.unwrap_or(3.6e6)).max(Duration::from_millis(1));
        let mut ping = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
        let close = Message::Close(Some(CloseFrame {
            code: CloseCode::Normal,
            reason: "".into(),
        }));
        loop {
            tokio::select! {
                frame = stream.next() => match frame {
                    Some(Ok(Message::Close(frame))) => {
                        let code = frame.as_ref().map(|f| u16::from(f.code));
                        if let Some(code) = code.filter(|c| *c != 1000 && *c != 1005) {
                            let reason = frame.as_ref().map(|f| f.reason.to_string()).filter(|r| !r.is_empty()).map(|r| format!(" {r}")).unwrap_or_default();
                            let _ = msg_tx.send(Err(format!("{nsid} closed: {code}{reason}")));
                        }
                        return;
                    }
                    Some(Ok(data)) => match decode_frame(&data, &nsid) {
                        Ok(Some(message)) => { let _ = msg_tx.send(Ok(message)); }
                        Ok(None) => {}
                        Err(error) => {
                            let _ = msg_tx.send(Err(error));
                            let _ = sink.send(close).await;
                            return;
                        }
                    },
                    Some(Err(_)) => { let _ = msg_tx.send(Err(format!("{nsid}: WebSocket error"))); return; }
                    None => { let _ = msg_tx.send(Err(format!("{nsid} closed: 1006"))); return; }
                },
                command = cmd_rx.recv() => match command {
                    Some(SubCmd::Send(data, ack)) => {
                        // sink.send waits for the socket to take the frame: the JS host's HIGH_WATER wait
                        let sent = match encode_frame(&data, framing.as_deref()) {
                            Ok(frame) => sink.send(frame).await.map_err(|_| format!("{nsid} is closed")),
                            Err(e) => Err(e),
                        };
                        let _ = ack.send(sent);
                    }
                    Some(SubCmd::Close) | None => { let _ = sink.send(close).await; return; }
                },
                _ = ping.tick(), if ping_ms.is_some() => {
                    if let Ok(frame) = encode_frame(&ping_data, framing.as_deref()) {
                        let _ = sink.send(frame).await;
                    }
                }
                _ = signal.aborted() => { let _ = sink.send(close).await; return; }
            }
        }
    });
    Sub { messages, commands }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cbor_round_trips_the_atproto_data_model() {
        let link = "bafyreie5737gdxlw5i64vzichcalba3z2v5n6icifvx5xytvske7mr3hpm";
        let value = json!({ "seq": 7, "neg": -3, "ok": true, "none": null, "blocks": { "$bytes": "AQID" }, "cid": { "$link": link }, "ops": [{ "path": "a/b" }] });
        let bytes = encode_cbor(&[value.clone()]).unwrap();
        assert_eq!(decode_cbor(&bytes).unwrap(), vec![value]);
        assert_eq!(
            encode_cbor(&[json!(1.5)]).unwrap_err(),
            "CBOR: DAG-CBOR in atproto has no floats"
        );
        // ints past 2^32 take the 8-byte head (as the JS host)
        let big = json!({ "time": 1_731_000_000_000u64, "neg": -5_000_000_000i64 });
        assert_eq!(decode_cbor(&encode_cbor(&[big.clone()]).unwrap()).unwrap(), vec![big]);
        // a 1 MB byte string
        let mb = json!({ "$bytes": STANDARD_NO_PAD.encode(vec![7u8; 1 << 20]) });
        assert_eq!(decode_cbor(&encode_cbor(&[mb.clone()]).unwrap()).unwrap(), vec![mb]);
    }

    #[test]
    fn cbor_nesting_is_bounded() {
        // ~1 MB of 0x81 (a one-item array, nested) errors instead of overflowing the stack
        assert_eq!(
            decode_cbor(&vec![0x81; 1 << 20]).unwrap_err(),
            "CBOR: nested deeper than 128"
        );
        let deep = |n: usize| [vec![0x81; n], vec![0x00]].concat();
        assert_eq!(decode_cbor(&deep(MAX_DEPTH)).unwrap().len(), 1);
        assert!(decode_cbor(&deep(MAX_DEPTH + 1)).is_err());
        // maps (0xa1 0x00 = { "0": ... }) and tags nest too
        assert!(decode_cbor(&[0xa1u8, 0x00].repeat(200)).is_err());
        assert!(decode_cbor(&vec![0xd8, 0x2a].repeat(200)).is_err());
        assert_eq!(decode_cbor(&[0x5a, 0xff, 0xff, 0xff, 0xff]).unwrap_err(), "CBOR: truncated");
    }

    #[test]
    fn frames_decode_and_encode_as_the_js_host() {
        let frame =
            encode_cbor(&[json!({ "op": 1, "t": "#commit" }), json!({ "seq": 1 })]).unwrap();
        assert_eq!(
            decode_frame(&Message::binary(frame), "com.atproto.sync.subscribeRepos").unwrap(),
            Some(json!({ "$type": "com.atproto.sync.subscribeRepos#commit", "seq": 1 }))
        );
        let error = encode_cbor(&[
            json!({ "op": -1 }),
            json!({ "error": "FutureCursor", "message": "too new" }),
        ])
        .unwrap();
        assert_eq!(
            decode_frame(&Message::binary(error), "x").unwrap_err(),
            "FutureCursor: too new"
        );
        let Message::Binary(out) = encode_frame(
            &json!({ "$type": "a.b.c#response", "status": 200 }),
            Some("cbor"),
        )
        .unwrap() else {
            panic!()
        };
        assert_eq!(
            decode_cbor(&out).unwrap(),
            vec![
                json!({ "op": 1, "t": "#response" }),
                json!({ "status": 200 })
            ]
        );
        // no "#": no t in the header (not the last character of $type)
        let Message::Binary(out) =
            encode_frame(&json!({ "$type": "a.b.c", "x": 1 }), Some("cbor")).unwrap()
        else {
            panic!()
        };
        assert_eq!(decode_cbor(&out).unwrap()[0], json!({ "op": 1 }));
        assert_eq!(
            decode_frame(&Message::binary(out), "a.b.c").unwrap(),
            Some(json!({ "$type": "a.b.c", "x": 1 }))
        );
        assert_eq!(
            encode_frame(&json!("hi"), None).unwrap(),
            Message::text("hi")
        );
        assert_eq!(
            encode_frame(&json!({ "a": 1 }), None).unwrap(),
            Message::text("{\"a\":1}")
        );
        assert_eq!(
            ws_url(
                "https://relay.example/",
                "a.b.c",
                Some(json!({ "cursor": 5, "q": "a b" }).as_object().unwrap())
            )
            .unwrap(),
            "wss://relay.example/xrpc/a.b.c?cursor=5&q=a+b"
        );
    }

    #[tokio::test]
    async fn subscribe_reads_cbor_frames_sends_in_order_and_ends_on_close() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = tokio::spawn(async move {
            let (tcp, _) = listener.accept().await.unwrap();
            let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap();
            for seq in 1..=2 {
                ws.send(Message::binary(
                    encode_cbor(&[json!({ "op": 1, "t": "#commit" }), json!({ "seq": seq })])
                        .unwrap(),
                ))
                .await
                .unwrap();
            }
            let mut got = vec![];
            while got.len() < 2 {
                if let Some(Ok(m)) = ws.next().await {
                    got.push(m);
                }
            }
            ws.close(None).await.unwrap();
            got
        });
        let mut sub = subscribe(
            format!("ws://127.0.0.1:{port}/xrpc/a.b.c"),
            "a.b.c".into(),
            Some("cbor".into()),
            None,
            vec![],
            Signal::new(),
        );
        assert_eq!(
            sub.messages.recv().await.unwrap().unwrap(),
            json!({ "$type": "a.b.c#commit", "seq": 1 })
        );
        for data in [json!("hello"), json!({ "$type": "a.b.c#ack", "seq": 1 })] {
            let (ack, acked) = tokio::sync::oneshot::channel();
            sub.commands.send(SubCmd::Send(data, ack)).unwrap();
            acked.await.unwrap().unwrap();
        }
        assert_eq!(sub.messages.recv().await.unwrap().unwrap()["seq"], 2);
        assert!(
            sub.messages.recv().await.is_none(),
            "a clean close ends the stream"
        );
        let got = server.await.unwrap();
        assert_eq!(got[0], Message::text("hello"));
        let Message::Binary(b) = &got[1] else {
            panic!()
        };
        assert_eq!(decode_cbor(b).unwrap()[1], json!({ "seq": 1 }));
    }

    #[tokio::test]
    async fn subscribe_sends_the_headers_it_is_given() {
        use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = tokio::spawn(async move {
            let (tcp, _) = listener.accept().await.unwrap();
            let mut seen = None;
            let mut ws = tokio_tungstenite::accept_hdr_async(tcp, |req: &Request, res: Response| {
                seen = req.headers().get("authorization").map(|v| v.to_str().unwrap().to_string());
                Ok(res)
            })
            .await
            .unwrap();
            ws.close(None).await.unwrap();
            seen
        });
        let mut sub = subscribe(
            format!("ws://127.0.0.1:{port}/xrpc/a.b.c"),
            "a.b.c".into(),
            None,
            None,
            vec![("authorization".into(), "Bearer jwt1".into())],
            Signal::new(),
        );
        assert!(sub.messages.recv().await.is_none());
        assert_eq!(server.await.unwrap().as_deref(), Some("Bearer jwt1"));
    }
}