Skip to content

Commit

Permalink
feat: use safe ser/deser for tfhe-rs ciphertexts
Browse files Browse the repository at this point in the history
A future commit will add support for keys too when available in tfhe-rs.

Note that safe ser/deser will support versioning of ciphertexts and keys
automatically in the upcoming 0.8 release.

Use FheUint2 for random bool generation for faster performance.
  • Loading branch information
dartdart26 committed Sep 20, 2024
1 parent a1c1a29 commit 14068b5
Show file tree
Hide file tree
Showing 17 changed files with 120 additions and 115 deletions.
13 changes: 7 additions & 6 deletions fhevm-engine/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions fhevm-engine/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ bincode = "1.3.3"
sha3 = "0.10.8"
anyhow = "1.0.86"
daggy = "0.8.0"
serde = "1.0.210"

[profile.dev.package.tfhe]
overflow-checks = false
Expand Down
4 changes: 2 additions & 2 deletions fhevm-engine/coprocessor/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ version = "0.1.0"
edition = "2021"

[target.'cfg(target_arch = "x86_64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
[target.'cfg(target_arch = "aarch64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }

[dependencies]
# Common dependencies
Expand Down
9 changes: 5 additions & 4 deletions fhevm-engine/coprocessor/src/tests/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ use crate::{
utils::{default_api_key, default_tenant_id, setup_test_app},
},
};
use fhevm_engine_common::utils::safe_serialize;
use tonic::metadata::MetadataValue;

#[tokio::test]
Expand Down Expand Up @@ -47,7 +48,7 @@ async fn test_coprocessor_input_errors() -> Result<(), Box<dyn std::error::Error
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();

let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

let mut input_ciphertexts = Vec::new();
for _ in 0..12 {
Expand Down Expand Up @@ -90,7 +91,7 @@ async fn test_coprocessor_input_errors() -> Result<(), Box<dyn std::error::Error
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();

let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

let mut input_ciphertexts = Vec::new();
input_ciphertexts.push(InputToUpload {
Expand Down Expand Up @@ -127,7 +128,7 @@ async fn test_coprocessor_input_errors() -> Result<(), Box<dyn std::error::Error
let the_list = builder
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();
let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

let mut input_ciphertexts = Vec::new();
input_ciphertexts.push(InputToUpload {
Expand Down Expand Up @@ -292,7 +293,7 @@ async fn test_coprocessor_computation_errors() -> Result<(), Box<dyn std::error:
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();

let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

let mut input_ciphertexts = Vec::new();
input_ciphertexts.push(InputToUpload {
Expand Down
5 changes: 3 additions & 2 deletions fhevm-engine/coprocessor/src/tests/inputs.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use std::str::FromStr;

use fhevm_engine_common::utils::safe_serialize;
use tfhe::integer::{bigint::StaticUnsignedBigInt, U256};
use tonic::metadata::MetadataValue;

Expand Down Expand Up @@ -55,7 +56,7 @@ async fn test_fhe_inputs() -> Result<(), Box<dyn std::error::Error>> {
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();

let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

println!("Encrypting inputs...");
let mut input_request = tonic::Request::new(InputUploadBatch {
Expand Down Expand Up @@ -144,7 +145,7 @@ async fn custom_insert_inputs() -> Result<(), Box<dyn std::error::Error>> {
.build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
.unwrap();

let serialized = bincode::serialize(&the_list).unwrap();
let serialized = safe_serialize(&the_list);

println!("Encrypting inputs...");
let mut input_request = tonic::Request::new(InputUploadBatch {
Expand Down
38 changes: 19 additions & 19 deletions fhevm-engine/coprocessor/src/tests/random.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,17 +117,17 @@ async fn test_fhe_random_basic() -> Result<(), Box<dyn std::error::Error>> {
let resp = decrypt_ciphertexts(&pool, 1, decrypt_request).await?;
let expected: Vec<DecryptionResult> = vec![
DecryptionResult { value: "true".to_string(), output_type: 0 },
DecryptionResult { value: "15".to_string(), output_type: 1 },
DecryptionResult { value: "255".to_string(), output_type: 2 },
DecryptionResult { value: "35839".to_string(), output_type: 3 },
DecryptionResult { value: "3850341375".to_string(), output_type: 4 },
DecryptionResult { value: "17749940961552600063".to_string(), output_type: 5 },
DecryptionResult { value: "338317369763394761689683669875179883519".to_string(), output_type: 6 },
DecryptionResult { value: "728577578444838901693251259047493597479225494527".to_string(), output_type: 7 },
DecryptionResult { value: "74507219184508265163276797316037056014989661721330011251665343119622538234879".to_string(), output_type: 8 },
DecryptionResult { value: "2504114756871049901977831169567626719242268143006537293974082655527485165166706069255105841561002259297707106091992898062733977791194469887698664398228479".to_string(), output_type: 9 },
DecryptionResult { value: "161403477149019279849094421471438276759470342665592799183135241729901366432142182887626249781582277693598739699657901471347278400933758701565974722374779502889670960728481539900065496353261065895631087051001039006688379936272227778512119687343404647253114276811856825699481770620842978083170204991519490804735".to_string(), output_type: 10 },
DecryptionResult { value: "25936938232639367084138738776796005240569894972685748071387345538441824752721115183420331032123662095068585061486170621701914537473863580716907485015792656836012274813782383618707951652014429130281418579493619200757130996710486367616379470398907811991574381901820339230001839609009629272311033481447460170915425494881184165217672126153498808951018389035017882126611344150334135611931908252734913533251704739949869925607308666786131997107573881460338809242245923254917133620850600852429219039727909517483206781101690821959886535327935681900259567708853284514443037134192173440669143533573266158029589611848512003083263".to_string(), output_type: 11 }
DecryptionResult { value: "6".to_string(), output_type: 1 },
DecryptionResult { value: "6".to_string(), output_type: 2 },
DecryptionResult { value: "23046".to_string(), output_type: 3 },
DecryptionResult { value: "2257672710".to_string(), output_type: 4 },
DecryptionResult { value: "12138718414261803526".to_string(), output_type: 5 },
DecryptionResult { value: "130536719590611940049803920731387550214".to_string(), output_type: 6 },
DecryptionResult { value: "971176705489787087023559718483701127113677560326".to_string(), output_type: 7 },
DecryptionResult { value: "62210255757460412253332620363065848989112923584999887570035464828426661222918".to_string(), output_type: 8 },
DecryptionResult { value: "167958935840398111366003661819132943572579228212385323643009044778284654758971531763634195717060767316412295162146605242695852136468800900790045270694406".to_string(), output_type: 9 },
DecryptionResult { value: "127460563385689404084570635453516642982330737396307363709535669246693726363369279326274116849562765049033667934125131507607869225026009107310544028242879211116101076829363291657387574479716476869613221980036198477470920343187777849916436388023322996436007563319615378730113313056846971613305517149919649028614".to_string(), output_type: 10 },
DecryptionResult { value: "29687326363179539154232170826093317060572491263948154715413122357200687474061448043555291795321984983113829977114301561317315809196828773909981565653610082891472340553741585442577497506409472143098823132371629384036451019214072899732235656145602725111017828708028912154841404994944466545632048686969494346234325709069045453046020648098209481065154942201598888424765642988091655940417557742117518483932517015160272576663001732809302519121630949039706341063098676812339442637939392896074884484156187775746589025164758187166306751922076107008755031211360389068550389609734783888124482836062055425119177200121882346609158".to_string(), output_type: 11 }
];

println!("results: {:#?}", resp);
Expand Down Expand Up @@ -187,14 +187,14 @@ async fn test_fhe_random_bounded() -> Result<(), Box<dyn std::error::Error>> {
];
let results = [
"true",
"3",
"127",
"3071",
"629115903",
"3914882906270436351",
"83105594572690914092152714301353724927",
"363202169112113172142330050868422842565242358783",
"16611174565850167451491304811693102088354669388509729231936551115665973414911",
"2",
"6",
"6662",
"110189062",
"2915346377407027718",
"45466127860377324183960268873445497350",
"240425886824335627921717302125559617285711288838",
"4314211138802314541547127858721895062477931252179605550306672824470096402950",
];

for (idx, the_type) in supported_types().iter().enumerate() {
Expand Down
4 changes: 2 additions & 2 deletions fhevm-engine/executor/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ version = "0.1.0"
edition = "2021"

[target.'cfg(target_arch = "x86_64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
[target.'cfg(target_arch = "aarch64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }

[dependencies]
clap.workspace = true
Expand Down
11 changes: 5 additions & 6 deletions fhevm-engine/executor/tests/sync_compute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use executor::server::executor::{
use executor::server::executor::{sync_input::Input, SyncInput};
use executor::server::SyncComputeError;
use fhevm_engine_common::types::{SupportedFheCiphertexts, HANDLE_LEN};
use fhevm_engine_common::utils::safe_serialize;
use tfhe::zk::ZkComputeLoad;
use tfhe::ProvenCompactCiphertextList;
use utils::get_test;
Expand All @@ -21,13 +22,12 @@ async fn get_input_ciphertext() {
.await
.unwrap();
let mut builder = ProvenCompactCiphertextList::builder(&test.keys.compact_public_key);
let list = bincode::serialize(
let list = safe_serialize(
&builder
.push(10_u8)
.build_with_proof_packed(&test.keys.public_params, &[], ZkComputeLoad::Proof)
.unwrap(),
)
.unwrap();
);
// TODO: tests for all types and avoiding passing in 2 as an identifier for FheUint8.
let input_handle = test.input_handle(&list, 0, 2);
let sync_input = SyncInput {
Expand Down Expand Up @@ -136,13 +136,12 @@ async fn compute_on_compact_and_serialized_ciphertexts() {
.await
.unwrap();
let mut builder_input = ProvenCompactCiphertextList::builder(&test.keys.compact_public_key);
let compact_list = bincode::serialize(
let compact_list = safe_serialize(
&builder_input
.push(10_u16)
.build_with_proof_packed(&test.keys.public_params, &[], ZkComputeLoad::Proof)
.unwrap(),
)
.unwrap();
);
let mut builder_cts = ProvenCompactCiphertextList::builder(&test.keys.compact_public_key);
let list = builder_cts
.push(11_u16)
Expand Down
5 changes: 3 additions & 2 deletions fhevm-engine/fhevm-engine-common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ version = "0.1.0"
edition = "2021"

[target.'cfg(target_arch = "x86_64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
[target.'cfg(target_arch = "aarch64")'.dependencies]
tfhe = { version = "0.8.0-alpha.6", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
tfhe = { version = "0.8.0-alpha.8", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }

[dependencies]
sha3.workspace = true
Expand All @@ -17,6 +17,7 @@ hex = "0.4"
bigdecimal = "0.4.5"
rand_chacha = "0.3.1"
rand = "0.8.5"
serde.workspace = true

[[bin]]
name = "generate-keys"
Expand Down
1 change: 1 addition & 0 deletions fhevm-engine/fhevm-engine-common/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod keys;
pub mod tfhe_ops;
pub mod types;
pub mod utils;
56 changes: 22 additions & 34 deletions fhevm-engine/fhevm-engine-common/src/tfhe_ops.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,18 @@
use crate::types::{
is_ebytes_type, FheOperationType, FhevmError, SupportedFheCiphertexts, SupportedFheOperations,
use crate::{
types::{
is_ebytes_type, FheOperationType, FhevmError, SupportedFheCiphertexts,
SupportedFheOperations,
},
utils::safe_deserialize,
};
use tfhe::{
integer::{bigint::StaticUnsignedBigInt, U256},
prelude::{
CastInto, FheEq, FheMax, FheMin, FheOrd, FheTryTrivialEncrypt, IfThenElse, RotateLeft,
RotateRight,
},
FheBool, FheUint1024, FheUint128, FheUint16, FheUint160, FheUint2048, FheUint256, FheUint32,
FheUint4, FheUint512, FheUint64, FheUint8, Seed,
FheBool, FheUint1024, FheUint128, FheUint16, FheUint160, FheUint2, FheUint2048, FheUint256,
FheUint32, FheUint4, FheUint512, FheUint64, FheUint8, Seed,
};

pub fn deserialize_fhe_ciphertext(
Expand All @@ -17,63 +21,51 @@ pub fn deserialize_fhe_ciphertext(
) -> Result<SupportedFheCiphertexts, FhevmError> {
match input_type {
0 => {
let v: tfhe::FheBool = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheBool = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheBool(v))
}
1 => {
let v: tfhe::FheUint4 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint4 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint4(v))
}
2 => {
let v: tfhe::FheUint8 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint8 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint8(v))
}
3 => {
let v: tfhe::FheUint16 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint16 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint16(v))
}
4 => {
let v: tfhe::FheUint32 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint32 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint32(v))
}
5 => {
let v: tfhe::FheUint64 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint64 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint64(v))
}
6 => {
let v: tfhe::FheUint128 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint128 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint128(v))
}
7 => {
let v: tfhe::FheUint160 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint160 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint160(v))
}
8 => {
let v: tfhe::FheUint256 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint256 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheUint256(v))
}
9 => {
let v: tfhe::FheUint512 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint512 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheBytes64(v))
}
10 => {
let v: tfhe::FheUint1024 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint1024 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheBytes128(v))
}
11 => {
let v: tfhe::FheUint2048 = bincode::deserialize(input_bytes)
.map_err(|e| FhevmError::DeserializationError(e))?;
let v: tfhe::FheUint2048 = safe_deserialize(input_bytes)?;
Ok(SupportedFheCiphertexts::FheBytes256(v))
}
_ => {
Expand Down Expand Up @@ -239,11 +231,7 @@ pub fn try_expand_ciphertext_list(
) -> Result<Vec<SupportedFheCiphertexts>, FhevmError> {
let mut res = Vec::new();

let the_list: tfhe::ProvenCompactCiphertextList = bincode::deserialize(input_ciphertext)
.map_err(|e| {
let err: Box<(dyn std::error::Error + Send + Sync)> = e;
FhevmError::DeserializationError(err)
})?;
let the_list: tfhe::ProvenCompactCiphertextList = safe_deserialize(input_ciphertext)?;

// TODO: we can do better and avoid cloning
tfhe::set_server_key(server_key.clone());
Expand Down Expand Up @@ -3209,7 +3197,7 @@ pub fn generate_random_number(
let subtract_from = 255;
match the_type {
0 => {
let num = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(seed), 1);
let num = FheUint2::generate_oblivious_pseudo_random_bounded(Seed(seed), 1);
SupportedFheCiphertexts::FheBool(num.gt(0))
}
1 => {
Expand Down
Loading

0 comments on commit 14068b5

Please sign in to comment.