Skip to main content

prolly/prolly/
sync.rs

1//! Store-to-store synchronization helpers for content-addressed tree nodes.
2
3use super::cid::Cid;
4use super::error::Error;
5use super::node::Node;
6use super::tree::Tree;
7use serde::{Deserialize, Serialize};
8use std::collections::{BTreeMap, VecDeque};
9
10/// Current in-memory format version for portable tree snapshot bundles.
11pub const SNAPSHOT_BUNDLE_FORMAT_VERSION: u32 = 1;
12
13const SNAPSHOT_BUNDLE_BYTES_VERSION: u32 = SNAPSHOT_BUNDLE_FORMAT_VERSION;
14
15/// Dry-run plan for making a destination store able to read one source tree.
16///
17/// `required_cids` are all node CIDs reachable from the tree root. `missing_cids`
18/// are the subset not currently present in the destination store. All CIDs are
19/// sorted by raw CID bytes for deterministic transport planning.
20#[derive(Clone, Debug, Default, PartialEq, Eq)]
21pub struct MissingNodePlan {
22    /// All reachable node CIDs required by the tree.
23    pub required_cids: Vec<Cid>,
24    /// Number of reachable nodes required by the tree.
25    pub required_nodes: usize,
26    /// Serialized bytes for all reachable nodes in the source tree.
27    pub required_bytes: usize,
28    /// Required node CIDs not present in the destination store.
29    pub missing_cids: Vec<Cid>,
30    /// Number of missing destination nodes.
31    pub missing_nodes: usize,
32    /// Serialized bytes for missing nodes as read from the source store.
33    pub missing_bytes: usize,
34}
35
36impl MissingNodePlan {
37    /// Whether the destination already has every required node.
38    pub fn is_empty(&self) -> bool {
39        self.missing_cids.is_empty()
40    }
41
42    /// Return all required CIDs in deterministic byte order.
43    pub fn required_cids(&self) -> &[Cid] {
44        &self.required_cids
45    }
46
47    /// Return missing CIDs in deterministic byte order.
48    pub fn missing_cids(&self) -> &[Cid] {
49        &self.missing_cids
50    }
51}
52
53/// Result of copying missing nodes from one store to another.
54#[derive(Clone, Debug, Default, PartialEq, Eq)]
55pub struct MissingNodeCopy {
56    /// Dry-run plan used for this copy.
57    pub plan: MissingNodePlan,
58    /// Number of nodes written to the destination store.
59    pub copied_nodes: usize,
60    /// Serialized bytes written to the destination store.
61    pub copied_bytes: usize,
62}
63
64/// One content-addressed node included in a portable tree snapshot bundle.
65#[derive(Clone, Debug, PartialEq, Eq)]
66pub struct SnapshotBundleNode {
67    /// CID bytes the node is stored under.
68    pub cid: Cid,
69    /// Serialized node bytes whose SHA-256 CID must equal `cid`.
70    pub bytes: Vec<u8>,
71}
72
73/// Self-contained transport bundle for one tree and its reachable node bytes.
74///
75/// The bundle is intended for import/export between stores, processes, and
76/// language bindings. `nodes` should contain exactly the node CIDs reachable
77/// from `tree.root`, sorted by raw CID bytes for deterministic transport.
78#[derive(Clone, Debug, PartialEq)]
79pub struct SnapshotBundle {
80    /// Bundle schema version. Currently always `1`.
81    pub format_version: u32,
82    /// Tree handle the imported store will be able to read.
83    pub tree: Tree,
84    /// Reachable serialized nodes for the tree.
85    pub nodes: Vec<SnapshotBundleNode>,
86}
87
88/// Compact metadata for a validated snapshot bundle.
89#[derive(Clone, Debug, PartialEq, Eq)]
90pub struct SnapshotBundleSummary {
91    /// Bundle schema version.
92    pub format_version: u32,
93    /// Tree root CID, or `None` for an empty tree.
94    pub root: Option<Cid>,
95    /// Number of unique node CIDs included in the bundle.
96    pub node_count: usize,
97    /// Total serialized bytes across unique bundled nodes.
98    pub byte_count: usize,
99    /// Smallest serialized node payload in the bundle.
100    pub min_node_bytes: usize,
101    /// Largest serialized node payload in the bundle.
102    pub max_node_bytes: usize,
103}
104
105/// Result of verifying a snapshot bundle as a self-contained tree.
106#[derive(Clone, Debug, PartialEq, Eq)]
107pub struct SnapshotBundleVerification {
108    /// True when the bundle has exactly the reachable node set for its tree.
109    pub valid: bool,
110    /// Validated bundle metadata.
111    pub summary: SnapshotBundleSummary,
112    /// Number of reachable CIDs discovered from the tree root.
113    pub reachable_nodes: usize,
114    /// Serialized bytes for reachable bundled nodes.
115    pub reachable_bytes: usize,
116    /// Reachable CIDs absent from the bundle.
117    pub missing_cids: Vec<Cid>,
118    /// Bundled CIDs not reachable from the tree root.
119    pub extra_cids: Vec<Cid>,
120}
121
122#[derive(Serialize, Deserialize)]
123struct SnapshotBundleWire {
124    version: u32,
125    tree: Tree,
126    nodes: Vec<SnapshotBundleNodeWire>,
127}
128
129#[derive(Serialize, Deserialize)]
130struct SnapshotBundleNodeWire {
131    cid: Vec<u8>,
132    bytes: Vec<u8>,
133}
134
135impl SnapshotBundle {
136    /// Create a versioned snapshot bundle from a tree and reachable node bytes.
137    pub fn new(tree: Tree, nodes: Vec<SnapshotBundleNode>) -> Self {
138        Self {
139            format_version: SNAPSHOT_BUNDLE_FORMAT_VERSION,
140            tree,
141            nodes,
142        }
143    }
144
145    /// Number of serialized nodes in the bundle.
146    pub fn node_count(&self) -> usize {
147        self.nodes.len()
148    }
149
150    /// Total serialized node bytes in the bundle.
151    pub fn byte_count(&self) -> usize {
152        self.nodes.iter().map(|node| node.bytes.len()).sum()
153    }
154
155    /// Return the SHA-256 digest of this bundle's canonical byte encoding.
156    ///
157    /// The digest is stable for semantically equivalent bundles: node entries
158    /// are canonicalized by CID before encoding, so caller-side ordering does
159    /// not affect the result.
160    pub fn digest(&self) -> Result<Cid, Error> {
161        self.to_bytes().map(|bytes| Cid::from_bytes(&bytes))
162    }
163
164    /// Return validated metadata for this bundle without importing it.
165    ///
166    /// The summary canonicalizes node order, deduplicates identical repeated
167    /// nodes, and verifies each node byte payload against its CID.
168    pub fn summary(&self) -> Result<SnapshotBundleSummary, Error> {
169        self.validate_format_version()?;
170        let nodes = canonical_snapshot_nodes(&self.nodes)?;
171        Ok(snapshot_bundle_summary(self, &nodes))
172    }
173
174    /// Verify that this bundle is complete and contains no unreachable nodes.
175    ///
176    /// This check is read-only: it validates version, canonicalizes nodes,
177    /// verifies every node byte payload by CID, decodes reachable nodes from
178    /// the supplied bytes, and compares the reachable CID set with the bundled
179    /// CID set.
180    pub fn verify(&self) -> Result<SnapshotBundleVerification, Error> {
181        self.validate_format_version()?;
182        let nodes = canonical_snapshot_nodes(&self.nodes)?;
183        let summary = snapshot_bundle_summary(self, &nodes);
184        let nodes_by_cid = nodes
185            .iter()
186            .map(|node| (node.cid.as_bytes().to_vec(), node.bytes.as_slice()))
187            .collect::<BTreeMap<_, _>>();
188        let reachability = reachable_snapshot_nodes(&self.tree, &nodes_by_cid)?;
189        let provided_cids = nodes_by_cid.keys().cloned().collect::<Vec<_>>();
190
191        let missing_cids = reachability
192            .reachable_cids
193            .iter()
194            .filter(|cid| !nodes_by_cid.contains_key(*cid))
195            .cloned()
196            .map(cid_from_wire_bytes)
197            .collect::<Result<Vec<_>, Error>>()?;
198        let extra_cids = provided_cids
199            .iter()
200            .filter(|cid| !reachability.reachable_cids.contains(*cid))
201            .cloned()
202            .map(cid_from_wire_bytes)
203            .collect::<Result<Vec<_>, Error>>()?;
204
205        Ok(SnapshotBundleVerification {
206            valid: missing_cids.is_empty() && extra_cids.is_empty(),
207            summary,
208            reachable_nodes: reachability.reachable_cids.len(),
209            reachable_bytes: reachability.reachable_bytes,
210            missing_cids,
211            extra_cids,
212        })
213    }
214
215    /// Validate that the bundle version is supported by this crate.
216    pub fn validate_format_version(&self) -> Result<(), Error> {
217        if self.format_version == SNAPSHOT_BUNDLE_FORMAT_VERSION {
218            Ok(())
219        } else {
220            Err(Error::InvalidSnapshotBundle(format!(
221                "unsupported format version {}",
222                self.format_version
223            )))
224        }
225    }
226
227    /// Serialize this bundle as deterministic, versioned bytes.
228    ///
229    /// The encoded form canonicalizes node order by CID bytes and deduplicates
230    /// repeated identical nodes. It rejects unsupported bundle versions,
231    /// malformed CIDs, and node bytes whose content hash does not match the
232    /// supplied CID.
233    pub fn to_bytes(&self) -> Result<Vec<u8>, Error> {
234        self.validate_format_version()?;
235        let nodes = canonical_snapshot_nodes(&self.nodes)?;
236        let wire = SnapshotBundleWire {
237            version: SNAPSHOT_BUNDLE_BYTES_VERSION,
238            tree: self.tree.clone(),
239            nodes: nodes
240                .into_iter()
241                .map(|node| SnapshotBundleNodeWire {
242                    cid: node.cid.as_bytes().to_vec(),
243                    bytes: node.bytes,
244                })
245                .collect(),
246        };
247        serde_cbor::ser::to_vec_packed(&wire).map_err(|err| Error::Serialize(err.to_string()))
248    }
249
250    /// Decode a deterministic, versioned snapshot bundle byte payload.
251    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
252        let wire: SnapshotBundleWire =
253            serde_cbor::from_slice(bytes).map_err(snapshot_bundle_deserialize)?;
254        if wire.version != SNAPSHOT_BUNDLE_BYTES_VERSION {
255            return Err(Error::InvalidSnapshotBundle(format!(
256                "unsupported bytes version {}",
257                wire.version
258            )));
259        }
260        let nodes = wire
261            .nodes
262            .into_iter()
263            .map(|node| {
264                let cid = cid_from_wire_bytes(node.cid)?;
265                verify_node_bytes(&cid, &node.bytes)?;
266                Ok(SnapshotBundleNode {
267                    cid,
268                    bytes: node.bytes,
269                })
270            })
271            .collect::<Result<Vec<_>, Error>>()?;
272        Ok(Self {
273            format_version: SNAPSHOT_BUNDLE_FORMAT_VERSION,
274            tree: wire.tree,
275            nodes: canonical_snapshot_nodes(&nodes)?,
276        })
277    }
278}
279
280struct SnapshotBundleReachability {
281    reachable_cids: Vec<Vec<u8>>,
282    reachable_bytes: usize,
283}
284
285fn snapshot_bundle_summary(
286    bundle: &SnapshotBundle,
287    nodes: &[SnapshotBundleNode],
288) -> SnapshotBundleSummary {
289    let byte_count = nodes.iter().map(|node| node.bytes.len()).sum();
290    let min_node_bytes = nodes
291        .iter()
292        .map(|node| node.bytes.len())
293        .min()
294        .unwrap_or_default();
295    let max_node_bytes = nodes
296        .iter()
297        .map(|node| node.bytes.len())
298        .max()
299        .unwrap_or_default();
300
301    SnapshotBundleSummary {
302        format_version: bundle.format_version,
303        root: bundle.tree.root.clone(),
304        node_count: nodes.len(),
305        byte_count,
306        min_node_bytes,
307        max_node_bytes,
308    }
309}
310
311fn reachable_snapshot_nodes(
312    tree: &Tree,
313    nodes_by_cid: &BTreeMap<Vec<u8>, &[u8]>,
314) -> Result<SnapshotBundleReachability, Error> {
315    let mut seen = BTreeMap::<Vec<u8>, ()>::new();
316    let mut missing = BTreeMap::<Vec<u8>, ()>::new();
317    let mut frontier = VecDeque::new();
318    let mut reachable_bytes = 0usize;
319
320    if let Some(root) = &tree.root {
321        frontier.push_back(root.as_bytes().to_vec());
322    }
323
324    while let Some(cid) = frontier.pop_front() {
325        if seen.contains_key(&cid) {
326            continue;
327        }
328        seen.insert(cid.clone(), ());
329
330        let Some(bytes) = nodes_by_cid.get(&cid) else {
331            missing.insert(cid, ());
332            continue;
333        };
334
335        let node = Node::from_bytes(bytes)?;
336        if node.keys.len() != node.vals.len() {
337            return Err(Error::InvalidNode);
338        }
339        reachable_bytes += bytes.len();
340
341        if !node.leaf {
342            for child in &node.vals {
343                let cid = child
344                    .as_slice()
345                    .try_into()
346                    .map(Cid)
347                    .map_err(|_| Error::InvalidNode)?;
348                if !seen.contains_key(cid.as_bytes()) {
349                    frontier.push_back(cid.as_bytes().to_vec());
350                }
351            }
352        }
353    }
354
355    let mut reachable_cids = seen.into_keys().collect::<Vec<_>>();
356    for cid in missing.into_keys() {
357        if !reachable_cids.contains(&cid) {
358            reachable_cids.push(cid);
359        }
360    }
361    reachable_cids.sort();
362
363    Ok(SnapshotBundleReachability {
364        reachable_cids,
365        reachable_bytes,
366    })
367}
368
369pub(crate) fn verify_node_bytes(expected: &Cid, bytes: &[u8]) -> Result<(), Error> {
370    let actual = Cid::from_bytes(bytes);
371    if &actual == expected {
372        Ok(())
373    } else {
374        Err(Error::CidMismatch {
375            expected: expected.clone(),
376            actual,
377        })
378    }
379}
380
381fn canonical_snapshot_nodes(
382    nodes: &[SnapshotBundleNode],
383) -> Result<Vec<SnapshotBundleNode>, Error> {
384    let mut by_cid: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
385    for node in nodes {
386        verify_node_bytes(&node.cid, &node.bytes)?;
387        let cid = node.cid.as_bytes().to_vec();
388        if let Some(existing) = by_cid.get(&cid) {
389            if existing != &node.bytes {
390                return Err(Error::InvalidSnapshotBundle(format!(
391                    "bundle contains conflicting duplicate node CID {}",
392                    hex_bytes(&cid)
393                )));
394            }
395            continue;
396        }
397        by_cid.insert(cid, node.bytes.clone());
398    }
399    by_cid
400        .into_iter()
401        .map(|(cid, bytes)| {
402            Ok(SnapshotBundleNode {
403                cid: cid_from_wire_bytes(cid)?,
404                bytes,
405            })
406        })
407        .collect()
408}
409
410fn cid_from_wire_bytes(bytes: Vec<u8>) -> Result<Cid, Error> {
411    let cid: [u8; 32] = bytes.try_into().map_err(|bytes: Vec<u8>| {
412        Error::InvalidSnapshotBundle(format!("CID must be exactly 32 bytes, got {}", bytes.len()))
413    })?;
414    Ok(Cid(cid))
415}
416
417fn snapshot_bundle_deserialize(error: serde_cbor::Error) -> Error {
418    Error::InvalidSnapshotBundle(format!("could not decode bundle bytes: {error}"))
419}
420
421fn hex_bytes(bytes: &[u8]) -> String {
422    bytes.iter().map(|byte| format!("{byte:02x}")).collect()
423}