Merge pull request #83 from ralexstokes/address-option

error on extra input when deserializing `Union[None]` for some SSZ union type
ralexstokes · Jul 8, 2023 · 197c460 · 197c460
2 parents 885a887 + 39d1bb7
commit 197c460
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/ssz-rs-derive/src/lib.rs b/ssz-rs-derive/src/lib.rs
@@ -285,7 +285,15 @@ fn derive_deserialize_impl(data: &Data) -> TokenStream {
                         }
                         Fields::Unit => {
                             quote_spanned! { variant.span() =>
-                                0 => Ok(Self::None),
+                                0 => {
+                                    if encoding.len() != 1 {
+                                        return Err(DeserializeError::AdditionalInput {
+                                            provided: encoding.len(),
+                                            expected: 1,
+                                        })
+                                    }
+                                    Ok(Self::None)
+                                },
                             }
                         }
                         _ => unreachable!(),

diff --git a/ssz-rs/src/union.rs b/ssz-rs/src/union.rs
@@ -12,6 +12,7 @@ use crate::{
 ///     None,
 ///     Some(T),
 /// }
+/// The SSZ schema for this value would be `Union[None, T]`.
 impl<T: SimpleSerialize> Sized for Option<T> {
     fn is_variable_size() -> bool {
         true
@@ -49,7 +50,15 @@ where
 
         // SAFETY: index is safe because encoding is not empty; qed
         match encoding[0] {
-            0 => Ok(None),
+            0 => {
+                if encoding.len() != 1 {
+                    return Err(DeserializeError::AdditionalInput {
+                        provided: encoding.len(),
+                        expected: 1,
+                    })
+                }
+                Ok(None)
+            }
             1 => {
                 // SAFETY: index is safe because encoding is not empty; qed
                 let inner = T::deserialize(&encoding[1..])?;
@@ -170,6 +179,17 @@ mod tests {
         assert_eq!(x, recovered);
     }
 
+    #[test]
+    fn test_options_with_extra_input() {
+        let buffer = vec![0u8, 123, 234];
+        let result = Option::<u8>::deserialize(&buffer);
+        assert!(matches!(result, Err(DeserializeError::AdditionalInput { .. })));
+
+        let buffer = vec![0u8, 123, 234];
+        let result = AnotherOption::deserialize(&buffer);
+        assert!(matches!(result, Err(DeserializeError::AdditionalInput { .. })));
+    }
+
     #[test]
     fn test_another_option() {
         let mut x = AnotherOption::A(12u8);