//! mst node decoding, shared by the full-repo walk here and the sparse scanner in
//! `sparse_mst`.
use std::borrow::Borrow;
use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::ops::Deref;
use std::sync::Arc;
use bytes::Bytes;
use cid::Cid as IpldCid;
use jacquard_common::types::crypto::{KeyCodec, PublicKey};
use jacquard_common::types::string::Did;
use jacquard_common::types::tid::Tid;
use miette::{IntoDiagnostic, Result, WrapErr};
use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use smol_str::SmolStr;
use crate::car::BlockMap;
pub(crate) use crate::car::compute_cid;
/// signed atproto repository commit object stored as the root block of a repo car.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct AtpCommit {
pub did: Did,
pub version: i64,
pub data: IpldCid,
pub rev: Tid,
pub prev: Option<IpldCid>,
#[serde(
with = "jacquard_common::serde_bytes_helper",
skip_serializing_if = "Bytes::is_empty"
)]
pub sig: Bytes,
}
impl AtpCommit {
pub(crate) fn from_cbor(bytes: &[u8]) -> Result<Self> {
serde_ipld_dagcbor::from_slice(bytes).into_diagnostic()
}
pub(crate) fn to_cbor(&self) -> Result<Vec<u8>> {
serde_ipld_dagcbor::to_vec(self).into_diagnostic()
}
pub(crate) fn to_cid(&self) -> Result<IpldCid> {
Ok(compute_cid(&self.to_cbor()?))
}
pub(crate) fn unsigned_bytes(&self) -> Result<Vec<u8>> {
Self {
sig: Bytes::new(),
..self.clone()
}
.to_cbor()
}
pub(crate) fn verify(&self, pubkey: &PublicKey<'_>) -> Result<()> {
let unsigned = self.unsigned_bytes()?;
match pubkey.codec {
KeyCodec::Secp256k1 => {
use k256::ecdsa::{Signature, VerifyingKey, signature::Verifier};
let vk = VerifyingKey::from_sec1_bytes(pubkey.bytes.as_ref()).into_diagnostic()?;
let sig = Signature::from_slice(self.sig.as_ref()).into_diagnostic()?;
vk.verify(&unsigned, &sig).into_diagnostic()
}
KeyCodec::P256 => {
use p256::ecdsa::{Signature, VerifyingKey, signature::Verifier};
let vk = VerifyingKey::from_sec1_bytes(pubkey.bytes.as_ref()).into_diagnostic()?;
let sig = Signature::from_slice(self.sig.as_ref()).into_diagnostic()?;
vk.verify(&unsigned, &sig).into_diagnostic()
}
other => miette::bail!("unsupported signing key codec: {other:?}"),
}
}
#[cfg(test)]
pub(crate) fn sign_k256(
did: Did,
data: IpldCid,
rev: Tid,
prev: Option<IpldCid>,
key: &k256::ecdsa::SigningKey,
) -> Self {
use k256::ecdsa::signature::Signer;
let mut commit = Self {
did,
version: 3,
data,
rev,
prev,
sig: Bytes::new(),
};
let unsigned = commit.unsigned_bytes().expect("unsigned commit encodes");
let sig: k256::ecdsa::Signature = key.sign(&unsigned);
commit.sig = Bytes::copy_from_slice(&sig.to_bytes());
commit
}
}
/// mst layer (depth from the leaves) of a key: `floor(leading_zero_bits(sha256(key)) / 2)`.
pub(crate) fn layer_for_key(key: &str) -> usize {
let hash = Sha256::digest(key.as_bytes());
let mut zeros = 0usize;
for &byte in hash.iter() {
if byte == 0 {
zeros += 8;
} else {
zeros += byte.leading_zeros() as usize;
break;
}
}
zeros / 2
}
/// inline capacity for mst keys (`collection/rkey`). 100% of keys across the benchmark
/// corpus are <= 62 bytes (`size_of::<MstKey>() == 64`), so normal keys never touch the heap.
const INLINE_KEY_CAP: usize = 62;
#[derive(Clone)]
pub(crate) enum MstKey {
Inline { len: u8, buf: [u8; INLINE_KEY_CAP] },
Heap(Arc<str>),
}
impl MstKey {
#[inline]
pub(crate) fn new(s: impl AsRef<str>) -> Self {
let s = s.as_ref();
let bytes = s.as_bytes();
if bytes.len() <= INLINE_KEY_CAP {
let mut buf = [0u8; INLINE_KEY_CAP];
buf[..bytes.len()].copy_from_slice(bytes);
Self::Inline {
len: bytes.len() as u8,
buf,
}
} else {
Self::Heap(Arc::from(s))
}
}
#[inline]
pub(crate) fn as_bytes(&self) -> &[u8] {
match self {
Self::Inline { len, buf } => &buf[..usize::from(*len)],
Self::Heap(s) => s.as_bytes(),
}
}
#[inline]
pub(crate) fn as_str(&self) -> &str {
match self {
Self::Inline { len, buf } => {
std::str::from_utf8(&buf[..usize::from(*len)]).expect("MstKey holds valid utf8")
}
Self::Heap(s) => s,
}
}
}
impl Deref for MstKey {
type Target = str;
#[inline]
fn deref(&self) -> &str {
self.as_str()
}
}
impl AsRef<str> for MstKey {
#[inline]
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl AsRef<[u8]> for MstKey {
#[inline]
fn as_ref(&self) -> &[u8] {
self.as_bytes()
}
}
impl Borrow<str> for MstKey {
#[inline]
fn borrow(&self) -> &str {
self.as_str()
}
}
impl PartialEq for MstKey {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.as_bytes() == other.as_bytes()
}
}
impl Eq for MstKey {}
impl PartialOrd for MstKey {
#[inline]
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for MstKey {
#[inline]
fn cmp(&self, other: &Self) -> Ordering {
self.as_bytes().cmp(other.as_bytes())
}
}
impl Hash for MstKey {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
self.as_str().hash(state);
}
}
impl fmt::Debug for MstKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(self.as_str(), f)
}
}
impl fmt::Display for MstKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(self.as_str(), f)
}
}
impl PartialEq<str> for MstKey {
#[inline]
fn eq(&self, other: &str) -> bool {
self.as_bytes() == other.as_bytes()
}
}
impl PartialEq<&str> for MstKey {
#[inline]
fn eq(&self, other: &&str) -> bool {
self.as_bytes() == other.as_bytes()
}
}
impl PartialEq<SmolStr> for MstKey {
#[inline]
fn eq(&self, other: &SmolStr) -> bool {
self.as_bytes() == other.as_bytes()
}
}
#[cfg(test)]
impl From<SmolStr> for MstKey {
#[inline]
fn from(s: SmolStr) -> Self {
Self::new(s.as_str())
}
}
/// where [`walk`] reads nodes from. a read hands back owned bytes and can fail, so a car
/// that isn't held in memory can serve the walk too.
pub(crate) trait BlockSource {
fn read(&self, cid: &IpldCid) -> Result<Option<Bytes>>;
}
impl<M: BlockMap> BlockSource for M {
fn read(&self, cid: &IpldCid) -> Result<Option<Bytes>> {
Ok(self.get(cid).cloned())
}
}
#[derive(Serialize, Deserialize)]
pub(crate) struct RawTreeEntry<'a> {
#[serde(rename = "k", with = "serde_bytes", borrow)]
pub(crate) key_suffix: &'a [u8],
#[serde(rename = "p")]
pub(crate) prefix_len: u8,
#[serde(rename = "t")]
pub(crate) tree: Option<IpldCid>,
#[serde(rename = "v")]
pub(crate) value: IpldCid,
}
#[derive(Serialize)]
pub(crate) struct RawNodeData<'a> {
#[serde(rename = "e")]
pub(crate) entries: Vec<RawTreeEntry<'a>>,
#[serde(rename = "l")]
pub(crate) left: Option<IpldCid>,
}
#[derive(Debug, Clone)]
pub(crate) enum FlatEntry {
Tree { cid: IpldCid },
Leaf { key: MstKey, cid: IpldCid },
}
/// a node's subtrees and leaves in key order, with full keys rebuilt from prefix compression.
pub(crate) fn decode_node(bytes: &[u8]) -> Result<Vec<FlatEntry>> {
let mut entries = Vec::new();
decode_node_into(bytes, &mut entries)?;
Ok(entries)
}
/// appends a decoded node's entries onto `entries`, reusing its allocation when empty.
pub(crate) fn decode_node_into(bytes: &[u8], entries: &mut Vec<FlatEntry>) -> Result<()> {
let mut de = serde_ipld_dagcbor::de::Deserializer::from_slice(bytes);
NodeSeed(entries).deserialize(&mut de).into_diagnostic()?;
de.end().into_diagnostic()?;
Ok(())
}
struct NodeSeed<'a>(&'a mut Vec<FlatEntry>);
impl<'de> DeserializeSeed<'de> for NodeSeed<'_> {
type Value = ();
fn deserialize<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_map(self)
}
}
impl<'de> Visitor<'de> for NodeSeed<'_> {
type Value = ();
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("an mst node map")
}
fn visit_map<A>(self, mut map: A) -> std::result::Result<Self::Value, A::Error>
where
A: MapAccess<'de>,
{
let base = self.0.len();
let mut left: Option<IpldCid> = None;
let mut seen_entries = false;
while let Some(field) = map.next_key::<&str>()? {
match field {
"l" => {
left = map.next_value()?;
}
"e" => {
if seen_entries {
return Err(serde::de::Error::duplicate_field("e"));
}
seen_entries = true;
map.next_value_seed(EntriesSeed(&mut *self.0))?;
}
_ => {
map.next_value::<IgnoredAny>()?;
}
}
}
if !seen_entries {
return Err(serde::de::Error::missing_field("e"));
}
if let Some(cid) = left {
self.0.insert(base, FlatEntry::Tree { cid });
}
Ok(())
}
}
struct EntriesSeed<'a>(&'a mut Vec<FlatEntry>);
impl<'de> DeserializeSeed<'de> for EntriesSeed<'_> {
type Value = ();
fn deserialize<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_seq(self)
}
}
impl<'de> Visitor<'de> for EntriesSeed<'_> {
type Value = ();
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("an mst entry list")
}
fn visit_seq<A>(self, mut seq: A) -> std::result::Result<Self::Value, A::Error>
where
A: SeqAccess<'de>,
{
let count = seq.size_hint().unwrap_or(0);
if count > 0 {
self.0.reserve(count * 2 + 1);
}
// each key is stored as a prefix length into the previous key plus a suffix
let mut key = KeyBuilder::new();
let mut idx = 0usize;
while let Some(entry) = seq.next_element::<RawTreeEntry<'de>>()? {
let leaf_key = key
.advance(idx > 0, usize::from(entry.prefix_len), entry.key_suffix)
.map_err(serde::de::Error::custom)?;
self.0.push(FlatEntry::Leaf {
key: leaf_key,
cid: entry.value,
});
if let Some(cid) = entry.tree {
self.0.push(FlatEntry::Tree { cid });
}
idx += 1;
}
Ok(())
}
}
enum KeyBuilder {
Inline {
len: usize,
buf: [u8; INLINE_KEY_CAP],
},
Heap(String),
}
impl KeyBuilder {
#[inline]
fn new() -> Self {
Self::Inline {
len: 0,
buf: [0u8; INLINE_KEY_CAP],
}
}
#[inline]
fn advance(
&mut self,
check_order: bool,
prefix_len: usize,
suffix_bytes: &[u8],
) -> std::result::Result<MstKey, String> {
let suffix = std::str::from_utf8(suffix_bytes)
.map_err(|err| format!("invalid utf8 in mst key suffix: {err}"))?;
let prev_bytes = match self {
Self::Inline { len, buf } => &buf[..*len],
Self::Heap(s) => s.as_bytes(),
};
if !is_utf8_char_boundary(prev_bytes, prefix_len) {
return Err(format!("invalid mst key prefix length {prefix_len}"));
}
if check_order && suffix.as_bytes() <= &prev_bytes[prefix_len..] {
let prev = std::str::from_utf8(prev_bytes).expect("KeyBuilder holds valid utf8");
let cur = format!("{}{suffix}", &prev[..prefix_len]);
return Err(format!("mst key {cur} does not sort after {prev}"));
}
let new_len = prefix_len + suffix.len();
match self {
Self::Inline { len, buf } if new_len <= INLINE_KEY_CAP => {
buf[prefix_len..new_len].copy_from_slice(suffix.as_bytes());
*len = new_len;
Ok(MstKey::Inline {
len: new_len as u8,
buf: *buf,
})
}
Self::Inline { buf, .. } => {
let mut s = String::with_capacity(new_len.max(128));
s.push_str(
std::str::from_utf8(&buf[..prefix_len]).expect("KeyBuilder holds valid utf8"),
);
s.push_str(suffix);
let key = MstKey::Heap(Arc::from(s.as_str()));
*self = Self::Heap(s);
Ok(key)
}
Self::Heap(s) => {
s.truncate(prefix_len);
s.push_str(suffix);
Ok(MstKey::new(s.as_str()))
}
}
}
}
#[inline]
fn is_utf8_char_boundary(bytes: &[u8], index: usize) -> bool {
if index == bytes.len() {
return true;
}
match bytes.get(index) {
None => false,
Some(&b) => (b as i8) >= -0x40,
}
}
/// every leaf of the mst rooted at `root`, in key order.
///
/// the tree comes from an untrusted car, so anything a valid mst cannot contain is an error:
/// a missing node, a node reached twice (which also ends cid cycles), or keys that are not
/// strictly increasing. the walk keeps its own stack, so tree depth cannot overflow the
/// thread's.
// backfill walks with `walk`, so only tests and the benches collect every leaf at once
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) fn leaves(blocks: &impl BlockSource, root: IpldCid) -> Result<Vec<(MstKey, IpldCid)>> {
walk(blocks, root).collect()
}
/// [`leaves`] one at a time, holding only the path being walked rather than every leaf. it
/// ends after the first error.
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) fn walk<B: BlockSource + ?Sized>(blocks: &B, root: IpldCid) -> Walk<'_, B> {
Walk {
blocks,
visited: HashSet::with_hasher(ahash::RandomState::new()),
stack: vec![FlatEntry::Tree { cid: root }],
last: None,
}
}
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) struct Walk<'a, B: ?Sized> {
blocks: &'a B,
visited: HashSet<IpldCid, ahash::RandomState>,
stack: Vec<FlatEntry>,
last: Option<MstKey>,
}
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
impl<B: BlockSource + ?Sized> Walk<'_, B> {
fn step(&mut self) -> Result<Option<(MstKey, IpldCid)>> {
while let Some(entry) = self.stack.pop() {
match entry {
FlatEntry::Leaf { key, cid } => {
if let Some(previous) = self.last.as_ref().filter(|previous| **previous >= key)
{
miette::bail!("mst key {key} does not sort after {previous}");
}
self.last = Some(key.clone());
return Ok(Some((key, cid)));
}
FlatEntry::Tree { cid } => {
if !self.visited.insert(cid) {
miette::bail!("mst node {cid} is reachable more than once");
}
let bytes = self
.blocks
.read(&cid)?
.ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
let start = self.stack.len();
decode_node_into(&bytes, &mut self.stack)
.wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
self.stack[start..].reverse();
}
}
}
Ok(None)
}
}
impl<B: BlockSource + ?Sized> Iterator for Walk<'_, B> {
type Item = Result<(MstKey, IpldCid)>;
fn next(&mut self) -> Option<Self::Item> {
let step = self.step();
if step.is_err() {
self.stack.clear();
}
step.transpose()
}
}
/// shared prefix in bytes, ending on a char boundary and capped at `u8::MAX`.
fn common_prefix_len(a: &str, b: &str) -> usize {
let mut len = 0usize;
for (ca, cb) in a.chars().zip(b.chars()) {
if ca != cb || len + ca.len_utf8() > usize::from(u8::MAX) {
break;
}
len += ca.len_utf8();
}
len
}
fn encode_node(
left: Option<IpldCid>,
entries: &[(&str, IpldCid, Option<IpldCid>)],
) -> Result<(IpldCid, Bytes)> {
let mut last = "";
let raw_entries = entries
.iter()
.map(|&(key, value, tree)| {
let prefix_len = common_prefix_len(last, key);
last = key;
RawTreeEntry {
key_suffix: &key.as_bytes()[prefix_len..],
prefix_len: prefix_len as u8,
tree,
value,
}
})
.collect();
let node = RawNodeData {
entries: raw_entries,
left,
};
let bytes = Bytes::from(serde_ipld_dagcbor::to_vec(&node).into_diagnostic()?);
Ok((compute_cid(&bytes), bytes))
}
/// sha-256 has 256 bits (128 2-bit chunks), so no valid mst can be deeper than this.
const MAX_MST_DEPTH: usize = 128;
#[derive(Clone)]
enum FlatNode {
Tree { cid: IpldCid, layer: usize },
Leaf { key: MstKey, cid: IpldCid, layer: usize },
}
impl FlatNode {
fn layer(&self) -> usize {
match self {
Self::Tree { layer, .. } | Self::Leaf { layer, .. } => *layer,
}
}
}
/// builds an mst from `(key, record_cid)` pairs, returning `(root_cid, nodes_in_preorder)`.
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) fn build_tree<K: AsRef<str>>(
leaves: impl IntoIterator<Item = (K, IpldCid)>,
) -> Result<(IpldCid, Vec<(IpldCid, Bytes)>)> {
let mut items: Vec<FlatNode> = Vec::new();
let mut sorted = true;
for (key, cid) in leaves {
let key_str = key.as_ref();
validate_mst_key(key_str)?;
let mst_key = MstKey::new(key_str);
if let Some(FlatNode::Leaf { key: prev, .. }) = items.last()
&& *prev >= mst_key
{
sorted = false;
}
items.push(FlatNode::Leaf {
key: mst_key,
cid,
layer: layer_for_key(key_str),
});
}
if !sorted {
items.sort_by(|a, b| match (a, b) {
(FlatNode::Leaf { key: ka, .. }, FlatNode::Leaf { key: kb, .. }) => ka.cmp(kb),
_ => unreachable!(),
});
items.dedup_by(|a, b| match (a, b) {
(
FlatNode::Leaf { key: ka, cid: ca, .. },
FlatNode::Leaf { key: kb, cid: cb, .. },
) if ka == kb => {
*cb = *ca;
true
}
_ => false,
});
}
if items.is_empty() {
let (cid, bytes) = encode_node(None, &[])?;
return Ok((cid, vec![(cid, bytes)]));
}
let top_layer = items.iter().map(FlatNode::layer).max().unwrap_or(0);
let mut nodes = Vec::new();
let root = build_layer(&items, top_layer, None, Some(&mut nodes))?;
Ok((root, nodes))
}
fn build_layer(
slice: &[FlatNode],
layer: usize,
blocks: Option<&dyn BlockSource>,
mut out: Option<&mut Vec<(IpldCid, Bytes)>>,
) -> Result<IpldCid> {
if let [FlatNode::Tree { cid, layer: l }] = slice
&& *l == layer
{
return Ok(*cid);
}
let mut expanded: Vec<FlatNode>;
let mut slice = slice;
while slice
.iter()
.any(|n| matches!(n, FlatNode::Tree { layer: l, .. } if *l == layer))
{
let blocks = blocks.expect("tree node requires block source");
let mut next = Vec::with_capacity(slice.len() + 4);
for node in slice {
if let FlatNode::Tree { cid, layer: l } = node
&& *l == layer
{
expand_tree(blocks, *cid, *l, &mut next)?;
} else {
next.push(node.clone());
}
}
expanded = next;
slice = &expanded;
}
let slot = out.as_deref_mut().map(|v| {
let s = v.len();
v.push((IpldCid::default(), Bytes::new()));
s
});
let mut start = 0usize;
let mut left: Option<IpldCid> = None;
let mut entries: Vec<(&str, IpldCid, Option<IpldCid>)> = Vec::new();
for (i, node) in slice.iter().enumerate() {
if let FlatNode::Leaf { key, cid, layer: l } = node
&& *l == layer
{
if start < i {
let child = build_layer(&slice[start..i], layer - 1, blocks, out.as_deref_mut())?;
if let Some(last) = entries.last_mut() {
last.2 = Some(child);
} else {
left = Some(child);
}
}
entries.push((key.as_str(), *cid, None));
start = i + 1;
}
}
if start < slice.len() {
let child = build_layer(&slice[start..], layer - 1, blocks, out.as_deref_mut())?;
if let Some(last) = entries.last_mut() {
last.2 = Some(child);
} else {
left = Some(child);
}
}
let (node_cid, node_bytes) = encode_node(left, &entries)?;
if let (Some(out), Some(slot)) = (out, slot) {
out[slot] = (node_cid, node_bytes);
}
Ok(node_cid)
}
fn expand_tree(
blocks: &dyn BlockSource,
cid: IpldCid,
layer: usize,
out: &mut Vec<FlatNode>,
) -> Result<()> {
let bytes = blocks
.read(&cid)?
.ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
let flat =
decode_node(&bytes).wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
let child_layer = layer.saturating_sub(1);
for entry in flat {
match entry {
FlatEntry::Leaf { key, cid } => out.push(FlatNode::Leaf {
layer: layer_for_key(&key),
key,
cid,
}),
FlatEntry::Tree { cid } => out.push(FlatNode::Tree {
cid,
layer: child_layer,
}),
}
}
Ok(())
}
fn find_layer(blocks: &impl BlockSource, cid: IpldCid, depth: usize) -> Result<usize> {
if depth > MAX_MST_DEPTH {
miette::bail!("mst depth exceeds maximum ({MAX_MST_DEPTH})");
}
let bytes = blocks
.read(&cid)?
.ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
let flat =
decode_node(&bytes).wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
for entry in &flat {
if let FlatEntry::Leaf { key, .. } = entry {
return Ok(layer_for_key(key));
}
}
for entry in flat {
if let FlatEntry::Tree { cid: child } = entry {
return Ok(find_layer(blocks, child, depth + 1)? + 1);
}
}
Ok(0)
}
fn validate_mst_key(key: &str) -> Result<()> {
if key.is_empty() || key.len() > 830 {
miette::bail!("invalid mst key length: {}", key.len());
}
let Some((coll, rkey)) = key.split_once('/') else {
miette::bail!("mst key missing '/': {key}");
};
if coll.is_empty() || rkey.is_empty() || rkey.contains('/') {
miette::bail!("invalid mst key format: {key}");
}
if !key
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b':' | b'~' | b'-' | b'/'))
{
miette::bail!("invalid characters in mst key: {key}");
}
Ok(())
}
/// lazy in-memory mst for inverting firehose commit operations against a commit car.
pub(crate) struct MstNode {
root: IpldCid,
nodes: Option<Vec<FlatNode>>,
dirty: bool,
}
impl MstNode {
pub(crate) fn load(root: IpldCid) -> Self {
Self {
root,
nodes: None,
dirty: false,
}
}
fn seek(
&mut self,
blocks: &impl BlockSource,
key: &str,
split: bool,
) -> Result<std::result::Result<usize, usize>> {
let key_layer = layer_for_key(key);
if self.nodes.is_none() {
let layer = find_layer(blocks, self.root, 0)?;
self.nodes = Some(vec![FlatNode::Tree {
cid: self.root,
layer,
}]);
}
let nodes = self.nodes.as_mut().unwrap();
for _ in 0..=MAX_MST_DEPTH {
let i = nodes
.iter()
.position(|n| matches!(n, FlatNode::Leaf { key: k, .. } if k.as_str() >= key))
.unwrap_or(nodes.len());
if let Some(FlatNode::Leaf { key: k, .. }) = nodes.get(i)
&& k.as_str() == key
{
return Ok(Ok(i));
}
if i > 0
&& let FlatNode::Tree { cid, layer } = nodes[i - 1]
&& (split || layer >= key_layer)
{
let mut children = Vec::new();
expand_tree(blocks, cid, layer, &mut children)?;
nodes.splice((i - 1)..i, children);
} else {
return Ok(Err(i));
}
}
miette::bail!("mst depth exceeds maximum ({MAX_MST_DEPTH})")
}
pub(crate) fn get(&mut self, blocks: &impl BlockSource, key: &str) -> Result<Option<IpldCid>> {
validate_mst_key(key)?;
match self.seek(blocks, key, false)? {
Ok(i) => {
let FlatNode::Leaf { cid, .. } = self.nodes.as_ref().unwrap()[i] else {
unreachable!()
};
Ok(Some(cid))
}
Err(_) => Ok(None),
}
}
pub(crate) fn add(&mut self, blocks: &impl BlockSource, key: &str, cid: IpldCid) -> Result<()> {
validate_mst_key(key)?;
let layer = layer_for_key(key);
let leaf = FlatNode::Leaf {
key: MstKey::new(key),
cid,
layer,
};
match self.seek(blocks, key, true)? {
Ok(i) => self.nodes.as_mut().unwrap()[i] = leaf,
Err(i) => self.nodes.as_mut().unwrap().insert(i, leaf),
}
self.dirty = true;
Ok(())
}
pub(crate) fn delete(&mut self, blocks: &impl BlockSource, key: &str) -> Result<IpldCid> {
validate_mst_key(key)?;
match self.seek(blocks, key, false)? {
Ok(i) => {
let FlatNode::Leaf { cid, .. } = self.nodes.as_mut().unwrap().remove(i) else {
unreachable!()
};
self.dirty = true;
Ok(cid)
}
Err(_) => miette::bail!("mst key {key} not found"),
}
}
pub(crate) fn root_cid(&mut self, blocks: &impl BlockSource) -> Result<IpldCid> {
if !self.dirty {
return Ok(self.root);
}
let nodes = self.nodes.as_mut().unwrap();
let top_layer = loop {
let Some(max_layer) = nodes.iter().map(FlatNode::layer).max() else {
let (cid, _) = encode_node(None, &[])?;
self.root = cid;
self.dirty = false;
return Ok(cid);
};
if let Some(i) = nodes
.iter()
.position(|n| matches!(n, FlatNode::Tree { layer, .. } if *layer == max_layer))
{
let FlatNode::Tree { cid, layer } = nodes[i] else {
unreachable!()
};
let mut children = Vec::new();
expand_tree(blocks, cid, layer, &mut children)?;
nodes.splice(i..=i, children);
continue;
}
break max_layer;
};
let cid = build_layer(nodes, top_layer, Some(blocks), None)?;
self.root = cid;
self.dirty = false;
Ok(cid)
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::car::CarBlocks;
use cid::multihash::Multihash;
use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256};
pub(crate) fn cid(byte: u8) -> IpldCid {
let mh = Multihash::<64>::wrap(SHA2_256, &[byte; 32]).unwrap();
IpldCid::new_v1(DAG_CBOR, mh)
}
/// builds a node from full keys, prefix-compressing them like the wire format.
pub(crate) fn node<'a>(
entries: Vec<(&'a str, IpldCid, Option<IpldCid>)>,
left: Option<IpldCid>,
) -> RawNodeData<'a> {
let mut last = "";
let entries = entries
.into_iter()
.map(|(key, value, tree)| {
let prefix_len = common_prefix_len(last, key);
last = key;
RawTreeEntry {
key_suffix: &key.as_bytes()[prefix_len..],
prefix_len: prefix_len as u8,
tree,
value,
}
})
.collect();
RawNodeData { entries, left }
}
fn blocks(nodes: Vec<(IpldCid, RawNodeData<'_>)>) -> CarBlocks {
nodes
.into_iter()
.map(|(cid, node)| (cid, serde_ipld_dagcbor::to_vec(&node).unwrap().into()))
.collect()
}
fn keys(leaves: &[(MstKey, IpldCid)]) -> Vec<&str> {
leaves.iter().map(|(key, _)| key.as_str()).collect()
}
#[test]
fn leaves_match_jacquard_walk() {
let mut input: Vec<(MstKey, IpldCid)> = (0..2000)
.map(|i| {
let collection = [
"app.bsky.feed.like",
"app.bsky.feed.post",
"sh.tangled.repo",
][i % 3];
(
MstKey::new(format!("{collection}/{i:013}")),
cid((i % 251) as u8),
)
})
.collect();
let (root, nodes) = build_tree(input.iter().map(|(k, c)| (k.as_str(), *c))).unwrap();
assert!(nodes.len() > 1, "test tree should span several nodes");
let built_map: std::collections::BTreeMap<IpldCid, Bytes> =
nodes.iter().cloned().collect();
let ours = leaves(&nodes.into_iter().collect::<CarBlocks>(), root).unwrap();
input.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(ours, input);
// Deleting and re-adding keys through MstNode reproduces the exact root at each step.
let mut mst = MstNode::load(root);
let mut remaining = input.clone();
for i in (0..input.len()).step_by(97).rev() {
let (key, expected_cid) = remaining.remove(i);
assert_eq!(mst.get(&built_map, &key).unwrap(), Some(expected_cid));
assert_eq!(mst.delete(&built_map, &key).unwrap(), expected_cid);
assert_eq!(mst.get(&built_map, &key).unwrap(), None);
let (expected_root, _) =
build_tree(remaining.iter().map(|(k, c)| (k.as_str(), *c))).unwrap();
assert_eq!(mst.root_cid(&built_map).unwrap(), expected_root);
}
for (key, expected_cid) in &input {
if mst.get(&built_map, key).unwrap().is_none() {
mst.add(&built_map, key, *expected_cid).unwrap();
}
}
assert_eq!(mst.root_cid(&built_map).unwrap(), root);
}
#[test]
fn walks_an_empty_tree() {
let root = cid(1);
assert!(
leaves(&blocks(vec![(root, node(vec![], None))]), root)
.unwrap()
.is_empty()
);
}
#[test]
fn rejects_a_missing_node() {
let (root, missing) = (cid(1), cid(2));
let tree = blocks(vec![(
root,
node(vec![("a/1", cid(9), Some(missing))], None),
)]);
let err = leaves(&tree, root).unwrap_err();
assert!(err.to_string().contains("missing from car"), "{err}");
}
#[test]
fn walk_hands_out_leaves_up_to_a_missing_node_and_nothing_after() {
let (root, missing) = (cid(1), cid(2));
let tree = blocks(vec![(
root,
node(
vec![("a/1", cid(9), Some(missing)), ("b/1", cid(9), None)],
None,
),
)]);
let mut walked = walk(&tree, root);
assert_eq!(
walked.next().unwrap().unwrap(),
(MstKey::new("a/1"), cid(9))
);
let err = walked.next().unwrap().unwrap_err();
assert!(err.to_string().contains("missing from car"), "{err}");
// b/1 is still on the stack, but nothing after an error counts
assert!(walked.next().is_none());
}
#[test]
fn rejects_a_subtree_reachable_twice() {
let (root, shared) = (cid(1), cid(2));
let tree = blocks(vec![
(
root,
node(
vec![("a/1", cid(9), Some(shared)), ("c/1", cid(9), Some(shared))],
None,
),
),
(shared, node(vec![("b/1", cid(9), None)], None)),
]);
let err = leaves(&tree, root).unwrap_err();
assert!(
err.to_string().contains("reachable more than once"),
"{err}"
);
}
#[test]
fn stops_on_a_node_that_is_its_own_left_subtree() {
// a car that is not cid-verified can claim any cid for a block
let root = cid(1);
let tree = blocks(vec![(root, node(vec![("a/1", cid(9), None)], Some(root)))]);
let err = leaves(&tree, root).unwrap_err();
assert!(
err.to_string().contains("reachable more than once"),
"{err}"
);
}
#[test]
fn rejects_keys_out_of_order() {
let root = cid(1);
let tree = blocks(vec![(
root,
node(vec![("b/1", cid(9), None), ("a/1", cid(9), None)], None),
)]);
let err = leaves(&tree, root).unwrap_err();
assert!(
format!("{err:?}").contains("does not sort after"),
"{err:?}"
);
}
#[test]
fn rejects_a_duplicate_key_across_subtrees() {
let (root, right) = (cid(1), cid(2));
let tree = blocks(vec![
(root, node(vec![("a/1", cid(9), Some(right))], None)),
(right, node(vec![("a/1", cid(8), None)], None)),
]);
let err = leaves(&tree, root).unwrap_err();
assert!(err.to_string().contains("does not sort after"), "{err}");
}
#[test]
fn rejects_a_prefix_that_splits_a_char() {
let root = cid(1);
let mut bad = node(vec![("é/1", cid(9), None), ("é/2", cid(9), None)], None);
bad.entries[1].prefix_len = 1;
let err = leaves(&blocks(vec![(root, bad)]), root).unwrap_err();
assert!(
format!("{err:?}").contains("invalid mst key prefix length 1"),
"{err:?}"
);
}
#[test]
fn rebuilds_prefix_compressed_keys() {
let long_prefix = "io.atcr.hold.extremely.long.collection.name.exceeding.inline/";
let long_a = format!("{long_prefix}3k2b1f5moe22a");
let long_b = format!("{long_prefix}3k2b1f5moe22b");
assert!(long_a.len() > INLINE_KEY_CAP);
let root = cid(1);
let tree = blocks(vec![(
root,
node(
vec![
("app.bsky.feed.like/1", cid(9), None),
("app.bsky.feed.post/1", cid(9), None),
("app.bsky.feed.post/2", cid(9), None),
(long_a.as_str(), cid(9), None),
(long_b.as_str(), cid(9), None),
("io.atcr.hold.short/1", cid(9), None),
// both start with byte 0xc3, so they share no whole-char prefix
("è/1", cid(9), None),
("é/1", cid(9), None),
],
None,
),
)]);
assert_eq!(
keys(&leaves(&tree, root).unwrap()),
[
"app.bsky.feed.like/1",
"app.bsky.feed.post/1",
"app.bsky.feed.post/2",
long_a.as_str(),
long_b.as_str(),
"io.atcr.hold.short/1",
"è/1",
"é/1"
]
);
}
/// a chain of `depth` nodes that only point left, ending in one leaf, for checking that a
/// walk does not recurse per level. returns the top cid and every encoded node.
pub(crate) fn left_chain(depth: usize, key: &str) -> (IpldCid, Vec<(IpldCid, bytes::Bytes)>) {
let mut child_bytes =
serde_ipld_dagcbor::to_vec(&node(vec![(key, cid(9), None)], None)).unwrap();
let mut child = compute_cid(&child_bytes);
let mut nodes = vec![(child, child_bytes.into())];
for _ in 0..depth {
child_bytes = serde_ipld_dagcbor::to_vec(&node(vec![], Some(child))).unwrap();
child = compute_cid(&child_bytes);
nodes.push((child, child_bytes.into()));
}
(child, nodes)
}
#[test]
fn walks_a_tree_deeper_than_the_thread_stack() {
let (top, nodes) = left_chain(100_000, "a/1");
assert_eq!(
keys(&leaves(&nodes.into_iter().collect::<CarBlocks>(), top).unwrap()),
["a/1"]
);
}
#[test]
fn commit_cbor_and_signature_verification_roundtrip() {
let did = Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").unwrap();
let rev = Tid::new("3jzfcijpj2z2a").unwrap();
let key = k256::ecdsa::SigningKey::from_bytes(&[7u8; 32].into()).unwrap();
let pubkey = PublicKey {
codec: KeyCodec::Secp256k1,
bytes: std::borrow::Cow::Owned(
key.verifying_key()
.to_encoded_point(true)
.as_bytes()
.to_vec(),
),
};
let signed = AtpCommit::sign_k256(did, cid(1), rev, Some(cid(2)), &key);
let cbor = signed.to_cbor().unwrap();
let decoded = AtpCommit::from_cbor(&cbor).unwrap();
assert_eq!(decoded, signed);
assert_eq!(decoded.to_cid().unwrap(), compute_cid(&cbor));
decoded.verify(&pubkey).unwrap();
let mut tampered = decoded.clone();
tampered.version = 2;
assert!(tampered.verify(&pubkey).is_err());
}
}