github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-09 14:47:56 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat(integer): plug metadata into lower level ZK APIs

Browse Source
This commit is contained in:
Arthur Meyre
2024-09-06 18:44:31 +02:00
parent aee4c1ed18
commit 0612ef5be5
10 changed files with 136 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tfhe/benches/integer/zk_pke.rs
View File

@@ -4,6 +4,7 @@
mod utilities;
use criterion::{criterion_group, criterion_main, Criterion};
use rand::prelude::*;
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::Path;
@@ -49,6 +50,11 @@ fn pke_zk_proof(c: &mut Criterion) {
let _casting_key =
KeySwitchingKey::new((&compact_private_key, None), (&cks, &sks), _param_casting);
// We have a use case with 320 bits of metadata
let mut metadata = [0u8; (320 / u8::BITS) as usize];
let mut rng = rand::thread_rng();
metadata.fill_with(|| rng.gen());
for bits in [640usize, 1280, 4096] {
assert_eq!(bits % 64, 0);
// Packing, so we take the message and carry modulus to compute our block count
@@ -77,7 +83,7 @@ fn pke_zk_proof(c: &mut Criterion) {
b.iter(|| {
let _ct1 = tfhe::integer::ProvenCompactCiphertextList::builder(&pk)
.extend(messages.iter().copied())
.build_with_proof_packed(public_params, compute_load)
.build_with_proof_packed(public_params, &metadata, compute_load)
.unwrap();
})
});
@@ -129,6 +135,11 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
let casting_key =
KeySwitchingKey::new((&compact_private_key, None), (&cks, &sks), param_casting);
// We have a use case with 320 bits of metadata
let mut metadata = [0u8; (320 / u8::BITS) as usize];
let mut rng = rand::thread_rng();
metadata.fill_with(|| rng.gen());
for bits in [640usize, 1280, 4096] {
assert_eq!(bits % 64, 0);
// Packing, so we take the message and carry modulus to compute our block count
@@ -184,7 +195,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
println!("Generating proven ciphertext ({zk_load})... ");
let ct1 = tfhe::integer::ProvenCompactCiphertextList::builder(&pk)
.extend(messages.iter().copied())
.build_with_proof_packed(public_params, compute_load)
.build_with_proof_packed(public_params, &metadata, compute_load)
.unwrap();
let proven_ciphertext_list_serialized = bincode::serialize(&ct1).unwrap();
@@ -231,7 +242,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
bench_group.bench_function(&bench_id_verify, |b| {
b.iter(|| {
let _ret = ct1.verify(public_params, &pk);
let _ret = ct1.verify(public_params, &pk, &metadata);
});
});
@@ -241,6 +252,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
.verify_and_expand(
public_params,
&pk,
&metadata,
IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(&sks),
IntegerCompactCiphertextListCastingMode::CastIfNecessary(
casting_key.as_view(),

11
tfhe/c_api_tests/test_high_level_zk.c
View File

@@ -35,6 +35,9 @@ int main(void) {
status = compact_pke_crs_public_params(crs, &public_params);
assert(status == 0);
#define METADATA_LEN 5
uint8_t metadata[METADATA_LEN] = {'c', '-', 'a', 'p', 'i'};
ClientKey *client_key;
ServerKey *server_key;
status = generate_keys(config, &client_key, &server_key);
@@ -47,8 +50,6 @@ int main(void) {
status = compact_public_key_new(client_key, &pk);
assert(status == 0);
// Then, we create the compact list
ProvenCompactCiphertextList *compact_list = NULL;
{
@@ -69,8 +70,8 @@ int main(void) {
status = compact_ciphertext_list_builder_push_u2(builder, 3);
assert(status == 0);
status = compact_ciphertext_list_builder_build_with_proof_packed(builder, public_params,
ZkComputeLoadProof, &compact_list);
status = compact_ciphertext_list_builder_build_with_proof_packed(
builder, public_params, metadata, METADATA_LEN, ZkComputeLoadProof, &compact_list);
assert(status == 0);
// Don't forget to destroy the builder
@@ -85,7 +86,7 @@ int main(void) {
{
CompactCiphertextListExpander *expander = NULL;
status = proven_compact_ciphertext_list_verify_and_expand(compact_list, public_params, pk,
&expander);
metadata, METADATA_LEN, &expander);
assert(status == 0);
status = compact_ciphertext_list_expander_get_fhe_uint32(expander, 0, &a);

19
tfhe/docs/guides/zk-pok.md
View File

@@ -19,8 +19,7 @@ use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
pub fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut rng = thread_rng();
let params =
tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let params = tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let config = tfhe::ConfigBuilder::with_custom_parameters(params);
let client_key = tfhe::ClientKey::generate(config.clone());
@@ -29,21 +28,24 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
let public_zk_params = crs.public_params();
let server_key = tfhe::ServerKey::new(&client_key);
let public_key = tfhe::CompactPublicKey::try_new(&client_key).unwrap();
// This can be left empty, but if provided allows to tie the proof to arbitrary data
let metadata = [b'T', b'F', b'H', b'E', b'-', b'r', b's'];
let clear_a = rng.gen::<u64>();
let clear_b = rng.gen::<u64>();
let proven_compact_list = tfhe::ProvenCompactCiphertextList::builder(&public_key)
.push(clear_a)
.push(clear_b)
.build_with_proof_packed(public_zk_params, ZkComputeLoad::Proof)?;
.build_with_proof_packed(public_zk_params, &metadata, ZkComputeLoad::Proof)?;
// Server side
let result = {
set_server_key(server_key);
// Verify the ciphertexts
let mut expander = proven_compact_list.verify_and_expand(public_zk_params, &public_key)?;
let mut expander =
proven_compact_list.verify_and_expand(public_zk_params, &public_key, &metadata)?;
let a: tfhe::FheUint64 = expander.get(0).unwrap()?;
let b: tfhe::FheUint64 = expander.get(1).unwrap()?;
@@ -101,6 +103,8 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
let public_zk_params = crs.public_params();
let server_key = tfhe::ServerKey::new(&client_key);
let public_key = tfhe::CompactPublicKey::try_new(&client_key).unwrap();
// This can be left empty, but if provided allows to tie the proof to arbitrary data
let metadata = [b'T', b'F', b'H', b'E', b'-', b'r', b's'];
let clear_a = rng.gen::<u64>();
let clear_b = rng.gen::<u64>();
@@ -108,14 +112,15 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
let proven_compact_list = tfhe::ProvenCompactCiphertextList::builder(&public_key)
.push(clear_a)
.push(clear_b)
.build_with_proof_packed(public_zk_params, ZkComputeLoad::Verify)?;
.build_with_proof_packed(public_zk_params, &metadata, ZkComputeLoad::Verify)?;
// Server side
let result = {
set_server_key(server_key);
// Verify the ciphertexts
let mut expander = proven_compact_list.verify_and_expand(public_zk_params, &public_key)?;
let mut expander =
proven_compact_list.verify_and_expand(public_zk_params, &public_key, &metadata)?;
let a: tfhe::FheUint64 = expander.get(0).unwrap()?;
let b: tfhe::FheUint64 = expander.get(1).unwrap()?;

23
tfhe/src/c_api/high_level_api/compact_list.rs
View File

@@ -78,15 +78,25 @@ pub unsafe extern "C" fn compact_ciphertext_list_builder_build_packed(
pub unsafe extern "C" fn compact_ciphertext_list_builder_build_with_proof_packed(
builder: *const CompactCiphertextListBuilder,
public_params: *const CompactPkePublicParams,
metadata: *const u8,
metadata_len: usize,
compute_load: ZkComputeLoad,
list: *mut *mut ProvenCompactCiphertextList,
) -> c_int {
catch_panic(|| {
let builder = get_ref_checked(builder).unwrap();
let public_params = get_ref_checked(public_params).unwrap();
let metadata = if metadata.is_null() {
&[]
} else {
let _metadata_check_ptr = get_ref_checked(metadata).unwrap();
core::slice::from_raw_parts(metadata, metadata_len)
};
let inner = builder
.0
.build_with_proof_packed(&public_params.0, compute_load.into())
.build_with_proof_packed(&public_params.0, metadata, compute_load.into())
.unwrap();
*list = Box::into_raw(Box::new(ProvenCompactCiphertextList(inner)));
@@ -173,6 +183,8 @@ pub unsafe extern "C" fn proven_compact_ciphertext_list_verify_and_expand(
compact_list: *const ProvenCompactCiphertextList,
public_params: *const CompactPkePublicParams,
public_key: *const CompactPublicKey,
metadata: *const u8,
metadata_len: usize,
expander: *mut *mut CompactCiphertextListExpander,
) -> c_int {
catch_panic(|| {
@@ -180,9 +192,16 @@ pub unsafe extern "C" fn proven_compact_ciphertext_list_verify_and_expand(
let public_params = get_ref_checked(public_params).unwrap();
let public_key = get_ref_checked(public_key).unwrap();
let metadata = if metadata.is_null() {
&[]
} else {
let _metadata_check_ptr = get_ref_checked(metadata).unwrap();
core::slice::from_raw_parts(metadata, metadata_len)
};
let inner = list
.0
.verify_and_expand(&public_params.0, &public_key.0)
.verify_and_expand(&public_params.0, &public_key.0, metadata)
.unwrap();
*expander = Box::into_raw(Box::new(CompactCiphertextListExpander(inner)));

25
tfhe/src/high_level_api/compact_list.rs
View File

@@ -233,6 +233,7 @@ mod zk {
&self,
public_params: &CompactPkePublicParams,
pk: &CompactPublicKey,
metadata: &[u8],
) -> crate::Result<CompactCiphertextListExpander> {
// For WASM
if !self.inner.is_packed() && !self.inner.needs_casting() {
@@ -241,6 +242,7 @@ mod zk {
inner: self.inner.verify_and_expand(
public_params,
&pk.key.key,
metadata,
IntegerCompactCiphertextListUnpackingMode::NoUnpacking,
IntegerCompactCiphertextListCastingMode::NoCasting,
)?,
@@ -274,7 +276,13 @@ mod zk {
};
self.inner
.verify_and_expand(public_params, &pk.key.key, unpacking_mode, casting_mode)
.verify_and_expand(
public_params,
&pk.key.key,
metadata,
unpacking_mode,
casting_mode,
)
.map(|expander| CompactCiphertextListExpander {
inner: expander,
tag: self.tag.clone(),
@@ -405,10 +413,11 @@ impl CompactCiphertextListBuilder {
pub fn build_with_proof_packed(
&self,
public_params: &CompactPkePublicParams,
metadata: &[u8],
compute_load: ZkComputeLoad,
) -> crate::Result<ProvenCompactCiphertextList> {
self.inner
.build_with_proof_packed(public_params, compute_load)
.build_with_proof_packed(public_params, metadata, compute_load)
.map(|proved_list| ProvenCompactCiphertextList {
inner: proved_list,
tag: self.tag.clone(),
@@ -496,19 +505,21 @@ mod tests {
// Intentionally low to that we test when multiple lists and proofs are needed
let crs = CompactPkeCrs::from_config(config, 32).unwrap();
let metadata = [b'h', b'l', b'a', b'p', b'i'];
let compact_list = ProvenCompactCiphertextList::builder(&pk)
.push(17u32)
.push(-1i64)
.push(false)
.push_with_num_bits(3u32, 2)
.unwrap()
.build_with_proof_packed(crs.public_params(), ZkComputeLoad::Proof)
.build_with_proof_packed(crs.public_params(), &metadata, ZkComputeLoad::Proof)
.unwrap();
let serialized = bincode::serialize(&compact_list).unwrap();
let compact_list: ProvenCompactCiphertextList = bincode::deserialize(&serialized).unwrap();
let expander = compact_list
.verify_and_expand(crs.public_params(), &pk)
.verify_and_expand(crs.public_params(), &pk, &metadata)
.unwrap();
{
@@ -563,19 +574,21 @@ mod tests {
// Intentionally low to that we test when multiple lists and proofs are needed
let crs = CompactPkeCrs::from_config(config, 32).unwrap();
let metadata = [b'h', b'l', b'a', b'p', b'i'];
let compact_list = ProvenCompactCiphertextList::builder(&pk)
.push(17u32)
.push(-1i64)
.push(false)
.push_with_num_bits(3u32, 2)
.unwrap()
.build_with_proof_packed(crs.public_params(), ZkComputeLoad::Proof)
.build_with_proof_packed(crs.public_params(), &metadata, ZkComputeLoad::Proof)
.unwrap();
let serialized = bincode::serialize(&compact_list).unwrap();
let compact_list: ProvenCompactCiphertextList = bincode::deserialize(&serialized).unwrap();
let expander = compact_list
.verify_and_expand(crs.public_params(), &pk)
.verify_and_expand(crs.public_params(), &pk, &metadata)
.unwrap();
{

10
tfhe/src/high_level_api/tests/tags_on_entities.rs
View File

@@ -34,6 +34,8 @@ fn test_tag_propagation_zk_pok() {
ConfigBuilder::with_custom_parameters(PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64).build();
let crs = crate::zk::CompactPkeCrs::from_config(config, (2 * 32) + (2 * 64) + 2).unwrap();
let metadata = [b'h', b'l', b'a', b'p', b'i'];
let mut cks = ClientKey::generate(config);
let tag_value = random();
cks.tag_mut().set_u64(tag_value);
@@ -55,14 +57,18 @@ fn test_tag_propagation_zk_pok() {
.push(i64::MIN)
.push(false)
.push(true)
.build_with_proof_packed(crs.public_params(), crate::zk::ZkComputeLoad::Proof)
.build_with_proof_packed(
crs.public_params(),
&metadata,
crate::zk::ZkComputeLoad::Proof,
)
.unwrap();
let list_packed: ProvenCompactCiphertextList = serialize_then_deserialize(list_packed);
assert_eq!(list_packed.tag(), cks.tag());
let expander = list_packed
.verify_and_expand(crs.public_params(), &cpk)
.verify_and_expand(crs.public_params(), &cpk, &metadata)
.unwrap();
{

27
tfhe/src/integer/ciphertext/compact_list.rs
View File

@@ -193,11 +193,13 @@ impl CompactCiphertextListBuilder {
pub fn build_with_proof(
&self,
public_params: &CompactPkePublicParams,
metadata: &[u8],
load: ZkComputeLoad,
) -> crate::Result<ProvenCompactCiphertextList> {
let ct_list = self.pk.key.encrypt_and_prove_slice(
self.messages.as_slice(),
public_params,
metadata,
load,
self.pk.key.parameters.message_modulus.0 as u64,
)?;
@@ -211,6 +213,7 @@ impl CompactCiphertextListBuilder {
pub fn build_with_proof_packed(
&self,
public_params: &CompactPkePublicParams,
metadata: &[u8],
load: ZkComputeLoad,
) -> crate::Result<ProvenCompactCiphertextList> {
if self.pk.key.parameters.carry_modulus.0 < self.pk.key.parameters.message_modulus.0 {
@@ -230,6 +233,7 @@ impl CompactCiphertextListBuilder {
let ct_list = self.pk.key.encrypt_and_prove_slice(
packed_messages.as_slice(),
public_params,
metadata,
load,
msg_mod * msg_mod,
)?;
@@ -568,14 +572,17 @@ impl ProvenCompactCiphertextList {
&self,
public_params: &CompactPkePublicParams,
public_key: &CompactPublicKey,
metadata: &[u8],
) -> ZkVerificationOutCome {
self.ct_list.verify(public_params, &public_key.key)
self.ct_list
.verify(public_params, &public_key.key, metadata)
}
pub fn verify_and_expand(
&self,
public_params: &CompactPkePublicParams,
public_key: &CompactPublicKey,
metadata: &[u8],
unpacking_mode: IntegerCompactCiphertextListUnpackingMode<'_>,
casting_mode: IntegerCompactCiphertextListCastingMode<'_>,
) -> crate::Result<CompactCiphertextListExpander> {
@@ -592,9 +599,12 @@ impl ProvenCompactCiphertextList {
)));
}
let expanded_blocks =
self.ct_list
.verify_and_expand(public_params, &public_key.key, casting_mode.into())?;
let expanded_blocks = self.ct_list.verify_and_expand(
public_params,
&public_key.key,
metadata,
casting_mode.into(),
)?;
let expanded_blocks = if is_packed {
match unpacking_mode {
@@ -643,6 +653,10 @@ impl ProvenCompactCiphertextList {
self.ct_list.proof_size()
}
pub fn message_modulus(&self) -> MessageModulus {
self.ct_list.message_modulus()
}
pub fn len(&self) -> usize {
self.info.len()
}
@@ -679,6 +693,8 @@ mod tests {
let ksk_params = PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let fhe_params = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let metadata = [b'i', b'n', b't', b'e', b'g', b'e', b'r'];
let num_blocks = 4usize;
let modulus = (pke_params.message_modulus.0 as u64)
.checked_pow(num_blocks as u32)
@@ -697,13 +713,14 @@ mod tests {
let proven_ct = CompactCiphertextList::builder(&pk)
.extend_with_num_blocks(msgs.iter().copied(), num_blocks)
.build_with_proof_packed(crs.public_params(), ZkComputeLoad::Proof)
.build_with_proof_packed(crs.public_params(), &metadata, ZkComputeLoad::Proof)
.unwrap();
let expander = proven_ct
.verify_and_expand(
crs.public_params(),
&pk,
&metadata,
IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(&sk),
IntegerCompactCiphertextListCastingMode::CastIfNecessary(ksk.as_view()),
)

6
tfhe/src/js_on_wasm_api/js_high_level_api/integers.rs
View File

@@ -779,11 +779,12 @@ impl ProvenCompactCiphertextList {
&self,
public_params: &CompactPkePublicParams,
public_key: &TfheCompactPublicKey,
metadata: &[u8],
) -> Result<CompactCiphertextListExpander, JsError> {
catch_panic_result(|| {
let inner = self
.0
.verify_and_expand(&public_params.0, &public_key.0)
.verify_and_expand(&public_params.0, &public_key.0, metadata)
.map_err(into_js_error)?;
Ok(CompactCiphertextListExpander(inner))
})
@@ -1009,11 +1010,12 @@ impl CompactCiphertextListBuilder {
pub fn build_with_proof_packed(
&self,
public_params: &CompactPkePublicParams,
metadata: &[u8],
compute_load: ZkComputeLoad,
) -> Result<ProvenCompactCiphertextList, JsError> {
catch_panic_result(|| {
self.0
.build_with_proof_packed(&public_params.0, compute_load.into())
.build_with_proof_packed(&public_params.0, metadata, compute_load.into())
.map_err(into_js_error)
.map(ProvenCompactCiphertextList)
})

10
tfhe/src/js_on_wasm_api/js_high_level_api/zk.rs
View File

@@ -26,7 +26,10 @@ pub struct CompactPkeCrs(pub(crate) crate::core_crypto::entities::CompactPkeCrs)
#[wasm_bindgen]
pub struct CompactPkePublicParams(pub(crate) crate::zk::CompactPkePublicParams);
#[allow(clippy::use_self, reason = "wasm bindgen is fragile and prefers the actual type vs. Self")]
#[allow(
clippy::use_self,
reason = "wasm bindgen is fragile and prefers the actual type vs. Self"
)]
#[wasm_bindgen]
impl CompactPkePublicParams {
#[wasm_bindgen]
@@ -73,7 +76,10 @@ impl CompactPkePublicParams {
}
}
#[allow(clippy::use_self, reason = "wasm bindgen is fragile and prefers the actual type vs. Self")]
#[allow(
clippy::use_self,
reason = "wasm bindgen is fragile and prefers the actual type vs. Self"
)]
#[wasm_bindgen]
impl CompactPkeCrs {
#[wasm_bindgen]

23
tfhe/web_wasm_parallel_tests/worker.js
View File

@@ -378,6 +378,10 @@ async function compactPublicKeyZeroKnowledge() {
console.timeEnd("CRS generation");
let public_params = crs.public_params();
// 320 bits is a use case we have, 8 bits per byte
const metadata = new Uint8Array(320 / 8);
crypto.getRandomValues(metadata);
{
let input = generateRandomBigInt(64);
let start = performance.now();
@@ -386,6 +390,7 @@ async function compactPublicKeyZeroKnowledge() {
builder.push_u64(input);
let list = builder.build_with_proof_packed(
public_params,
metadata,
ZkComputeLoad.Proof,
);
let end = performance.now();
@@ -399,7 +404,11 @@ async function compactPublicKeyZeroKnowledge() {
console.log("CompactCiphertextList size:", serialized.length);
let deserialized = ProvenCompactCiphertextList.deserialize(serialized);
let expander = deserialized.verify_and_expand(public_params, publicKey);
let expander = deserialized.verify_and_expand(
public_params,
publicKey,
metadata,
);
assert_eq(expander.get_uint64(0).decrypt(clientKey), input);
}
@@ -418,6 +427,7 @@ async function compactPublicKeyZeroKnowledge() {
}
let encrypted = builder.build_with_proof_packed(
public_params,
metadata,
ZkComputeLoad.Proof,
);
let end = performance.now();
@@ -427,7 +437,11 @@ async function compactPublicKeyZeroKnowledge() {
" ms",
);
let expander = encrypted.verify_and_expand(public_params, publicKey);
let expander = encrypted.verify_and_expand(
public_params,
publicKey,
metadata,
);
assert_eq(expander.get_uint64(0).decrypt(clientKey), inputs[0]);
@@ -651,6 +665,10 @@ async function compactPublicKeyZeroKnowledgeBench() {
let crs = CompactPkeCrs.from_config(config, encrypt_count * 64);
console.timeEnd("CRS generation");
// 320 bits is a use case we have, 8 bits per byte
const metadata = new Uint8Array(320 / 8);
crypto.getRandomValues(metadata);
let public_params = crs.public_params();
let inputs = Array.from(Array(encrypt_count).keys()).map((_) => U64_MAX);
for (const loadChoice of load_choices) {
@@ -666,6 +684,7 @@ async function compactPublicKeyZeroKnowledgeBench() {
const start = performance.now();
let list = compact_list_builder.build_with_proof_packed(
public_params,
metadata,
loadChoice,
);
const end = performance.now();