didbot_data/mst/

read.rs

//! Reading a node back out of a block.
//!
//! [`Tree`](super::Tree) builds nodes and hands out their blocks. This is the
//! other direction, and it is the one that has to be suspicious: a node that
//! arrives in a CAR file or a firehose frame was encoded by somebody else, and
//! the encoding has several ways to say something that is not a node.
//!
//! Everything structural is checked here, so that a [`Node`] that exists is
//! one a tree can be walked over:
//!
//! * the map is exactly `l` and `e`, and each entry exactly `p`, `k`, `v`
//!   and `t`. A spare field is a different block with a different CID, so a
//!   node that carried one and was re-encoded would change its own name.
//! * a prefix length that reaches past the previous key is refused, and so
//!   is one that is *shorter* than it could be — if the byte after the shared
//!   part matches, the encoder should have shared it. Both spellings name the
//!   same key and only one of them is the block that key belongs in.
//! * keys strictly ascend. A node whose entries are out of order has no
//!   defined lookup.
//! * every key in one node sits at the same layer, which is what a layer
//!   *is*. A node mixing heights is a tree no other implementation would have
//!   built.
//!
//! What is deliberately not checked is anything needing another block: that a
//! subtree is one layer down, that it holds keys in the right range, that the
//! links resolve at all. A node is one block, and those are claims about a
//! tree — see [`Store`](super::Store), which checks them as it walks.

use crate::cid::Cid;
use crate::dag_cbor::{self, CborError};
use crate::value::Value;

use super::{key_height, Entry, Node};

/// Why a block is not a Merkle search tree node.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum NodeError {
    /// The block is not DAG-CBOR at all.
    #[error("a node block is not dag-cbor: {0}")]
    Cbor(#[from] CborError),
    /// The block decoded, and to something that is not a node.
    #[error("entry {at}: {why}")]
    Shape {
        /// Which entry the problem is in, or `0` for the node itself.
        at: usize,
        /// What was wrong.
        why: String,
    },
    /// A key that is not UTF-8.
    ///
    /// A repository path is a collection NSID and a record key, both of which
    /// are ASCII by rule. The tree is keyed on bytes and this crate is happy
    /// to *compare* bytes, but a path it cannot name is one it cannot report
    /// in a diff or an operation, so it is refused at the door rather than
    /// lost in the middle.
    #[error("entry {at}: the key is not valid utf-8")]
    KeyNotUtf8 {
        /// Which entry the key is in.
        at: usize,
    },
}

impl NodeError {
    /// A shape error at one entry.
    fn shape(at: usize, why: impl Into<String>) -> Self {
        Self::Shape {
            at,
            why: why.into(),
        }
    }
}

/// A node's contents in the form the tree algorithms want them.
///
/// The encoding stores each key against the previous one, which is compact
/// and useless to walk over. This is the same node with the keys rebuilt and
/// the subtrees interleaved the way the protocol describes them:
///
/// ```text
/// keys:          (0,    1,    2,    3)
/// values:        (0,    1,    2,    3)
/// subtrees:   (0,    1,    2,    3,    4)
/// ```
///
/// There is always exactly one more subtree than there are keys: the one
/// before the first key is the node's `l`, and each entry's `t` is the one
/// after it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Parts {
    /// The full keys, ascending.
    pub keys: Vec<String>,
    /// The record each key names.
    pub values: Vec<Cid>,
    /// The subtrees between and around the keys; always `keys.len() + 1`.
    pub subtrees: Vec<Option<Cid>>,
}

impl Node {
    /// Decodes one node out of its block.
    ///
    /// ```
    /// use std::collections::BTreeMap;
    /// use didbot_data::mst::{Node, Tree};
    ///
    /// let tree = Tree::build(&BTreeMap::new());
    /// let block = &tree.blocks()[tree.root()];
    /// assert_eq!(Node::decode(block).unwrap(), Node::default());
    /// ```
    pub fn decode(block: &[u8]) -> Result<Self, NodeError> {
        Self::from_value(&dag_cbor::decode(block)?)
    }

    /// Reads a node out of an already-decoded value.
    pub fn from_value(value: &Value) -> Result<Self, NodeError> {
        let Value::Map(fields) = value else {
            return Err(NodeError::shape(0, "a node is a map"));
        };
        if fields.len() != 2 || !fields.contains_key("l") || !fields.contains_key("e") {
            return Err(NodeError::shape(0, "a node is exactly the fields l and e"));
        }
        let left = link_or_null(&fields["l"]).ok_or_else(|| {
            NodeError::shape(0, "l is a link to the subtree left of every entry, or null")
        })?;
        let Value::List(entries) = &fields["e"] else {
            return Err(NodeError::shape(0, "e is a list of entries"));
        };

        let mut node = Node {
            left,
            entries: Vec::with_capacity(entries.len()),
        };
        let mut previous: Vec<u8> = Vec::new();
        let mut layer: Option<usize> = None;
        for (at, entry) in entries.iter().enumerate() {
            let key = read_entry(at, entry, &previous, &mut node)?;
            let height = key_height(&key);
            match layer {
                None => layer = Some(height),
                Some(layer) if layer == height => {}
                Some(layer) => {
                    return Err(NodeError::shape(
                        at,
                        format!("{key:?} is at layer {height}, and this node is layer {layer}"),
                    ))
                }
            }
            previous = key.into_bytes();
        }
        Ok(node)
    }

    /// The node with its keys rebuilt; see [`Parts`].
    ///
    /// Every key was checked to be UTF-8 when the node was decoded, and a
    /// node this crate built has UTF-8 keys because it was handed strings. A
    /// node assembled by hand out of arbitrary bytes is the one case where
    /// this is lossy, and it is the case [`Node::decode`] refuses.
    #[must_use]
    pub fn parts(&self) -> Parts {
        let mut parts = Parts {
            keys: Vec::with_capacity(self.entries.len()),
            values: Vec::with_capacity(self.entries.len()),
            subtrees: Vec::with_capacity(self.entries.len() + 1),
        };
        parts.subtrees.push(self.left.clone());
        let mut previous: Vec<u8> = Vec::new();
        for entry in &self.entries {
            let mut key = previous[..entry.prefix.min(previous.len())].to_vec();
            key.extend_from_slice(&entry.suffix);
            previous = key.clone();
            parts.keys.push(String::from_utf8_lossy(&key).into_owned());
            parts.values.push(entry.value.clone());
            parts.subtrees.push(entry.tree.clone());
        }
        parts
    }

    /// Builds a node from full keys, values and subtrees.
    ///
    /// The inverse of [`Node::parts`], and the only way to make a node
    /// without going through [`Tree::build`](super::Tree::build): the shared
    /// prefixes are computed here, so a caller cannot hold a node whose
    /// encoding disagrees with its keys.
    ///
    /// # Panics
    ///
    /// If there is not exactly one more subtree than there are keys, or if
    /// the keys and values are of different lengths. Both are the caller
    /// building something that is not a node.
    #[must_use]
    pub fn from_parts(parts: &Parts) -> Self {
        assert_eq!(
            parts.keys.len(),
            parts.values.len(),
            "a node names one record per key"
        );
        assert_eq!(
            parts.subtrees.len(),
            parts.keys.len() + 1,
            "a node has one more subtree than it has keys"
        );
        let mut previous = "";
        let entries = parts
            .keys
            .iter()
            .zip(&parts.values)
            .zip(&parts.subtrees[1..])
            .map(|((key, value), tree)| {
                let prefix = super::common_prefix_len(previous, key);
                previous = key;
                Entry {
                    prefix,
                    suffix: key.as_bytes()[prefix..].to_vec(),
                    value: value.clone(),
                    tree: tree.clone(),
                }
            })
            .collect();
        Self {
            left: parts.subtrees[0].clone(),
            entries,
        }
    }

    /// The block this node encodes to, and the CID naming it.
    #[must_use]
    pub fn block(&self) -> (Cid, Vec<u8>) {
        let block = dag_cbor::encode(&self.to_value());
        (Cid::of_dag_cbor(&block), block)
    }

    /// Which layer this node is on, where that can be told from the node
    /// alone.
    ///
    /// A node with keys is at the layer its keys are at — they all share one,
    /// which [`Node::decode`] checks. A node with no keys and no left subtree
    /// is the empty root, which is layer zero by definition. A node with no
    /// keys and a left subtree is the empty intermediate the layering rule
    /// produces, and it has no height of its own: it is one below whatever
    /// pointed at it, and nothing inside it says what that was.
    #[must_use]
    pub fn layer(&self) -> Option<usize> {
        // A node's first entry shares nothing with a predecessor it does not
        // have, so its suffix is a whole key — the one key readable without
        // rebuilding any other.
        if let Some(entry) = self.entries.first() {
            return Some(key_height(&String::from_utf8_lossy(&entry.suffix)));
        }
        if self.left.is_none() {
            return Some(0);
        }
        None
    }

    /// Whether this is the empty node: no keys, and nothing below it.
    ///
    /// This is the node an account's first commit signs over, and it is also
    /// what a subtree pointer of `null` means — which is why an empty node is
    /// never *stored* as a subtree. A tree carrying one would have two
    /// spellings of "nothing here".
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty() && self.left.is_none()
    }
}

/// Reads one entry, appends it to `node`, and answers its full key.
fn read_entry(
    at: usize,
    entry: &Value,
    previous: &[u8],
    node: &mut Node,
) -> Result<String, NodeError> {
    let Value::Map(fields) = entry else {
        return Err(NodeError::shape(at, "an entry is a map"));
    };
    if fields.len() != 4
        || !["p", "k", "v", "t"]
            .iter()
            .all(|key| fields.contains_key(*key))
    {
        return Err(NodeError::shape(
            at,
            "an entry is exactly the fields p, k, v and t",
        ));
    }
    let Value::Integer(prefix) = fields["p"] else {
        return Err(NodeError::shape(at, "p is the shared prefix length"));
    };
    let prefix = usize::try_from(prefix)
        .map_err(|_| NodeError::shape(at, "p is a length and cannot be negative"))?;
    let Value::Bytes(suffix) = &fields["k"] else {
        return Err(NodeError::shape(at, "k is the rest of the key, as bytes"));
    };
    let Value::Link(value) = &fields["v"] else {
        return Err(NodeError::shape(at, "v is a link to the record"));
    };
    let tree = link_or_null(&fields["t"])
        .ok_or_else(|| NodeError::shape(at, "t is a link to a subtree, or null"))?;

    if prefix > previous.len() {
        return Err(NodeError::shape(
            at,
            format!(
                "p is {prefix}, and the previous key is {} bytes",
                previous.len()
            ),
        ));
    }
    // The encoder shares every byte it can, so a shared byte left unshared is
    // a second spelling of the same key — a different block for a node that
    // holds the same thing.
    if previous.get(prefix).is_some() && previous.get(prefix) == suffix.first() {
        return Err(NodeError::shape(
            at,
            "p is shorter than the prefix these keys share",
        ));
    }
    let mut key = previous[..prefix].to_vec();
    key.extend_from_slice(suffix);
    if key.as_slice() <= previous {
        return Err(NodeError::shape(at, "keys must ascend within a node"));
    }
    let key = String::from_utf8(key).map_err(|_| NodeError::KeyNotUtf8 { at })?;

    node.entries.push(Entry {
        prefix,
        suffix: suffix.clone(),
        value: value.clone(),
        tree,
    });
    Ok(key)
}

/// Reads a field that is a link or `null`, and refuses anything else.
///
/// The outer `Option` is "this was not one of those two"; the inner one is
/// the null. A field that is simply absent never reaches here — the field
/// count is checked first.
fn link_or_null(value: &Value) -> Option<Option<Cid>> {
    match value {
        Value::Null => Some(None),
        Value::Link(cid) => Some(Some(cid.clone())),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::*;
    use crate::mst::Tree;

    /// Every node of a built tree decodes back to the node it was built as.
    /// This is the claim, and it is made over a tree deep enough to have
    /// entries with shared prefixes, subtrees, and an empty intermediate.
    #[test]
    fn every_node_of_a_tree_decodes_to_itself() {
        let keys: Vec<String> = (0..300)
            .map(|i| format!("app.bsky.feed.post/3jt6wh3ck{i:03}"))
            .collect();
        let leaves: BTreeMap<String, Cid> = keys
            .iter()
            .map(|key| (key.clone(), Cid::of_dag_cbor(key.as_bytes())))
            .collect();
        let tree = Tree::build(&leaves);
        for (cid, block) in tree.blocks() {
            let node = Node::decode(block).expect("a node this crate built");
            assert_eq!(&node, &tree.nodes()[cid]);
            assert_eq!(node.block().0, *cid);
            assert_eq!(Node::from_parts(&node.parts()), node);
        }
    }

    /// A node re-encoded from its parts is the same block, which is what
    /// makes `from_parts` usable for anything that rewrites a tree.
    #[test]
    fn parts_round_trip_through_the_encoding() {
        let mut parts = Parts {
            keys: vec!["app.bsky.feed.post/aaaa".to_owned()],
            values: vec![Cid::of_dag_cbor(b"a record")],
            subtrees: vec![None, None],
        };
        parts.keys.push("app.bsky.feed.post/aaab".to_owned());
        parts.values.push(Cid::of_dag_cbor(b"another record"));
        parts.subtrees.push(None);
        let node = Node::from_parts(&parts);
        assert_eq!(node.parts(), parts);
        assert_eq!(Node::decode(&node.block().1).unwrap(), node);
        // The second entry shares all but its last byte.
        assert_eq!(node.entries[1].prefix, 22);
    }

    /// A prefix length that reaches past the previous key would index out of
    /// bounds in a reader that trusted it.
    #[test]
    fn a_prefix_past_the_previous_key_is_refused() {
        let block = node_block(&[("k", 40, b"x")]);
        assert!(matches!(
            Node::decode(&block),
            Err(NodeError::Shape { at: 0, .. })
        ));
    }

    /// Two spellings of one key: `p` may not be shorter than the prefix the
    /// keys actually share, because the tree is content-addressed and a
    /// second spelling is a second CID for one node.
    #[test]
    fn a_prefix_shorter_than_the_shared_one_is_refused() {
        // "aa" then "ab": the shared prefix is one byte, so p=1 k="b".
        let good = node_block(&[("", 0, b"aa"), ("", 1, b"b")]);
        assert!(Node::decode(&good).is_ok());
        let sloppy = node_block(&[("", 0, b"aa"), ("", 0, b"ab")]);
        assert!(matches!(
            Node::decode(&sloppy),
            Err(NodeError::Shape { at: 1, .. })
        ));
    }

    /// Entries out of order have no defined lookup, so the node is not one.
    #[test]
    fn keys_must_ascend() {
        let descending = node_block(&[("", 0, b"bb"), ("", 0, b"aa")]);
        assert!(matches!(
            Node::decode(&descending),
            Err(NodeError::Shape { at: 1, .. })
        ));
        let repeated = node_block(&[("", 0, b"aa"), ("", 0, b"aa")]);
        assert!(Node::decode(&repeated).is_err());
    }

    /// A field the protocol does not define changes the block's CID, so a
    /// node carrying one cannot be re-encoded as itself.
    #[test]
    fn a_spare_field_is_not_a_node() {
        let mut map = BTreeMap::from([
            ("l".to_owned(), Value::Null),
            ("e".to_owned(), Value::List(Vec::new())),
        ]);
        assert!(Node::from_value(&Value::Map(map.clone())).is_ok());
        map.insert("extra".to_owned(), Value::Integer(1));
        assert!(Node::from_value(&Value::Map(map)).is_err());
    }

    /// Every key in a node is at the node's layer. A node mixing them is a
    /// tree built by rules nobody else uses.
    #[test]
    fn a_node_holds_one_layer() {
        // "asdf" is at layer 0 and "blue" is at layer 1.
        let mixed = node_block(&[("", 0, b"asdf"), ("", 0, b"blue")]);
        assert!(matches!(
            Node::decode(&mixed),
            Err(NodeError::Shape { at: 1, .. })
        ));
    }

    /// Encodes a node from raw entry fields, including ones a correct encoder
    /// would never write. The first tuple element is unused padding so the
    /// cases read as (comment, p, k).
    fn node_block(entries: &[(&str, i64, &[u8])]) -> Vec<u8> {
        let value = Value::Map(BTreeMap::from([
            ("l".to_owned(), Value::Null),
            (
                "e".to_owned(),
                Value::List(
                    entries
                        .iter()
                        .map(|(_, prefix, suffix)| {
                            Value::Map(BTreeMap::from([
                                ("p".to_owned(), Value::Integer(*prefix)),
                                ("k".to_owned(), Value::Bytes((*suffix).to_vec())),
                                ("v".to_owned(), Value::Link(Cid::of_dag_cbor(b"a record"))),
                                ("t".to_owned(), Value::Null),
                            ]))
                        })
                        .collect(),
                ),
            ),
        ]));
        dag_cbor::encode(&value)
    }
}