maarcadetweet/crates/pds-server/tests/firehose_integration.rs

//! Integration tests for `com.atproto.sync.subscribeRepos`.
//!
//! Same contract as the other integration suites in this crate: they talk to a
//! PDS listening on `127.0.0.1:2583` and **fail open** — if nothing answers
//! `/healthz`, the test prints a note and returns green rather than failing a
//! developer's `cargo test` on a machine with no server running. Start the
//! server (`./target/debug/pds-server` with `.env` sourced) to actually
//! exercise them.
//!
//! ## Why the frames are decoded by hand here
//!
//! `pds-server` is a binary, so a test crate cannot import its `dag_cbor`
//! module — and that is a feature, not a limitation. These tests are the
//! *consumer* side of the wire contract, and a consumer that reuses the
//! producer's encoder proves nothing: it would happily agree with a frame
//! that no other implementation can read. The decoder below is written from
//! the spec (tag 42, identity prefix, length-first map keys) and knows
//! nothing about how the server produced the bytes.

use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::time::Duration;

use futures::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;

const PDS_URL: &str = "http://127.0.0.1:2583";
const PDS_WS: &str = "ws://127.0.0.1:2583";

// -- harness ---------------------------------------------------------------

fn http() -> reqwest::Client {
    reqwest::Client::builder()
        .timeout(Duration::from_secs(5))
        .build()
        .unwrap()
}

async fn wait_for_pds() -> bool {
    let c = http();
    for _ in 0..20 {
        if let Ok(r) = c.get(format!("{PDS_URL}/healthz")).send().await {
            if r.status().is_success() {
                return true;
            }
        }
        tokio::time::sleep(Duration::from_millis(250)).await;
    }
    false
}

async fn fresh_user(prefix: &str) -> (reqwest::Client, String, String) {
    let c = http();
    let handle = format!(
        "{}_{}.maarcadetweet.local",
        prefix,
        uuid::Uuid::new_v4().simple()
    );
    let acc: Value = c
        .post(format!("{PDS_URL}/xrpc/com.atproto.server.createAccount"))
        .json(&json!({"handle": handle, "password": "hunter2hunter2"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let did = acc["did"].as_str().expect("createAccount did").to_string();
    let jwt = acc["access_jwt"].as_str().expect("access_jwt").to_string();
    (c, did, jwt)
}

async fn create_post(c: &reqwest::Client, did: &str, jwt: &str, text: &str) -> Value {
    c.post(format!("{PDS_URL}/xrpc/com.atproto.repo.createRecord"))
        .bearer_auth(jwt)
        .json(&json!({
            "repo": did,
            "collection": "app.twi.post",
            "record": { "text": text, "createdAt": "2026-09-10T12:00:00Z" },
        }))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap()
}

type Socket = tokio_tungstenite::WebSocketStream<
    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>;

async fn subscribe(cursor: Option<i64>) -> Socket {
    let url = match cursor {
        Some(c) => format!("{PDS_WS}/xrpc/com.atproto.sync.subscribeRepos?cursor={c}"),
        None => format!("{PDS_WS}/xrpc/com.atproto.sync.subscribeRepos"),
    };
    let (socket, _resp) = tokio_tungstenite::connect_async(&url)
        .await
        .unwrap_or_else(|e| panic!("subscribeRepos connect to {url}: {e}"));
    socket
}

/// Read the next **binary** message, or `None` on timeout.
///
/// Text messages would be a protocol violation on this endpoint and are
/// asserted against rather than skipped.
async fn next_frame(socket: &mut Socket) -> Option<Vec<u8>> {
    let deadline = Duration::from_secs(10);
    loop {
        match tokio::time::timeout(deadline, socket.next()).await {
            Err(_) => return None,
            Ok(None) => return None,
            Ok(Some(Ok(Message::Binary(b)))) => return Some(b),
            Ok(Some(Ok(Message::Ping(_)))) | Ok(Some(Ok(Message::Pong(_)))) => continue,
            Ok(Some(Ok(Message::Close(_)))) => return None,
            Ok(Some(Ok(other))) => panic!("subscribeRepos sent a non-binary frame: {other:?}"),
            Ok(Some(Err(e))) => panic!("subscribeRepos socket error: {e}"),
        }
    }
}

/// Give the server a moment to finish `on_upgrade` and actually subscribe to
/// the broadcast channel before we trigger a write.
///
/// The TCP handshake completing does not mean the handler has run. Without
/// this the test would occasionally write before the subscription exists and
/// then wait for a frame that was published to nobody. (The *cursor* replay
/// path is the real fix for that race in production; the live-only test is
/// deliberately testing the raceable path, so it waits.)
async fn settle() {
    tokio::time::sleep(Duration::from_millis(300)).await;
}

// -- an independent DAG-CBOR reader ----------------------------------------

#[derive(Debug, Clone, PartialEq)]
enum Cbor {
    Null,
    Bool(bool),
    Int(i64),
    Bytes(Vec<u8>),
    Text(String),
    Array(Vec<Cbor>),
    Map(BTreeMap<String, Cbor>),
    /// tag(42) + bytes(0x00 || cid) — the binary CID is kept raw and
    /// re-encoded to a `bafy…` string on demand, so the test never depends on
    /// the server's own CID formatting.
    Link(Vec<u8>),
}

impl Cbor {
    fn get(&self, key: &str) -> &Cbor {
        match self {
            Cbor::Map(m) => m
                .get(key)
                .unwrap_or_else(|| panic!("missing key `{key}` in {self:?}")),
            other => panic!("not a map: {other:?}"),
        }
    }
    fn opt(&self, key: &str) -> Option<&Cbor> {
        match self {
            Cbor::Map(m) => m.get(key),
            _ => None,
        }
    }
    fn int(&self) -> i64 {
        match self {
            Cbor::Int(i) => *i,
            other => panic!("not an int: {other:?}"),
        }
    }
    fn text(&self) -> &str {
        match self {
            Cbor::Text(s) => s,
            other => panic!("not text: {other:?}"),
        }
    }
    fn bool(&self) -> bool {
        match self {
            Cbor::Bool(b) => *b,
            other => panic!("not a bool: {other:?}"),
        }
    }
    fn bytes(&self) -> &[u8] {
        match self {
            Cbor::Bytes(b) => b,
            other => panic!("not bytes: {other:?}"),
        }
    }
    fn array(&self) -> &[Cbor] {
        match self {
            Cbor::Array(a) => a,
            other => panic!("not an array: {other:?}"),
        }
    }
    /// The link's CID rendered as a base32 `bafy…` string, for comparison
    /// against what the XRPC JSON responses return.
    fn link_cid(&self) -> String {
        match self {
            Cbor::Link(raw) => cid::Cid::read_bytes(&raw[..])
                .expect("tag-42 payload must be a valid binary CID")
                .to_string(),
            other => panic!("not a link: {other:?}"),
        }
    }
}

fn read_head(b: &[u8], p: usize) -> (u8, u64, usize) {
    let first = b[p];
    let major = first >> 5;
    let low = first & 0x1f;
    let (arg, extra) = match low {
        0..=23 => (low as u64, 0usize),
        24 => (b[p + 1] as u64, 1),
        25 => (u16::from_be_bytes([b[p + 1], b[p + 2]]) as u64, 2),
        26 => (
            u32::from_be_bytes([b[p + 1], b[p + 2], b[p + 3], b[p + 4]]) as u64,
            4,
        ),
        27 => {
            let mut n = 0u64;
            for i in 0..8 {
                n = (n << 8) | b[p + 1 + i] as u64;
            }
            (n, 8)
        }
        other => panic!("indefinite or reserved CBOR head 0x{other:02x} — illegal in DAG-CBOR"),
    };
    (major, arg, p + 1 + extra)
}

fn decode_at(b: &[u8], p: usize) -> (Cbor, usize) {
    let (major, arg, mut p) = read_head(b, p);
    match major {
        0 => (Cbor::Int(arg as i64), p),
        1 => (Cbor::Int(-(arg as i64) - 1), p),
        2 => {
            let end = p + arg as usize;
            (Cbor::Bytes(b[p..end].to_vec()), end)
        }
        3 => {
            let end = p + arg as usize;
            (
                Cbor::Text(std::str::from_utf8(&b[p..end]).unwrap().to_string()),
                end,
            )
        }
        4 => {
            let mut items = Vec::new();
            for _ in 0..arg {
                let (v, next) = decode_at(b, p);
                items.push(v);
                p = next;
            }
            (Cbor::Array(items), p)
        }
        5 => {
            let mut m = BTreeMap::new();
            let mut prev_key: Option<String> = None;
            for _ in 0..arg {
                let (k, next) = decode_at(b, p);
                p = next;
                let key = k.text().to_string();
                // Canonical DAG-CBOR order: shorter keys first, then
                // bytewise. Asserted here because a consumer written against
                // a strict codec (cborg's `dag-cbor` decoder, for one) will
                // reject a frame whose keys are out of order.
                if let Some(prev) = &prev_key {
                    let ordered = (prev.len(), prev.as_bytes()) < (key.len(), key.as_bytes());
                    assert!(ordered, "map keys out of canonical order: {prev:?} then {key:?}");
                }
                prev_key = Some(key.clone());
                let (v, next) = decode_at(b, p);
                p = next;
                m.insert(key, v);
            }
            (Cbor::Map(m), p)
        }
        6 => {
            assert_eq!(arg, 42, "DAG-CBOR permits only tag 42");
            let (inner, next) = decode_at(b, p);
            let raw = match inner {
                Cbor::Bytes(v) => v,
                other => panic!("tag 42 must wrap bytes, got {other:?}"),
            };
            assert_eq!(
                raw.first(),
                Some(&0x00),
                "a binary CID link must start with the 0x00 multibase identity prefix"
            );
            (Cbor::Link(raw[1..].to_vec()), next)
        }
        7 => match arg {
            20 => (Cbor::Bool(false), p),
            21 => (Cbor::Bool(true), p),
            22 => (Cbor::Null, p),
            other => panic!("unsupported CBOR simple value {other}"),
        },
        other => panic!("unsupported CBOR major type {other}"),
    }
}

/// Split one binary frame into its header and body values, and assert that
/// the two together consume the whole message — a frame with trailing bytes
/// would silently desynchronise a streaming consumer.
fn parse_frame(bytes: &[u8]) -> (Cbor, Cbor) {
    let (header, after_header) = decode_at(bytes, 0);
    let (body, end) = decode_at(bytes, after_header);
    assert_eq!(end, bytes.len(), "frame must be exactly two CBOR values");
    (header, body)
}

/// Minimal CAR v1 reader: returns the root CIDs and the block CIDs, both as
/// `bafy…` strings.
fn parse_car(bytes: &[u8]) -> (Vec<String>, Vec<String>) {
    fn varint(b: &[u8], p: &mut usize) -> u64 {
        let (mut v, mut shift) = (0u64, 0u32);
        loop {
            let byte = b[*p];
            *p += 1;
            v |= ((byte & 0x7f) as u64) << shift;
            if byte & 0x80 == 0 {
                return v;
            }
            shift += 7;
        }
    }
    let mut p = 0usize;
    let header_len = varint(bytes, &mut p) as usize;
    let header = &bytes[p..p + header_len];
    p += header_len;

    // The CAR header is a CBOR map; walk it with the same head reader. Note
    // that this server's CAR header tags its roots *without* the 0x00
    // identity prefix (a documented deviation in `car.rs`), so the roots are
    // read as plain tagged byte strings rather than through `decode_at`.
    let mut hp = 0usize;
    let (major, n, next) = read_head(header, hp);
    assert_eq!(major, 5, "CAR header must be a map");
    hp = next;
    let mut roots = Vec::new();
    for _ in 0..n {
        let (m, len, next) = read_head(header, hp);
        assert_eq!(m, 3);
        hp = next;
        let key = std::str::from_utf8(&header[hp..hp + len as usize]).unwrap().to_string();
        hp += len as usize;
        if key == "version" {
            let (m, v, next) = read_head(header, hp);
            assert_eq!(m, 0);
            assert_eq!(v, 1, "CAR must be v1");
            hp = next;
        } else if key == "roots" {
            let (m, count, next) = read_head(header, hp);
            assert_eq!(m, 4);
            hp = next;
            for _ in 0..count {
                let (m, tag, next) = read_head(header, hp);
                assert_eq!((m, tag), (6, 42), "root must be a tag-42 link");
                hp = next;
                let (m, len, next) = read_head(header, hp);
                assert_eq!(m, 2);
                hp = next;
                // A DAG-CBOR link wraps `0x00 || <binary CID>`; the 0x00 is
                // the multibase identity prefix and is not part of the CID.
                // Assert on it rather than skipping it silently — this
                // reader stands in for a foreign consumer, and dropping the
                // check would let the header regress unnoticed.
                let raw = &header[hp..hp + len as usize];
                assert_eq!(
                    raw.first(),
                    Some(&0x00),
                    "CAR root link must carry the multibase identity prefix"
                );
                roots.push(cid::Cid::read_bytes(&raw[1..]).unwrap().to_string());
                hp += len as usize;
            }
        } else {
            panic!("unexpected CAR header key {key}");
        }
    }

    let mut blocks = Vec::new();
    while p < bytes.len() {
        let section_len = varint(bytes, &mut p) as usize;
        let section = &bytes[p..p + section_len];
        let cid = cid::Cid::read_bytes(section).unwrap();
        blocks.push(cid.to_string());
        p += section_len;
    }
    (roots, blocks)
}

// -- tests -----------------------------------------------------------------

/// A live subscriber receives a `#commit` frame for a record created after it
/// connected, and every field of that frame says what it should.
#[tokio::test]
async fn live_subscriber_receives_a_commit_frame() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("fhlive").await;

    let mut socket = subscribe(None).await;
    settle().await;

    let created = create_post(&c, &did, &jwt, "hello firehose").await;
    let record_cid = created["cid"].as_str().expect("createRecord cid").to_string();
    let commit_cid = created["commit"]["cid"]
        .as_str()
        .expect("createRecord commit.cid")
        .to_string();
    let commit_rev = created["commit"]["rev"].as_str().unwrap().to_string();
    let uri = created["uri"].as_str().unwrap().to_string();
    let rkey = uri.rsplit('/').next().unwrap().to_string();

    // The account was created moments ago and has never written before, so
    // the first frame we see for it is this post's. Other accounts may be
    // writing concurrently, so filter by DID rather than taking frame 1.
    let (header, body) = loop {
        let bytes = next_frame(&mut socket)
            .await
            .expect("expected a #commit frame within the timeout");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) == Some("#commit")
            && body.get("repo").text() == did
        {
            break (header, body);
        }
    };

    assert_eq!(header.get("op").int(), 1, "regular frames carry op = 1");
    assert_eq!(header.get("t").text(), "#commit");

    assert!(body.get("seq").int() > 0, "seq must be a real cursor value");
    assert_eq!(body.get("repo").text(), did);
    assert_eq!(body.get("rev").text(), commit_rev);
    assert_eq!(body.get("commit").link_cid(), commit_cid);
    assert!(!body.get("rebase").bool());
    assert!(!body.get("tooBig").bool());
    assert!(body.get("blobs").array().is_empty());
    // First-ever commit on a brand new repo — there is no predecessor.
    assert_eq!(*body.get("since"), Cbor::Null, "first commit has since = null");
    // RFC 3339 with a Z offset.
    let time = body.get("time").text();
    assert!(
        chrono_like_rfc3339(time),
        "time must be RFC 3339, got {time}"
    );

    // ops
    let ops = body.get("ops").array();
    assert_eq!(ops.len(), 1, "one record write means one op");
    assert_eq!(ops[0].get("action").text(), "create");
    assert_eq!(ops[0].get("path").text(), format!("app.twi.post/{rkey}"));
    assert_eq!(ops[0].get("cid").link_cid(), record_cid);

    // blocks: a CAR rooted at the commit, containing the commit block and the
    // new record block.
    let (roots, blocks) = parse_car(body.get("blocks").bytes());
    assert_eq!(roots, vec![commit_cid.clone()], "CAR root is the commit");
    assert!(
        blocks.contains(&commit_cid),
        "CAR must carry the commit block itself; got {blocks:?}"
    );
    assert!(
        blocks.contains(&record_cid),
        "CAR must carry the new record block; got {blocks:?}"
    );

    let _ = socket.send(Message::Close(None)).await;
}

/// The same event, fetched again from the durable log with a cursor, is
/// byte-identical to the live frame.
///
/// Byte-identity is the strong form of the claim and the one that matters: a
/// consumer that deduplicates by hashing frames, or that verifies a signature
/// over them, must not see two different representations of one event.
#[tokio::test]
async fn cursor_replay_reproduces_the_live_frame_exactly() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("fhreplay").await;

    let mut live = subscribe(None).await;
    settle().await;
    create_post(&c, &did, &jwt, "replay me").await;

    let live_bytes = loop {
        let bytes = next_frame(&mut live)
            .await
            .expect("expected a live #commit frame");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) == Some("#commit")
            && body.get("repo").text() == did
        {
            break bytes;
        }
    };
    let (_h, live_body) = parse_frame(&live_bytes);
    let seq = live_body.get("seq").int();
    let _ = live.send(Message::Close(None)).await;

    // Reconnect asking for everything after the event *before* ours, so the
    // replay's first matching frame is the one we just saw.
    let mut replayed = subscribe(Some(seq - 1)).await;
    let replay_bytes = loop {
        let bytes = next_frame(&mut replayed)
            .await
            .expect("expected the event to come back from the replay");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) == Some("#commit")
            && body.get("seq").int() == seq
        {
            break bytes;
        }
    };

    assert_eq!(
        replay_bytes, live_bytes,
        "a replayed frame must be byte-identical to the live one"
    );
    let _ = replayed.send(Message::Close(None)).await;
}

/// Handing over from replay to live loses nothing and duplicates nothing.
///
/// Connect with a cursor at the current head (so the replay is empty), then
/// write twice: both events must arrive, in order, exactly once each.
#[tokio::test]
async fn replay_to_live_handover_has_no_gap_and_no_duplicate() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("fhhandover").await;

    // Establish where the log currently ends by writing one event and reading
    // its seq off the live stream.
    let mut probe = subscribe(None).await;
    settle().await;
    create_post(&c, &did, &jwt, "probe").await;
    let head_seq = loop {
        let bytes = next_frame(&mut probe).await.expect("probe frame");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) == Some("#commit")
            && body.get("repo").text() == did
        {
            break body.get("seq").int();
        }
    };
    let _ = probe.send(Message::Close(None)).await;

    // Now reconnect at that exact cursor: nothing to replay, straight to live.
    let mut socket = subscribe(Some(head_seq)).await;
    settle().await;

    create_post(&c, &did, &jwt, "after handover one").await;
    create_post(&c, &did, &jwt, "after handover two").await;

    let mut seen: Vec<i64> = Vec::new();
    while seen.len() < 2 {
        let bytes = next_frame(&mut socket)
            .await
            .expect("expected both post-handover frames");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) != Some("#commit") {
            continue;
        }
        let seq = body.get("seq").int();
        assert!(
            seq > head_seq,
            "the cursor said we already had seq {head_seq}; got {seq} again"
        );
        if body.get("repo").text() == did {
            assert!(!seen.contains(&seq), "event {seq} delivered twice");
            seen.push(seq);
        }
    }
    assert_eq!(seen.len(), 2);
    assert!(seen[0] < seen[1], "events must arrive in seq order: {seen:?}");
    let _ = socket.send(Message::Close(None)).await;
}

/// A cursor past the end of the log is a terminal error frame, not silence.
#[tokio::test]
async fn future_cursor_gets_an_error_frame() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let mut socket = subscribe(Some(i64::MAX / 2)).await;
    let bytes = next_frame(&mut socket)
        .await
        .expect("expected an error frame for a future cursor");
    let (header, body) = parse_frame(&bytes);
    assert_eq!(header.get("op").int(), -1, "error frames carry op = -1");
    assert!(
        header.opt("t").is_none(),
        "an error header has no `t`, only `op`"
    );
    assert_eq!(body.get("error").text(), "FutureCursor");
    assert!(
        !body.get("message").text().is_empty(),
        "the error should say what went wrong"
    );
}

/// A record deleted through `deleteRecord` produces a `delete` op with a null
/// CID — the one op shape that is not a link.
#[tokio::test]
async fn delete_produces_a_delete_op_with_a_null_cid() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("fhdelete").await;
    let created = create_post(&c, &did, &jwt, "to be deleted").await;
    let rkey = created["uri"]
        .as_str()
        .unwrap()
        .rsplit('/')
        .next()
        .unwrap()
        .to_string();

    let mut socket = subscribe(None).await;
    settle().await;

    let resp = c
        .post(format!("{PDS_URL}/xrpc/com.atproto.repo.deleteRecord"))
        .bearer_auth(&jwt)
        .json(&json!({
            "repo": did,
            "collection": "app.twi.post",
            "rkey": rkey,
        }))
        .send()
        .await
        .unwrap();
    assert!(resp.status().is_success(), "deleteRecord: {:?}", resp.status());

    let (_header, body) = loop {
        let bytes = next_frame(&mut socket).await.expect("expected a delete frame");
        let (header, body) = parse_frame(&bytes);
        if header.opt("t").map(|t| t.text()) == Some("#commit")
            && body.get("repo").text() == did
        {
            break (header, body);
        }
    };

    let ops = body.get("ops").array();
    assert_eq!(ops.len(), 1);
    assert_eq!(ops[0].get("action").text(), "delete");
    assert_eq!(ops[0].get("path").text(), format!("app.twi.post/{rkey}"));
    assert_eq!(*ops[0].get("cid"), Cbor::Null, "a delete has no resulting CID");

    // The second commit on this repo, so `since` is the previous revision.
    assert!(
        matches!(body.get("since"), Cbor::Text(_)),
        "a follow-up commit must name its predecessor's rev, got {:?}",
        body.get("since")
    );
    let _ = socket.send(Message::Close(None)).await;
}

/// Shape check for the frame's `time`: RFC 3339, UTC, with a `Z` suffix.
fn chrono_like_rfc3339(s: &str) -> bool {
    chrono::DateTime::parse_from_rfc3339(s).is_ok() && s.ends_with('Z')
}