From 39d1bb70119973f4737ea743471175445b7e41c2 Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Tue, 27 Jun 2023 11:03:43 -0600
Subject: [PATCH] error on extra input when deserializing `Union[None]` for
 some SSZ union type

---
 ssz-rs-derive/src/lib.rs | 10 +++++++++-
 ssz-rs/src/union.rs      | 22 +++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/ssz-rs-derive/src/lib.rs b/ssz-rs-derive/src/lib.rs
index f413897e..1fe70901 100644
--- a/ssz-rs-derive/src/lib.rs
+++ b/ssz-rs-derive/src/lib.rs
@@ -304,7 +304,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
index 7bee0afe..5f84897b 100644
--- 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);