Merge branch 'main' into davidk/fhe-rand-generation

LICENSE

@@ -0,0 +1,33 @@
+BSD 3-Clause Clear License
+
+Copyright © 2024 ZAMA.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this
+list of conditions and the following disclaimer in the documentation and/or other
+materials provided with the distribution.
+
+3. Neither the name of ZAMA nor the names of its contributors may be used to endorse
+or promote products derived from this software without specific prior written permission.
+
+NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE*.
+THIS SOFTWARE IS PROVIDED BY THE ZAMA AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
+ZAMA OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
+OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+*In addition to the rights carried by this license, ZAMA grants to the user a non-exclusive, 
+free and non-commercial license on all patents filed in its name relating to the open-source 
+code (the "Patents") for the sole purpose of evaluation, development, research, prototyping 
+and experimentation.

fhevm-engine/coprocessor/src/tests/mod.rs

@@ -8,7 +8,7 @@ use crate::server::coprocessor::{
     TrivialEncryptRequestSingle,
 };
 use tonic::metadata::MetadataValue;
-use utils::{decrypt_ciphertexts, default_api_key, wait_until_all_ciphertexts_computed};
+use utils::{decrypt_ciphertexts, default_api_key, random_handle, wait_until_all_ciphertexts_computed};
 
 mod errors;
 mod inputs;
@@ -29,6 +29,11 @@ async fn test_smoke() -> Result<(), Box<dyn std::error::Error>> {
     let api_key_header = format!("bearer {}", default_api_key());
     let ct_type = 4; // i32
 
+    let h1 = random_handle().to_be_bytes();
+    let h2 = random_handle().to_be_bytes();
+    let h3 = random_handle().to_be_bytes();
+    let h4 = random_handle().to_be_bytes();
+
     // encrypt two ciphertexts
     {
         let mut encrypt_request = tonic::Request::new(TrivialEncryptBatch {
@@ -59,10 +64,10 @@ async fn test_smoke() -> Result<(), Box<dyn std::error::Error>> {
             computations: vec![
                 AsyncComputation {
                     operation: FheOperation::FheAdd.into(),
-                    output_handle: vec![0x0a, 0xbf],
+                    output_handle: h3.to_vec(),
                     inputs: vec![
                         AsyncComputationInput {
-                            input: Some(Input::InputHandle(vec![0x0a, 0xbe])),
+                            input: Some(Input::InputHandle(h4.to_vec())),
                         },
                         AsyncComputationInput {
                             input: Some(Input::Scalar(vec![0x00, 0x10])),
@@ -71,13 +76,13 @@ async fn test_smoke() -> Result<(), Box<dyn std::error::Error>> {
                 },
                 AsyncComputation {
                     operation: FheOperation::FheAdd.into(),
-                    output_handle: vec![0x0a, 0xbe],
+                    output_handle: h4.to_vec(),
                     inputs: vec![
                         AsyncComputationInput {
-                            input: Some(Input::InputHandle(vec![0x0a, 0xbc])),
+                            input: Some(Input::InputHandle(h1.to_vec())),
                         },
                         AsyncComputationInput {
-                            input: Some(Input::InputHandle(vec![0x0a, 0xbd])),
+                            input: Some(Input::InputHandle(h2.to_vec())),
                         },
                     ],
                 },
@@ -136,4 +141,4 @@ async fn test_custom_function() -> Result<(), Box<dyn std::error::Error>> {
     println!("{:#?}", res);
 
     Ok(())
-}
+}

fhevm-engine/coprocessor/src/tests/operators.rs

@@ -4,7 +4,7 @@ use crate::server::coprocessor::{
     AsyncComputation, AsyncComputeRequest, TrivialEncryptBatch, TrivialEncryptRequestSingle,
 };
 use crate::tests::utils::{
-    decrypt_ciphertexts, random_handle_start, wait_until_all_ciphertexts_computed,
+    decrypt_ciphertexts, random_handle, random_handle_start, wait_until_all_ciphertexts_computed,
 };
 use crate::{
     server::coprocessor::{async_computation_input::Input, AsyncComputationInput},
@@ -85,7 +85,7 @@ async fn test_fhe_binary_operands() -> Result<(), Box<dyn std::error::Error>> {
         .await?;
     let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
 
-    let mut handle_counter: u64 = random_handle_start();
+    let mut handle_counter: u64 = random_handle();
     let mut next_handle = || {
         let out: u64 = handle_counter;
         handle_counter += 1;
@@ -221,7 +221,7 @@ async fn test_fhe_unary_operands() -> Result<(), Box<dyn std::error::Error>> {
         .await?;
     let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
 
-    let mut handle_counter: u64 = random_handle_start();
+    let mut handle_counter: u64 = random_handle();
     let mut next_handle = || {
         let out: u64 = handle_counter;
         handle_counter += 1;
@@ -332,7 +332,7 @@ async fn test_fhe_casts() -> Result<(), Box<dyn std::error::Error>> {
         .await?;
     let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
 
-    let mut handle_counter = random_handle_start();
+    let mut handle_counter = random_handle();
     let mut next_handle = || {
         let out: u64 = handle_counter;
         handle_counter += 1;
@@ -465,7 +465,7 @@ async fn test_fhe_if_then_else() -> Result<(), Box<dyn std::error::Error>> {
         .await?;
     let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
 
-    let mut handle_counter = random_handle_start();
+    let mut handle_counter = random_handle();
     let mut next_handle = || {
         let out: u64 = handle_counter;
         handle_counter += 1;

fhevm-engine/coprocessor/src/tests/utils.rs

@@ -1,9 +1,10 @@
 use crate::cli::Args;
 use fhevm_engine_common::tfhe_ops::{current_ciphertext_version, deserialize_fhe_ciphertext};
-use rand::RngCore;
+use rand::{Rnd, RngCore};
 use std::collections::BTreeMap;
 use std::sync::atomic::{AtomicU16, Ordering};
 use testcontainers::{core::WaitFor, runners::AsyncRunner, GenericImage, ImageExt};
+use tokio::sync::watch::Receiver;
 
 pub struct TestInstance {
     // just to destroy container
@@ -41,32 +42,81 @@ pub fn default_tenant_id() -> i32 {
     1
 }
 
+pub fn random_handle() -> u64 {
+    rand::thread_rng().gen()
+}
+
 pub async fn setup_test_app() -> Result<TestInstance, Box<dyn std::error::Error>> {
     if std::env::var("COPROCESSOR_TEST_LOCALHOST").is_ok() {
         setup_test_app_existing_localhost().await
+    } else if std::env::var("COPROCESSOR_TEST_LOCAL_DB").is_ok() {
+        setup_test_app_existing_db().await
     } else {
         setup_test_app_custom_docker().await
     }
 }
 
-pub async fn setup_test_app_existing_localhost() -> Result<TestInstance, Box<dyn std::error::Error>>
-{
+pub async fn setup_test_app_existing_localhost() -> Result<TestInstance, Box<dyn std::error::Error>> {
+    const LOCAL_DB_URL: &str = "postgresql://postgres:[email protected]:5432/coprocessor";
     Ok(TestInstance {
         _container: None,
         app_close_channel: None,
         app_url: "http://127.0.0.1:50051".to_string(),
-        db_url: "postgresql://postgres:[email protected]:5432/coprocessor".to_string(),
+        db_url: LOCAL_DB_URL.to_string(),
     })
 }
 
-pub async fn setup_test_app_custom_docker() -> Result<TestInstance, Box<dyn std::error::Error>> {
+async fn setup_test_app_existing_db() -> Result<TestInstance, Box<dyn std::error::Error>> {
+    let app_port = get_app_port();
+    let (app_close_channel, rx) = tokio::sync::watch::channel(false);
+    start_coprocessor(rx, app_port, &LOCAL_DB_URL).await;
+    Ok(TestInstance {
+        _container: None,
+        app_close_channel: Some(app_close_channel),
+        app_url: format!("http://127.0.0.1:{app_port}"),
+        db_url: LOCAL_DB_URL.to_string(),
+    })
+}
+
+async fn start_coprocessor(rx: Receiver<bool>, app_port: u16, db_url: &str) {
+    let args: Args = Args {
+        run_bg_worker: true,
+        run_server: true,
+        generate_fhe_keys: false,
+        server_maximum_ciphertexts_to_schedule: 5000,
+        work_items_batch_size: 40,
+        tenant_key_cache_size: 4,
+        coprocessor_fhe_threads: 4,
+        maximum_handles_per_input: 255,
+        tokio_threads: 2,
+        pg_pool_max_connections: 2,
+        server_addr: format!("127.0.0.1:{app_port}"),
+        database_url: Some(db_url.to_string()),
+        maximimum_compact_inputs_upload: 10,
+        coprocessor_private_key: "./coprocessor.key".to_string(),
+    };
+
+    std::thread::spawn(move || {
+        crate::start_runtime(args, Some(rx));
+    });
+
+    // wait until app port is opened
+    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+}
+
+fn get_app_port() -> u16 {
     static PORT_COUNTER: AtomicU16 = AtomicU16::new(10000);
 
     let app_port = PORT_COUNTER.fetch_add(1, Ordering::SeqCst);
     // wrap around, if we ever have that many tests?
     if app_port >= 50000 {
         PORT_COUNTER.store(10000, Ordering::SeqCst);
     }
+    app_port
+}
+
+async fn setup_test_app_custom_docker() -> Result<TestInstance, Box<dyn std::error::Error>> {
+    let app_port = get_app_port();
 
     let container = GenericImage::new("postgres", "15.7")
         .with_wait_for(WaitFor::message_on_stderr(
@@ -102,30 +152,7 @@ pub async fn setup_test_app_custom_docker() -> Result<TestInstance, Box<dyn std:
     println!("DB prepared");
 
     let (app_close_channel, rx) = tokio::sync::watch::channel(false);
-    let args: Args = Args {
-        run_bg_worker: true,
-        run_server: true,
-        generate_fhe_keys: false,
-        server_maximum_ciphertexts_to_schedule: 5000,
-        work_items_batch_size: 40,
-        tenant_key_cache_size: 4,
-        coprocessor_fhe_threads: 4,
-        maximum_handles_per_input: 255,
-        tokio_threads: 2,
-        pg_pool_max_connections: 2,
-        server_addr: format!("127.0.0.1:{app_port}"),
-        database_url: Some(db_url.clone()),
-        maximimum_compact_inputs_upload: 10,
-        coprocessor_private_key: "./coprocessor.key".to_string(),
-    };
-
-    std::thread::spawn(move || {
-        crate::start_runtime(args, Some(rx));
-    });
-
-    // wait until app port is opened
-    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
-
+    start_coprocessor(rx, app_port, &db_url).await;
     Ok(TestInstance {
         _container: Some(container),
         app_close_channel: Some(app_close_channel),