count type table cost

dfinity · Feb 16, 2024 · d81d25c · d81d25c
1 parent ab91b4c
commit d81d25c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rust/candid/src/de.rs b/rust/candid/src/de.rs
@@ -29,13 +29,14 @@ pub struct IDLDeserialize<'de> {
 impl<'de> IDLDeserialize<'de> {
     /// Create a new deserializer with IDL binary message.
     pub fn new(bytes: &'de [u8]) -> Result<Self> {
-        let de = Deserializer::from_bytes(bytes).with_context(|| {
+        let mut de = Deserializer::from_bytes(bytes).with_context(|| {
             if bytes.len() <= 500 {
                 format!("Cannot parse header {}", &hex::encode(bytes))
             } else {
                 "Cannot parse header".to_string()
             }
         })?;
+        de.add_cost(de.input.position() as usize * 2)?;
         Ok(IDLDeserialize { de })
     }
     /// Create a new deserializer with IDL binary message. The config is used to adjust some parameters in the deserializer.
@@ -49,6 +50,7 @@ impl<'de> IDLDeserialize<'de> {
         })?;
         de.decoding_cost = config.decoding_cost;
         de.full_error_message = config.full_error_message;
+        de.add_cost(de.input.position() as usize * 2)?;
         Ok(IDLDeserialize { de })
     }
     /// Deserialize one value from deserializer.