From 95734b695c39bf8c3aa3636d009ddb51744ef191 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Tue, 25 Feb 2025 10:15:27 +0100
Subject: [PATCH] fixed potential buffer overflow & added
 vec_znx_normalize_tmp_bytes

---
 base2k/src/vec_znx.rs | 9 +++++++++
 base2k/src/vmp.rs     | 1 +
 2 files changed, 10 insertions(+)

diff --git a/base2k/src/vec_znx.rs b/base2k/src/vec_znx.rs
index 540d120..88791c4 100644
--- a/base2k/src/vec_znx.rs
+++ b/base2k/src/vec_znx.rs
@@ -469,6 +469,7 @@ pub fn rsh<T: VecZnxCommon>(log_base2k: usize, a: &mut T, k: usize, carry: &mut
 pub trait VecZnxCommon: VecZnxApi + Infos {}
 
 pub trait VecZnxOps {
+
     /// Allocates a new [VecZnx].
     ///
     /// # Arguments
@@ -480,6 +481,8 @@ pub trait VecZnxOps {
     /// a new [VecZnx] through [VecZnx::from_bytes].
     fn bytes_of_vec_znx(&self, cols: usize) -> usize;
 
+    fn vec_znx_normalize_tmp_bytes(&self) -> usize;
+
     /// c <- a + b.
     fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
         &self,
@@ -557,6 +560,12 @@ impl VecZnxOps for Module {
         self.n() * cols * 8
     }
 
+    fn vec_znx_normalize_tmp_bytes(&self) -> usize{
+        unsafe{
+            vec_znx::vec_znx_normalize_tmp_bytes(self.0) as usize
+        }   
+    }
+
     // c <- a + b
     fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
         &self,
diff --git a/base2k/src/vmp.rs b/base2k/src/vmp.rs
index 39be0ee..d572e11 100644
--- a/base2k/src/vmp.rs
+++ b/base2k/src/vmp.rs
@@ -439,6 +439,7 @@ impl VmpPMatOps for Module {
     }
 
     fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, buf: &mut [u8]) {
+        debug_assert!(a.len() == b.cols() * self.n());
         unsafe {
             vmp::vmp_prepare_row(
                 self.0,