From c10f1def705f900569e32f2ac936b7ef8e238cf2 Mon Sep 17 00:00:00 2001
From: Thomas Montaigu <thomas.montaigu@zama.ai>
Date: Thu, 23 Oct 2025 15:06:53 +0200
Subject: [PATCH] fix: Tag propagation in XofKeySet

---
 tfhe/src/high_level_api/xof_key_set.rs | 32 +++++++++++++++++++-------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/tfhe/src/high_level_api/xof_key_set.rs b/tfhe/src/high_level_api/xof_key_set.rs
index e12d8fcbf..d0210a49a 100644
--- a/tfhe/src/high_level_api/xof_key_set.rs
+++ b/tfhe/src/high_level_api/xof_key_set.rs
@@ -27,7 +27,7 @@ use crate::shortint::server_key::{
 use crate::{
     integer, shortint, CompactPublicKey, CompressedCompactPublicKey,
     CompressedReRandomizationKeySwitchingKey, CompressedServerKey, ReRandomizationKeySwitchingKey,
-    ServerKey, Tag,
+    ServerKey,
 };
 use aligned_vec::ABox;
 use serde::{Deserialize, Serialize};
@@ -75,6 +75,7 @@ mod cfg_test_imports {
     pub(super) use tfhe_csprng::seeders::Seed;
 
     pub(super) use crate::high_level_api::keys::ReRandomizationKeyGenerationInfo;
+    pub(super) use crate::Tag;
 }
 #[cfg(test)]
 use cfg_test_imports::*;
@@ -193,7 +194,7 @@ impl CompressedXofKeySet {
             noise_squashing_key,
             noise_squashing_compression_key,
             integer_cpk_re_rand_ksk,
-            Tag::default(),
+            self.compressed_server_key.tag,
         );
 
         Ok(XofKeySet {
@@ -253,6 +254,7 @@ impl CompressedCompactPublicKey {
     fn generate_with_pre_seeded_generator<Gen>(
         private_key: &CompactPrivateKey,
         generator: &mut EncryptionRandomGenerator<Gen>,
+        tag: Tag,
     ) -> Self
     where
         Gen: ByteRandomGenerator,
@@ -280,7 +282,7 @@ impl CompressedCompactPublicKey {
                     private_key.0.key.parameters(),
                 ),
             ),
-            Tag::default(),
+            tag,
         )
     }
 
@@ -307,7 +309,7 @@ impl CompressedCompactPublicKey {
 
         let shortint_pk = shortint::CompactPublicKey::from_raw_parts(pk, shortint_cpk.parameters);
         let integer_pk = integer::CompactPublicKey::from_raw_parts(shortint_pk);
-        CompactPublicKey::from_raw_parts(integer_pk, Tag::default())
+        CompactPublicKey::from_raw_parts(integer_pk, self.tag.clone())
     }
 }
 
@@ -1437,6 +1439,7 @@ mod test {
                 CompressedCompactPublicKey::generate_with_pre_seeded_generator(
                     dedicated_pk_key,
                     &mut encryption_rand_gen,
+                    ck.tag.clone(),
                 );
 
             let glwe_secret_key = match &shortint_client_key.atomic_pattern {
@@ -1620,7 +1623,7 @@ mod test {
                 noise_squashing_bs_key,
                 noise_squashing_compression_key,
                 cpk_re_randomization_key_switching_key_material,
-                Tag::default(),
+                ck.tag.clone(),
             );
 
             Ok(Self {
@@ -1651,7 +1654,8 @@ mod test {
             .enable_ciphertext_re_randomization(re_rand_ksk_params)
             .build();
 
-        let cks = ClientKey::generate(config);
+        let mut cks = ClientKey::generate(config);
+        cks.tag_mut().set_data(b"classic 2_2");
 
         let mut seeder = new_seeder();
         let pub_seed = XofSeed::new_u128(
@@ -1664,6 +1668,8 @@ mod test {
         );
 
         let compressed_key_set = CompressedXofKeySet::with_seed(pub_seed, priv_seed, &cks).unwrap();
+        assert_eq!(cks.tag(), compressed_key_set.compressed_public_key.tag());
+        assert_eq!(cks.tag(), compressed_key_set.compressed_server_key.tag());
         test_xof_key_set(&compressed_key_set, config, &cks);
     }
 
@@ -1687,7 +1693,8 @@ mod test {
             .enable_ciphertext_re_randomization(re_rand_ksk_params)
             .build();
 
-        let cks = ClientKey::generate(config);
+        let mut cks = ClientKey::generate(config);
+        cks.tag_mut().set_data(b"ks32 big pke");
 
         let mut seeder = new_seeder();
         let pub_seed = XofSeed::new_u128(
@@ -1700,6 +1707,8 @@ mod test {
         );
 
         let compressed_key_set = CompressedXofKeySet::with_seed(pub_seed, priv_seed, &cks).unwrap();
+        assert_eq!(cks.tag(), compressed_key_set.compressed_public_key.tag());
+        assert_eq!(cks.tag(), compressed_key_set.compressed_server_key.tag());
         test_xof_key_set(&compressed_key_set, config, &cks);
     }
 
@@ -1723,7 +1732,8 @@ mod test {
             .enable_ciphertext_re_randomization(re_rand_ksk_params)
             .build();
 
-        let cks = ClientKey::generate(config);
+        let mut cks = ClientKey::generate(config);
+        cks.tag_mut().set_data(b"ks32 small pke");
 
         let mut seeder = new_seeder();
         let pub_seed = XofSeed::new_u128(
@@ -1736,6 +1746,8 @@ mod test {
         );
 
         let compressed_key_set = CompressedXofKeySet::with_seed(pub_seed, priv_seed, &cks).unwrap();
+        assert_eq!(cks.tag(), compressed_key_set.compressed_public_key.tag());
+        assert_eq!(cks.tag(), compressed_key_set.compressed_server_key.tag());
         test_xof_key_set(&compressed_key_set, config, &cks);
     }
 
@@ -1752,6 +1764,8 @@ mod test {
             crate::safe_serialization::safe_deserialize(data.as_slice(), compressed_size_limit)
                 .unwrap();
 
+        let expected_pk_tag = compressed_key_set.compressed_public_key.tag().clone();
+        let expected_sk_tag = compressed_key_set.compressed_server_key.tag().clone();
         let key_set = compressed_key_set.decompress().unwrap();
 
         let size_limit = 1 << 32;
@@ -1761,6 +1775,8 @@ mod test {
             crate::safe_serialization::safe_deserialize(data.as_slice(), size_limit).unwrap();
 
         let (pk, sk) = key_set.into_raw_parts();
+        assert_eq!(pk.tag(), &expected_pk_tag);
+        assert_eq!(sk.tag(), &expected_sk_tag);
 
         assert!(sk.is_conformant(&config.into()));