chadtmiller.com/atproto
atproto's primitives in Rust, small and free of I/O
atproto 0.1.1 Latest 1 package
0
0
//! The Merkle search tree a repo's records hang from: keys
//! (`collection/rkey`) to the CIDs of their records, in a tree whose shape
//! follows from the keys alone, so that the same records always hash to
//! the same root.
//!
//! A tree is built whole from its keys ([`Tree::build`]) rather than
//! changed in place. That is all a repo needs while it is small; changing
//! it in place is for when it is not.
use crate::cid::{Cid, DAG_CBOR};
use crate::drisl::{self, Value};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MstError(pub String);
impl fmt::Display for MstError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for MstError {}
fn err<T>(what: impl Into<String>) -> Result<T, MstError> {
Err(MstError(what.into()))
}
/// The layer a key sits at: the leading zero bits of its SHA-256, in
/// pairs, so that each layer holds about a quarter of the one below.
pub fn depth(key: &[u8]) -> u32 {
let hash = Sha256::digest(key);
let mut zeros = 0;
for byte in hash {
zeros += byte.leading_zeros();
if byte != 0 {
break;
}
}
zeros / 2
}
/// How many bytes `a` and `b` start with in common.
pub fn common_prefix(a: &[u8], b: &[u8]) -> usize {
a.iter().zip(b).take_while(|(x, y)| x == y).count()
}
/// A part of a node, in key order: a subtree of the keys between its
/// neighbours, or a key and its record's CID.
#[derive(Debug, Clone, PartialEq)]
pub enum Entry {
Tree(Node),
Leaf { key: String, value: Cid },
}
/// A node of the tree, with its bytes and CID worked out once.
#[derive(Debug, Clone, PartialEq)]
pub struct Node {
entries: Vec<Entry>,
bytes: Vec<u8>,
cid: Cid,
}
impl Node {
fn new(entries: Vec<Entry>) -> Node {
let bytes = drisl::encode(&Node::value(&entries));
Node {
cid: Cid::of_cbor(&bytes),
entries,
bytes,
}
}
/// The node as DRISL has it: `l`, its leftmost subtree, then `e`, its
/// keys, each with its prefix shared with the key before it elided, and
/// `t`, the subtree to its right.
fn value(entries: &[Entry]) -> Value {
let link = |e: Option<&Entry>| match e {
Some(Entry::Tree(n)) => Value::Link(n.cid),
_ => Value::Null,
};
let mut e = Vec::new();
let mut last: &[u8] = &[];
for (i, entry) in entries.iter().enumerate() {
if let Entry::Leaf { key, value } = entry {
let p = common_prefix(last, key.as_bytes());
e.push(Value::Map(BTreeMap::from([
("p".into(), Value::Integer(p as i64)),
("k".into(), Value::Bytes(key.as_bytes()[p..].to_vec())),
("v".into(), Value::Link(*value)),
("t".into(), link(entries.get(i + 1))),
])));
last = key.as_bytes();
}
}
Value::Map(BTreeMap::from([
("l".into(), link(entries.first())),
("e".into(), Value::Array(e)),
]))
}
pub fn cid(&self) -> Cid {
self.cid
}
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
pub fn entries(&self) -> &[Entry] {
&self.entries
}
/// The index of the first key at or after `key`, or past the end.
fn index_at_or_after(&self, key: &str) -> usize {
self.entries
.iter()
.position(|e| matches!(e, Entry::Leaf { key: k, .. } if k.as_str() >= key))
.unwrap_or(self.entries.len())
}
fn tree_at(&self, index: usize) -> Option<&Node> {
match self.entries.get(index) {
Some(Entry::Tree(n)) => Some(n),
_ => None,
}
}
fn before(&self, index: usize) -> Option<&Node> {
index.checked_sub(1).and_then(|i| self.tree_at(i))
}
fn add_to(&self, blocks: &mut Blocks) {
blocks.insert(self.cid, self.bytes.clone());
}
// What a reader needs to see `key`'s place in the tree, as the
// reference implementation gives it, so that the firehose's proofs are
// the same: the nodes down to where the key is or would be, and down
// the edges of the subtrees either side of it.
fn proof_for_key(&self, key: &str, blocks: &mut Blocks) {
let index = self.index_at_or_after(key);
let found = matches!(self.entries.get(index), Some(Entry::Leaf { key: k, .. }) if k == key);
if !found {
match self.before(index) {
Some(prev) => prev.proof_for_key(key, blocks),
None => return,
}
}
self.add_to(blocks);
}
fn proof_for_left(&self, key: &str, blocks: &mut Blocks) {
let index = self.index_at_or_after(key);
if let Some(prev) = self.before(index) {
prev.proof_for_left(key, blocks);
}
self.add_to(blocks);
}
fn proof_for_right(&self, key: &str, blocks: &mut Blocks) {
let index = self.index_at_or_after(key);
let at = match self.entries.get(index) {
Some(e) => Some((index, e)),
None => index.checked_sub(1).map(|i| (i, &self.entries[i])),
};
// Its neighbour from `index`, not from where it was found, as the
// reference has it: past the end, that is the key itself, and so no
// subtree.
let next = match at {
Some((_, Entry::Tree(n))) => Some(n),
Some((_, Entry::Leaf { key: k, .. })) if k == key => self.tree_at(index + 1),
Some((_, Entry::Leaf { .. })) => self.before(index),
None => None,
};
if let Some(n) = next {
n.proof_for_right(key, blocks);
}
self.add_to(blocks);
}
fn walk<'a>(&'a self, nodes: &mut Vec<&'a Node>, leaves: &mut Vec<(&'a str, Cid)>) {
nodes.push(self);
for e in &self.entries {
match e {
Entry::Tree(n) => n.walk(nodes, leaves),
Entry::Leaf { key, value } => leaves.push((key, *value)),
}
}
}
}
/// Blocks by CID: a node's, a record's, a commit's.
pub type Blocks = BTreeMap<Cid, Vec<u8>>;
/// A whole tree.
#[derive(Debug, Clone, PartialEq)]
pub struct Tree {
root: Node,
}
impl Tree {
/// The tree of `records`, keys to their records' CIDs.
pub fn build(records: &BTreeMap<String, Cid>) -> Tree {
let items: Vec<(&str, Cid, u32)> = records
.iter()
.map(|(k, v)| (k.as_str(), *v, depth(k.as_bytes())))
.collect();
let top = items.iter().map(|i| i.2).max().unwrap_or(0);
Tree {
root: build(&items, top),
}
}
/// The root's CID: a commit's `data`.
pub fn root(&self) -> Cid {
self.root.cid
}
pub fn root_node(&self) -> &Node {
&self.root
}
/// Every node, parents before children, in key order.
pub fn nodes(&self) -> Vec<&Node> {
let (mut nodes, mut leaves) = (Vec::new(), Vec::new());
self.root.walk(&mut nodes, &mut leaves);
nodes
}
/// Every key and its record's CID, in key order.
pub fn leaves(&self) -> Vec<(&str, Cid)> {
let (mut nodes, mut leaves) = (Vec::new(), Vec::new());
self.root.walk(&mut nodes, &mut leaves);
leaves
}
/// Every node's block.
pub fn blocks(&self) -> Blocks {
self.nodes()
.into_iter()
.map(|n| (n.cid, n.bytes.clone()))
.collect()
}
/// The nodes of this tree that `before` does not have: what a commit
/// from `before` to this adds.
pub fn new_since(&self, before: &Tree) -> Blocks {
let had: BTreeSet<Cid> = before.nodes().into_iter().map(|n| n.cid).collect();
self.nodes()
.into_iter()
.filter(|n| !had.contains(&n.cid))
.map(|n| (n.cid, n.bytes.clone()))
.collect()
}
/// The nodes a reader needs, beside the new ones, to undo a change to
/// `key` and find the tree before: the nodes down to its place, and
/// down the edges of its neighbours.
pub fn covering_proof(&self, key: &str) -> Blocks {
let mut blocks = Blocks::new();
self.root.proof_for_key(key, &mut blocks);
self.root.proof_for_left(key, &mut blocks);
self.root.proof_for_right(key, &mut blocks);
blocks
}
/// The tree whose root is `root`, read from `blocks`, every node
/// checked: keys in order, each at its depth, and the whole the one
/// tree its keys make. A node missing from `blocks` fails it.
pub fn read(root: Cid, blocks: &Blocks) -> Result<Tree, MstError> {
let mut records = BTreeMap::new();
read_node(root, blocks, &mut records, &mut BTreeSet::new(), 0)?;
let tree = Tree::build(&records);
if tree.root() != root {
return err("the tree is not the one its keys make");
}
Ok(tree)
}
}
fn build(items: &[(&str, Cid, u32)], height: u32) -> Node {
let mut entries = Vec::new();
let mut run = 0;
for (i, &(key, value, d)) in items.iter().enumerate() {
if d == height {
if run < i {
entries.push(Entry::Tree(build(&items[run..i], height - 1)));
}
entries.push(Entry::Leaf {
key: key.to_string(),
value,
});
run = i + 1;
}
}
if run < items.len() {
entries.push(Entry::Tree(build(&items[run..], height - 1)));
}
Node::new(entries)
}
/// How deep a tree read from elsewhere may go.
const MAX_LAYERS: usize = 64;
/// Reads the node `cid` and the tree under it into `records`. A node
/// reached twice is refused: no tree reaches one node from two places, and
/// blocks that do could make the walk take time exponential in depth.
fn read_node(
cid: Cid,
blocks: &Blocks,
records: &mut BTreeMap<String, Cid>,
seen: &mut BTreeSet<Cid>,
layer: usize,
) -> Result<(), MstError> {
if layer > MAX_LAYERS {
return err("the tree is too deep");
}
if !seen.insert(cid) {
return err(format!("node {cid} is reached twice"));
}
if cid.codec() != DAG_CBOR {
return err(format!("node {cid} is not DAG-CBOR"));
}
let Some(bytes) = blocks.get(&cid) else {
return err(format!("node {cid} is missing"));
};
let value = drisl::decode(bytes).or_else(|e| err(format!("node {cid}: {e}")))?;
let bad = || MstError(format!("node {cid} is not a tree node"));
let Value::Map(map) = value else {
return Err(bad());
};
let (Some(Value::Array(entries)), Some(left), 2) = (map.get("e"), map.get("l"), map.len())
else {
return Err(bad());
};
let subtree = |v: &Value, records: &mut BTreeMap<String, Cid>, seen: &mut BTreeSet<Cid>| match v
{
Value::Null => Ok(()),
Value::Link(c) => read_node(*c, blocks, records, seen, layer + 1),
_ => Err(bad()),
};
subtree(left, records, seen)?;
let mut last: Vec<u8> = Vec::new();
for e in entries {
let Value::Map(e) = e else {
return Err(bad());
};
let (Some(Value::Integer(p)), Some(Value::Bytes(k)), Some(Value::Link(v)), Some(t), 4) =
(e.get("p"), e.get("k"), e.get("v"), e.get("t"), e.len())
else {
return Err(bad());
};
let p = usize::try_from(*p).map_err(|_| bad())?;
if p > last.len() {
return Err(bad());
}
let key = [&last[..p], k].concat();
let key = String::from_utf8(key).map_err(|_| bad())?;
records.insert(key.clone(), *v);
subtree(t, records, seen)?;
last = key.into_bytes();
}
Ok(())
}