feat(tools): Concrete Rust keygen with initial secret keys

tools/concrete-rust/src/keygen.rs

@@ -1,5 +1,5 @@
 use clap::{Arg, Command};
-use concrete_rust::generate_keyset;
+use concrete_rust::{generate_keyset, read_secret_key_from_file};
 
 pub fn main() {
     let matches = Command::new("concrete-keygen")
@@ -31,6 +31,14 @@ pub fn main() {
                 .default_value("0")
                 .long("enc-seed"),
         )
+        .arg(
+            Arg::new("initial-secret-key")
+                .help("Initial secret key for key generation")
+                .required(false)
+                .long("sk")
+                .num_args(1)
+                .action(clap::ArgAction::Append),
+        )
         .get_matches();
 
     let keyset_info_path = matches.get_one::<String>("keyset-info").unwrap();
@@ -45,12 +53,27 @@ pub fn main() {
         .unwrap()
         .parse::<u128>()
         .expect("Invalid encryption seed");
+    let initial_secret_keys_paths: Vec<String> = matches
+        .get_many::<String>("initial-secret-key")
+        .unwrap_or_default()
+        .map(|s| s.to_string())
+        .collect();
+    let initial_secret_keys: std::collections::HashMap<
+        u32,
+        capnp::message::Reader<capnp::serialize::OwnedSegments>,
+    > = initial_secret_keys_paths
+        .iter()
+        .map(|s| {
+            let (id, reader) = read_secret_key_from_file(s);
+            (id, reader)
+        })
+        .collect();
 
     generate_keyset(
         keyset_info_path,
         sec_seed,
         enc_seed,
         keyset_path,
-        std::collections::HashMap::new(),
+        &initial_secret_keys,
     );
 }

tools/concrete-rust/src/lib.rs

@@ -26,7 +26,7 @@ use std::u128;
 
 use crate::concrete_protocol_capnp::*;
 use capnp::message::{HeapAllocator, ReaderOptions};
-use capnp::serialize;
+use capnp::serialize::{self, OwnedSegments};
 use concrete_csprng::generators::SoftwareRandomGenerator;
 use tfhe::core_crypto::prelude::{
     generate_lwe_keyswitch_key, par_generate_circuit_bootstrap_lwe_pfpksk_list,
@@ -426,36 +426,59 @@ fn generate_keyset_message(
 ///
 /// The function generates secret keys, bootstrap keys, keyswitch keys, and packing keyswitch keys based on the keyset information.
 /// It uses initial secret keys instead of generating them if they are provided.
+///
 /// # Arguments
 ///
 /// * `keyset_info_path`: The path to the keyset information file.
 /// * `secret_seed`: The seed for the secret random generator.
 /// * `enc_seed`: The seed for the encryption random generator.
 /// * `keyset_path`: The path to the output keyset file.
-/// * `init_lwe_secret_keys`: A map of initial secret keys to use instead of generating them.
+/// * `init_secret_keys`: A map of initial secret keys to use instead of generating them.
 ///
 /// # Panics
-/// This function panics if it fails to open the keyset info file or create the keyset file.
+///  This function panics if it fails to open the keyset info file or create the keyset file.
 pub fn generate_keyset(
     keyset_info_path: &str,
     secret_seed: u128,
     enc_seed: u128,
     keyset_path: &str,
-    init_lwe_secret_keys: std::collections::HashMap<u32, lwe_secret_key::Reader>,
+    init_secret_keys: &std::collections::HashMap<u32, capnp::message::Reader<OwnedSegments>>,
 ) {
     // read keyset info from input file
     let file = std::fs::File::open(keyset_info_path).expect("Failed to open keyset info file");
     let reader =
         serialize::read_message(std::io::BufReader::new(file), ReaderOptions::new()).unwrap();
     let key_set_info = reader.get_root::<keyset_info::Reader>().unwrap();
     // keygen
+    let mut init_lwe_secret_keys = init_secret_keys
+        .iter()
+        .map(|(k, v)| (*k, v.get_root::<lwe_secret_key::Reader>().unwrap().clone()))
+        .collect::<std::collections::HashMap<u32, lwe_secret_key::Reader>>();
     let builder = generate_keyset_message(
         key_set_info,
         secret_seed,
         enc_seed,
-        &mut init_lwe_secret_keys.clone(),
+        &mut init_lwe_secret_keys,
     );
     // write keyset to output file
     let output = std::fs::File::create(keyset_path).expect("Failed to create keyset file");
     serialize::write_message(output, &builder).unwrap();
 }
+
+/// Reads a secret key from a file and returns the key ID and Cap'n Proto reader.
+///
+/// # Arguments
+///
+/// * `path`: The path to the secret key file.
+///
+/// # Returns
+///
+/// A tuple containing the key ID and Cap'n Proto reader for the secret key.
+pub fn read_secret_key_from_file(path: &str) -> (u32, capnp::message::Reader<OwnedSegments>) {
+    let file = std::fs::File::open(path).expect("Failed to open secret key file");
+    let reader =
+        serialize::read_message(std::io::BufReader::new(file), ReaderOptions::new()).unwrap();
+    let key = reader.get_root::<lwe_secret_key::Reader>().unwrap();
+    let id = key.get_info().unwrap().get_id();
+    (id, reader)
+}