From d98033c71df9dfecb5253d60d61dcf73cfb581e9 Mon Sep 17 00:00:00 2001
From: Nicolas Sarlin <nicolas.sarlin@zama.ai>
Date: Thu, 4 Dec 2025 11:19:29 +0100
Subject: [PATCH] fix(integer): check overflows when computing expected list
 size

---
 tfhe/src/integer/ciphertext/compact_list.rs | 30 +++++++++++++++------
 tfhe/src/integer/ciphertext/utils.rs        | 15 +++++++++++
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/tfhe/src/integer/ciphertext/compact_list.rs b/tfhe/src/integer/ciphertext/compact_list.rs
index 2feafd9e6..c6defbdd8 100644
--- a/tfhe/src/integer/ciphertext/compact_list.rs
+++ b/tfhe/src/integer/ciphertext/compact_list.rs
@@ -430,10 +430,10 @@ impl ParameterSetConformant for CompactCiphertextList {
             return false;
         }
 
-        let total_expected_num_blocks: usize = info
-            .iter()
-            .map(|a| a.num_blocks(self.message_modulus()))
-            .sum();
+        let Ok(total_expected_num_blocks) = self.expected_num_block() else {
+            // Return early if the sum overflowed
+            return false;
+        };
 
         let total_expected_lwe_count = if is_packed {
             total_expected_num_blocks.div_ceil(2)
@@ -1001,6 +1001,13 @@ impl CompactCiphertextList {
     pub fn message_modulus(&self) -> MessageModulus {
         self.ct_list.message_modulus
     }
+
+    /// Computes the expected number of blocks based on the `info` metadata.
+    ///
+    /// Returns an error if the sum overflows
+    fn expected_num_block(&self) -> Result<usize, ()> {
+        DataKind::total_block_count(&self.info, self.message_modulus())
+    }
 }
 
 #[cfg(feature = "zk-pok")]
@@ -1127,6 +1134,13 @@ impl ProvenCompactCiphertextList {
     pub fn get_kind_of(&self, index: usize) -> Option<DataKind> {
         self.info.get(index).copied()
     }
+
+    /// Computes the expected number of blocks based on the `info` metadata.
+    ///
+    /// Returns an error if the sum overflows
+    fn expected_num_block(&self) -> Result<usize, ()> {
+        DataKind::total_block_count(&self.info, self.message_modulus())
+    }
 }
 
 #[cfg(feature = "zk-pok")]
@@ -1233,10 +1247,10 @@ impl ParameterSetConformant for ProvenCompactCiphertextList {
             return false;
         }
 
-        let total_expected_num_blocks: usize = info
-            .iter()
-            .map(|a| a.num_blocks(self.message_modulus()))
-            .sum();
+        let Ok(total_expected_num_blocks) = self.expected_num_block() else {
+            // Return early if the sum overflowed
+            return false;
+        };
 
         let total_expected_lwe_count = if is_packed {
             total_expected_num_blocks.div_ceil(2)
diff --git a/tfhe/src/integer/ciphertext/utils.rs b/tfhe/src/integer/ciphertext/utils.rs
index 8d0273eb3..57171e8a3 100644
--- a/tfhe/src/integer/ciphertext/utils.rs
+++ b/tfhe/src/integer/ciphertext/utils.rs
@@ -30,6 +30,21 @@ impl DataKind {
             }
         }
     }
+
+    pub(crate) fn total_block_count(
+        info: &[Self],
+        message_modulus: MessageModulus,
+    ) -> Result<usize, ()> {
+        if message_modulus.0 == 0 {
+            return Err(());
+        }
+
+        info.iter()
+            .try_fold(0usize, |acc, &x| {
+                acc.checked_add(x.num_blocks(message_modulus))
+            })
+            .ok_or(())
+    }
 }
 
 pub trait Expandable: Sized {