Binary operators test harness

fhevm-engine/Cargo.toml

@@ -22,6 +22,8 @@ clap = { version = "4.5", features = ["derive"] }
 lru = "0.12.3"
 bincode = "1.3.3"
 hex = "0.4"
+strum = { version = "0.26", features = ["derive"] }
+bigdecimal = "0.4"
 
 [dev-dependencies]
 testcontainers = "0.21"

fhevm-engine/src/server.rs

@@ -1,12 +1,10 @@
-use coprocessor::DebugDecryptResponse;
-use sqlx::query;
-use tfhe::prelude::FheTryTrivialEncrypt;
-use tfhe::FheUint32;
+use coprocessor::{DebugDecryptResponse, DebugDecryptResponseSingle};
+use sqlx::{query, Acquire};
 use tonic::transport::Server;
 use crate::db_queries::{check_if_api_key_is_valid, check_if_ciphertexts_exist_in_db};
 use crate::utils::sort_computations_by_dependencies;
-use crate::types::{CoprocessorError, SupportedFheCiphertexts};
-use crate::tfhe_ops::{self, check_fhe_operand_types, current_ciphertext_version};
+use crate::types::CoprocessorError;
+use crate::tfhe_ops::{self, check_fhe_operand_types, current_ciphertext_version, debug_trivial_encrypt_le_bytes};
 use crate::server::coprocessor::GenericResponse;
 
 pub mod coprocessor {
@@ -63,20 +61,40 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
 
         let public_key = public_key.pop().unwrap();
 
-        let value_to_encrypt = req.original_value as u32;
-        let handle = req.handle.clone();
-        let (db_type, db_bytes) = tokio::task::spawn_blocking(move || {
+        let cloned = req.values.clone();
+        let out_cts = tokio::task::spawn_blocking(move || {
             let server_key: tfhe::ServerKey = bincode::deserialize(&public_key.sks_key).unwrap();
             tfhe::set_server_key(server_key);
-            let encrypted = FheUint32::try_encrypt_trivial(value_to_encrypt).unwrap();
-            SupportedFheCiphertexts::FheUint32(encrypted).serialize()
+
+            // single threaded implementation as this is debug function and it is simple to implement
+            let mut res: Vec<(String, i16, Vec<u8>)> = Vec::with_capacity(cloned.len());
+            for v in cloned {
+                let ct = debug_trivial_encrypt_le_bytes(v.output_type as i16, &v.le_value);
+                let (ct_type, ct_bytes) = ct.serialize();
+                res.push((
+                    v.handle,
+                    ct_type,
+                    ct_bytes
+                ));
+            }
+
+            res
         }).await.unwrap();
 
-        sqlx::query!("
-          INSERT INTO ciphertexts(tenant_id, handle, ciphertext, ciphertext_version, ciphertext_type)
-          VALUES ($1, $2, $3, $4, $5)
-        ", tenant_id, handle, db_bytes, current_ciphertext_version(), db_type)
-        .execute(&self.pool).await.map_err(Into::<CoprocessorError>::into)?;
+        let mut conn = self.pool.acquire().await.map_err(Into::<CoprocessorError>::into)?;
+        let mut trx = conn.begin().await.map_err(Into::<CoprocessorError>::into)?;
+
+        for (handle, db_type, db_bytes) in out_cts {
+            sqlx::query!("
+                    INSERT INTO ciphertexts(tenant_id, handle, ciphertext, ciphertext_version, ciphertext_type)
+                    VALUES ($1, $2, $3, $4, $5)
+                ",
+                tenant_id, handle, db_bytes, current_ciphertext_version(), db_type as i16
+            )
+            .execute(trx.as_mut()).await.map_err(Into::<CoprocessorError>::into)?;
+        }
+
+        trx.commit().await.map_err(Into::<CoprocessorError>::into)?;
 
         return Ok(tonic::Response::new(GenericResponse { response_code: 0 }));
     }
@@ -103,32 +121,38 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
 
         assert_eq!(priv_key.len(), 1);
 
-        let mut cts = sqlx::query!("
-          SELECT ciphertext, ciphertext_type
+        let cts = sqlx::query!("
+          SELECT ciphertext, ciphertext_type, handle
           FROM ciphertexts
           WHERE tenant_id = $1
-          AND handle = $2
+          AND handle = ANY($2::TEXT[])
           AND ciphertext_version = $3
-        ", tenant_id, &req.handle, current_ciphertext_version())
+        ", tenant_id, &req.handles, current_ciphertext_version())
         .fetch_all(&self.pool)
         .await.map_err(Into::<CoprocessorError>::into)?;
 
         if cts.is_empty() {
             return Err(tonic::Status::not_found("ciphertext not found"));
         }
 
-        assert_eq!(cts.len(), 1);
-
         let priv_key = priv_key.pop().unwrap().cks_key.unwrap();
-        let ciphertext = cts.pop().unwrap();
 
-        let value = tokio::task::spawn_blocking(move || {
+        let values = tokio::task::spawn_blocking(move || {
             let client_key: tfhe::ClientKey = bincode::deserialize(&priv_key).unwrap();
-            let deserialized = tfhe_ops::deserialize_fhe_ciphertext(ciphertext.ciphertext_type, &ciphertext.ciphertext).unwrap();
-            deserialized.decrypt(&client_key)
+
+            let mut decrypted: Vec<DebugDecryptResponseSingle> = Vec::with_capacity(cts.len());
+            for ct in cts {
+                let deserialized = tfhe_ops::deserialize_fhe_ciphertext(ct.ciphertext_type, &ct.ciphertext).unwrap();
+                decrypted.push(DebugDecryptResponseSingle {
+                    output_type: ct.ciphertext_type as i32,
+                    value: deserialized.decrypt(&client_key),
+                });
+            }
+
+            decrypted
         }).await.unwrap();
 
-        return Ok(tonic::Response::new(DebugDecryptResponse { value }));
+        return Ok(tonic::Response::new(DebugDecryptResponse { values }));
     }
 
     async fn upload_ciphertexts(

fhevm-engine/src/tests/mod.rs

@@ -1,35 +1,39 @@
+use std::str::FromStr;
+
 use tonic::metadata::MetadataValue;
+use utils::default_api_key;
 use crate::server::coprocessor::fhevm_coprocessor_client::FhevmCoprocessorClient;
-use crate::server::coprocessor::{AsyncComputeRequest, FheOperation, DebugEncryptRequest, DebugDecryptRequest, AsyncComputation};
+use crate::server::coprocessor::{AsyncComputation, AsyncComputeRequest, DebugDecryptRequest, DebugEncryptRequest, DebugEncryptRequestSingle, FheOperation};
 
 mod utils;
+mod operators;
 
 #[tokio::test]
 async fn test_smoke() -> Result<(), Box<dyn std::error::Error>> {
     let app = utils::setup_test_app().await?;
 
     let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
 
-    let api_key = "Bearer a1503fb6-d79b-4e9e-826d-44cf262f3e05";
+    let api_key_header = format!("Bearer {}", default_api_key());
+    let ct_type = 4; // i32
 
-    // ciphertext A
+    // encrypt two ciphertexts
     {
         let mut encrypt_request = tonic::Request::new(DebugEncryptRequest {
-            handle: "0x0abc".to_string(),
-            original_value: 123,
-        });
-        encrypt_request.metadata_mut().append("authorization", MetadataValue::from_static(api_key));
-        let resp = client.debug_encrypt_ciphertext(encrypt_request).await?;
-        println!("encryption request: {:?}", resp);
-    }
-
-    // ciphertext B
-    {
-        let mut encrypt_request = tonic::Request::new(DebugEncryptRequest {
-            handle: "0x0abd".to_string(),
-            original_value: 124,
+            values: vec![
+                DebugEncryptRequestSingle {
+                    handle: "0x0abc".to_string(),
+                    le_value: vec![123],
+                    output_type: ct_type,
+                },
+                DebugEncryptRequestSingle {
+                    handle: "0x0abd".to_string(),
+                    le_value: vec![124],
+                    output_type: ct_type,
+                },
+            ],
         });
-        encrypt_request.metadata_mut().append("authorization", MetadataValue::from_static(api_key));
+        encrypt_request.metadata_mut().append("authorization", MetadataValue::from_str(&api_key_header).unwrap());
         let resp = client.debug_encrypt_ciphertext(encrypt_request).await?;
         println!("encryption request: {:?}", resp);
     }
@@ -58,34 +62,32 @@ async fn test_smoke() -> Result<(), Box<dyn std::error::Error>> {
                 },
             ]
         });
-        compute_request.metadata_mut().append("authorization", MetadataValue::from_static(api_key));
+        compute_request.metadata_mut().append("authorization", MetadataValue::from_str(&api_key_header).unwrap());
         let resp = client.async_compute(compute_request).await?;
         println!("compute request: {:?}", resp);
     }
 
     println!("sleeping for computation to complete...");
     tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
 
-    // decrypt first
-    {
-        let mut decrypt_request = tonic::Request::new(DebugDecryptRequest {
-            handle: "0x0abe".to_string()
-        });
-        decrypt_request.metadata_mut().append("authorization", MetadataValue::from_static(api_key));
-        let resp = client.debug_decrypt_ciphertext(decrypt_request).await?;
-        println!("decrypt request: {:?}", resp);
-        assert_eq!(resp.get_ref().value, "247");
-    }
-
-    // decrypt second
+    // decrypt values
     {
         let mut decrypt_request = tonic::Request::new(DebugDecryptRequest {
-            handle: "0x0abf".to_string()
+            handles: vec![
+                "0x0abe".to_string(),
+                "0x0abf".to_string(),
+            ],
         });
-        decrypt_request.metadata_mut().append("authorization", MetadataValue::from_static(api_key));
+        decrypt_request.metadata_mut().append("authorization", MetadataValue::from_str(&api_key_header).unwrap());
         let resp = client.debug_decrypt_ciphertext(decrypt_request).await?;
         println!("decrypt request: {:?}", resp);
-        assert_eq!(resp.get_ref().value, "263");
+        assert_eq!(resp.get_ref().values.len(), 2);
+        // first value
+        assert_eq!(resp.get_ref().values[0].value, "247");
+        assert_eq!(resp.get_ref().values[0].output_type, ct_type);
+        // second value
+        assert_eq!(resp.get_ref().values[1].value, "263");
+        assert_eq!(resp.get_ref().values[1].output_type, ct_type);
     }
 
     Ok(())