rkyv_test/collections/btree_map/
validation.rs

1//! Validation implementation for BTreeMap.
2
3use super::{
4    ArchivedBTreeMap, ClassifiedNode, InnerNode, InnerNodeEntry, LeafNode, LeafNodeEntry, Node,
5    NodeHeader, MIN_ENTRIES_PER_INNER_NODE, MIN_ENTRIES_PER_LEAF_NODE,
6};
7use crate::{
8    rel_ptr::RelPtr,
9    validation::{ArchiveContext, LayoutRaw},
10    Archived, Fallible,
11};
12use bytecheck::{CheckBytes, Error};
13use core::{
14    alloc::Layout,
15    convert::{Infallible, TryFrom},
16    fmt,
17    hint::unreachable_unchecked,
18    ptr,
19};
20
21impl<K, C> CheckBytes<C> for InnerNodeEntry<K>
22where
23    K: CheckBytes<C>,
24    C: ArchiveContext + ?Sized,
25    C::Error: Error,
26{
27    type Error = K::Error;
28
29    #[inline]
30    unsafe fn check_bytes<'a>(
31        value: *const Self,
32        context: &mut C,
33    ) -> Result<&'a Self, Self::Error> {
34        RelPtr::manual_check_bytes(ptr::addr_of!((*value).ptr), context)
35            .unwrap_or_else(|_| core::hint::unreachable_unchecked());
36        K::check_bytes(ptr::addr_of!((*value).key), context)?;
37
38        Ok(&*value)
39    }
40}
41
42/// An error that can occur while checking a leaf node entry.
43#[derive(Debug)]
44pub enum LeafNodeEntryError<K, V> {
45    /// An error occurred while checking the entry's key.
46    KeyCheckError(K),
47    /// An error occurred while checking the entry's value.
48    ValueCheckError(V),
49}
50
51impl<K: fmt::Display, V: fmt::Display> fmt::Display for LeafNodeEntryError<K, V> {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        match self {
54            LeafNodeEntryError::KeyCheckError(e) => write!(f, "key check error: {}", e),
55            LeafNodeEntryError::ValueCheckError(e) => write!(f, "value check error: {}", e),
56        }
57    }
58}
59
60#[cfg(feature = "std")]
61const _: () = {
62    use std::error::Error;
63
64    impl<K: Error + 'static, V: Error + 'static> Error for LeafNodeEntryError<K, V> {
65        fn source(&self) -> Option<&(dyn Error + 'static)> {
66            match self {
67                Self::KeyCheckError(e) => Some(e as &dyn Error),
68                Self::ValueCheckError(e) => Some(e as &dyn Error),
69            }
70        }
71    }
72};
73
74impl<K, V, C> CheckBytes<C> for LeafNodeEntry<K, V>
75where
76    K: CheckBytes<C>,
77    V: CheckBytes<C>,
78    C: Fallible + ?Sized,
79    C::Error: Error,
80{
81    type Error = LeafNodeEntryError<K::Error, V::Error>;
82
83    #[inline]
84    unsafe fn check_bytes<'a>(
85        value: *const Self,
86        context: &mut C,
87    ) -> Result<&'a Self, Self::Error> {
88        K::check_bytes(ptr::addr_of!((*value).key), context)
89            .map_err(LeafNodeEntryError::KeyCheckError)?;
90        V::check_bytes(ptr::addr_of!((*value).value), context)
91            .map_err(LeafNodeEntryError::ValueCheckError)?;
92        Ok(&*value)
93    }
94}
95
96/// Errors that can occur while checking an archived B-tree.
97#[derive(Debug)]
98pub enum ArchivedBTreeMapError<K, V, C> {
99    /// An error occurred while checking the bytes of a key
100    KeyCheckError(K),
101    /// An error occurred while checking the bytes of a value
102    ValueCheckError(V),
103    /// The number of entries in the inner node is less than the minimum number of entries required
104    TooFewInnerNodeEntries(usize),
105    /// The number of entries in the leaf node is less than the minimum number of entries
106    TooFewLeafNodeEntries(usize),
107    /// An error occurred while checking the entries of an inner node
108    CheckInnerNodeEntryError {
109        /// The index of the inner node entry
110        index: usize,
111        /// The inner error that occurred
112        inner: K,
113    },
114    /// An error occurred while checking the entries of a leaf node
115    CheckLeafNodeEntryError {
116        /// The index of the leaf node entry
117        index: usize,
118        /// The inner error that occurred
119        inner: LeafNodeEntryError<K, V>,
120    },
121    /// The size of an inner node was invalid
122    InvalidNodeSize(usize),
123    /// The child of an inner node had a first key that did not match the inner node's key
124    MismatchedInnerChildKey,
125    /// The leaf level of the B-tree contained an inner node
126    InnerNodeInLeafLevel,
127    /// The leaves of the B-tree were not all located at the same depth
128    InvalidLeafNodeDepth {
129        /// The depth of the first leaf node in the tree
130        expected: usize,
131        /// The depth of the invalid leaf node
132        actual: usize,
133    },
134    /// A leaf node did not contain entries in sorted order
135    UnsortedLeafNodeEntries,
136    /// A leaf node is not linked after a node despite being the next leaf node
137    UnlinkedLeafNode,
138    /// A leaf node with lesser keys is linked after a leaf node with greater keys
139    UnsortedLeafNode,
140    /// The forward pointer of the last leaf did not have an offset of 0
141    LastLeafForwardPointerNotNull,
142    /// The number of entries the B-tree claims to have does not match the actual number of entries
143    LengthMismatch {
144        /// The number of entries the B-tree claims to have
145        expected: usize,
146        /// The actual number of entries in the B-tree
147        actual: usize,
148    },
149    /// The keys for an inner node were incorrect
150    IncorrectChildKey,
151    /// An context error occurred
152    ContextError(C),
153}
154
155impl<K, V, C> From<Infallible> for ArchivedBTreeMapError<K, V, C> {
156    fn from(_: Infallible) -> Self {
157        unsafe { core::hint::unreachable_unchecked() }
158    }
159}
160
161impl<K, V, C> fmt::Display for ArchivedBTreeMapError<K, V, C>
162where
163    K: fmt::Display,
164    V: fmt::Display,
165    C: fmt::Display,
166{
167    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
168        match self {
169            Self::KeyCheckError(e) => write!(f, "key check error: {}", e),
170            Self::ValueCheckError(e) => write!(f, "value check error: {}", e),
171            Self::TooFewInnerNodeEntries(n) => write!(
172                f,
173                "too few inner node entries (expected at least {}): {}",
174                MIN_ENTRIES_PER_INNER_NODE, n
175            ),
176            Self::TooFewLeafNodeEntries(n) => write!(
177                f,
178                "too few leaf node entries (expected at least {}): {}",
179                MIN_ENTRIES_PER_LEAF_NODE, n,
180            ),
181            Self::CheckInnerNodeEntryError { index, inner } => write!(
182                f,
183                "inner node entry check error: index {}, error {}",
184                index, inner
185            ),
186            Self::CheckLeafNodeEntryError { index, inner } => write!(
187                f,
188                "leaf node entry check error: index {}, error {}",
189                index, inner
190            ),
191            Self::InvalidNodeSize(n) => write!(f, "invalid node size: {}", n),
192            Self::MismatchedInnerChildKey => write!(f, "mismatched inner child key"),
193            Self::InnerNodeInLeafLevel => write!(f, "inner node in leaf level"),
194            Self::InvalidLeafNodeDepth { expected, actual } => write!(
195                f,
196                "expected leaf node depth {} but found leaf node depth {}",
197                expected, actual,
198            ),
199            Self::UnsortedLeafNodeEntries => write!(f, "leaf node contains keys in unsorted order"),
200            Self::UnlinkedLeafNode => write!(f, "leaf nodes are not linked in the sorted order"),
201            Self::UnsortedLeafNode => write!(f, "leaf nodes are not linked in sorted order"),
202            Self::LastLeafForwardPointerNotNull => {
203                write!(f, "the forward pointer of the last leaf was not null")
204            }
205            Self::LengthMismatch { expected, actual } => write!(
206                f,
207                "expected {} entries but there were actually {} entries",
208                expected, actual,
209            ),
210            Self::IncorrectChildKey => write!(f, "incorrect child key in inner node"),
211            Self::ContextError(e) => write!(f, "context error: {}", e),
212        }
213    }
214}
215
216#[cfg(feature = "std")]
217const _: () = {
218    use std::error::Error;
219
220    impl<K, V, C> Error for ArchivedBTreeMapError<K, V, C>
221    where
222        K: Error + 'static,
223        V: Error + 'static,
224        C: Error + 'static,
225    {
226        fn source(&self) -> Option<&(dyn Error + 'static)> {
227            match self {
228                Self::KeyCheckError(e) => Some(e as &dyn Error),
229                Self::ValueCheckError(e) => Some(e as &dyn Error),
230                Self::TooFewInnerNodeEntries(_) => None,
231                Self::TooFewLeafNodeEntries(_) => None,
232                Self::CheckInnerNodeEntryError { inner, .. } => Some(inner as &dyn Error),
233                Self::CheckLeafNodeEntryError { inner, .. } => Some(inner as &dyn Error),
234                Self::InvalidNodeSize(_) => None,
235                Self::MismatchedInnerChildKey => None,
236                Self::InnerNodeInLeafLevel => None,
237                Self::InvalidLeafNodeDepth { .. } => None,
238                Self::UnsortedLeafNodeEntries => None,
239                Self::UnlinkedLeafNode => None,
240                Self::UnsortedLeafNode => None,
241                Self::LastLeafForwardPointerNotNull => None,
242                Self::LengthMismatch { .. } => None,
243                Self::IncorrectChildKey => None,
244                Self::ContextError(e) => Some(e as &dyn Error),
245            }
246        }
247    }
248};
249
250impl<T> LayoutRaw for Node<[T]> {
251    fn layout_raw(value: *const Self) -> Layout {
252        let len = ptr_meta::metadata(value);
253        let result = Layout::new::<NodeHeader>()
254            .extend(Layout::array::<T>(len).unwrap())
255            .unwrap()
256            .0;
257        #[cfg(not(feature = "strict"))]
258        {
259            result
260        }
261        #[cfg(feature = "strict")]
262        {
263            result.pad_to_align()
264        }
265    }
266}
267
268type ABTMError<K, V, C> = ArchivedBTreeMapError<
269    <K as CheckBytes<C>>::Error,
270    <V as CheckBytes<C>>::Error,
271    <C as Fallible>::Error,
272>;
273
274impl NodeHeader {
275    #[inline]
276    unsafe fn manual_check_bytes<'a, K, V, C>(
277        value: *const Self,
278        context: &mut C,
279    ) -> Result<&'a Self, ABTMError<K, V, C>>
280    where
281        K: CheckBytes<C>,
282        V: CheckBytes<C>,
283        C: ArchiveContext + ?Sized,
284        C::Error: Error,
285    {
286        let raw_node = Self::manual_check_header(value, context)
287            .map_err(ArchivedBTreeMapError::ContextError)?;
288        Self::manual_check_contents::<K, V, C>(raw_node, context)?;
289
290        Ok(raw_node)
291    }
292
293    #[inline]
294    unsafe fn manual_check_header<'a, C>(
295        value: *const Self,
296        context: &mut C,
297    ) -> Result<&'a Self, C::Error>
298    where
299        C: ArchiveContext + ?Sized,
300        C::Error: Error,
301    {
302        CheckBytes::check_bytes(ptr::addr_of!((*value).meta), context).map_err(
303            // SAFETY: Infallible cannot exist
304            |_: Infallible| unreachable_unchecked(),
305        )?;
306        CheckBytes::check_bytes(ptr::addr_of!((*value).size), context).map_err(
307            // SAFETY: Infallible cannot exist
308            |_: Infallible| unreachable_unchecked(),
309        )?;
310        RelPtr::manual_check_bytes(ptr::addr_of!((*value).ptr), context).map_err(
311            // SAFETY: Infallible cannot exist
312            |_: Infallible| unreachable_unchecked(),
313        )?;
314
315        // All the fields have been checked and this is a valid RawNode
316        Ok(&*value)
317    }
318
319    #[inline]
320    unsafe fn manual_check_contents<K, V, C>(
321        raw_node: &Self,
322        context: &mut C,
323    ) -> Result<(), ABTMError<K, V, C>>
324    where
325        K: CheckBytes<C>,
326        V: CheckBytes<C>,
327        C: ArchiveContext + ?Sized,
328        C::Error: Error,
329    {
330        // Now that the fields have been checked, we can start checking the specific subtype
331        let root = (raw_node as *const Self).cast();
332        let size = from_archived!(raw_node.size) as usize;
333        let offset =
334            -isize::try_from(size).map_err(|_| ArchivedBTreeMapError::InvalidNodeSize(size))?;
335        let start = context
336            .check_ptr(root, offset, ())
337            .map_err(ArchivedBTreeMapError::ContextError)?;
338
339        // Push a new suffix range and check the inner or leaf part
340        let range = context
341            .push_suffix_subtree_range(start, root)
342            .map_err(ArchivedBTreeMapError::ContextError)?;
343        if raw_node.is_inner() {
344            InnerNode::manual_check_bytes::<V, C>(raw_node.classify_inner_ptr::<K>(), context)?;
345        } else {
346            CheckBytes::check_bytes(raw_node.classify_leaf_ptr::<K, V>(), context)?;
347        }
348        context
349            .pop_suffix_range(range)
350            .map_err(ArchivedBTreeMapError::ContextError)?;
351
352        Ok(())
353    }
354}
355
356impl<K> InnerNode<K> {
357    #[allow(clippy::type_complexity)]
358    fn verify_integrity<'a, V, C>(
359        &'a self,
360    ) -> Result<&K, ArchivedBTreeMapError<K::Error, V::Error, C::Error>>
361    where
362        K: CheckBytes<C> + PartialEq,
363        V: CheckBytes<C> + 'a,
364        C: Fallible + ?Sized,
365    {
366        for entry in self.tail.iter() {
367            let child = unsafe { &*entry.ptr.as_ptr() }.classify::<K, V>();
368            let first_key = match child {
369                ClassifiedNode::Inner(c) => c.verify_integrity::<V, C>()?,
370                ClassifiedNode::Leaf(c) => &c.tail[0].key,
371            };
372            if !entry.key.eq(first_key) {
373                return Err(ArchivedBTreeMapError::IncorrectChildKey);
374            }
375        }
376
377        let least_child = unsafe { &*self.header.ptr.as_ptr() }.classify::<K, V>();
378        let first_key = match least_child {
379            ClassifiedNode::Inner(c) => c.verify_integrity::<V, C>()?,
380            ClassifiedNode::Leaf(c) => &c.tail[0].key,
381        };
382
383        Ok(first_key)
384    }
385
386    #[inline]
387    unsafe fn manual_check_bytes<'a, V, C>(
388        value: *const Self,
389        context: &mut C,
390    ) -> Result<&'a Self, ABTMError<K, V, C>>
391    where
392        K: CheckBytes<C>,
393        V: CheckBytes<C>,
394        C: ArchiveContext + ?Sized,
395        C::Error: Error,
396    {
397        // meta, size, and ptr have already been checked by the check_bytes for RawNode
398        let len = ptr_meta::metadata(value);
399
400        // Each inner node actually contains one more entry that the length indicates (the least
401        // child pointer)
402        if len + 1 < MIN_ENTRIES_PER_INNER_NODE {
403            return Err(ArchivedBTreeMapError::TooFewInnerNodeEntries(len + 1));
404        }
405
406        // The subtree range has already been set up for us so we can just check our tail
407        let tail_ptr = ptr::addr_of!((*value).tail) as *const InnerNodeEntry<K>;
408        for index in (0..len).rev() {
409            CheckBytes::check_bytes(tail_ptr.add(index), context).map_err(|inner| {
410                ArchivedBTreeMapError::CheckInnerNodeEntryError { index, inner }
411            })?;
412        }
413
414        Ok(&*value)
415    }
416}
417
418impl<K, V, C> CheckBytes<C> for LeafNode<K, V>
419where
420    K: CheckBytes<C>,
421    V: CheckBytes<C>,
422    C: ArchiveContext + ?Sized,
423    C::Error: Error,
424{
425    type Error = ArchivedBTreeMapError<K::Error, V::Error, C::Error>;
426
427    #[inline]
428    unsafe fn check_bytes<'a>(
429        value: *const Self,
430        context: &mut C,
431    ) -> Result<&'a Self, Self::Error> {
432        // meta, size, and ptr have already been checked by the check_bytes for RawNode
433        let len = ptr_meta::metadata(value);
434
435        if len < MIN_ENTRIES_PER_LEAF_NODE {
436            return Err(ArchivedBTreeMapError::TooFewLeafNodeEntries(len));
437        }
438
439        // The subtree range has already been set up for us so we can just check our tail
440        let tail_ptr = ptr::addr_of!((*value).tail) as *const LeafNodeEntry<K, V>;
441        for index in (0..len).rev() {
442            CheckBytes::check_bytes(tail_ptr.add(index), context)
443                .map_err(|inner| ArchivedBTreeMapError::CheckLeafNodeEntryError { index, inner })?;
444        }
445
446        Ok(&*value)
447    }
448}
449
450#[cfg(feature = "alloc")]
451const _: () = {
452    #[cfg(not(feature = "std"))]
453    use alloc::collections::VecDeque;
454    #[cfg(feature = "std")]
455    use std::collections::VecDeque;
456
457    impl<K, V, C> CheckBytes<C> for ArchivedBTreeMap<K, V>
458    where
459        K: CheckBytes<C> + Ord,
460        V: CheckBytes<C>,
461        C: ArchiveContext + ?Sized,
462        C::Error: Error,
463    {
464        type Error = ArchivedBTreeMapError<K::Error, V::Error, C::Error>;
465
466        unsafe fn check_bytes<'a>(
467            value: *const Self,
468            context: &mut C,
469        ) -> Result<&'a Self, Self::Error> {
470            let len = from_archived!(*Archived::<usize>::check_bytes(
471                ptr::addr_of!((*value).len),
472                context,
473            )?) as usize;
474
475            if len > 0 {
476                let root_rel_ptr =
477                    RelPtr::manual_check_bytes(ptr::addr_of!((*value).root), context)?;
478
479                // Walk all the inner nodes, claim their memory, and check their contents
480                let mut nodes = VecDeque::new();
481                let root_ptr = context
482                    .check_subtree_rel_ptr(root_rel_ptr)
483                    .map_err(ArchivedBTreeMapError::ContextError)?;
484
485                // Before checking all the nodes, we have to push an additional prefix subtree with
486                // the root node
487                // Otherwise, when the suffix subtree of the root node is popped it will remove any
488                // trailing suffix space that should be checked by subsequent fields
489                let root = NodeHeader::manual_check_header(root_ptr, context)
490                    .map_err(ArchivedBTreeMapError::ContextError)?;
491                // To push the subtree, we need to know the real size of the root node
492                // Since the header is checked, we can just classify the pointer and use layout_raw
493                let root_layout = if root.is_inner() {
494                    Node::layout_raw(root.classify_inner_ptr::<K>())
495                } else {
496                    Node::layout_raw(root.classify_leaf_ptr::<K, V>())
497                };
498                let nodes_range = context
499                    .push_prefix_subtree_range(
500                        root_ptr.cast(),
501                        root_ptr.cast::<u8>().add(root_layout.size()),
502                    )
503                    .map_err(ArchivedBTreeMapError::ContextError)?;
504
505                // Now we're finally ready to check node subtrees
506                NodeHeader::manual_check_contents::<K, V, C>(root, context)?;
507
508                nodes.push_back((root, 0));
509
510                while let Some(&(node, depth)) = nodes.front() {
511                    // Break when a leaf is found
512                    if !node.is_inner() {
513                        break;
514                    }
515                    nodes.pop_front();
516                    let inner = node.classify_inner::<K>();
517
518                    let child_ptr = context
519                        .check_subtree_rel_ptr(&inner.header.ptr)
520                        .map_err(ArchivedBTreeMapError::ContextError)?;
521                    let child = NodeHeader::manual_check_bytes::<K, V, C>(child_ptr, context)?;
522                    nodes.push_back((child, depth + 1));
523
524                    // The invariant that this node contains keys less than the first key of this node will
525                    // be checked when we iterate through the leaf nodes in order and check ordering
526                    for entry in inner.tail.iter() {
527                        let child_ptr = context
528                            .check_subtree_rel_ptr(&entry.ptr)
529                            .map_err(ArchivedBTreeMapError::ContextError)?;
530                        let child = NodeHeader::manual_check_bytes::<K, V, C>(child_ptr, context)?;
531                        nodes.push_back((child, depth + 1));
532                    }
533                }
534
535                // We're done checking node subtrees now
536                context
537                    .pop_prefix_range(nodes_range)
538                    .map_err(ArchivedBTreeMapError::ContextError)?;
539
540                // The remaining nodes must all be leaf nodes
541                let mut entry_count = 0;
542                for (i, (node, depth)) in nodes.iter().enumerate() {
543                    if !node.is_leaf() {
544                        return Err(ArchivedBTreeMapError::InnerNodeInLeafLevel);
545                    }
546                    let leaf = node.classify_leaf::<K, V>();
547
548                    // Leaf nodes must all be the same depth
549                    let expected_depth = nodes.front().unwrap().1;
550                    if *depth != expected_depth {
551                        return Err(ArchivedBTreeMapError::InvalidLeafNodeDepth {
552                            expected: expected_depth,
553                            actual: *depth,
554                        });
555                    }
556
557                    // They must contain entries in sorted order
558                    for (prev, next) in leaf.tail.iter().zip(leaf.tail.iter().skip(1)) {
559                        if next.key < prev.key {
560                            return Err(ArchivedBTreeMapError::UnsortedLeafNodeEntries);
561                        }
562                    }
563
564                    // And they must link together in sorted order
565                    if i < nodes.len() - 1 {
566                        let next_ptr = context
567                            .check_rel_ptr(&leaf.header.ptr)
568                            .map_err(ArchivedBTreeMapError::ContextError)?;
569                        let next_node = nodes[i + 1].0.classify_leaf();
570
571                        if next_ptr != (next_node as *const LeafNode<K, V>).cast() {
572                            return Err(ArchivedBTreeMapError::UnlinkedLeafNode);
573                        }
574                        if next_node.tail[0].key < leaf.tail[leaf.tail.len() - 1].key {
575                            return Err(ArchivedBTreeMapError::UnsortedLeafNode);
576                        }
577                    } else {
578                        // The last node must have a null pointer forward
579                        if !leaf.header.ptr.is_null() {
580                            return Err(ArchivedBTreeMapError::LastLeafForwardPointerNotNull);
581                        }
582                    }
583
584                    // Keep track of the number of entries found
585                    entry_count += leaf.tail.len();
586                }
587
588                // Make sure that the number of entries matches the length
589                if entry_count != len {
590                    return Err(ArchivedBTreeMapError::LengthMismatch {
591                        expected: len,
592                        actual: entry_count,
593                    });
594                }
595
596                // Make sure that inner nodes are constructed appropriately
597                if root.is_inner() {
598                    root.classify_inner::<K>().verify_integrity::<V, C>()?;
599                }
600            }
601
602            Ok(&*value)
603        }
604    }
605};
rkyv_test/collections/btree_map/validation.rs

rkyv_test/collections/btree_map/
validation.rs
