maarcadetweet/crates/appview/src/cbor.rs

//! A minimal, allocation-honest CBOR reader — just enough to decode the
//! frames of `com.atproto.sync.subscribeRepos` and the header of a CAR
//! file.
//!
//! ## Why not `ciborium`?
//!
//! Two reasons, and both come from what the wire format actually is.
//!
//! 1. **Two values per message.** A subscribeRepos frame is *two*
//!    DAG-CBOR values written back to back (header, then body) inside
//!    one WebSocket binary message. A `serde`-shaped reader gives us
//!    "decode one value from this slice" and no cursor we can resume
//!    from, so we would have to guess where the header ended.
//! 2. **Tag 42.** DAG-CBOR encodes a CID link as `tag(42) +
//!    bytes(<cid>)`. `ciborium`'s `serde` mapping turns tags into a
//!    private newtype dance that does not survive a round trip through
//!    `serde_json::Value`, which is the shape the rest of the AppView
//!    speaks.
//!
//! So this module decodes CBOR into its own small [`Cbor`] tree with an
//! explicit byte offset, which makes "read the header, then read the
//! body from where the header stopped" a two-line function.
//!
//! ## What it deliberately does not do
//!
//! No indefinite-length items (DAG-CBOR forbids them; we reject them
//! rather than guess), no half floats beyond a plain `f64` widening, no
//! canonicalisation checks. It is a *reader* for input we already
//! decided to trust at the transport layer, with hard length checks so
//! a malformed frame returns `Err` instead of panicking.
//!
//! Record blocks inside the CAR payload are NOT decoded with this
//! module: they are written by `ciborium::into_writer(&serde_json::Value)`
//! on the PDS side (see `pds-server/src/routes/repo.rs`), which means
//! CIDs inside a record are plain strings, not tag-42 links. Their
//! inverse is `ciborium::from_reader::<serde_json::Value, _>`, and
//! that is what [`crate::pds_firehose`] uses for them.

use anyhow::{anyhow, bail, Result};
use cid::Cid;

/// A decoded CBOR value.
///
/// `Nint` carries the already-negated value (CBOR stores `-1 - n`), so
/// callers never have to remember the bias.
#[derive(Debug, Clone, PartialEq)]
pub enum Cbor {
    Uint(u64),
    Nint(i64),
    Bytes(Vec<u8>),
    Text(String),
    Array(Vec<Cbor>),
    /// Kept as an ordered key/value list rather than a map: DAG-CBOR
    /// keys are text and already canonically ordered, and a `Vec` keeps
    /// the decoder free of hashing while the maps we read have a
    /// handful of entries at most.
    Map(Vec<(Cbor, Cbor)>),
    Tag(u64, Box<Cbor>),
    Bool(bool),
    Null,
    Undefined,
    Float(f64),
}

impl Cbor {
    /// Look up a text key in a map. `None` for non-maps and misses.
    pub fn get(&self, key: &str) -> Option<&Cbor> {
        match self {
            Cbor::Map(entries) => entries.iter().find_map(|(k, v)| match k {
                Cbor::Text(s) if s == key => Some(v),
                _ => None,
            }),
            _ => None,
        }
    }

    /// The value as a signed integer, accepting both CBOR integer
    /// majors. `None` when the value is not an integer, or when an
    /// unsigned value exceeds `i64::MAX` (which cannot happen for a
    /// `seq`, but silently wrapping would be worse than a miss).
    pub fn as_i64(&self) -> Option<i64> {
        match self {
            Cbor::Uint(n) => i64::try_from(*n).ok(),
            Cbor::Nint(n) => Some(*n),
            _ => None,
        }
    }

    pub fn as_str(&self) -> Option<&str> {
        match self {
            Cbor::Text(s) => Some(s.as_str()),
            _ => None,
        }
    }

    pub fn as_bytes(&self) -> Option<&[u8]> {
        match self {
            Cbor::Bytes(b) => Some(b.as_slice()),
            _ => None,
        }
    }

    pub fn as_bool(&self) -> Option<bool> {
        match self {
            Cbor::Bool(b) => Some(*b),
            _ => None,
        }
    }

    pub fn as_array(&self) -> Option<&[Cbor]> {
        match self {
            Cbor::Array(items) => Some(items.as_slice()),
            _ => None,
        }
    }

    pub fn is_null(&self) -> bool {
        matches!(self, Cbor::Null | Cbor::Undefined)
    }

    /// Decode a DAG-CBOR CID link: `tag(42) + bytes(...)`.
    ///
    /// The spec prefixes the CID bytes with a single `0x00` (the
    /// multibase "identity" marker), because a CID inside a byte string
    /// has no textual multibase prefix to carry. **This repository's own
    /// writer omits that byte** — `pds-server/src/car.rs` writes
    /// `cbor_bytes(&cid.to_bytes())` — so we accept both spellings: a
    /// leading `0x00` is skipped, anything else is parsed as-is. Being
    /// lenient here costs nothing (a real CIDv1 never starts with
    /// `0x00`, and CIDv0 starts with `0x12`) and it means the AppView
    /// keeps working whether the PDS follows the spec or the house
    /// convention.
    pub fn as_cid(&self) -> Option<Cid> {
        let inner = match self {
            Cbor::Tag(42, inner) => inner.as_ref(),
            // A bare byte string where a link is expected: CARs written by
            // older builds of this PDS tagged the root CID without the
            // `0x00` identity prefix, and some encoders drop the tag
            // entirely. Both still have to parse.
            Cbor::Bytes(_) => self,
            // A record block written from `serde_json::Value` spells a
            // CID as a plain string — accept that too, so callers do
            // not need a second code path for the block contents.
            Cbor::Text(s) => return s.parse::<Cid>().ok(),
            _ => return None,
        };
        let raw = inner.as_bytes()?;
        let raw = match raw.first() {
            Some(0x00) => &raw[1..],
            _ => raw,
        };
        Cid::read_bytes(raw).ok()
    }
}

/// Decode exactly one CBOR value starting at `pos`.
///
/// Returns the value and the offset just past it, so a caller can read
/// the next value from the same buffer — which is precisely what a
/// two-value subscribeRepos frame needs.
pub fn decode_at(bytes: &[u8], pos: usize) -> Result<(Cbor, usize)> {
    let (major, arg, mut p) = read_head(bytes, pos)?;
    match major {
        0 => Ok((Cbor::Uint(arg), p)),
        1 => {
            // CBOR negative integers store `-1 - n`. Values below
            // `i64::MIN` cannot occur in anything we consume, and
            // wrapping them would produce a positive number, so bail.
            let n = i64::try_from(arg)
                .map_err(|_| anyhow!("negative integer out of i64 range"))?;
            Ok((Cbor::Nint(-1 - n), p))
        }
        2 => {
            let end = checked_end(bytes, p, arg, "byte string")?;
            let v = bytes[p..end].to_vec();
            Ok((Cbor::Bytes(v), end))
        }
        3 => {
            let end = checked_end(bytes, p, arg, "text string")?;
            let s = std::str::from_utf8(&bytes[p..end])
                .map_err(|e| anyhow!("invalid UTF-8 in CBOR text: {e}"))?
                .to_string();
            Ok((Cbor::Text(s), end))
        }
        4 => {
            let mut items = Vec::with_capacity(sane_capacity(arg));
            for _ in 0..arg {
                let (v, next) = decode_at(bytes, p)?;
                items.push(v);
                p = next;
            }
            Ok((Cbor::Array(items), p))
        }
        5 => {
            let mut entries = Vec::with_capacity(sane_capacity(arg));
            for _ in 0..arg {
                let (k, next) = decode_at(bytes, p)?;
                let (v, next) = decode_at(bytes, next)?;
                entries.push((k, v));
                p = next;
            }
            Ok((Cbor::Map(entries), p))
        }
        6 => {
            let (inner, next) = decode_at(bytes, p)?;
            Ok((Cbor::Tag(arg, Box::new(inner)), next))
        }
        7 => match arg {
            20 => Ok((Cbor::Bool(false), p)),
            21 => Ok((Cbor::Bool(true), p)),
            22 => Ok((Cbor::Null, p)),
            23 => Ok((Cbor::Undefined, p)),
            // Floats arrive as the raw bit pattern in `arg`; the width
            // is implied by the additional-information byte, which
            // `read_head` has already consumed. We only ever see f64 in
            // practice (DAG-CBOR requires it), so the narrower widths
            // are decoded for completeness rather than need.
            _ => Ok((Cbor::Float(f64::from_bits(arg)), p)),
        },
        other => bail!("unsupported CBOR major type {other}"),
    }
}

/// Decode a single CBOR value that must span the whole buffer.
pub fn decode(bytes: &[u8]) -> Result<Cbor> {
    let (v, end) = decode_at(bytes, 0)?;
    if end != bytes.len() {
        bail!("trailing bytes after CBOR value ({} left)", bytes.len() - end);
    }
    Ok(v)
}

/// Read a CBOR head: major type plus its argument.
///
/// Indefinite-length encodings (`additional information == 31`) are
/// rejected: DAG-CBOR forbids them, and accepting them would mean
/// implementing break-stop scanning for input that should never carry
/// it.
fn read_head(bytes: &[u8], pos: usize) -> Result<(u8, u64, usize)> {
    let first = *bytes
        .get(pos)
        .ok_or_else(|| anyhow!("CBOR read past end of input at {pos}"))?;
    let major = first >> 5;
    let low = first & 0x1f;
    let (arg, extra) = match low {
        0..=23 => (low as u64, 0usize),
        24 => (read_uint(bytes, pos + 1, 1)?, 1),
        25 => (read_uint(bytes, pos + 1, 2)?, 2),
        26 => (read_uint(bytes, pos + 1, 4)?, 4),
        27 => (read_uint(bytes, pos + 1, 8)?, 8),
        31 => bail!("indefinite-length CBOR item is not valid DAG-CBOR"),
        other => bail!("reserved CBOR additional information {other}"),
    };
    // For major 7 the "argument" of a float is the raw bit pattern, and
    // f32/f16 need widening before `f64::from_bits` makes sense.
    let arg = match (major, low) {
        (7, 25) => f64::from(half_to_f32(arg as u16)).to_bits(),
        (7, 26) => f64::from(f32::from_bits(arg as u32)).to_bits(),
        _ => arg,
    };
    Ok((major, arg, pos + 1 + extra))
}

fn read_uint(bytes: &[u8], pos: usize, len: usize) -> Result<u64> {
    if pos + len > bytes.len() {
        bail!("truncated CBOR integer of {len} byte(s) at {pos}");
    }
    let mut n: u64 = 0;
    for b in &bytes[pos..pos + len] {
        n = (n << 8) | u64::from(*b);
    }
    Ok(n)
}

/// IEEE-754 half → single. Only reached for `f16` inputs, which nothing
/// in this protocol emits; kept so a stray value decodes instead of
/// erroring out mid-frame.
fn half_to_f32(bits: u16) -> f32 {
    let sign = ((bits >> 15) & 1) as u32;
    let exp = ((bits >> 10) & 0x1f) as u32;
    let frac = (bits & 0x3ff) as u32;
    let out = match exp {
        0 if frac == 0 => sign << 31,
        0 => {
            // Subnormal: renormalise.
            let mut e = -1i32;
            let mut f = frac;
            while f & 0x400 == 0 {
                f <<= 1;
                e -= 1;
            }
            let exp32 = (127 - 15 + e) as u32;
            (sign << 31) | (exp32 << 23) | ((f & 0x3ff) << 13)
        }
        0x1f => (sign << 31) | (0xff << 23) | (frac << 13),
        _ => (sign << 31) | ((exp + 127 - 15) << 23) | (frac << 13),
    };
    f32::from_bits(out)
}

/// Bounds-check a string/bytes payload before slicing it.
fn checked_end(bytes: &[u8], pos: usize, len: u64, what: &str) -> Result<usize> {
    let len = usize::try_from(len).map_err(|_| anyhow!("{what} length overflows usize"))?;
    let end = pos
        .checked_add(len)
        .ok_or_else(|| anyhow!("{what} length overflows"))?;
    if end > bytes.len() {
        bail!("{what} of {len} byte(s) exceeds input at offset {pos}");
    }
    Ok(end)
}

/// Cap the pre-allocation a declared array/map length can trigger. A
/// corrupt frame claiming `map(2^40)` must not make us reserve 40 GiB
/// before the first missing byte errors out; the collection still grows
/// naturally for genuinely large inputs.
fn sane_capacity(declared: u64) -> usize {
    usize::try_from(declared.min(1024)).unwrap_or(0)
}

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

    /// Build `{"op": 1, "t": "#commit"}` by hand — the exact bytes the
    /// PDS writes for a regular frame header.
    fn commit_header_bytes() -> Vec<u8> {
        let mut v = vec![0xA2]; // map(2)
        v.push(0x62); // text(2)
        v.extend_from_slice(b"op");
        v.push(0x01); // uint 1
        v.push(0x61); // text(1)
        v.extend_from_slice(b"t");
        v.push(0x67); // text(7)
        v.extend_from_slice(b"#commit");
        v
    }

    #[test]
    fn decodes_a_frame_header() {
        let v = decode(&commit_header_bytes()).unwrap();
        assert_eq!(v.get("op").unwrap().as_i64(), Some(1));
        assert_eq!(v.get("t").unwrap().as_str(), Some("#commit"));
        assert!(v.get("missing").is_none());
    }

    #[test]
    fn decodes_negative_op_of_an_error_header() {
        // {"op": -1} → map(1), text(2)"op", nint(0) = -1
        let bytes = vec![0xA1, 0x62, b'o', b'p', 0x20];
        let v = decode(&bytes).unwrap();
        assert_eq!(v.get("op").unwrap().as_i64(), Some(-1));
    }

    #[test]
    fn decodes_multi_byte_integers() {
        // uint16 300, uint32 70000, uint64 2^33, nint -300
        assert_eq!(decode(&[0x19, 0x01, 0x2C]).unwrap().as_i64(), Some(300));
        assert_eq!(
            decode(&[0x1A, 0x00, 0x01, 0x11, 0x70]).unwrap().as_i64(),
            Some(70000)
        );
        assert_eq!(
            decode(&[0x1B, 0, 0, 0, 2, 0, 0, 0, 0]).unwrap().as_i64(),
            Some(8_589_934_592)
        );
        assert_eq!(decode(&[0x39, 0x01, 0x2B]).unwrap().as_i64(), Some(-300));
    }

    #[test]
    fn decode_at_reads_two_values_back_to_back() {
        // This is the whole reason the module exists: a frame is header
        // + body concatenated with no separator.
        let mut buf = commit_header_bytes();
        let body_start = buf.len();
        buf.extend_from_slice(&[0xA1, 0x63, b's', b'e', b'q', 0x18, 0x2A]); // {"seq": 42}
        let (header, next) = decode_at(&buf, 0).unwrap();
        assert_eq!(next, body_start);
        assert_eq!(header.get("t").unwrap().as_str(), Some("#commit"));
        let (body, end) = decode_at(&buf, next).unwrap();
        assert_eq!(end, buf.len());
        assert_eq!(body.get("seq").unwrap().as_i64(), Some(42));
    }

    #[test]
    fn decodes_tag_42_cid_link_both_spellings() {
        let cid = cid_for_cbor(b"a block").unwrap();
        let raw = cid.to_bytes();

        // House spelling: tag(42) + bytes(<cid>) with no 0x00 prefix.
        let mut bare = vec![0xD8, 42]; // tag(42) via 1-byte extension
        bare.push(0x58); // bytes, 1-byte length
        bare.push(raw.len() as u8);
        bare.extend_from_slice(&raw);
        assert_eq!(decode(&bare).unwrap().as_cid(), Some(cid));

        // Spec spelling: the same, with the identity multibase prefix.
        let mut prefixed = vec![0xD8, 42, 0x58, (raw.len() + 1) as u8, 0x00];
        prefixed.extend_from_slice(&raw);
        assert_eq!(decode(&prefixed).unwrap().as_cid(), Some(cid));

        // And the string spelling records use.
        let text = {
            let s = cid.to_string();
            let mut v = vec![0x78, s.len() as u8];
            v.extend_from_slice(s.as_bytes());
            v
        };
        assert_eq!(decode(&text).unwrap().as_cid(), Some(cid));
    }

    #[test]
    fn decodes_simple_values_and_containers() {
        // [true, false, null] → array(3)
        let v = decode(&[0x83, 0xF5, 0xF4, 0xF6]).unwrap();
        let items = v.as_array().unwrap();
        assert_eq!(items[0].as_bool(), Some(true));
        assert_eq!(items[1].as_bool(), Some(false));
        assert!(items[2].is_null());
    }

    #[test]
    fn rejects_truncated_and_indefinite_input() {
        // text(7) claiming 7 bytes but carrying 2.
        assert!(decode(&[0x67, b'a', b'b']).is_err());
        // Indefinite-length array.
        assert!(decode(&[0x9F, 0x01, 0xFF]).is_err());
        // Trailing garbage after a complete value.
        assert!(decode(&[0x01, 0x02]).is_err());
        // Empty input.
        assert!(decode(&[]).is_err());
    }

    #[test]
    fn oversized_declared_length_errors_instead_of_allocating() {
        // map(2^32) with nothing behind it. Must return Err quickly
        // rather than trying to reserve the declared capacity.
        let bytes = vec![0xBA, 0xFF, 0xFF, 0xFF, 0xFF];
        assert!(decode(&bytes).is_err());
    }
}