From 03956a9a24abd06146c5fa37effd85955e4cc57c Mon Sep 17 00:00:00 2001
From: Nicolas Sarlin <nicolas.sarlin@zama.ai>
Date: Fri, 13 Dec 2024 10:56:49 +0100
Subject: [PATCH] chore(zk): check that k <= d for zk crs

---
 tfhe-zk-pok/src/proofs/mod.rs    | 1 +
 tfhe-zk-pok/src/proofs/pke_v2.rs | 8 +++++++-
 tfhe/src/zk/mod.rs               | 4 ++++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tfhe-zk-pok/src/proofs/mod.rs b/tfhe-zk-pok/src/proofs/mod.rs
index f3cd1b5305..fceb97ca22 100644
--- a/tfhe-zk-pok/src/proofs/mod.rs
+++ b/tfhe-zk-pok/src/proofs/mod.rs
@@ -153,6 +153,7 @@ fn assert_pke_proof_preconditions(
     big_d: usize,
     big_d_max: usize,
 ) {
+    assert!(k_max <= d);
     assert_eq!(c1.len(), d);
     assert_eq!(e1.len(), d);
 
diff --git a/tfhe-zk-pok/src/proofs/pke_v2.rs b/tfhe-zk-pok/src/proofs/pke_v2.rs
index 4b3d4d374e..0d743fb2d4 100644
--- a/tfhe-zk-pok/src/proofs/pke_v2.rs
+++ b/tfhe-zk-pok/src/proofs/pke_v2.rs
@@ -497,6 +497,12 @@ pub fn compute_crs_params(
     msbs_zero_padding_bit_count: u64,
     bound_type: Bound,
 ) -> (usize, usize, u128, usize) {
+    assert!(
+        k <= d,
+        "Invalid parameters for zk_pok, the maximum number of messages k should be smaller \
+than the lwe dimension d. Please pick a smaller k: k = {k}, d = {d}"
+    );
+
     let mut B_bound_squared = {
         (match bound_type {
             // GHL factor is 9.75, 9.75**2 = 95.0625
@@ -526,7 +532,7 @@ Please select a smaller B, d and/or k"
     // safely used for this
     assert!(
         m_bound <= 64,
-        "Invalid parameters for zk_pok, w e only support 64 bits integer. \
+        "Invalid parameters for zk_pok, we only support 64 bits integer. \
 The computed m parameter is {m_bound} > 64. Please select a smaller B, d and/or k"
     );
 
diff --git a/tfhe/src/zk/mod.rs b/tfhe/src/zk/mod.rs
index e039f0655e..760489d6e7 100644
--- a/tfhe/src/zk/mod.rs
+++ b/tfhe/src/zk/mod.rs
@@ -298,6 +298,10 @@ impl CompactPkeCrs {
         Scalar: UnsignedInteger + CastInto<u64> + Debug,
         NoiseDistribution: BoundedDistribution<Scalar::Signed>,
     {
+        if max_num_cleartext.0 > lwe_dim.0 {
+            return Err("Maximum number of cleartexts is greater than the lwe dimension".into());
+        }
+
         let noise_bound = match zk_scheme {
             CompactPkeZkScheme::V1 => Self::compute_bound_v1(noise_distribution)?,
             CompactPkeZkScheme::V2 => Self::compute_bound_v2(noise_distribution)?,