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

tfhe/benches/integer/zk_pke.rs

@@ -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(),

tfhe/c_api_tests/test_high_level_zk.c

@@ -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);

tfhe/docs/guides/zk-pok.md

@@ -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,21 +103,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::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()?;
 

tfhe/src/c_api/high_level_api/compact_list.rs

@@ -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,16 +183,25 @@ 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(|| {
         let list = get_ref_checked(compact_list).unwrap();
         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)));

tfhe/src/high_level_api/compact_list.rs

@@ -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();
 
         {

tfhe/src/high_level_api/tests/tags_on_entities.rs

@@ -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();
 
     {