diff --git a/coprocessor/fhevm-engine/fhevm-engine-common/src/types.rs b/coprocessor/fhevm-engine/fhevm-engine-common/src/types.rs index d8080cc8c7..6feefb836c 100644 --- a/coprocessor/fhevm-engine/fhevm-engine-common/src/types.rs +++ b/coprocessor/fhevm-engine/fhevm-engine-common/src/types.rs @@ -23,6 +23,12 @@ pub enum FhevmError { DeserializationError(Box), CiphertextExpansionError(tfhe::Error), ReRandomisationError(tfhe::Error), + CiphertextCompressionError(tfhe::Error), + CiphertextCompressionRequiresEmptyCarries, + CiphertextCompressionPanic { + message: String, + }, + CannotCompressScalar, CiphertextExpansionUnsupportedCiphertextKind(tfhe::FheTypes), FheOperationOnlyOneOperandCanBeScalar { fhe_operation: i32, @@ -153,6 +159,21 @@ impl std::fmt::Display for FhevmError { Self::ReRandomisationError(e) => { write!(f, "error re-randomising ciphertext: {:?}", e) } + Self::CiphertextCompressionError(e) => { + write!(f, "error compressing ciphertext: {:?}", e) + } + Self::CiphertextCompressionRequiresEmptyCarries => { + write!( + f, + "cannot compress ciphertext because block carries are not empty" + ) + } + Self::CiphertextCompressionPanic { message } => { + write!(f, "panic while compressing ciphertext: {}", message) + } + Self::CannotCompressScalar => { + write!(f, "cannot compress scalar input") + } Self::CiphertextExpansionUnsupportedCiphertextKind(e) => { write!( f, @@ -322,6 +343,18 @@ impl std::fmt::Display for FhevmError { } } +// TFHE panics with both &str and String payloads depending on call site. +// Normalize to a stable String so callers can log and map consistently. +fn panic_payload_to_string(payload: Box) -> String { + if let Some(message) = payload.downcast_ref::<&str>() { + (*message).to_string() + } else if let Some(message) = payload.downcast_ref::() { + message.clone() + } else { + "unknown panic payload".to_string() + } +} + #[derive(Clone)] pub enum SupportedFheCiphertexts { FheBool(tfhe::FheBool), @@ -522,10 +555,12 @@ impl SupportedFheCiphertexts { } } - pub fn compress(&self) -> (i16, Vec) { - let type_num = self.type_num(); + pub fn compress(&self) -> std::result::Result, FhevmError> { let mut builder = CompressedCiphertextListBuilder::new(); match self { + SupportedFheCiphertexts::Scalar(_) => { + return Err(FhevmError::CannotCompressScalar); + } SupportedFheCiphertexts::FheBool(c) => builder.push(c.clone()), SupportedFheCiphertexts::FheUint4(c) => builder.push(c.clone()), SupportedFheCiphertexts::FheUint8(c) => builder.push(c.clone()), @@ -538,13 +573,21 @@ impl SupportedFheCiphertexts { SupportedFheCiphertexts::FheBytes64(c) => builder.push(c.clone()), SupportedFheCiphertexts::FheBytes128(c) => builder.push(c.clone()), SupportedFheCiphertexts::FheBytes256(c) => builder.push(c.clone()), - SupportedFheCiphertexts::Scalar(_) => { - // TODO: Need to fix that, scalars are not ciphertexts. - panic!("cannot compress a scalar"); + }; + let list = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| builder.build())) + { + Ok(Ok(list)) => list, + Ok(Err(error)) => return Err(FhevmError::CiphertextCompressionError(error)), + Err(panic_payload) => { + let message = panic_payload_to_string(panic_payload); + if message == "Ciphertexts must have empty carries to be compressed" { + return Err(FhevmError::CiphertextCompressionRequiresEmptyCarries); + } + + return Err(FhevmError::CiphertextCompressionPanic { message }); } }; - let list = builder.build().expect("ciphertext compression"); - (type_num, safe_serialize(&list)) + Ok(safe_serialize(&list)) } #[cfg(feature = "gpu")] @@ -1009,3 +1052,15 @@ pub type BlockchainProvider = FillProvider< >, RootProvider, >; + +#[cfg(test)] +mod tests { + use super::{FhevmError, SupportedFheCiphertexts}; + + #[test] + fn compress_scalar_returns_error() { + let scalar = SupportedFheCiphertexts::Scalar(vec![1, 2, 3]); + let compressed = scalar.compress(); + assert!(matches!(compressed, Err(FhevmError::CannotCompressScalar))); + } +} diff --git a/coprocessor/fhevm-engine/scheduler/src/dfg/scheduler.rs b/coprocessor/fhevm-engine/scheduler/src/dfg/scheduler.rs index 5c03f93c55..1363fb66ba 100644 --- a/coprocessor/fhevm-engine/scheduler/src/dfg/scheduler.rs +++ b/coprocessor/fhevm-engine/scheduler/src/dfg/scheduler.rs @@ -554,14 +554,26 @@ fn run_computation( compress_span.set_attribute(KeyValue::new("ct_type", inputs[0].type_name())); compress_span.set_attribute(KeyValue::new("operation", "FheGetCiphertext")); - let (ct_type, ct_bytes) = inputs[0].compress(); - compress_span.set_attribute(KeyValue::new("compressed_size", ct_bytes.len() as i64)); - compress_span.end(); - - ( - graph_node_index, - Ok((inputs[0].clone(), Some((ct_type, ct_bytes)))), - ) + let ct_type = inputs[0].type_num(); + let compressed = inputs[0].compress(); + match compressed { + Ok(ct_bytes) => { + compress_span + .set_attribute(KeyValue::new("compressed_size", ct_bytes.len() as i64)); + compress_span.end(); + ( + graph_node_index, + Ok((inputs[0].clone(), Some((ct_type, ct_bytes)))), + ) + } + Err(error) => { + compress_span.set_status(Status::Error { + description: error.to_string().into(), + }); + compress_span.end(); + (graph_node_index, Err(error.into())) + } + } } Ok(fhe_op) => { let op_name = fhe_op.as_str_name(); @@ -586,13 +598,25 @@ fn run_computation( telemetry::set_txn_id(&mut compress_span, transaction_id); compress_span.set_attribute(KeyValue::new("ct_type", result.type_name())); compress_span.set_attribute(KeyValue::new("operation", op_name)); - - let (ct_type, ct_bytes) = result.compress(); - compress_span - .set_attribute(KeyValue::new("compressed_size", ct_bytes.len() as i64)); - compress_span.end(); - - (graph_node_index, Ok((result, Some((ct_type, ct_bytes))))) + let ct_type = result.type_num(); + let compressed = result.compress(); + match compressed { + Ok(ct_bytes) => { + compress_span.set_attribute(KeyValue::new( + "compressed_size", + ct_bytes.len() as i64, + )); + compress_span.end(); + (graph_node_index, Ok((result, Some((ct_type, ct_bytes))))) + } + Err(error) => { + compress_span.set_status(Status::Error { + description: error.to_string().into(), + }); + compress_span.end(); + (graph_node_index, Err(error.into())) + } + } } else { (graph_node_index, Ok((result, None))) } diff --git a/coprocessor/fhevm-engine/tfhe-worker/src/server.rs b/coprocessor/fhevm-engine/tfhe-worker/src/server.rs index f28afabb2f..ed99e25aa1 100644 --- a/coprocessor/fhevm-engine/tfhe-worker/src/server.rs +++ b/coprocessor/fhevm-engine/tfhe-worker/src/server.rs @@ -518,7 +518,8 @@ impl CoprocessorService { let server_key_clone = server_key.clone(); let (handle, serialized_ct, serialized_type) = spawn_blocking(move || { tfhe::set_server_key(server_key_clone); - let (serialized_type, serialized_ct) = the_ct.compress(); + let serialized_type = the_ct.type_num(); + let serialized_ct = the_ct.compress().unwrap(); let mut handle_hash = Keccak256::new(); handle_hash.update(&blob_hash_clone); handle_hash.update([ct_idx as u8]); @@ -764,7 +765,8 @@ impl CoprocessorService { let ct = trivial_encrypt_be_bytes(v.output_type as i16, &v.be_value); span.end(); let mut span = inner_tracer.child_span("compress_ciphertext"); - let (ct_type, ct_bytes) = ct.compress(); + let ct_type = ct.type_num(); + let ct_bytes = ct.compress().unwrap(); span.end(); res.push((v.handle, ct_type, ct_bytes)); } diff --git a/coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs b/coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs index c4a394cc9e..0c56dce3db 100644 --- a/coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs +++ b/coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs @@ -493,7 +493,8 @@ fn create_ciphertext( // Add the full 256bit hash as re-randomization metadata, NOT the // truncated hash of the handle the_ct.add_re_randomization_metadata(&handle); - let (serialized_type, compressed) = the_ct.compress(); + let serialized_type = the_ct.type_num(); + let compressed = the_ct.compress()?; // idx cast to u8 must succeed because we don't allow // more handles than u8 size