Fixed CostModel range check for different versions of plutus

Author: vsubhuman

rust/src/plutus.rs

@@ -200,8 +200,6 @@ impl ConstrPlutusData {
     }
 }
 
-const COST_MODEL_OP_COUNT: usize = 166;
-
 #[wasm_bindgen]
 #[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
 pub struct CostModel(Vec<Int>);
@@ -210,26 +208,35 @@ to_from_bytes!(CostModel);
 
 #[wasm_bindgen]
 impl CostModel {
+
+    #[deprecated(since = "11.0.0", note="Use `.new_with_length`")]
     pub fn new() -> Self {
-        let mut costs = Vec::with_capacity(COST_MODEL_OP_COUNT);
-        for _ in 0 .. COST_MODEL_OP_COUNT {
+        Self::new_with_length(TxBuilderConstants::PLUTUS_V1_OP_COUNT)
+    }
+
+    /// Creates a new CostModel instance with the specified underlying length
+    pub fn new_with_length(len: usize) -> Self {
+        let mut costs = Vec::with_capacity(len);
+        for _ in 0 .. len {
             costs.push(Int::new_i32(0));
         }
         Self(costs)
     }
 
     pub fn set(&mut self, operation: usize, cost: &Int) -> Result<Int, JsError> {
-        if operation >= COST_MODEL_OP_COUNT {
-            return Err(JsError::from_str(&format!("CostModel operation {} out of bounds. Max is {}", operation, COST_MODEL_OP_COUNT)));
+        let max = self.0.len();
+        if operation >= max {
+            return Err(JsError::from_str(&format!("CostModel operation {} out of bounds. Max is {}", operation, max)));
         }
         let old = self.0[operation].clone();
         self.0[operation] = cost.clone();
         Ok(old)
     }
 
     pub fn get(&self, operation: usize) -> Result<Int, JsError> {
-        if operation >= COST_MODEL_OP_COUNT {
-            return Err(JsError::from_str(&format!("CostModel operation {} out of bounds. Max is {}", operation, COST_MODEL_OP_COUNT)));
+        let max = self.0.len();
+        if operation >= max {
+            return Err(JsError::from_str(&format!("CostModel operation {} out of bounds. Max is {}", operation, max)));
         }
         Ok(self.0[operation].clone())
     }
@@ -794,6 +801,7 @@ impl Strings {
 // Serialization
 
 use std::io::{SeekFrom};
+use crate::tx_builder_constants::TxBuilderConstants;
 
 
 impl cbor_event::se::Serialize for PlutusScript {
@@ -898,7 +906,7 @@ impl Deserialize for ConstrPlutusData {
 
 impl cbor_event::se::Serialize for CostModel {
     fn serialize<'se, W: Write>(&self, serializer: &'se mut Serializer<W>) -> cbor_event::Result<&'se mut Serializer<W>> {
-        serializer.write_array(cbor_event::Len::Len(COST_MODEL_OP_COUNT as u64))?;
+        serializer.write_array(cbor_event::Len::Len(self.0.len() as u64))?;
         for cost in &self.0 {
             cost.serialize(serializer)?;
         }
@@ -918,10 +926,12 @@ impl Deserialize for CostModel {
                 }
                 arr.push(Int::deserialize(raw)?);
             }
-            if arr.len() != COST_MODEL_OP_COUNT {
+            let min = TxBuilderConstants::PLUTUS_V1_OP_COUNT;
+            let max = TxBuilderConstants::PLUTUS_V2_OP_COUNT;
+            if arr.len() != min && arr.len() != max {
                 return Err(DeserializeFailure::OutOfRange{
-                    min: COST_MODEL_OP_COUNT,
-                    max: COST_MODEL_OP_COUNT,
+                    min,
+                    max,
                     found: arr.len()
                 }.into());
             }
@@ -1574,4 +1584,10 @@ mod tests {
             Language::new_plutus_v2(),
         );
     }
+
+    #[test]
+    fn test_cost_model_roundtrip() {
+        let costmodels = TxBuilderConstants::plutus_vasil_cost_models();
+        assert_eq!(costmodels, Costmdls::from_bytes(costmodels.to_bytes()).unwrap());
+    }
 }

rust/src/tx_builder.rs

@@ -1593,7 +1593,6 @@ impl TransactionBuilder {
 
 #[cfg(test)]
 mod tests {
-    use itertools::Itertools;
     use super::*;
     use fees::*;
     use crate::fakes::{fake_base_address, fake_bytes_32, fake_key_hash, fake_tx_hash, fake_tx_input, fake_tx_input2, fake_value, fake_value2};
@@ -5901,7 +5900,7 @@ mod tests {
         let fake_out = TransactionOutput::new(&collateral_return_address, &possible_ret);
         let min_ada = min_ada_for_output(&fake_out, &tx_builder.config.utxo_cost()).unwrap();
 
-        let mut total_collateral_value = to_bignum(collateral_input_value).checked_sub(&min_ada).unwrap();
+        let total_collateral_value = to_bignum(collateral_input_value).checked_sub(&min_ada).unwrap();
 
         tx_builder.set_total_collateral_and_return(
             &total_collateral_value,
@@ -5923,8 +5922,6 @@ mod tests {
         let mut tx_builder = create_reallistic_tx_builder();
         tx_builder.set_fee(&to_bignum(123456));
 
-        let masset = fake_multiasset(123);
-
         let mut inp = TxInputsBuilder::new();
         let collateral_input_value = 2_000_000;
         inp.add_input(

rust/src/tx_builder_constants.rs

@@ -12,6 +12,9 @@ pub struct TxBuilderConstants();
 #[wasm_bindgen]
 impl TxBuilderConstants {
 
+    pub const PLUTUS_V1_OP_COUNT: usize = 166;
+    pub const PLUTUS_V2_OP_COUNT: usize = 175;
+
     pub fn plutus_default_cost_models() -> Costmdls {
         TxBuilderConstants::plutus_alonzo_cost_models()
     }