Fix AnyMap for TypeIds with 128 bit (#1311)

* Fix AnyMap for TypeIds with 128 bit * make const * added test, removed static_assertions
AFLplusplus · Jun 9, 2023 · 62b1bde · 62b1bde
1 parent fa63493
commit 62b1bde
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 15 deletions.
diff --git a/libafl/Cargo.toml b/libafl/Cargo.toml
@@ -67,7 +67,6 @@ serde = { version = "1.0", default-features = false, features = ["alloc", "deriv
 erased-serde = { version = "0.3.21", default-features = false, features = ["alloc"] } # erased serde
 postcard = { version = "1.0", features = ["alloc"] } # no_std compatible serde serialization fromat
 bincode = {version = "1.3", optional = true }
-static_assertions = "1.1.0"
 c2rust-bitfields = { version = "0.17", features = ["no_std"] }
 num_enum = { version = "0.5.7", default-features = false }
 typed-builder = "0.14" # Implement the builder pattern at compiletime

diff --git a/libafl/src/bolts/anymap.rs b/libafl/src/bolts/anymap.rs
@@ -3,6 +3,7 @@
 use alloc::boxed::Box;
 use core::{
     any::{Any, TypeId},
+    mem::size_of,
     ptr::addr_of,
 };
 
@@ -43,21 +44,41 @@ macro_rules! impl_asany {
 ///
 /// # Note
 /// Probably not safe for future compilers, fine for now.
+/// The size changed in later rust versions, see <https://github.com/rust-lang/compiler-team/issues/608>
+#[inline]
 #[must_use]
-pub fn pack_type_id(id: u64) -> TypeId {
-    assert_eq_size!(TypeId, u64);
-    unsafe { *(addr_of!(id) as *const TypeId) }
+pub const fn pack_type_id(id: u128) -> TypeId {
+    match size_of::<TypeId>() {
+        8 => {
+            let id_64 = id as u64;
+            unsafe { *(addr_of!(id_64) as *const TypeId) }
+        }
+        16 => unsafe { *(addr_of!(id) as *const TypeId) },
+        _ => {
+            // TypeId size of this size is not yet supported"
+            panic!("Unsupported size for TypeId");
+        }
+    }
 }
 
 /// Unpack a `type_id` to an `u64`
 /// Opposite of [`pack_type_id(id)`].
 ///
 /// # Note
 /// Probably not safe for future compilers, fine for now.
+/// The size changed in later rust versions, see <https://github.com/rust-lang/compiler-team/issues/608>
+#[inline]
 #[must_use]
-pub fn unpack_type_id(id: TypeId) -> u64 {
-    assert_eq_size!(TypeId, u64);
-    unsafe { *(addr_of!(id) as *const u64) }
+pub const fn unpack_type_id(id: TypeId) -> u128 {
+    #[allow(clippy::cast_ptr_alignment)] // we never actually cast to u128 if the type is u64.
+    match size_of::<TypeId>() {
+        8 => unsafe { *(addr_of!(id) as *const u64) as u128 },
+        16 => unsafe { *(addr_of!(id) as *const u128) },
+        _ => {
+            // TypeId size of this size is not yet supported"
+            panic!("Unsupported size for TypeId");
+        }
+    }
 }
 
 /// Create `AnyMap` and `NamedAnyMap` for a given trait
@@ -432,3 +453,22 @@ macro_rules! create_anymap_for_trait {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use core::any::TypeId;
+
+    use super::{pack_type_id, unpack_type_id};
+
+    #[test]
+    fn test_type_id() {
+        let type_id_u64 = unpack_type_id(TypeId::of::<u64>());
+        let type_id_u128 = unpack_type_id(TypeId::of::<u128>());
+
+        assert_eq!(pack_type_id(type_id_u64), TypeId::of::<u64>());
+        assert_eq!(pack_type_id(type_id_u128), TypeId::of::<u128>());
+
+        assert_ne!(pack_type_id(type_id_u64), TypeId::of::<u128>());
+        assert_ne!(pack_type_id(type_id_u128), TypeId::of::<u64>());
+    }
+}
diff --git a/libafl/src/bolts/serdeany.rs b/libafl/src/bolts/serdeany.rs
@@ -100,7 +100,7 @@ macro_rules! create_serde_registry_for_trait {
                 where
                     V: serde::de::SeqAccess<'de>,
                 {
-                    let id: u64 = visitor.next_element()?.unwrap();
+                    let id: u128 = visitor.next_element()?.unwrap();
                     let cb = unsafe {
                         *REGISTRY
                             .deserializers
@@ -117,7 +117,7 @@ macro_rules! create_serde_registry_for_trait {
 
             #[allow(unused_qualifications)]
             struct Registry {
-                deserializers: Option<HashMap<u64, DeserializeCallback<dyn $trait_name>>>,
+                deserializers: Option<HashMap<u128, DeserializeCallback<dyn $trait_name>>>,
                 finalized: bool,
             }
 
@@ -174,7 +174,7 @@ macro_rules! create_serde_registry_for_trait {
             /// in the registry
             #[derive(Debug, Serialize, Deserialize)]
             pub struct SerdeAnyMap {
-                map: HashMap<u64, Box<dyn $trait_name>>,
+                map: HashMap<u128, Box<dyn $trait_name>>,
             }
 
             // Cloning by serializing and deserializing. It ain't fast, but it's honest work.
@@ -301,7 +301,7 @@ macro_rules! create_serde_registry_for_trait {
             #[allow(unused_qualifications)]
             #[derive(Debug, Serialize, Deserialize)]
             pub struct NamedSerdeAnyMap {
-                map: HashMap<u64, HashMap<u64, Box<dyn $trait_name>>>,
+                map: HashMap<u128, HashMap<u64, Box<dyn $trait_name>>>,
             }
 
             // Cloning by serializing and deserializing. It ain't fast, but it's honest work.
@@ -467,8 +467,8 @@ macro_rules! create_serde_registry_for_trait {
                 pub fn all_typeids(
                     &self,
                 ) -> core::iter::Map<
-                    Keys<'_, u64, HashMap<u64, Box<dyn $trait_name>>>,
-                    fn(&u64) -> TypeId,
+                    Keys<'_, u128, HashMap<u64, Box<dyn $trait_name>>>,
+                    fn(&u128) -> TypeId,
                 > {
                     self.map.keys().map(|x| pack_type_id(*x))
                 }

diff --git a/libafl/src/lib.rs b/libafl/src/lib.rs
@@ -79,8 +79,6 @@ extern crate std;
 #[macro_use]
 #[doc(hidden)]
 pub extern crate alloc;
-#[macro_use]
-extern crate static_assertions;
 #[cfg(feature = "ctor")]
 #[doc(hidden)]
 pub use ctor::ctor;