chore: avoid re-serialization of public key in verification (#513)

Author: kc1212

core/service/src/cryptography/mod.rs

@@ -11,8 +11,6 @@ pub mod encryption;
 #[allow(deprecated)]
 pub mod signcryption;
 
-use crate::cryptography::encryption::UnifiedPublicEncKey;
-use crate::cryptography::internal_crypto_types::LegacySerialization;
 #[cfg(any(feature = "non-wasm", test))]
 use crate::cryptography::signatures::PrivateSigKey;
 #[cfg(any(feature = "non-wasm", test))]
@@ -28,17 +26,17 @@ pub(crate) fn compute_external_user_decrypt_signature(
     server_sk: &PrivateSigKey,
     payload: &UserDecryptionResponsePayload,
     eip712_domain: &Eip712Domain,
-    user_pk: &UnifiedPublicEncKey,
+    user_pk_buf: &[u8],
     extra_data: &[u8],
 ) -> anyhow::Result<Vec<u8>> {
-    let message = compute_user_decrypt_message(payload, user_pk, extra_data)?;
+    let message = compute_user_decrypt_message(payload, user_pk_buf, extra_data)?;
     tracing::debug!("Computing signature for UserDecryptResponseVerification");
     compute_eip712_signature(server_sk, &message, eip712_domain)
 }
 
 pub(crate) fn compute_user_decrypt_message(
     payload: &UserDecryptionResponsePayload,
-    user_pk: &UnifiedPublicEncKey,
+    user_pk_buf: &[u8],
     extra_data: &[u8],
 ) -> anyhow::Result<UserDecryptResponseVerification> {
     // convert external_handles back to 256-bit bytes32 to be signed
@@ -62,21 +60,14 @@ pub(crate) fn compute_user_decrypt_message(
 
     let user_decrypted_share_buf = bc2wrap::serialize(payload)?;
 
-    // LEGACY CODE: we used to only support ML-KEM1024 encoded with bincode
-    // the solidity structure to sign with EIP-712
-    // note that the JS client must also use the same encoding to verify the result
-    let user_pk_buf = user_pk
-        .to_legacy_bytes()
-        .map_err(|e| anyhow::anyhow!("serialization error: {e}"))?;
-
     tracing::info!(
         "Computed UserDecryptResponseVerification for handles {:?} and extra data \"{}\".",
         external_handles_bytes32,
         hex::encode(extra_data)
     );
 
     Ok(UserDecryptResponseVerification {
-        publicKey: user_pk_buf.into(),
+        publicKey: user_pk_buf.to_vec().into(),
         ctHandles: external_handles_bytes32,
         userDecryptedShare: user_decrypted_share_buf.into(),
         extraData: extra_data.to_vec().into(),

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

@@ -8,6 +8,7 @@ use crate::cryptography::attestation::SecurityModuleProxy;
 use crate::cryptography::compute_external_user_decrypt_signature;
 use crate::cryptography::decompression;
 use crate::cryptography::encryption::UnifiedPublicEncKey;
+use crate::cryptography::internal_crypto_types::LegacySerialization;
 use crate::cryptography::signatures::{PrivateSigKey, PublicSigKey, Signature};
 use crate::cryptography::signcryption::SigncryptFHEPlaintext;
 use crate::cryptography::signcryption::UnifiedSigncryptionKey;
@@ -506,7 +507,7 @@ pub async fn async_user_decrypt<
     rng: &mut (impl CryptoRng + RngCore),
     typed_ciphertexts: &[TypedCiphertext],
     req_digest: &[u8],
-    client_enc_key: &UnifiedPublicEncKey,
+    client_enc_key_bytes: &[u8],
     client_address: &alloy_primitives::Address,
     server_verf_key: Vec<u8>,
     domain: &alloy_sol_types::Eip712Domain,
@@ -518,6 +519,10 @@ pub async fn async_user_decrypt<
         metrics_names::{OP_USER_DECRYPT_INNER, TAG_TFHE_TYPE},
     };
 
+    // LEGACY CODE: see `from_legacy_bytes` for docs
+    let client_enc_key = UnifiedPublicEncKey::from_legacy_bytes(client_enc_key_bytes)
+        .map_err(|e| anyhow::anyhow!("Error deserializing UnifiedPublicEncKey: {e}"))?;
+
     let mut all_signcrypted_cts = vec![];
     for typed_ciphertext in typed_ciphertexts {
         let mut inner_timer = metrics::METRICS
@@ -538,7 +543,7 @@ pub async fn async_user_decrypt<
             fhe_type,
             ct_format,
             req_digest,
-            client_enc_key,
+            &client_enc_key,
             client_address.as_ref(),
         )?;
         all_signcrypted_cts.push(TypedSigncryptedCiphertext {
@@ -562,7 +567,7 @@ pub async fn async_user_decrypt<
         sig_key,
         &payload,
         domain,
-        client_enc_key,
+        client_enc_key_bytes,
         extra_data,
     )?;
 

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

@@ -47,7 +47,7 @@ pub async fn user_decrypt_impl<
     let (
         typed_ciphertexts,
         link,
-        client_enc_key,
+        client_enc_key_bytes,
         client_address,
         request_id,
         key_id,
@@ -152,7 +152,7 @@ pub async fn user_decrypt_impl<
                 &mut rng,
                 &typed_ciphertexts,
                 &link,
-                &client_enc_key,
+                &client_enc_key_bytes,
                 &client_address,
                 server_verf_key,
                 &domain,

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

@@ -52,6 +52,7 @@ use crate::{
     anyhow_error_and_log,
     cryptography::{
         compute_external_user_decrypt_signature,
+        encryption::UnifiedPublicEncKey,
         error::CryptographyError,
         internal_crypto_types::LegacySerialization,
         signcryption::{SigncryptFHEPlaintext, UnifiedSigncryptionKeyOwned},
@@ -169,6 +170,12 @@ impl<
     ///
     /// This function does not perform user decryption in a background thread.
     /// The return type should be [UserDecryptCallValues] except the final item in the tuple
+    ///
+    /// Note that the argument `client_enc_key_bytes` must be the original
+    /// bytes that was provided by the user, it should not go through any re-serialization.
+    /// The same information is in `signcryption_key.receiver_enc_key` but in there it is
+    /// already serialized, that's why the original bytes representation is still needed
+    /// for signing later on.
     #[allow(clippy::too_many_arguments)]
     async fn inner_user_decrypt(
         req_id: &RequestId,
@@ -179,6 +186,7 @@ impl<
         typed_ciphertexts: Vec<TypedCiphertext>,
         link: Vec<u8>,
         signcryption_key: Arc<UnifiedSigncryptionKeyOwned>,
+        client_enc_key_bytes_orig: Vec<u8>,
         fhe_keys: OwnedRwLockReadGuard<
             HashMap<(RequestId, EpochId), ThresholdFheKeys>,
             ThresholdFheKeys,
@@ -384,7 +392,7 @@ impl<
             &signcryption_key.signing_key,
             &payload,
             domain,
-            &signcryption_key.receiver_enc_key,
+            &client_enc_key_bytes_orig,
             &extra_data,
         )?;
         Ok((payload, external_signature, extra_data))
@@ -456,7 +464,7 @@ impl<
         let (
             typed_ciphertexts,
             link,
-            client_enc_key,
+            client_enc_key_bytes_orig, // the original bytes for the encryption key
             client_address,
             req_id,
             key_id,
@@ -506,6 +514,15 @@ impl<
             )
         })?)
         .clone();
+        let client_enc_key = UnifiedPublicEncKey::from_legacy_bytes(&client_enc_key_bytes_orig)
+            .map_err(|e| {
+                MetricedError::new(
+                    OP_USER_DECRYPT_REQUEST,
+                    Some(req_id),
+                    anyhow::anyhow!("Error deserializing UnifiedPublicEncKey: {e}"),
+                    tonic::Code::Internal,
+                )
+            })?;
         let signcryption_key = Arc::new(UnifiedSigncryptionKeyOwned::new(
             sk,
             client_enc_key,
@@ -543,6 +560,7 @@ impl<
                 typed_ciphertexts,
                 link,
                 signcryption_key,
+                client_enc_key_bytes_orig,
                 fhe_keys_rlock,
                 dec_mode,
                 &domain,

core/service/src/engine/validation_non_wasm.rs

@@ -186,7 +186,7 @@ pub(crate) fn validate_user_decrypt_req(
     (
         Vec<TypedCiphertext>,
         Vec<u8>,
-        UnifiedPublicEncKey,
+        Vec<u8>,
         alloy_primitives::Address,
         RequestId,
         KeyId,
@@ -214,7 +214,7 @@ fn unpack_user_decrypt_req(
     (
         Vec<TypedCiphertext>,
         Vec<u8>,
-        UnifiedPublicEncKey,
+        Vec<u8>,
         alloy_primitives::Address,
         RequestId,
         KeyId,
@@ -261,15 +261,17 @@ fn unpack_user_decrypt_req(
     };
 
     let (link, _) = req.compute_link_checked()?;
-    let client_enc_key = UnifiedPublicEncKey::from_legacy_bytes(&req.enc_key).map_err(|e| {
+    // Deserialize to validate the enc_key bytes, but don't return the typed key —
+    // callers use raw bytes for EIP-712 and deserialize at point-of-use for crypto.
+    let _client_enc_key = UnifiedPublicEncKey::from_legacy_bytes(&req.enc_key).map_err(|e| {
         Into::<Box<dyn std::error::Error + Send + Sync>>::into(anyhow::anyhow!(
             "Error deserializing UnifiedPublicEncKey from UserDecryptionRequest: {e}"
         ))
     })?;
     Ok((
         req.typed_ciphertexts.clone(),
         link,
-        client_enc_key,
+        req.enc_key.clone(),
         client_verf_key,
         request_id,
         key_id,

core/service/src/engine/validation_wasm.rs

@@ -13,8 +13,6 @@ use crate::{
     client::user_decryption_wasm::{ParsedUserDecryptionRequest, compute_link},
     cryptography::{
         compute_user_decrypt_message,
-        encryption::UnifiedPublicEncKey,
-        internal_crypto_types::LegacySerialization,
         signatures::{
             PublicSigKey, Signature, internal_verify_sig, recover_address_from_ext_signature,
         },
@@ -82,12 +80,8 @@ pub(crate) fn check_ext_user_decryption_signature(
     eip712_domain: &Eip712Domain,
     expected_addr: &alloy_primitives::Address,
 ) -> anyhow::Result<()> {
-    // NOTE: we need to support legacy user_pk, so try to deserialize MlKem1024 encoded with bincode first
-    let unified_pk = UnifiedPublicEncKey::from_legacy_bytes(request.enc_key()).map_err(|e| {
-        anyhow_error_and_log(format!("Error deserializing UnifiedPublicEncKey: {e}"))
-    })?;
     let extra_data = request.extra_data();
-    let message = compute_user_decrypt_message(payload, &unified_pk, extra_data)?;
+    let message = compute_user_decrypt_message(payload, request.enc_key(), extra_data)?;
     tracing::debug!(
         "Verifying external user decryption signature for UserDecryptResponseVerification"
     );
@@ -593,7 +587,7 @@ mod tests {
             &sk0,
             &payload,
             &domain,
-            &eph_client_pk,
+            request.enc_key(),
             request.extra_data(),
         )
         .unwrap();
@@ -621,7 +615,7 @@ mod tests {
                 &sk_bad,
                 &payload,
                 &domain,
-                &eph_client_pk,
+                request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -747,7 +741,7 @@ mod tests {
             &sk0,
             &pivot_resp,
             &dummy_domain,
-            &eph_client_pk,
+            client_request.enc_key(),
             &extra_data,
         )
         .unwrap();
@@ -1024,7 +1018,7 @@ mod tests {
                 &sk1,
                 &payload0,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -1053,7 +1047,7 @@ mod tests {
                 &sk2,
                 &payload,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -1082,7 +1076,7 @@ mod tests {
                 &sk3,
                 &payload,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -1111,7 +1105,7 @@ mod tests {
                 &sk4,
                 &payload,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -1249,7 +1243,7 @@ mod tests {
                     &sk3,
                     &payload,
                     &dummy_domain,
-                    &eph_client_pk,
+                    client_request.enc_key(),
                     &[],
                 )
                 .unwrap();
@@ -1399,7 +1393,7 @@ mod tests {
                 &sk1,
                 &payload0,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();
@@ -1428,7 +1422,7 @@ mod tests {
                 &sk2,
                 &payload,
                 &dummy_domain,
-                &eph_client_pk,
+                client_request.enc_key(),
                 &[],
             )
             .unwrap();