feat: make sure keygen generates homPRF key

ai-docs/ARCHITECTURE.md

@@ -129,10 +129,19 @@ The primary service is `CoreServiceEndpoint`. Its RPCs group into:
   preprocessing), `KeyGen`, and `NewEpoch` for key rotation. Multiple keyset
   configurations are supported (standard, decompression-only, compressed
   variants).
+  Standard threshold keygen persists a dedicated OPRF LWE secret-key share in
+  each party's private key material and includes the corresponding OPRF server
+  key in the generated TFHE server key. Legacy private keysets that predate this
+  field are upgraded with the OPRF share absent; `UseExisting` keygen generates
+  and persists a fresh OPRF share for such legacy material before regenerating
+  public keys.
 - **Decryption** — `PublicDecrypt` (returns plaintext) and `UserDecrypt`
   (user-initiated, EIP-712 authenticated).
 - **CRS** — `CRSGen` for ZK-proof common reference strings.
-- **Reshare** — `Reshare` to rotate parties / refresh secret shares.
+- **Reshare** — `Reshare` to rotate parties / refresh secret shares. When
+  resharing legacy key material that has no dedicated OPRF secret-key share,
+  the OPRF sub-protocol is skipped and the reshared private keyset keeps that
+  field absent.
 - **Session management** — creation, result retrieval, and cleanup for
   long-running threshold sessions.
 

core/experiments/src/choreography/tfhe_rs/grpc.rs

@@ -2593,9 +2593,11 @@ where
             let public_key_set_decompressed = key_ref.pub_keyset_decompressed.clone();
             let old_private_key_set = key_ref.priv_keyset.as_ref().clone();
             let params = key_ref.as_ref().params;
+            let oprf_key_present = old_private_key_set.oprf_secret_key_share.is_some();
 
             //Perform preprocessing
-            let num_needed_preproc = ResharePreprocRequired::new(num_parties, params);
+            let num_needed_preproc =
+                ResharePreprocRequired::new(num_parties, params, oprf_key_present);
 
             let (mut preprocessing_64, mut preprocessing_128, sessions) = match session_type {
                 SessionType::Small => {
@@ -2678,6 +2680,7 @@ where
                 &mut preprocessing_64,
                 &mut Some(old_private_key_set),
                 params,
+                oprf_key_present,
             )
             .await
             .unwrap();

core/service/Cargo.toml

@@ -103,6 +103,7 @@ tfhe = { workspace = true, features = [
     "experimental-force_fft_algo_dif4",
     "extended-types",
 ] }
+tfhe-csprng.workspace = true
 tfhe-versionable.workspace = true
 thiserror.workspace = true
 thread-handles = { workspace = true, optional = true }

core/service/src/client/key_gen.rs

@@ -509,4 +509,61 @@ pub(crate) mod tests {
             }
         }
     }
+
+    /// Verifies that the dedicated OPRF server key embedded in `server_key` is
+    /// consistent with the OPRF private key in `client_key` by running the
+    /// encrypted PRF for several seeds and comparing each output to the
+    /// independently computed cleartext reference.
+    pub(crate) fn check_oprf_correctness(
+        server_key: &tfhe::ServerKey,
+        client_key: &tfhe::ClientKey,
+    ) {
+        use threshold_execution::tfhe_internals::test_feature::assert_oprf_matches_plaintext;
+
+        #[cfg(not(feature = "slow_tests"))]
+        const NUM_SEEDS: u128 = 2;
+        #[cfg(feature = "slow_tests")]
+        const NUM_SEEDS: u128 = 50;
+
+        let (integer_server_key, _, _, _, _, _, _, oprf_server_key, _) =
+            server_key.clone().into_raw_parts();
+        let Some(oprf_server_key) = oprf_server_key else {
+            panic!("expected oprf_server_key")
+        };
+        let target_shortint_server_key = integer_server_key.into_raw_parts();
+
+        let (
+            integer_client_key,
+            _compact_client_key,
+            _compression_key,
+            _noise_squashing_key,
+            _noise_squashing_compression_key,
+            _rerand_parameters,
+            oprf_private_key,
+            _tag,
+        ) = client_key.clone().into_raw_parts();
+        let Some(oprf_private_key) = oprf_private_key else {
+            panic!("expected oprf_private_key")
+        };
+
+        let shortint_ck = tfhe::shortint::ClientKey {
+            atomic_pattern: integer_client_key.into_raw_parts().atomic_pattern,
+        };
+
+        let prf_lwe_sk = match oprf_private_key.into_raw_parts().into_raw_parts() {
+            tfhe::shortint::oprf::AtomicPatternOprfPrivateKey::Standard(sk) => sk,
+            tfhe::shortint::oprf::AtomicPatternOprfPrivateKey::KeySwitch32(_) => {
+                panic!("Unsupported AtomicPatternOprfPrivateKey::KeySwitch32")
+            }
+        };
+        let oprf_server_key = oprf_server_key.into_raw_parts();
+
+        assert_oprf_matches_plaintext(
+            &shortint_ck,
+            &target_shortint_server_key,
+            &oprf_server_key,
+            &prf_lwe_sk,
+            NUM_SEEDS,
+        );
+    }
 }

core/service/src/client/tests/centralized/key_gen_tests.rs

@@ -363,6 +363,13 @@ pub async fn run_key_gen_centralized(
         assert_eq!(&tag, public_key.tag());
         assert_eq!(&tag, server_key.tag());
 
+        let (_, _, _, _, _, _, _, oprf_key, _) = server_key.clone().into_raw_parts();
+        assert!(
+            oprf_key.is_some(),
+            "centralized full keygen must embed a dedicated OPRF server key"
+        );
+
+        crate::client::key_gen::tests::check_oprf_correctness(&server_key, &client_key);
         crate::client::key_gen::tests::check_conformance(server_key, client_key);
     };
 

core/service/src/client/tests/threshold/key_gen_tests.rs

@@ -151,6 +151,7 @@ impl TestKeyGenResult {
             }
         };
 
+        crate::client::key_gen::tests::check_oprf_correctness(&server_key, client_key);
         check_conformance(server_key.clone(), client_key.clone());
 
         let pt1 = 27u8;
@@ -677,6 +678,7 @@ pub(crate) async fn preproc_and_keygen(
         assert_eq!(&tag, client_key.tag());
         assert_eq!(&tag, public_key.tag());
         assert_eq!(&tag, server_key.tag());
+        crate::client::key_gen::tests::check_oprf_correctness(&server_key, &client_key);
         crate::client::key_gen::tests::check_conformance(server_key, client_key);
     }
 
@@ -1087,6 +1089,7 @@ async fn poll_key_gen_preproc_result(
     resp_response_vec
 }
 
+#[expect(clippy::type_complexity)]
 #[cfg(any(feature = "slow_tests", feature = "insecure"))]
 fn try_reconstruct_shares(
     param: DKGParams,
@@ -1097,6 +1100,7 @@ fn try_reconstruct_shares(
     tfhe::core_crypto::prelude::GlweSecretKeyOwned<u64>,
     tfhe::core_crypto::prelude::GlweSecretKeyOwned<u128>,
     Option<NoiseSquashingCompressionPrivateKey>,
+    Option<tfhe::core_crypto::prelude::LweSecretKeyOwned<u64>>,
 ) {
     use tfhe::core_crypto::prelude::GlweSecretKeyOwned;
     use threshold_execution::tfhe_internals::{
@@ -1219,11 +1223,33 @@ fn try_reconstruct_shares(
             None
         };
 
+    let oprf_lwe_shares = all_threshold_fhe_keys
+        .iter()
+        .filter_map(|(k, v)| {
+            v.private_keys
+                .oprf_secret_key_share
+                .clone()
+                .map(|share| (*k, share.convert_to_z64().data))
+        })
+        .collect::<HashMap<_, _>>();
+    let oprf_lwe_secret_key = if oprf_lwe_shares.len() == all_threshold_fhe_keys.len() {
+        Some(
+            tfhe::core_crypto::prelude::LweSecretKeyOwned::from_container(reconstruct_bit_vec(
+                oprf_lwe_shares,
+                param_handle.lwe_dimension().0,
+                threshold,
+            )),
+        )
+    } else {
+        None
+    };
+
     (
         lwe_secret_key,
         glwe_sk,
         sns_glwe_sk,
         sns_compression_private_key,
+        oprf_lwe_secret_key,
     )
 }
 
@@ -1348,7 +1374,7 @@ pub(crate) async fn verify_keygen_responses(
     }
 
     let threshold = total_num_parties.div_ceil(3) - 1;
-    let (lwe_sk, glwe_sk, sns_glwe_sk, sns_compression_sk) = try_reconstruct_shares(
+    let (lwe_sk, glwe_sk, sns_glwe_sk, sns_compression_sk, oprf_lwe_sk) = try_reconstruct_shares(
         internal_client.params,
         threshold,
         all_threshold_fhe_keys.clone(),
@@ -1363,6 +1389,7 @@ pub(crate) async fn verify_keygen_responses(
         None,
         Some(sns_glwe_sk),
         sns_compression_sk,
+        oprf_lwe_sk,
     )
     .unwrap();
 
@@ -1880,6 +1907,11 @@ async fn secure_threshold_compressed_keygen_from_existing() -> anyhow::Result<()
                 CryptoMaterialReader::read_from_storage(storage, &keygen_id_2).await?;
 
             let (pk, server_key) = compressed_keyset.decompress().unwrap().into_raw_parts();
+            let (_, _, _, _, _, _, _, oprf_key, _) = server_key.clone().into_raw_parts();
+            assert!(
+                oprf_key.is_some(),
+                "Party {party_id}: compressed UseExisting keygen must embed a dedicated OPRF key"
+            );
             assert_eq!(
                 pk.tag(),
                 &expected_tag,
@@ -1930,6 +1962,27 @@ async fn secure_threshold_compressed_keygen_from_existing() -> anyhow::Result<()
                     &PrivDataType::FheKeyInfo.to_string(),
                 )
                 .await?;
+            let threshold_keys_old: crate::engine::threshold::service::ThresholdFheKeys =
+                read_versioned_at_request_and_epoch_id(
+                    &priv_storage,
+                    &keygen_id_1,
+                    &DEFAULT_EPOCH_ID,
+                    &PrivDataType::FheKeyInfo.to_string(),
+                )
+                .await?;
+            let old_oprf_share = &threshold_keys_old
+                .private_keys
+                .as_ref()
+                .oprf_secret_key_share;
+            let new_oprf_share = &threshold_keys.private_keys.as_ref().oprf_secret_key_share;
+            assert!(
+                old_oprf_share.is_some(),
+                "Party {party_id}: fresh keygen must persist the dedicated OPRF private share"
+            );
+            assert_eq!(
+                old_oprf_share, new_oprf_share,
+                "Party {party_id}: UseExisting keygen must reuse the persisted OPRF private share"
+            );
             match &threshold_keys.meta_data {
                 KeyGenMetadata::Current(inner) => {
                     let signed_pk_digest = inner

core/service/src/client/tests/threshold/mpc_epoch_tests.rs

@@ -290,6 +290,7 @@ pub(crate) async fn new_epoch_with_reshare_and_crs(
             let PrivateKeySet {
                 lwe_encryption_secret_key_share,
                 lwe_compute_secret_key_share,
+                oprf_secret_key_share,
                 glwe_secret_key_share,
                 glwe_secret_key_share_sns_as_lwe,
                 glwe_secret_key_share_compression,
@@ -300,6 +301,7 @@ pub(crate) async fn new_epoch_with_reshare_and_crs(
             let PrivateKeySet {
                 lwe_encryption_secret_key_share: reshared_lwe_encryption_secret_key_share,
                 lwe_compute_secret_key_share: reshared_lwe_compute_secret_key_share,
+                oprf_secret_key_share: reshared_oprf_secret_key_share,
                 glwe_secret_key_share: reshared_glwe_secret_key_share,
                 glwe_secret_key_share_sns_as_lwe: reshared_glwe_secret_key_share_sns_as_lwe,
                 glwe_secret_key_share_compression: reshared_glwe_secret_key_share_compression,
@@ -318,6 +320,9 @@ pub(crate) async fn new_epoch_with_reshare_and_crs(
                 lwe_compute_secret_key_share,
                 reshared_lwe_compute_secret_key_share
             );
+            if oprf_secret_key_share.is_some() {
+                assert_ne!(oprf_secret_key_share, reshared_oprf_secret_key_share);
+            }
             assert_ne!(glwe_secret_key_share, reshared_glwe_secret_key_share);
             if glwe_secret_key_share_sns_as_lwe.is_some() {
                 assert_ne!(

core/service/src/engine/centralized/central_kms.rs

@@ -1415,6 +1415,13 @@ pub(crate) mod tests {
             decompressed.is_ok(),
             "Compressed keyset should be decompressible"
         );
+
+        let (_pk, server_key) = decompressed.unwrap().into_raw_parts();
+        let (_, _, _, _, _, _, _, oprf_key, _) = server_key.into_raw_parts();
+        assert!(
+            oprf_key.is_some(),
+            "centralized full keygen must embed a dedicated OPRF server key"
+        );
     }
 
     #[tokio::test]

core/service/src/engine/threshold/service/epoch_manager.rs

@@ -509,11 +509,22 @@ impl<
                             .await?
                     }
                 };
+                // S1 has the previous epoch's private shares, so we read
+                // `oprf_key_present` from local state. The S2-only path in
+                // `reshare_as_set_2` has no private share and derives the same
+                // flag from the verified public `ServerKey` instead. Both
+                // derivations must yield the same value for the reshare
+                // sub-protocols to converge; this holds by construction
+                // because public and private OPRF material are produced
+                // together (legacy keysets predating the dedicated OPRF share
+                // have neither).
+                let oprf_key_present = private_keys.oprf_secret_key_share.is_some();
 
                 Reshare::reshare_sk_two_sets_as_s1(
                     &mut two_sets_session,
                     &mut private_keys,
                     key_info.key_parameters,
+                    oprf_key_present,
                 )
                 .await?;
                 keys_metadata.push(key_metadata);
@@ -831,9 +842,23 @@ impl<
 
             let mut new_private_keysets = Vec::new();
             let sessions_online = &mut (two_sets_session, session_online);
-            for key_info in verified_previous_epoch.keys_info.iter() {
-                let num_needed_preproc =
-                    ResharePreprocRequired::new(num_parties_set_1, key_info.key_parameters);
+            for (key_info, verified_material) in verified_previous_epoch
+                .keys_info
+                .iter()
+                .zip_eq(verified_fhe_public_materials.iter())
+            {
+                // S2 has no private share for the previous epoch, so unlike
+                // the S1 / both-sets paths (which read
+                // `private_keys.oprf_secret_key_share.is_some()`) we derive
+                // `oprf_key_present` from the verified public `ServerKey`.
+                // The protocol assumes both derivations yield the same value
+                // — see the comment in `reshare_as_set_1`.
+                let oprf_key_present = verified_material.has_oprf_key();
+                let num_needed_preproc = ResharePreprocRequired::new(
+                    num_parties_set_1,
+                    key_info.key_parameters,
+                    oprf_key_present,
+                );
 
                 let (mut correlated_randomness_z64, mut correlated_randomness_z128) =
                     Self::compute_s2_preproc(
@@ -848,6 +873,7 @@ impl<
                     &mut correlated_randomness_z128,
                     &mut correlated_randomness_z64,
                     key_info.key_parameters,
+                    oprf_key_present,
                 )
                 .await?;
 
@@ -962,9 +988,17 @@ impl<
                             .await?
                     }
                 };
-
-                let num_needed_preproc =
-                    ResharePreprocRequired::new(num_parties_set_1, key_info.key_parameters);
+                // Same as `reshare_as_set_1`: derived from local private
+                // state. The pure-S2 path in `reshare_as_set_2` derives the
+                // same flag from the verified public `ServerKey`; both must
+                // agree.
+                let oprf_key_present = private_keys.oprf_secret_key_share.is_some();
+
+                let num_needed_preproc = ResharePreprocRequired::new(
+                    num_parties_set_1,
+                    key_info.key_parameters,
+                    oprf_key_present,
+                );
                 let (mut correlated_randomness_z64, mut correlated_randomness_z128) =
                     Self::compute_s2_preproc(
                         &mut session_z64_set_2,
@@ -979,6 +1013,7 @@ impl<
                     &mut correlated_randomness_z64,
                     &mut private_keys,
                     key_info.key_parameters,
+                    oprf_key_present,
                 )
                 .await?;
                 new_private_keysets.push(new_private_keyset);