From c269fbc6eaf4e2361d0c2f1047814c1da3f0d433 Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Fri, 7 Mar 2025 11:29:42 +0100
Subject: [PATCH] chore: pre-generate keyswitching keys for shortint tests

- we run in a cross process race condition which fucks up the key file
- no rust crate seems to help and linux locks are just a fucking mess
- also avoid truncating file when we are going to write to it, get a lock
first
---
 .../examples/utilities/generates_test_keys.rs | 25 ++++++++++---------
 tfhe/src/keycache/mod.rs                      | 22 +++++++++++++---
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/tfhe/examples/utilities/generates_test_keys.rs b/tfhe/examples/utilities/generates_test_keys.rs
index 6edec4d797..4cc0aaa4e1 100644
--- a/tfhe/examples/utilities/generates_test_keys.rs
+++ b/tfhe/examples/utilities/generates_test_keys.rs
@@ -27,6 +27,16 @@ use tfhe::shortint::parameters::{
 };
 use tfhe::shortint::MultiBitPBSParameters;
 
+const KSK_PARAMS: [(
+    ClassicPBSParameters,
+    ClassicPBSParameters,
+    ShortintKeySwitchingParameters,
+); 1] = [(
+    V1_0_PARAM_MESSAGE_1_CARRY_1_KS_PBS_GAUSSIAN_2M128,
+    V1_0_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128,
+    V1_0_PARAM_KEYSWITCH_1_1_KS_PBS_TO_2_2_KS_PBS_GAUSSIAN_2M128,
+)];
+
 fn client_server_keys() {
     let matches = Command::new("test key gen")
         .arg(
@@ -45,6 +55,9 @@ fn client_server_keys() {
         )
         .get_matches();
 
+    // Always generate those as they may be used in the different cases
+    generate_ksk_keys(&KSK_PARAMS);
+
     // If set using the command line flag "--ladner-fischer" this algorithm will be used in
     // additions
     let multi_bit_only: bool = matches.get_flag("multi_bit_only");
@@ -81,18 +94,6 @@ fn client_server_keys() {
             generate_pbs_multi_bit_keys(&MULTI_BIT_PARAMS);
         }
 
-        const KSK_PARAMS: [(
-            ClassicPBSParameters,
-            ClassicPBSParameters,
-            ShortintKeySwitchingParameters,
-        ); 1] = [(
-            V1_0_PARAM_MESSAGE_1_CARRY_1_KS_PBS_GAUSSIAN_2M128,
-            V1_0_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128,
-            V1_0_PARAM_KEYSWITCH_1_1_KS_PBS_TO_2_2_KS_PBS_GAUSSIAN_2M128,
-        )];
-
-        generate_ksk_keys(&KSK_PARAMS);
-
         #[cfg(feature = "experimental")]
         {
             const WOPBS_PARAMS: [(ClassicPBSParameters, WopbsParameters); 1] = [(
diff --git a/tfhe/src/keycache/mod.rs b/tfhe/src/keycache/mod.rs
index 8797322bb9..8d06c41524 100644
--- a/tfhe/src/keycache/mod.rs
+++ b/tfhe/src/keycache/mod.rs
@@ -7,7 +7,7 @@ pub mod utils {
     use fs2::FileExt;
     use serde::de::DeserializeOwned;
     use serde::Serialize;
-    use std::fs::File;
+    use std::fs::{File, OpenOptions};
     use std::io::{BufReader, BufWriter};
     use std::ops::Deref;
     use std::path::PathBuf;
@@ -110,7 +110,10 @@ pub mod utils {
 
             if path_buf.exists() {
                 let file = File::open(&path_buf).unwrap();
-                // Lock for reading
+                // TODO Manage file locking for inter process stuff, unfortunately linux locks are a
+                // mess and nothing seems to work
+                //
+                // Lock for reading this only works for our process not inter process
                 fs2::FileExt::lock_shared(&file).unwrap();
                 let file_reader = BufReader::new(file);
                 bincode::deserialize_from::<_, (P, K)>(file_reader)
@@ -128,9 +131,20 @@ pub mod utils {
             path_buf.push(param.name());
             path_buf.set_extension("bin");
 
-            let file = File::create(&path_buf).unwrap();
-            // Lock for writing
+            let file = OpenOptions::new()
+                .create(true)
+                .write(true)
+                .truncate(false)
+                .open(&path_buf)
+                .unwrap();
+
+            // TODO Manage file locking for inter process stuff, unfortunately linux locks are a
+            // mess and nothing seems to work
+            //
+            // Lock for writing this only works for our process not inter process
             file.lock_exclusive().unwrap();
+            // Truncate manually
+            file.set_len(0).unwrap();
 
             let file_writer = BufWriter::new(file);
             bincode::serialize_into(file_writer, &(param, key)).unwrap();