Merge pull request #132 from zama-ai/davidk/allow-customizing-executo…

…r-keys-dir

feat: allow to customize key directory in the executor

fhevm-engine/executor/src/cli.rs

@@ -14,6 +14,11 @@ pub struct Args {
 
     #[arg(long, default_value = "127.0.0.1:50051")]
     pub server_addr: String,
+
+    /// directory for fhe keys, target directory expected to contain files named:
+    /// sks (server evaluation key), pks (compact public key), pp (public key params)
+    #[arg(long)]
+    pub fhe_keys_directory: String,
 }
 
 pub fn parse_args() -> Args {

fhevm-engine/executor/src/server.rs

@@ -32,10 +32,10 @@ thread_local! {
 }
 
 pub fn start(args: &crate::cli::Args) -> Result<()> {
-    let keys: FhevmKeys = SerializedFhevmKeys::load_from_disk().into();
+    let keys: FhevmKeys = SerializedFhevmKeys::load_from_disk(&args.fhe_keys_directory).into();
     SERVER_KEY.set(Some(keys.server_key.clone()));
     LOCAL_RAYON_THREADS.set(args.policy_fhe_compute_threads);
-    let executor = FhevmExecutorService::new();
+    let executor = FhevmExecutorService::new(keys.clone());
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .worker_threads(args.tokio_threads)
         .max_blocking_threads(args.fhe_compute_threads)
@@ -150,9 +150,9 @@ impl FhevmExecutor for FhevmExecutorService {
 }
 
 impl FhevmExecutorService {
-    fn new() -> Self {
+    fn new(keys: FhevmKeys) -> Self {
         FhevmExecutorService {
-            keys: SerializedFhevmKeys::load_from_disk().into(),
+            keys,
         }
     }
 
@@ -300,41 +300,6 @@ impl FhevmExecutorService {
     }
 }
 
-pub fn run_computation(
-    operation: i32,
-    inputs: Result<Vec<SupportedFheCiphertexts>, SyncComputeError>,
-    graph_node_index: usize,
-) -> Result<(usize, InMemoryCiphertext), SyncComputeError> {
-    let op = FheOperation::try_from(operation);
-    match inputs {
-        Ok(inputs) => match op {
-            Ok(FheOperation::FheGetCiphertext) => {
-                let res = InMemoryCiphertext {
-                    expanded: inputs[0].clone(),
-                    compressed: inputs[0].clone().compress().1,
-                };
-                Ok((graph_node_index, res))
-            }
-            Ok(_) => match perform_fhe_operation(operation as i16, &inputs) {
-                Ok(result) => {
-                    let res = InMemoryCiphertext {
-                        expanded: result.clone(),
-                        compressed: result.compress().1,
-                    };
-                    Ok((graph_node_index, res))
-                }
-                Err(_) => Err::<(usize, InMemoryCiphertext), SyncComputeError>(
-                    SyncComputeError::ComputationFailed,
-                ),
-            },
-            _ => Err::<(usize, InMemoryCiphertext), SyncComputeError>(
-                SyncComputeError::InvalidOperation,
-            ),
-        },
-        Err(_) => Err(SyncComputeError::ComputationFailed),
-    }
-}
-
 pub fn build_taskgraph_from_request(
     dfg: &mut DFGraph,
     req: &SyncComputeRequest,

fhevm-engine/executor/tests/utils.rs

@@ -19,10 +19,10 @@ pub struct TestInstance {
 impl TestInstance {
     pub async fn new() -> Self {
         // Get defaults by parsing a cmd line without any arguments.
-        let args = Args::parse_from(&["test"]);
+        let args = Args::parse_from(&["test", "--fhe-keys-directory", "../fhevm-keys"]);
 
         let instance = TestInstance {
-            keys: SerializedFhevmKeys::load_from_disk().into(),
+            keys: SerializedFhevmKeys::load_from_disk("../fhevm-keys").into(),
             server_addr: format!("http://{}", args.server_addr),
         };
 

fhevm-engine/fhevm-engine-common/src/keys.rs

@@ -28,6 +28,7 @@ pub const TFHE_KS_PARAMS: ShortintKeySwitchingParameters =
 
 pub const MAX_BITS_TO_PROVE: usize = 2048;
 
+#[derive(Clone)]
 pub struct FhevmKeys {
     pub server_key: ServerKey,
     pub client_key: Option<ClientKey>,
@@ -79,6 +80,7 @@ impl SerializedFhevmKeys {
     const PKS: &'static str = "../fhevm-keys/pks";
     const PUBLIC_PARAMS: &'static str = "../fhevm-keys/pp";
 
+    // generating keys is only for testing, so it is okay these are hardcoded
     pub fn save_to_disk(self) {
         println!("Creating directory {}", Self::DIRECTORY);
         std::fs::create_dir_all(Self::DIRECTORY).expect("create keys directory");
@@ -99,14 +101,15 @@ impl SerializedFhevmKeys {
             .expect("write public params");
     }
 
-    pub fn load_from_disk() -> Self {
-        let server_key = read(Self::SKS).expect("read server key");
-        let client_key = read(Self::CKS);
-        let compact_public_key = read(Self::PKS).expect("read compact public key");
-        let public_params = read(Self::PUBLIC_PARAMS).expect("read public params");
+    pub fn load_from_disk(keys_directory: &str) -> Self {
+        let keys_dir = std::path::Path::new(&keys_directory);
+        let server_key = read(keys_dir.join("sks")).expect("read server key");
+        let client_key = read(keys_dir.join("cks")).ok();
+        let compact_public_key = read(keys_dir.join("pks")).expect("read compact public key");
+        let public_params = read(keys_dir.join("pp")).expect("read public params");
         SerializedFhevmKeys {
             server_key,
-            client_key: client_key.ok(),
+            client_key,
             compact_public_key,
             public_params,
         }