From 7abc26ebc97562e1a9167b22895a36def074a119 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Wed, 20 Mar 2024 14:45:23 +0100 Subject: [PATCH] Change AnyMap API, add unsafe_ assert (#1958) * Change AnyMap API, add unsafe_ assert * clippy * Add anymap serialize test * Add test * fmt --- libafl/src/stages/tuneable.rs | 2 +- libafl/src/state/mod.rs | 4 +- libafl_bolts/src/serdeany.rs | 72 +++++++++++++++++++++++++++++------ 3 files changed, 63 insertions(+), 15 deletions(-) diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs index ed6f897130..c3fc25856b 100644 --- a/libafl/src/stages/tuneable.rs +++ b/libafl/src/stages/tuneable.rs @@ -478,7 +478,7 @@ where Z: Evaluator, Z::State: HasCorpus + HasRand + HasNamedMetadata, { - /// Creates a new tranforming mutational stage + /// Creates a new transforming mutational stage #[must_use] pub fn transforming(state: &mut Z::State, mutator: M, name: &str) -> Self { let _ = state.named_metadata_or_insert_with(name, TuneableMutationalStageMetadata::default); diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index 6a7450c186..ae5b5c759e 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -176,7 +176,7 @@ pub trait HasMetadata { where M: SerdeAny, { - self.metadata_map_mut().or_insert_with::(default) + self.metadata_map_mut().get_or_insert_with::(default) } /// Remove a metadata from the metadata map @@ -258,7 +258,7 @@ pub trait HasNamedMetadata { M: SerdeAny, { self.named_metadata_map_mut() - .or_insert_with::(name, default) + .get_or_insert_with::(name, default) } /// Check for a metadata diff --git a/libafl_bolts/src/serdeany.rs b/libafl_bolts/src/serdeany.rs index fa646c8fc9..356aeb8088 100644 --- a/libafl_bolts/src/serdeany.rs +++ b/libafl_bolts/src/serdeany.rs @@ -113,7 +113,7 @@ pub mod serdeany_registry { boxed::Box, string::{String, ToString}, }; - use core::{fmt, hash::BuildHasherDefault}; + use core::{any::TypeId, fmt, hash::BuildHasherDefault}; use hashbrown::{ hash_map::{Values, ValuesMut}, @@ -129,6 +129,9 @@ pub mod serdeany_registry { Error, }; + /// A [`HashMap`] that maps from [`TypeRepr`] to a deserializer and its [`TypeId`]. + type DeserializeCallbackMap = HashMap, TypeId)>; + /// Visitor object used internally for the [`crate::serdeany::SerdeAny`] registry. #[derive(Debug)] pub struct BoxDynVisitor {} @@ -146,12 +149,13 @@ pub mod serdeany_registry { { let id: TypeRepr = visitor.next_element()?.unwrap(); let cb = unsafe { - *REGISTRY + REGISTRY .deserializers .as_ref() .expect("Empty types registry") .get(&id) .expect("Cannot deserialize an unregistered type") + .0 }; let seed = DeserializeCallbackSeed:: { cb }; let obj: Self::Value = visitor.next_element_seed(seed)?.unwrap(); @@ -161,8 +165,7 @@ pub mod serdeany_registry { #[allow(unused_qualifications)] struct Registry { - deserializers: - Option>>, + deserializers: Option, finalized: bool, } @@ -175,9 +178,17 @@ pub mod serdeany_registry { assert!(!self.finalized, "Registry is already finalized!"); let deserializers = self.deserializers.get_or_insert_with(HashMap::default); - deserializers.insert(type_repr_owned::(), |de| { - Ok(Box::new(erased_serde::deserialize::(de)?)) - }); + let _entry = deserializers + .entry(type_repr_owned::()) + .or_insert_with(|| { + ( + |de| Ok(Box::new(erased_serde::deserialize::(de)?)), + TypeId::of::(), + ) + }); + + #[cfg(feature = "unsafe_stable_anymap")] + assert_eq!(_entry.1, TypeId::of::(), "Fatal safety error: TypeId of type {} is not equals to the deserializer's TypeId for this type! Two registered types have the same type_name!", type_repr::()); } pub fn finalize(&mut self) { @@ -360,15 +371,15 @@ pub mod serdeany_registry { } /// Gets a value by type, or inserts it using the given construction function `default` - pub fn or_insert_with(&mut self, default: impl FnOnce() -> T) -> &mut T + pub fn get_or_insert_with(&mut self, default: impl FnOnce() -> T) -> &mut T where T: SerdeAny, { - self.or_insert_with_boxed::(|| Box::new(default())) + self.get_or_insert_with_boxed::(|| Box::new(default())) } /// Gets a value by type, or inserts it using the given construction function `default` (returning a boxed value) - pub fn or_insert_with_boxed(&mut self, default: impl FnOnce() -> Box) -> &mut T + pub fn get_or_insert_with_boxed(&mut self, default: impl FnOnce() -> Box) -> &mut T where T: SerdeAny + 'static, { @@ -664,7 +675,7 @@ pub mod serdeany_registry { } /// Gets a value by name, or inserts it using the given construction function `default` - pub fn or_insert_with(&mut self, name: &str, default: impl FnOnce() -> T) -> &mut T + pub fn get_or_insert_with(&mut self, name: &str, default: impl FnOnce() -> T) -> &mut T where T: SerdeAny, { @@ -675,7 +686,7 @@ pub mod serdeany_registry { } /// Gets a value by name, or inserts it using the given construction function `default` (returning a boxed value) - pub fn or_insert_with_boxed( + pub fn get_or_insert_with_boxed( &mut self, name: &str, default: impl FnOnce() -> Box, @@ -874,6 +885,7 @@ macro_rules! impl_serdeany { /// /// # Safety /// This may never be called concurrently as it dereferences the `RegistryBuilder` without acquiring a lock. + #[allow(unused)] pub unsafe fn register() { $crate::serdeany::RegistryBuilder::register::<$struct_name>(); } @@ -882,3 +894,39 @@ macro_rules! impl_serdeany { $crate::create_register!($struct_name); }; } + +#[cfg(test)] +mod tests { + use serde::{Deserialize, Serialize}; + + use crate::serdeany::RegistryBuilder; + + #[derive(Debug, Serialize, Deserialize)] + struct MyType(u32); + impl_serdeany!(MyType); + + mod inner { + use serde::{Deserialize, Serialize}; + + #[derive(Debug, Serialize, Deserialize)] + pub(super) struct MyType(f32); + impl_serdeany!(MyType); + } + + #[test] + fn test_deserialize_serialize() { + unsafe { + RegistryBuilder::register::(); + RegistryBuilder::register::(); + } + + let val = MyType(1); + let serialized = postcard::to_allocvec(&val).unwrap(); + + assert_eq!( + postcard::from_bytes::(&serialized).unwrap().0, + val.0 + ); + assert!(postcard::from_bytes::(&serialized).is_err()); + } +}