process signage_point_index as u32 to prevent overflows

Chia-Network · Jan 27, 2025 · 4294ff4 · 4294ff4
1 parent 638943b
commit 4294ff4
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 31 deletions.
diff --git a/crates/chia-protocol/src/block_record.rs b/crates/chia-protocol/src/block_record.rs
@@ -165,7 +165,7 @@ impl BlockRecord {
         calculate_sp_iters(
             num_sps_sub_slot,
             self.sub_slot_iters,
-            self.signage_point_index,
+            self.signage_point_index as u32,
         )
     }
 
@@ -178,7 +178,7 @@ impl BlockRecord {
             num_sps_sub_slot,
             num_sp_intervals_extra,
             self.sub_slot_iters,
-            self.signage_point_index,
+            self.signage_point_index as u32,
             self.required_iters,
         )
     }

diff --git a/crates/chia-protocol/src/pot_iterations.rs b/crates/chia-protocol/src/pot_iterations.rs
@@ -8,12 +8,12 @@
 pub fn is_overflow_block(
     num_sps_sub_slot: u32,
     num_sp_intervals_extra: u8,
-    signage_point_index: u8,
+    signage_point_index: u32,
 ) -> pyo3::PyResult<bool> {
-    if signage_point_index as u32 >= num_sps_sub_slot {
+    if signage_point_index >= num_sps_sub_slot {
         return Err(pyo3::exceptions::PyValueError::new_err("SP index too high"));
     }
-    Ok(signage_point_index as u32 >= num_sps_sub_slot - num_sp_intervals_extra as u32)
+    Ok(signage_point_index >= num_sps_sub_slot - num_sp_intervals_extra as u32)
 }
 
 #[cfg(feature = "py-bindings")]
@@ -35,9 +35,9 @@ pub fn calculate_sp_interval_iters(
 pub fn calculate_sp_iters(
     num_sps_sub_slot: u32,
     sub_slot_iters: u64,
-    signage_point_index: u8,
+    signage_point_index: u32,
 ) -> pyo3::PyResult<u64> {
-    if signage_point_index as u32 >= num_sps_sub_slot {
+    if signage_point_index >= num_sps_sub_slot {
         return Err(pyo3::exceptions::PyValueError::new_err("SP index too high"));
     }
     Ok(calculate_sp_interval_iters(num_sps_sub_slot, sub_slot_iters)? * signage_point_index as u64)
@@ -49,7 +49,7 @@ pub fn calculate_ip_iters(
     num_sps_sub_slot: u32,
     num_sp_intervals_extra: u8,
     sub_slot_iters: u64,
-    signage_point_index: u8,
+    signage_point_index: u32,
     required_iters: u64,
 ) -> pyo3::PyResult<u64> {
     let sp_interval_iters = calculate_sp_interval_iters(num_sps_sub_slot, sub_slot_iters)?;
@@ -73,7 +73,6 @@ pub fn calculate_ip_iters(
 mod tests {
     use super::*;
     static NUM_SPS_SUB_SLOT: u32 = 32;
-    static NUM_SPS_SUB_SLOT_U8: u8 = 32;
     static NUM_SP_INTERVALS_EXTRA: u8 = 3;
 
     #[test]
@@ -127,27 +126,34 @@ mod tests {
         let sp_iters = sp_interval_iters * 13;
 
         // required_iters too high
-        // disabled this test as rusts typing enforces it already
-        // assert!(matches!(
-        //     calculate_ip_iters(
-        //         NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), sp_interval_iters
-        //     ),
-        //     Err(_)
-        // ));
-
-        // // required_iters too high
-        // assert!(matches!(
-        //     calculate_ip_iters(
-        //         NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), sp_interval_iters * 12
-        //     ),
-        //     Err(_)
-        // ));
-
-        // // required_iters too low (0)
-        // assert!(matches!(
-        //     calculate_ip_iters(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), 0),
-        //     Err(_)
-        // ));
+        assert!(calculate_ip_iters(
+            NUM_SPS_SUB_SLOT,
+            NUM_SP_INTERVALS_EXTRA,
+            ssi,
+            sp_interval_iters.try_into().unwrap(),
+            sp_interval_iters
+        )
+        .is_err());
+
+        // required_iters too high
+        assert!(calculate_ip_iters(
+            NUM_SPS_SUB_SLOT,
+            NUM_SP_INTERVALS_EXTRA,
+            ssi,
+            sp_interval_iters.try_into().unwrap(),
+            sp_interval_iters * 12
+        )
+        .is_err());
+
+        // required_iters too low (0)
+        assert!(calculate_ip_iters(
+            NUM_SPS_SUB_SLOT,
+            NUM_SP_INTERVALS_EXTRA,
+            ssi,
+            sp_interval_iters.try_into().unwrap(),
+            0
+        )
+        .is_err());
 
         let required_iters = sp_interval_iters - 1;
         let ip_iters = calculate_ip_iters(
@@ -198,7 +204,7 @@ mod tests {
             NUM_SPS_SUB_SLOT,
             NUM_SP_INTERVALS_EXTRA,
             ssi,
-            NUM_SPS_SUB_SLOT_U8 - 1_u8,
+            NUM_SPS_SUB_SLOT - 1_u32,
             required_iters,
         )
         .expect("valid");