maarcadetweet/crates/at-mst/src/node.rs

use anyhow::{anyhow, Result};
use cid::Cid;
use serde::{Deserialize, Serialize};

use at_crypto::cid::cid_for_cbor;

/// A single MST entry. The `key` is the **base64url-encoded** form of the
/// user-facing key string. The `tree` is the CID of the sub-tree immediately
/// to the right of this entry (i.e. the sub-tree that contains all keys
/// strictly between this entry's key and the next entry's key).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MstEntry {
    pub key: String,
    pub value: Cid,
    #[serde(rename = "t", skip_serializing_if = "Option::is_none")]
    pub tree: Option<Cid>,
}

impl MstEntry {
    pub fn new(encoded_key: impl Into<String>, value: Cid, tree: Option<Cid>) -> Self {
        Self {
            key: encoded_key.into(),
            value,
            tree,
        }
    }
}

/// Tag used to distinguish a node that only contains leaf entries (no sub-trees
/// pointing further down) from an inner node.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeKind {
    Leaf,
    Inner,
}

/// In-memory representation of an MST node.
#[derive(Debug, Clone)]
pub struct MstNode {
    pub left: Option<Cid>,
    pub entries: Vec<MstEntry>,
    pub cid: Cid,
}

impl MstNode {
    pub fn leaf(entries: Vec<MstEntry>, cid: Cid) -> Self {
        Self {
            left: None,
            entries,
            cid,
        }
    }

    pub fn kind(&self) -> NodeKind {
        if self.left.is_some() || self.entries.iter().any(|e| e.tree.is_some()) {
            NodeKind::Inner
        } else {
            NodeKind::Leaf
        }
    }

    pub fn is_leaf(&self) -> bool {
        self.kind() == NodeKind::Leaf
    }
}

// -- CBOR wire format ----------------------------------------------------
//
// Per the atproto MST spec (datamodel-repo#node-data), each MST
// node is a DAG-CBOR object:
//
//     {
//       "l": <CID> | null,       // left sub-tree (keys < first entry)
//       "e": [{                  // entries in sort order
//         "k": "<encoded>",      // see encode_key below
//         "v": <CID>,            // value block pointer
//         "t": <CID> | null      // right sub-tree for this entry
//       }, ...]
//     }
//
// Optional fields (`t`) are CBOR-omitted via `serde(skip_serializing_if)`.
// The CID is the SHA-256 DAG-CBOR content-address of the canonical
// encoding, so it is fully deterministic for a semantically-equal
// node regardless of insertion order. (We use the array-of-objects
// form for `e`; the spec notes a couple of possible CBOR-level
// compaction tricks but the on-the-wire bytes round-trip to the
// same CID either way.)

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct WireNode {
    #[serde(rename = "l", skip_serializing_if = "Option::is_none")]
    pub left: Option<Cid>,
    #[serde(rename = "e")]
    pub entries: Vec<WireEntry>,
}

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct WireEntry {
    #[serde(rename = "k")]
    pub key: String,
    #[serde(rename = "v")]
    pub value: Cid,
    #[serde(rename = "t", skip_serializing_if = "Option::is_none")]
    pub tree: Option<Cid>,
}

/// Encode the node `(left, entries)` to its canonical DAG-CBOR bytes.
pub(crate) fn encode_cbor(left: Option<&Cid>, entries: &[MstEntry]) -> Result<Vec<u8>> {
    let wire_entries: Vec<WireEntry> = entries
        .iter()
        .map(|e| WireEntry {
            key: e.key.clone(),
            value: e.value,
            tree: e.tree,
        })
        .collect();
    let node = WireNode {
        left: left.cloned(),
        entries: wire_entries,
    };
    let mut buf = Vec::new();
    ciborium::into_writer(&node, &mut buf)?;
    Ok(buf)
}

/// Decode a node from CBOR bytes. Returns `(left, entries, computed_cid)`.
/// `computed_cid` is the CID implied by the canonical encoding of `bytes`,
/// callers can verify it matches the CID used to fetch the block.
pub(crate) fn decode_cbor(bytes: &[u8]) -> Result<(Option<Cid>, Vec<MstEntry>, Cid)> {
    let wire: WireNode = ciborium::from_reader(bytes)
        .map_err(|e| anyhow!("failed to decode MST node CBOR: {e}"))?;
    let entries: Vec<MstEntry> = wire
        .entries
        .into_iter()
        .map(|we| MstEntry {
            key: we.key,
            value: we.value,
            tree: we.tree,
        })
        .collect();
    let cid = cid_for_cbor(bytes)?;
    Ok((wire.left, entries, cid))
}

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

    #[test]
    fn leaf_kind_detection() {
        let e = MstEntry::new("a", Cid::default(), None);
        // We can't easily build a real CID without a hash; this test is mainly
        // for the leaf/inner classification logic which only depends on the
        // Option<Cid> fields.
        let node = MstNode {
            left: None,
            entries: vec![e],
            cid: Cid::default(),
        };
        assert!(node.is_leaf());
    }
}