Skip to content

Commit

Permalink
#239 Fix BTreeMap validation erasing suffix ranges
Browse files Browse the repository at this point in the history
Non-compact BTreeMaps would accidentally erase the suffix subtree after
their node data. This effectively made any struct members after a
BTreeMap fail to validate with an invalid subtree pointer range. The fix
for this is to push an additional prefix subtree range for all nodes in
the BTreeMap before checking each one individually.
  • Loading branch information
djkoloski committed Jan 8, 2022
1 parent 0ac6c04 commit 23d4537
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 24 deletions.
80 changes: 73 additions & 7 deletions rkyv/src/collections/btree_map/validation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use bytecheck::{CheckBytes, Error};
use core::{
alloc::Layout,
convert::{Infallible, TryFrom},
hint::unreachable_unchecked,
fmt, ptr,
};

Expand Down Expand Up @@ -281,15 +282,55 @@ impl NodeHeader {
C: ArchiveContext + ?Sized,
C::Error: Error,
{
CheckBytes::check_bytes(ptr::addr_of!((*value).meta), context)?;
CheckBytes::check_bytes(ptr::addr_of!((*value).size), context)?;
RelPtr::manual_check_bytes(ptr::addr_of!((*value).ptr), context)?;
let raw_node = Self::manual_check_header(value, context)
.map_err(ArchivedBTreeMapError::ContextError)?;
Self::manual_check_contents::<K, V, C>(raw_node, context)?;

Ok(raw_node)
}

#[inline]
unsafe fn manual_check_header<'a, C>(
value: *const Self,
context: &mut C,
) -> Result<&'a Self, C::Error>
where
C: ArchiveContext + ?Sized,
C::Error: Error,
{
CheckBytes::check_bytes(ptr::addr_of!((*value).meta), context)
.map_err(|_: Infallible| unsafe {
// SAFETY: Infallible cannot exist
unreachable_unchecked()
})?;
CheckBytes::check_bytes(ptr::addr_of!((*value).size), context)
.map_err(|_: Infallible| unsafe {
// SAFETY: Infallible cannot exist
unreachable_unchecked()
})?;
RelPtr::manual_check_bytes(ptr::addr_of!((*value).ptr), context)
.map_err(|_: Infallible| unsafe {
// SAFETY: Infallible cannot exist
unreachable_unchecked()
})?;

// All the fields have been checked and this is a valid RawNode
let raw_node = &*value;
Ok(&*value)
}

#[inline]
unsafe fn manual_check_contents<K, V, C>(
raw_node: &Self,
context: &mut C,
) -> Result<(), ABTMError<K, V, C>>
where
K: CheckBytes<C>,
V: CheckBytes<C>,
C: ArchiveContext + ?Sized,
C::Error: Error,
{
// Now that the fields have been checked, we can start checking the specific subtype
let root = value.cast();
let root = (raw_node as *const Self).cast();
let size = from_archived!(raw_node.size) as usize;
let offset =
-isize::try_from(size).map_err(|_| ArchivedBTreeMapError::InvalidNodeSize(size))?;
Expand All @@ -310,7 +351,7 @@ impl NodeHeader {
.pop_suffix_range(range)
.map_err(ArchivedBTreeMapError::ContextError)?;

Ok(&*value)
Ok(())
}
}

Expand Down Expand Up @@ -442,7 +483,28 @@ const _: () = {
let root_ptr = context
.check_subtree_rel_ptr(root_rel_ptr)
.map_err(ArchivedBTreeMapError::ContextError)?;
let root = NodeHeader::manual_check_bytes::<K, V, C>(root_ptr, context)?;

// Before checking all the nodes, we have to push an additional prefix subtree with
// the root node
// Otherwise, when the suffix subtree of the root node is popped it will remove any
// trailing suffix space that should be checked by subsequent fields
let root = NodeHeader::manual_check_header(root_ptr, context)
.map_err(ArchivedBTreeMapError::ContextError)?;
// To push the subtree, we need to know the real size of the root node
// Since the header is checked, we can just classify the pointer and use layout_raw
let root_layout = if root.is_inner() {
Node::layout_raw(root.classify_inner_ptr::<K>())
} else {
Node::layout_raw(root.classify_leaf_ptr::<K, V>())
};
let nodes_range = context.push_prefix_subtree_range(
root_ptr.cast(),
root_ptr.cast::<u8>().add(root_layout.size())
).map_err(ArchivedBTreeMapError::ContextError)?;

// Now we're finally ready to check node subtrees
NodeHeader::manual_check_contents::<K, V, C>(root, context)?;

nodes.push_back((root, 0));

while let Some(&(node, depth)) = nodes.front() {
Expand Down Expand Up @@ -470,6 +532,10 @@ const _: () = {
}
}

// We're done checking node subtrees now
context.pop_prefix_range(nodes_range)
.map_err(ArchivedBTreeMapError::ContextError)?;

// The remaining nodes must all be leaf nodes
let mut entry_count = 0;
for (i, (node, depth)) in nodes.iter().enumerate() {
Expand Down
46 changes: 29 additions & 17 deletions rkyv_test/src/validation/test_alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ mod tests {
#[cfg(not(feature = "std"))]
use alloc::{
boxed::Box,
collections::BTreeMap,
rc::Rc,
string::{String, ToString},
vec,
Expand All @@ -12,10 +13,10 @@ mod tests {
use bytecheck::CheckBytes;
use rkyv::{
check_archived_root, check_archived_value, ser::Serializer, AlignedBytes, Archive,
Serialize,
Deserialize, Infallible, Serialize,
};
#[cfg(feature = "std")]
use std::rc::Rc;
use std::{collections::BTreeMap, rc::Rc};

#[cfg(feature = "wasm")]
use wasm_bindgen_test::*;
Expand Down Expand Up @@ -405,11 +406,6 @@ mod tests {
#[test]
#[cfg_attr(feature = "wasm", wasm_bindgen_test)]
fn check_b_tree() {
#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap;
#[cfg(feature = "std")]
use std::collections::BTreeMap;

let mut value = BTreeMap::new();
value.insert("foo".to_string(), 10);
value.insert("bar".to_string(), 20);
Expand All @@ -426,11 +422,6 @@ mod tests {
#[test]
#[cfg_attr(feature = "wasm", wasm_bindgen_test)]
fn check_empty_b_tree() {
#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap;
#[cfg(feature = "std")]
use std::collections::BTreeMap;

let value = BTreeMap::<u8, ()>::new();

let mut serializer = DefaultSerializer::default();
Expand All @@ -446,11 +437,6 @@ mod tests {
// This test creates structures too big to fit in 16-bit offsets
#[cfg(not(feature = "size_16"))]
fn check_b_tree_large() {
#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap;
#[cfg(feature = "std")]
use std::collections::BTreeMap;

let mut value = BTreeMap::new();
for i in 0..100_000 {
value.insert(i.to_string(), i);
Expand All @@ -463,6 +449,32 @@ mod tests {
check_archived_root::<BTreeMap<String, i32>>(buf.as_ref()).unwrap();
}

#[test]
#[cfg_attr(feature = "wasm", wasm_bindgen_test)]
fn b_tree_struct_member() {
#[derive(Archive, Serialize, Deserialize, Debug, Default)]
#[archive_attr(derive(CheckBytes))]
pub struct MyType {
pub some_list: BTreeMap<String, Vec<f32>>,
pub values: Vec<f32>,
}

let mut value = MyType::default();

value.some_list
.entry("Asdf".to_string())
.and_modify(|e| e.push(1.0))
.or_insert_with(|| vec![2.0]);

let mut serializer = DefaultSerializer::default();
serializer.serialize_value(&value).unwrap();
let buf = serializer.into_serializer().into_inner();
let data = buf.into_vec();

let value = check_archived_root::<MyType>(&data).unwrap();
let _: MyType = value.deserialize(&mut Infallible).unwrap();
}

#[test]
#[cfg_attr(feature = "wasm", wasm_bindgen_test)]
fn check_atomics_endianness() {
Expand Down

0 comments on commit 23d4537

Please sign in to comment.