fix(soundness): fix a soundness bug with the Maybe<T> type and the …

…`HasSchema` derive macro for generic types. (#456)

There was a bug in the `HasSchema` derive macro for generic types, where
the enum variant schema was returning the schema for a previously
computed instantiation of the generic type, instead of specializing
based on the Rust type ID.

This fixes that, fixing undefined behavior that may be triggered when
using `Maybe<T>` along with any other generic enums that have derived
`HasSchema`.

This should also fix #296, which I can't reproduce right now, but could
have been caused randomly by `Maybe::schema()` being unsound.

Fixes: #296

framework_crates/bones_schema/macros/src/lib.rs

@@ -254,12 +254,38 @@ pub fn derive_has_schema(input: TokenStream) -> TokenStream {
                     let name = v.name.to_string();
                     let variant_schema_name = format!("{}::{}", e.name, name);
                     let fields = parse_struct_fields(&v.contents);
+
+                    let register_schema = if input.generic_params().is_some() {
+                        quote! {
+                            static S: OnceLock<RwLock<HashMap<TypeId, &'static Schema>>> = OnceLock::new();
+                            let schema = {
+                                S.get_or_init(Default::default)
+                                    .read()
+                                    .get(&TypeId::of::<Self>())
+                                    .copied()
+                            };
+                            schema.unwrap_or_else(|| {
+                                let schema = compute_schema();
+
+                                S.get_or_init(Default::default)
+                                    .write()
+                                    .insert(TypeId::of::<Self>(), schema);
+
+                                schema
+                            })
+                        }
+                    } else {
+                        quote! {
+                            static S: ::std::sync::OnceLock<&'static #schema_mod::Schema> = ::std::sync::OnceLock::new();
+                            S.get_or_init(compute_schema)
+                        }
+                    };
+
                     variants.push(quote! {
                         #schema_mod::VariantInfo {
                             name: #name.into(),
                             schema: {
-                                static S: ::std::sync::OnceLock<&'static #schema_mod::Schema> = ::std::sync::OnceLock::new();
-                                S.get_or_init(|| {
+                                let compute_schema = || {
                                     #schema_mod::registry::SCHEMA_REGISTRY.register(#schema_mod::SchemaData {
                                         name: #variant_schema_name.into(),
                                         full_name: concat!(module_path!(), "::", #variant_schema_name).into(),
@@ -276,7 +302,8 @@ pub fn derive_has_schema(input: TokenStream) -> TokenStream {
                                         hash_fn: None,
                                         drop_fn: None,
                                     })
-                                })
+                                };
+                                #register_schema
                             }
                         }
                     })
@@ -325,7 +352,6 @@ pub fn derive_has_schema(input: TokenStream) -> TokenStream {
         quote! {
             unsafe impl #sync_send_generic_params #schema_mod::HasSchema for #name #generic_params {
                 fn schema() -> &'static #schema_mod::Schema {
-                    // TODO: use faster hashmap and rwlocks from bones_utils.
                     use ::std::sync::{OnceLock};
                     use ::std::any::TypeId;
                     use bones_utils::{HashMap, parking_lot::RwLock};

framework_crates/bones_schema/src/lib.rs

@@ -44,3 +44,125 @@ mod std_impls;
 /// Serde implementations for [`Schema`].
 #[cfg(feature = "serde")]
 pub mod ser_de;
+
+#[cfg(test)]
+mod test {
+    #[cfg(feature = "derive")]
+    mod derive_test {
+        #![allow(dead_code)]
+
+        use crate::prelude::*;
+
+        #[derive(HasSchema, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Default, Debug)]
+        #[schema_module(crate)]
+        #[repr(C, u8)]
+        pub enum Maybe<T> {
+            /// The value is not set.
+            #[default]
+            Unset,
+            /// The value is set.
+            Set(T),
+        }
+
+        #[derive(HasSchema, Clone, Copy, Debug, PartialEq, Eq, Default)]
+        #[schema_module(crate)]
+        #[repr(u8)]
+        pub enum E {
+            #[default]
+            None,
+            L,
+            R,
+            U,
+            D,
+            G,
+            S,
+        }
+
+        /// Represents the ball in the game
+        #[derive(HasSchema, Clone, Default)]
+        #[schema_module(crate)]
+        pub struct A {
+            pub c: Maybe<u32>,
+            pub d: SVec<u64>,
+            pub e: Maybe<u64>,
+            pub f: u32,
+            pub g: f32,
+            pub h: f32,
+            pub i: E,
+            pub j: u32,
+            pub k: u32,
+        }
+
+        #[derive(HasSchema, Clone, Default)]
+        #[schema_module(crate)]
+        #[repr(C)]
+        pub struct B {
+            pub c: Maybe<u32>,
+            pub d: SVec<u64>,
+            pub e: Maybe<u64>,
+            pub f: u32,
+            pub g: f32,
+            pub h: f32,
+            pub i: E,
+            pub j: u32,
+            pub k: u32,
+        }
+
+        #[derive(HasSchema, Clone, Default)]
+        #[schema_module(crate)]
+        pub struct C {
+            pub c: Maybe<u32>,
+            pub e: Maybe<u64>,
+        }
+
+        #[derive(HasSchema, Clone, Default)]
+        #[schema_module(crate)]
+        #[repr(C)]
+        pub struct D {
+            pub c: Maybe<u32>,
+            pub e: Maybe<u64>,
+        }
+
+        #[derive(HasSchema, Clone)]
+        #[schema(no_default)]
+        #[schema_module(crate)]
+        #[repr(C)]
+        struct F<T> {
+            a: bool,
+            b: T,
+        }
+
+        /// Makes sure that the layout reported in the schema for a generic type matches the layout
+        /// reported by Rust, for two different type parameters.
+        #[test]
+        fn generic_type_schema_layouts_match() {
+            assert_eq!(
+                Maybe::<u32>::schema().layout(),
+                std::alloc::Layout::new::<Maybe<u32>>()
+            );
+            assert_eq!(
+                Maybe::<u64>::schema().layout(),
+                std::alloc::Layout::new::<Maybe<u64>>()
+            );
+            assert_eq!(
+                F::<u64>::schema().layout(),
+                std::alloc::Layout::new::<F<u64>>()
+            );
+            assert_eq!(
+                F::<u32>::schema().layout(),
+                std::alloc::Layout::new::<F<u32>>()
+            );
+
+            // Check a normal enum too, just in case.
+            assert_eq!(E::schema().layout(), std::alloc::Layout::new::<E>());
+        }
+
+        // Makes sure that the layout reported for two structs, where the only difference between
+        // them is the `#[repr(C)]` annotation, matches.
+        #[test]
+        fn schema_layout_for_repr_c_matches_repr_rust() {
+            assert_eq!(A::schema().layout(), B::schema().layout());
+            assert_eq!(C::schema().layout(), D::schema().layout());
+        }
+    }
+}