prolly/prolly/content_graph/
walk.rs

1use super::kind::{ContentObjectKind, TypedContentObject, TypedContentRoot};
2use crate::prolly::cid::Cid;
3use crate::prolly::error::Error;
4use crate::prolly::node::Node;
5use crate::prolly::proximity::accelerator::catalog::{
6    CatalogAcceleratorKind, Manifest as CatalogManifest,
7};
8use crate::prolly::proximity::accelerator::composite::Manifest as CompositeManifest;
9use crate::prolly::proximity::accelerator::hnsw::storage::{GraphNode, Manifest as HnswManifest};
10use crate::prolly::proximity::accelerator::pq::Manifest as PqManifest;
11use crate::prolly::proximity::storage::quantized::ScalarQuantized;
12use crate::prolly::proximity::storage::vector::ExternalVector;
13use crate::prolly::proximity::storage::{Descriptor, PhysicalNodeKind, ProximityNode, VectorRef};
14use crate::prolly::store::Store;
15use std::collections::{BTreeMap, HashMap, HashSet};
16
17/// Hard limits for untrusted typed graph traversal.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct ContentGraphLimits {
20    pub max_objects: usize,
21    pub max_depth: usize,
22    pub max_bytes: usize,
23    pub max_references_per_object: usize,
24}
25
26impl Default for ContentGraphLimits {
27    fn default() -> Self {
28        Self {
29            max_objects: 10_000_000,
30            max_depth: 256,
31            max_bytes: usize::try_from(64_u64 * 1024 * 1024 * 1024).unwrap_or(usize::MAX),
32            max_references_per_object: 16_777_216,
33        }
34    }
35}
36
37/// Deterministic descendant-first traversal report.
38#[derive(Clone, Debug, Default, PartialEq, Eq)]
39pub struct ContentGraphWalk {
40    pub objects: Vec<TypedContentObject>,
41    pub total_bytes: usize,
42    pub maximum_depth: usize,
43    pub objects_by_kind: BTreeMap<ContentObjectKind, usize>,
44}
45
46pub fn walk_content_graph<S: Store>(
47    store: &S,
48    roots: &[TypedContentRoot],
49    limits: &ContentGraphLimits,
50) -> Result<ContentGraphWalk, Error> {
51    validate_limits(limits)?;
52    let mut ordered_roots = roots.to_vec();
53    ordered_roots.sort_by(compare_root);
54    ordered_roots.dedup();
55    let mut stack: Vec<_> = ordered_roots
56        .into_iter()
57        .rev()
58        .map(|root| Frame::Enter(root, 0))
59        .collect();
60    let mut visiting = HashSet::<Cid>::new();
61    let mut seen = HashSet::<Cid>::new();
62    let mut contexts = HashMap::<Cid, (ContentObjectKind, Option<u32>)>::new();
63    let mut pending = HashMap::<Cid, TypedContentObject>::new();
64    let mut walk = ContentGraphWalk::default();
65    let mut loaded_bytes = 0usize;
66
67    while let Some(frame) = stack.pop() {
68        match frame {
69            Frame::Enter(root, depth) => {
70                if let Some((kind, dimensions)) = contexts.get(&root.cid) {
71                    if *dimensions != root.dimensions || !compatible_kinds(*kind, root.kind) {
72                        return Err(invalid(
73                            "the same CID was referenced with conflicting type context",
74                        ));
75                    }
76                } else {
77                    contexts.insert(root.cid.clone(), (root.kind, root.dimensions));
78                    if contexts.len() > limits.max_objects {
79                        return Err(limit("objects", limits.max_objects, contexts.len()));
80                    }
81                }
82                if seen.contains(&root.cid) {
83                    continue;
84                }
85                if !visiting.insert(root.cid.clone()) {
86                    return Err(invalid("cycle detected in typed content graph"));
87                }
88                if depth > limits.max_depth {
89                    return Err(limit("depth", limits.max_depth, depth));
90                }
91                let bytes = load_content(store, &root.cid)?;
92                let next_bytes = loaded_bytes.saturating_add(bytes.len());
93                if next_bytes > limits.max_bytes {
94                    return Err(limit("bytes", limits.max_bytes, next_bytes));
95                }
96                loaded_bytes = next_bytes;
97                let (actual_kind, mut references) = references(&root, &bytes)?;
98                if references.len() > limits.max_references_per_object {
99                    return Err(limit(
100                        "references",
101                        limits.max_references_per_object,
102                        references.len(),
103                    ));
104                }
105                references.sort_by(compare_root);
106                references.dedup();
107                let object = TypedContentObject {
108                    root: TypedContentRoot {
109                        kind: actual_kind,
110                        cid: root.cid.clone(),
111                        dimensions: root.dimensions,
112                    },
113                    bytes,
114                    depth,
115                };
116                pending.insert(root.cid.clone(), object);
117                stack.push(Frame::Exit(root.cid));
118                for reference in references.into_iter().rev() {
119                    stack.push(Frame::Enter(reference, depth + 1));
120                }
121            }
122            Frame::Exit(cid) => {
123                visiting.remove(&cid);
124                if seen.insert(cid.clone()) {
125                    let object = pending.remove(&cid).expect("entered content object");
126                    walk.total_bytes = walk.total_bytes.saturating_add(object.bytes.len());
127                    walk.maximum_depth = walk.maximum_depth.max(object.depth);
128                    *walk.objects_by_kind.entry(object.root.kind).or_default() += 1;
129                    walk.objects.push(object);
130                }
131            }
132        }
133    }
134    Ok(walk)
135}
136
137/// Load one authenticated object and return its exact, deterministic typed
138/// reference list without recursively walking descendants.
139pub fn content_references<S: Store>(
140    store: &S,
141    root: &TypedContentRoot,
142) -> Result<Vec<TypedContentRoot>, Error> {
143    let bytes = load_content(store, &root.cid)?;
144    let (_, mut output) = references(root, &bytes)?;
145    output.sort_by(compare_root);
146    output.dedup();
147    Ok(output)
148}
149
150fn compatible_kinds(left: ContentObjectKind, right: ContentObjectKind) -> bool {
151    left == right || (is_proximity_node_kind(left) && is_proximity_node_kind(right))
152}
153
154fn is_proximity_node_kind(kind: ContentObjectKind) -> bool {
155    matches!(
156        kind,
157        ContentObjectKind::ProximityNode
158            | ContentObjectKind::OverflowDirectory
159            | ContentObjectKind::OverflowPage
160    )
161}
162
163enum Frame {
164    Enter(TypedContentRoot, usize),
165    Exit(Cid),
166}
167
168fn references(
169    root: &TypedContentRoot,
170    bytes: &[u8],
171) -> Result<(ContentObjectKind, Vec<TypedContentRoot>), Error> {
172    validate_root_context(root)?;
173    let mut output = Vec::new();
174    let actual = match root.kind {
175        ContentObjectKind::OrderedNode => {
176            let node = Node::from_bytes(bytes)?;
177            if node.keys.len() != node.vals.len() {
178                return Err(invalid("ordered node key/value count mismatch"));
179            }
180            if !node.leaf {
181                for value in node.vals {
182                    let cid = Cid(value
183                        .try_into()
184                        .map_err(|_| invalid("ordered internal value is not a CID"))?);
185                    output.push(TypedContentRoot::new(ContentObjectKind::OrderedNode, cid));
186                }
187            }
188            ContentObjectKind::OrderedNode
189        }
190        ContentObjectKind::ProximityDescriptor => {
191            let descriptor = Descriptor::decode(bytes)?;
192            if let Some(cid) = descriptor.directory.root {
193                output.push(TypedContentRoot::new(ContentObjectKind::OrderedNode, cid));
194            }
195            output.push(
196                TypedContentRoot::new(ContentObjectKind::ProximityNode, descriptor.proximity_root)
197                    .with_dimensions(descriptor.config.dimensions),
198            );
199            ContentObjectKind::ProximityDescriptor
200        }
201        ContentObjectKind::ProximityNode
202        | ContentObjectKind::OverflowDirectory
203        | ContentObjectKind::OverflowPage => {
204            let dimensions = root
205                .dimensions
206                .ok_or_else(|| invalid("PRXN traversal requires dimensions"))?;
207            let node = ProximityNode::decode(bytes, dimensions)?;
208            let actual = match node.kind {
209                PhysicalNodeKind::OverflowDirectory => ContentObjectKind::OverflowDirectory,
210                PhysicalNodeKind::OverflowPage => ContentObjectKind::OverflowPage,
211                PhysicalNodeKind::Leaf | PhysicalNodeKind::Route => {
212                    ContentObjectKind::ProximityNode
213                }
214            };
215            if root.kind != ContentObjectKind::ProximityNode && root.kind != actual {
216                return Err(invalid("PRXN physical kind disagrees with typed reference"));
217            }
218            if let Some(cid) = node.quantizer {
219                output.push(TypedContentRoot::new(
220                    ContentObjectKind::ScalarQuantization,
221                    cid,
222                ));
223            }
224            for entry in node.entries {
225                if let VectorRef::External(cid) = entry.vector {
226                    output.push(TypedContentRoot::new(
227                        ContentObjectKind::ExternalVector,
228                        cid,
229                    ));
230                }
231                if let Some(cid) = entry.child {
232                    output.push(
233                        TypedContentRoot::new(ContentObjectKind::ProximityNode, cid)
234                            .with_dimensions(dimensions),
235                    );
236                }
237            }
238            actual
239        }
240        ContentObjectKind::ExternalVector => {
241            ExternalVector::decode(bytes)?;
242            ContentObjectKind::ExternalVector
243        }
244        ContentObjectKind::ScalarQuantization => {
245            ScalarQuantized::decode(bytes)?;
246            ContentObjectKind::ScalarQuantization
247        }
248        ContentObjectKind::ProductQuantization => {
249            let manifest = PqManifest::decode(bytes)?;
250            manifest.config.validate(
251                manifest.dimensions,
252                usize::from(manifest.config.centroids_per_subquantizer),
253            )?;
254            output.push(TypedContentRoot::proximity_descriptor(manifest.source));
255            output.push(TypedContentRoot::new(
256                ContentObjectKind::OrderedNode,
257                manifest.code_root,
258            ));
259            ContentObjectKind::ProductQuantization
260        }
261        ContentObjectKind::HnswManifest => {
262            let manifest = HnswManifest::decode(bytes)?;
263            manifest.config.validate()?;
264            output.push(TypedContentRoot::proximity_descriptor(manifest.source));
265            output.push(TypedContentRoot::new(
266                ContentObjectKind::OrderedNode,
267                manifest.graph_root,
268            ));
269            ContentObjectKind::HnswManifest
270        }
271        ContentObjectKind::HnswPage => {
272            GraphNode::decode(bytes)?;
273            ContentObjectKind::HnswPage
274        }
275        ContentObjectKind::CompositeAccelerator => {
276            let manifest = CompositeManifest::decode(bytes)?;
277            output.push(TypedContentRoot::proximity_descriptor(
278                manifest.current_source,
279            ));
280            output.push(TypedContentRoot::proximity_descriptor(manifest.base_source));
281            output.push(TypedContentRoot::new(
282                match manifest.base_kind {
283                    crate::prolly::proximity::CompositeBaseKind::Hnsw => {
284                        ContentObjectKind::HnswManifest
285                    }
286                    crate::prolly::proximity::CompositeBaseKind::ProductQuantized => {
287                        ContentObjectKind::ProductQuantization
288                    }
289                },
290                manifest.base_manifest,
291            ));
292            if let Some(root) = manifest.delta_root {
293                output.push(TypedContentRoot::new(ContentObjectKind::OrderedNode, root));
294            }
295            if let Some(root) = manifest.shadow_root {
296                output.push(TypedContentRoot::new(ContentObjectKind::OrderedNode, root));
297            }
298            ContentObjectKind::CompositeAccelerator
299        }
300        ContentObjectKind::AcceleratorCatalog => {
301            let manifest = CatalogManifest::decode(bytes)?;
302            output.push(TypedContentRoot::proximity_descriptor(manifest.source));
303            for entry in manifest.entries {
304                output.push(TypedContentRoot::new(
305                    match entry.kind {
306                        CatalogAcceleratorKind::Hnsw => ContentObjectKind::HnswManifest,
307                        CatalogAcceleratorKind::ProductQuantized => {
308                            ContentObjectKind::ProductQuantization
309                        }
310                        CatalogAcceleratorKind::Composite => {
311                            ContentObjectKind::CompositeAccelerator
312                        }
313                    },
314                    entry.manifest,
315                ));
316            }
317            ContentObjectKind::AcceleratorCatalog
318        }
319    };
320    Ok((actual, output))
321}
322
323fn validate_root_context(root: &TypedContentRoot) -> Result<(), Error> {
324    if is_proximity_node_kind(root.kind) != root.dimensions.is_some() {
325        return Err(invalid(
326            "dimensions must be present exactly for PRXN-family objects",
327        ));
328    }
329    Ok(())
330}
331
332fn compare_root(left: &TypedContentRoot, right: &TypedContentRoot) -> std::cmp::Ordering {
333    left.kind
334        .cmp(&right.kind)
335        .then_with(|| left.cid.as_bytes().cmp(right.cid.as_bytes()))
336        .then_with(|| left.dimensions.cmp(&right.dimensions))
337}
338
339fn load_content<S: Store>(store: &S, cid: &Cid) -> Result<Vec<u8>, Error> {
340    let bytes = store
341        .get(cid.as_bytes())
342        .map_err(|error| Error::Store(Box::new(error)))?
343        .ok_or_else(|| Error::NotFound(cid.clone()))?;
344    let actual = Cid::from_bytes(&bytes);
345    if actual != *cid {
346        return Err(Error::CidMismatch {
347            expected: cid.clone(),
348            actual,
349        });
350    }
351    Ok(bytes)
352}
353
354fn validate_limits(limits: &ContentGraphLimits) -> Result<(), Error> {
355    if limits.max_objects == 0 || limits.max_bytes == 0 || limits.max_references_per_object == 0 {
356        return Err(invalid("content graph limits must be greater than zero"));
357    }
358    Ok(())
359}
360
361fn limit(resource: &'static str, limit: usize, actual: usize) -> Error {
362    Error::ContentGraphResourceLimitExceeded {
363        resource,
364        limit,
365        actual,
366    }
367}
368
369fn invalid(reason: impl Into<String>) -> Error {
370    Error::InvalidProximityObject {
371        kind: "content graph",
372        reason: reason.into(),
373    }
374}
prolly/prolly/content_graph/walk.rs

prolly/prolly/content_graph/
walk.rs
