Merge pull request #185 from sosthene-nitrokey/wide-recursive

Add faillible version of size_hint to properly handle recursive structures
rust-fuzz · Sep 20, 2024 · 1cc0e46 · 1cc0e46
2 parents 4254bd2 + b6991ce
commit 1cc0e46
Show file tree

Hide file tree

Showing 16 changed files with 385 additions and 81 deletions.
diff --git a/derive/src/lib.rs b/derive/src/lib.rs
@@ -364,27 +364,27 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
                 determine_field_constructor(f).map(|field_constructor| {
                     match field_constructor {
                         FieldConstructor::Default | FieldConstructor::Value(_) => {
-                            quote!((0, Some(0)))
+                            quote!(Ok((0, Some(0))))
                         }
                         FieldConstructor::Arbitrary => {
-                            quote! { <#ty as arbitrary::Arbitrary>::size_hint(depth) }
+                            quote! { <#ty as arbitrary::Arbitrary>::try_size_hint(depth) }
                         }
 
                         // Note that in this case it's hard to determine what size_hint must be, so size_of::<T>() is
                         // just an educated guess, although it's gonna be inaccurate for dynamically
                         // allocated types (Vec, HashMap, etc.).
                         FieldConstructor::With(_) => {
-                            quote! { (::core::mem::size_of::<#ty>(), None) }
+                            quote! { Ok((::core::mem::size_of::<#ty>(), None)) }
                         }
                     }
                 })
             })
             .collect::<Result<Vec<TokenStream>>>()
             .map(|hints| {
                 quote! {
-                    arbitrary::size_hint::and_all(&[
-                        #( #hints ),*
-                    ])
+                    Ok(arbitrary::size_hint::and_all(&[
+                        #( #hints? ),*
+                    ]))
                 }
             })
     };
@@ -393,7 +393,12 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
             quote! {
                 #[inline]
                 fn size_hint(depth: usize) -> (usize, ::core::option::Option<usize>) {
-                    arbitrary::size_hint::recursion_guard(depth, |depth| #hint)
+                    Self::try_size_hint(depth).unwrap_or_default()
+                }
+
+                #[inline]
+                fn try_size_hint(depth: usize) -> Result<(usize, ::core::option::Option<usize>), arbitrary::MaxRecursionReached> {
+                    arbitrary::size_hint::try_recursion_guard(depth, |depth| #hint)
                 }
             }
         })
@@ -413,14 +418,17 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
             .collect::<Result<Vec<TokenStream>>>()
             .map(|variants| {
                 quote! {
-                    #[inline]
                     fn size_hint(depth: usize) -> (usize, ::core::option::Option<usize>) {
-                        arbitrary::size_hint::and(
-                            <u32 as arbitrary::Arbitrary>::size_hint(depth),
-                            arbitrary::size_hint::recursion_guard(depth, |depth| {
-                                arbitrary::size_hint::or_all(&[ #( #variants ),* ])
-                            }),
-                        )
+                        Self::try_size_hint(depth).unwrap_or_default()
+                    }
+                    #[inline]
+                    fn try_size_hint(depth: usize) -> Result<(usize, ::core::option::Option<usize>), arbitrary::MaxRecursionReached> {
+                        Ok(arbitrary::size_hint::and(
+                            <u32 as arbitrary::Arbitrary>::try_size_hint(depth)?,
+                            arbitrary::size_hint::try_recursion_guard(depth, |depth| {
+                                Ok(arbitrary::size_hint::or_all(&[ #( #variants? ),* ]))
+                            })?,
+                        ))
                     }
                 }
             }),

diff --git a/src/foreign/alloc/borrow.rs b/src/foreign/alloc/borrow.rs
@@ -14,8 +14,13 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::recursion_guard(depth, |depth| {
-            <<A as ToOwned>::Owned as Arbitrary>::size_hint(depth)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
+        size_hint::try_recursion_guard(depth, |depth| {
+            <<A as ToOwned>::Owned as Arbitrary>::try_size_hint(depth)
         })
     }
 }
diff --git a/src/foreign/alloc/boxed.rs b/src/foreign/alloc/boxed.rs
@@ -13,7 +13,12 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
+        size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
     }
 }
 

diff --git a/src/foreign/alloc/rc.rs b/src/foreign/alloc/rc.rs
@@ -13,7 +13,12 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
+        size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
     }
 }
 

diff --git a/src/foreign/alloc/sync.rs b/src/foreign/alloc/sync.rs
@@ -13,7 +13,12 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
+        size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
     }
 }
 

diff --git a/src/foreign/core/array.rs b/src/foreign/core/array.rs
@@ -64,9 +64,13 @@ where
     }
 
     #[inline]
-    fn size_hint(d: usize) -> (usize, Option<usize>) {
-        size_hint::and_all(&array::from_fn::<_, N, _>(|_| {
-            <T as Arbitrary>::size_hint(d)
-        }))
+    fn size_hint(depth: usize) -> (usize, Option<usize>) {
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
+        let hint = <T as Arbitrary>::try_size_hint(depth)?;
+        Ok(size_hint::and_all(&array::from_fn::<_, N, _>(|_| hint)))
     }
 }
diff --git a/src/foreign/core/cell.rs b/src/foreign/core/cell.rs
@@ -1,5 +1,5 @@
 use {
-    crate::{Arbitrary, Result, Unstructured},
+    crate::{Arbitrary, MaxRecursionReached, Result, Unstructured},
     core::cell::{Cell, RefCell, UnsafeCell},
 };
 
@@ -13,7 +13,12 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        <A as Arbitrary<'a>>::size_hint(depth)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        <A as Arbitrary<'a>>::try_size_hint(depth)
     }
 }
 
@@ -27,7 +32,12 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        <A as Arbitrary<'a>>::size_hint(depth)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        <A as Arbitrary<'a>>::try_size_hint(depth)
     }
 }
 
@@ -41,6 +51,11 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        <A as Arbitrary<'a>>::size_hint(depth)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        <A as Arbitrary<'a>>::try_size_hint(depth)
     }
 }
diff --git a/src/foreign/core/num.rs b/src/foreign/core/num.rs
@@ -1,5 +1,5 @@
 use {
-    crate::{Arbitrary, Error, Result, Unstructured},
+    crate::{Arbitrary, Error, MaxRecursionReached, Result, Unstructured},
     core::{
         mem,
         num::{
@@ -131,6 +131,11 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        <A as Arbitrary<'a>>::size_hint(depth)
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        <A as Arbitrary<'a>>::try_size_hint(depth)
     }
 }
diff --git a/src/foreign/core/ops.rs b/src/foreign/core/ops.rs
@@ -1,5 +1,5 @@
 use {
-    crate::{size_hint, Arbitrary, Result, Unstructured},
+    crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured},
     core::{
         mem,
         ops::{Bound, Range, RangeBounds, RangeFrom, RangeInclusive, RangeTo, RangeToInclusive},
@@ -25,53 +25,57 @@ macro_rules! impl_range {
 
             #[inline]
             fn size_hint(depth: usize) -> (usize, Option<usize>) {
+                Self::try_size_hint(depth).unwrap_or_default()
+            }
+
+            #[inline]
+            fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
                 #[allow(clippy::redundant_closure_call)]
                 $size_hint_closure(depth)
             }
         }
     };
 }
-
 impl_range!(
     Range<A>,
     |r: &Range<A>| (r.start.clone(), r.end.clone()),
     (A, A),
     bounded_range(|(a, b)| a..b),
-    |depth| size_hint::and(
-        <A as Arbitrary>::size_hint(depth),
-        <A as Arbitrary>::size_hint(depth)
-    )
+    |depth| Ok(crate::size_hint::and(
+        <A as Arbitrary>::try_size_hint(depth)?,
+        <A as Arbitrary>::try_size_hint(depth)?,
+    ))
 );
 impl_range!(
     RangeFrom<A>,
     |r: &RangeFrom<A>| r.start.clone(),
     A,
     unbounded_range(|a| a..),
-    |depth| <A as Arbitrary>::size_hint(depth)
+    |depth| <A as Arbitrary>::try_size_hint(depth)
 );
 impl_range!(
     RangeInclusive<A>,
     |r: &RangeInclusive<A>| (r.start().clone(), r.end().clone()),
     (A, A),
     bounded_range(|(a, b)| a..=b),
-    |depth| size_hint::and(
-        <A as Arbitrary>::size_hint(depth),
-        <A as Arbitrary>::size_hint(depth)
-    )
+    |depth| Ok(crate::size_hint::and(
+        <A as Arbitrary>::try_size_hint(depth)?,
+        <A as Arbitrary>::try_size_hint(depth)?,
+    ))
 );
 impl_range!(
     RangeTo<A>,
     |r: &RangeTo<A>| r.end.clone(),
     A,
     unbounded_range(|b| ..b),
-    |depth| <A as Arbitrary>::size_hint(depth)
+    |depth| <A as Arbitrary>::try_size_hint(depth)
 );
 impl_range!(
     RangeToInclusive<A>,
     |r: &RangeToInclusive<A>| r.end.clone(),
     A,
     unbounded_range(|b| ..=b),
-    |depth| <A as Arbitrary>::size_hint(depth)
+    |depth| <A as Arbitrary>::try_size_hint(depth)
 );
 
 pub(crate) fn bounded_range<CB, I, R>(bounds: (I, I), cb: CB) -> R
@@ -110,9 +114,14 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::or(
-            size_hint::and((1, Some(1)), A::size_hint(depth)),
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        Ok(size_hint::or(
+            size_hint::and((1, Some(1)), A::try_size_hint(depth)?),
             (1, Some(1)),
-        )
+        ))
     }
 }
diff --git a/src/foreign/core/option.rs b/src/foreign/core/option.rs
@@ -1,4 +1,4 @@
-use crate::{size_hint, Arbitrary, Result, Unstructured};
+use crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured};
 
 impl<'a, A> Arbitrary<'a> for Option<A>
 where
@@ -14,9 +14,14 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::and(
-            <bool as Arbitrary>::size_hint(depth),
-            size_hint::or((0, Some(0)), <A as Arbitrary>::size_hint(depth)),
-        )
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        Ok(size_hint::and(
+            <bool as Arbitrary>::try_size_hint(depth)?,
+            size_hint::or((0, Some(0)), <A as Arbitrary>::try_size_hint(depth)?),
+        ))
     }
 }
diff --git a/src/foreign/core/result.rs b/src/foreign/core/result.rs
@@ -1,4 +1,4 @@
-use crate::{size_hint, Arbitrary, Error, Unstructured};
+use crate::{size_hint, Arbitrary, Error, MaxRecursionReached, Unstructured};
 
 impl<'a, T, E> Arbitrary<'a> for Result<T, E>
 where
@@ -15,12 +15,17 @@ where
 
     #[inline]
     fn size_hint(depth: usize) -> (usize, Option<usize>) {
-        size_hint::and(
+        Self::try_size_hint(depth).unwrap_or_default()
+    }
+
+    #[inline]
+    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+        Ok(size_hint::and(
             <bool as Arbitrary>::size_hint(depth),
             size_hint::or(
-                <T as Arbitrary>::size_hint(depth),
-                <E as Arbitrary>::size_hint(depth),
+                <T as Arbitrary>::try_size_hint(depth)?,
+                <E as Arbitrary>::try_size_hint(depth)?,
             ),
-        )
+        ))
     }
 }
diff --git a/src/foreign/core/tuple.rs b/src/foreign/core/tuple.rs
@@ -1,4 +1,4 @@
-use crate::{size_hint, Arbitrary, Result, Unstructured};
+use crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured};
 
 macro_rules! arbitrary_tuple {
     () => {};
@@ -23,10 +23,14 @@ macro_rules! arbitrary_tuple {
 
             #[inline]
             fn size_hint(depth: usize) -> (usize, Option<usize>) {
-                size_hint::and_all(&[
-                    <$last as Arbitrary>::size_hint(depth),
-                    $( <$xs as Arbitrary>::size_hint(depth) ),*
-                ])
+                Self::try_size_hint(depth).unwrap_or_default()
+            }
+            #[inline]
+            fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
+                Ok(size_hint::and_all(&[
+                    <$last as Arbitrary>::try_size_hint(depth)?,
+                    $( <$xs as Arbitrary>::try_size_hint(depth)?),*
+                ]))
             }
         }
     };