improve safety

Author: lukesandberg

turbopack/crates/turbo-tasks/src/id.rs

@@ -41,9 +41,8 @@ macro_rules! define_id {
             }
             /// Constructs a wrapper type from the numeric identifier.
             ///
-            /// # Safety
-            ///
-            /// The passed `id` must not be zero.
+            /// Returns `None` if the provided `id` is zero, otherwise returns
+            /// `Some(Self)` containing the wrapped non-zero identifier.
             pub fn new(id: $primitive) -> Option<Self> {
                 NonZero::<$primitive>::new(id).map(|id| Self{id})
             }
@@ -161,7 +160,7 @@ impl TaskId {
 }
 
 macro_rules! make_serializable {
-    ($ty:ty, $get_object:path, $visitor_name:ident) => {
+    ($ty:ty, $get_object:path, $validate_type_id:path, $visitor_name:ident) => {
         impl Serialize for $ty {
             fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
             where
@@ -193,8 +192,16 @@ macro_rules! make_serializable {
             where
                 E: serde::de::Error,
             {
-                Self::Value::new(v)
-                    .ok_or_else(|| E::unknown_variant(&format!("{v}"), &["a non zero u32"]))
+                match Self::Value::new(v) {
+                    Some(value) => {
+                        if let Some(error) = $validate_type_id(value) {
+                            Err(E::custom(error))
+                        } else {
+                            Ok(value)
+                        }
+                    }
+                    None => Err(E::unknown_variant(&format!("{v}"), &["a non zero u32"])),
+                }
             }
         }
 
@@ -209,10 +216,21 @@ macro_rules! make_serializable {
     };
 }
 
-make_serializable!(ValueTypeId, registry::get_value_type, ValueTypeVisitor);
-make_serializable!(TraitTypeId, registry::get_trait, TraitTypeVisitor);
+make_serializable!(
+    ValueTypeId,
+    registry::get_value_type,
+    registry::validate_value_type_id,
+    ValueTypeVisitor
+);
+make_serializable!(
+    TraitTypeId,
+    registry::get_trait,
+    registry::validate_trait_type_id,
+    TraitTypeVisitor
+);
 make_serializable!(
     FunctionId,
     registry::get_native_function,
+    registry::validate_function_id,
     FunctionTypeVisitor
 );

turbopack/crates/turbo-tasks/src/registry.rs

@@ -1,5 +1,6 @@
 use std::num::NonZeroU32;
 
+use anyhow::Error;
 use once_cell::sync::Lazy;
 use rustc_hash::{FxHashMap, FxHashSet};
 
@@ -52,6 +53,17 @@ pub fn get_function_id(func: &'static NativeFunction) -> FunctionId {
         .expect("function isn't registered")
 }
 
+pub fn validate_function_id(id: FunctionId) -> Option<Error> {
+    let len = FUNCTIONS.id_to_value.len();
+    if *id as usize <= len {
+        None
+    } else {
+        Some(anyhow::anyhow!(
+            "Invalid function type id, {id} expected a value <= {len}"
+        ))
+    }
+}
+
 struct Values {
     id_to_value: Box<[&'static ValueType]>,
     value_to_id: FxHashMap<&'static ValueType, ValueTypeId>,
@@ -100,6 +112,17 @@ pub fn get_value_type(id: ValueTypeId) -> &'static ValueType {
     VALUES.id_to_value[*id as usize - 1]
 }
 
+pub fn validate_value_type_id(id: ValueTypeId) -> Option<Error> {
+    let len = VALUES.id_to_value.len();
+    if *id as usize <= len {
+        None
+    } else {
+        Some(anyhow::anyhow!(
+            "Invalid value type id, {id} expected a value <= {len}"
+        ))
+    }
+}
+
 struct Traits {
     id_to_trait: Box<[&'static TraitType]>,
     trait_to_id: FxHashMap<&'static TraitType, TraitTypeId>,
@@ -144,3 +167,13 @@ pub fn get_trait_type_id(trait_type: &'static TraitType) -> TraitTypeId {
 pub fn get_trait(id: TraitTypeId) -> &'static TraitType {
     TRAITS.id_to_trait[*id as usize - 1]
 }
+pub fn validate_trait_type_id(id: TraitTypeId) -> Option<Error> {
+    let len = TRAITS.id_to_trait.len();
+    if *id as usize <= len {
+        None
+    } else {
+        Some(anyhow::anyhow!(
+            "Invalid trait type id, {id} expected a value <= {len}"
+        ))
+    }
+}