Validate struct #[canonical(prefix)] when decoding

FuelLabs · Aug 29, 2023 · 4af7c63 · 4af7c63
1 parent f96ef89
commit 4af7c63
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/fuel-derive/src/attribute.rs b/fuel-derive/src/attribute.rs
@@ -56,7 +56,7 @@ fn parse_attrs(s: &synstructure::Structure) -> HashMap<String, TokenStream> {
 
 /// Pop-level `canonical` attributes for a struct
 pub struct StructAttrs {
-    /// The struct is prefixed with the given byte.
+    /// The struct is prefixed with the given word.
     /// Useful with`#[canonical(inner_discriminant)]`.
     pub prefix: Option<TokenStream>,
 }

diff --git a/fuel-derive/src/deserialize.rs b/fuel-derive/src/deserialize.rs
@@ -40,10 +40,13 @@ fn deserialize_struct(s: &mut synstructure::Structure) -> TokenStream2 {
         }
     });
 
-    let remove_prefix = if StructAttrs::parse(s).prefix.is_some() {
-        quote! {
-            <u64 as ::fuel_types::canonical::Deserialize>::decode_static(buffer)?;
-        }
+    let remove_prefix = if let Some(expected_prefix) = StructAttrs::parse(s).prefix {
+        quote! {{
+            let prefix = <u64 as ::fuel_types::canonical::Deserialize>::decode_static(buffer)?;
+            if prefix.try_into() != Ok(#expected_prefix) {
+                return ::core::result::Result::Err(::fuel_types::canonical::Error::InvalidPrefix)
+            }
+        }}
     } else {
         quote! {}
     };

diff --git a/fuel-types/src/canonical.rs b/fuel-types/src/canonical.rs
@@ -23,6 +23,8 @@ pub enum Error {
     BufferIsTooShort,
     /// Got unknown enum's discriminant.
     UnknownDiscriminant,
+    /// Struct prefix was invalid.
+    InvalidPrefix,
     /// Wrong align.
     WrongAlign,
     /// Allocation too large to be correct.