1use 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
10pub const SNAPSHOT_BUNDLE_FORMAT_VERSION: u32 = 1;
12
13const SNAPSHOT_BUNDLE_BYTES_VERSION: u32 = SNAPSHOT_BUNDLE_FORMAT_VERSION;
14
15#[derive(Clone, Debug, Default, PartialEq, Eq)]
21pub struct MissingNodePlan {
22 pub required_cids: Vec<Cid>,
24 pub required_nodes: usize,
26 pub required_bytes: usize,
28 pub missing_cids: Vec<Cid>,
30 pub missing_nodes: usize,
32 pub missing_bytes: usize,
34}
35
36impl MissingNodePlan {
37 pub fn is_empty(&self) -> bool {
39 self.missing_cids.is_empty()
40 }
41
42 pub fn required_cids(&self) -> &[Cid] {
44 &self.required_cids
45 }
46
47 pub fn missing_cids(&self) -> &[Cid] {
49 &self.missing_cids
50 }
51}
52
53#[derive(Clone, Debug, Default, PartialEq, Eq)]
55pub struct MissingNodeCopy {
56 pub plan: MissingNodePlan,
58 pub copied_nodes: usize,
60 pub copied_bytes: usize,
62}
63
64#[derive(Clone, Debug, PartialEq, Eq)]
66pub struct SnapshotBundleNode {
67 pub cid: Cid,
69 pub bytes: Vec<u8>,
71}
72
73#[derive(Clone, Debug, PartialEq)]
79pub struct SnapshotBundle {
80 pub format_version: u32,
82 pub tree: Tree,
84 pub nodes: Vec<SnapshotBundleNode>,
86}
87
88#[derive(Clone, Debug, PartialEq, Eq)]
90pub struct SnapshotBundleSummary {
91 pub format_version: u32,
93 pub root: Option<Cid>,
95 pub node_count: usize,
97 pub byte_count: usize,
99 pub min_node_bytes: usize,
101 pub max_node_bytes: usize,
103}
104
105#[derive(Clone, Debug, PartialEq, Eq)]
107pub struct SnapshotBundleVerification {
108 pub valid: bool,
110 pub summary: SnapshotBundleSummary,
112 pub reachable_nodes: usize,
114 pub reachable_bytes: usize,
116 pub missing_cids: Vec<Cid>,
118 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 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 pub fn node_count(&self) -> usize {
147 self.nodes.len()
148 }
149
150 pub fn byte_count(&self) -> usize {
152 self.nodes.iter().map(|node| node.bytes.len()).sum()
153 }
154
155 pub fn digest(&self) -> Result<Cid, Error> {
161 self.to_bytes().map(|bytes| Cid::from_bytes(&bytes))
162 }
163
164 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 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 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 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 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}