From bd4f54be40ad0186dd2face6e3301e5ab1826475 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Tue, 16 Dec 2025 18:10:22 +0500 Subject: [PATCH 01/19] feat: add share encryption (merge commits) --- Cargo.lock | 1 + .../enclave_event/encryption_key_created.rs | 37 +++ crates/events/src/enclave_event/mod.rs | 7 +- .../enclave_event/threshold_share_created.rs | 23 +- crates/keyshare/Cargo.toml | 1 + .../keyshare/src/encryption_key_collector.rs | 133 +++++++++ crates/keyshare/src/lib.rs | 2 + crates/keyshare/src/threshold_keyshare.rs | 280 ++++++++++++++---- crates/net/src/document_publisher.rs | 76 ++++- crates/test-helpers/src/usecase_helpers.rs | 116 +++++--- crates/tests/tests/integration.rs | 15 + crates/trbfv/Cargo.toml | 3 +- crates/trbfv/src/helpers.rs | 30 +- crates/trbfv/src/shares/bfv_encrypted.rs | 258 ++++++++++++++++ crates/trbfv/src/shares/mod.rs | 4 +- crates/trbfv/src/shares/pvw.rs | 66 ----- crates/trbfv/tests/integration.rs | 13 +- 17 files changed, 877 insertions(+), 188 deletions(-) create mode 100644 crates/events/src/enclave_event/encryption_key_created.rs create mode 100644 crates/keyshare/src/encryption_key_collector.rs create mode 100644 crates/trbfv/src/shares/bfv_encrypted.rs delete mode 100644 crates/trbfv/src/shares/pvw.rs diff --git a/Cargo.lock b/Cargo.lock index 991041a2d5..2c9ec89ed9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3078,6 +3078,7 @@ dependencies = [ "e3-request", "e3-trbfv", "e3-utils", + "fhe", "fhe-traits", "ndarray", "rand 0.8.5", diff --git a/crates/events/src/enclave_event/encryption_key_created.rs b/crates/events/src/enclave_event/encryption_key_created.rs new file mode 100644 index 0000000000..e3deb25af7 --- /dev/null +++ b/crates/events/src/enclave_event/encryption_key_created.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use crate::E3id; +use actix::Message; +use derivative::Derivative; +use e3_utils::utility_types::ArcBytes; +use serde::{Deserialize, Serialize}; +use std::{ + fmt::{self, Display}, + sync::Arc, +}; + +#[derive(Derivative, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[derivative(Debug)] +pub struct EncryptionKey { + pub party_id: u64, + #[derivative(Debug(format_with = "e3_utils::formatters::hexf"))] + pub pk_bfv: ArcBytes, +} + +#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[rtype(result = "()")] +pub struct EncryptionKeyCreated { + pub e3_id: E3id, + pub key: Arc, + pub external: bool, +} + +impl Display for EncryptionKeyCreated { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} diff --git a/crates/events/src/enclave_event/mod.rs b/crates/events/src/enclave_event/mod.rs index 0f7f2b07ad..9965e80f26 100644 --- a/crates/events/src/enclave_event/mod.rs +++ b/crates/events/src/enclave_event/mod.rs @@ -19,6 +19,7 @@ mod die; mod e3_request_complete; mod e3_requested; mod enclave_error; +mod encryption_key_created; mod keyshare_created; mod operator_activation_changed; mod plaintext_aggregated; @@ -47,6 +48,7 @@ pub use die::*; pub use e3_request_complete::*; pub use e3_requested::*; pub use enclave_error::*; +pub use encryption_key_created::*; pub use keyshare_created::*; pub use operator_activation_changed::*; pub use plaintext_aggregated::*; @@ -112,6 +114,7 @@ pub enum EnclaveEventData { Shutdown(Shutdown), DocumentReceived(DocumentReceived), ThresholdShareCreated(ThresholdShareCreated), + EncryptionKeyCreated(EncryptionKeyCreated), /// This is a test event to use in testing TestEvent(TestEvent), } @@ -291,6 +294,7 @@ impl EnclaveEvent { EnclaveEventData::CommitteeFinalized(ref data) => Some(data.e3_id.clone()), EnclaveEventData::TicketGenerated(ref data) => Some(data.e3_id.clone()), EnclaveEventData::TicketSubmitted(ref data) => Some(data.e3_id.clone()), + EnclaveEventData::EncryptionKeyCreated(ref data) => Some(data.e3_id.clone()), _ => None, } } @@ -322,7 +326,8 @@ impl_into_event_data!( Shutdown, TestEvent, DocumentReceived, - ThresholdShareCreated + ThresholdShareCreated, + EncryptionKeyCreated ); impl TryFrom<&EnclaveEvent> for EnclaveError { diff --git a/crates/events/src/enclave_event/threshold_share_created.rs b/crates/events/src/enclave_event/threshold_share_created.rs index d07fbd84b1..51aac476c5 100644 --- a/crates/events/src/enclave_event/threshold_share_created.rs +++ b/crates/events/src/enclave_event/threshold_share_created.rs @@ -7,7 +7,7 @@ use crate::E3id; use actix::Message; use derivative::Derivative; -use e3_trbfv::shares::{PvwEncrypted, SharedSecret}; +use e3_trbfv::shares::BfvEncryptedShares; use e3_utils::utility_types::ArcBytes; use serde::{Deserialize, Serialize}; use std::{ @@ -15,22 +15,23 @@ use std::{ sync::Arc, }; -/// Type Representing Pvw encrypted bytes -pub type PvwBytes = Vec; - -/// PVW encrypted shares list for a party in the DKG +/// BFV-encrypted shares list for a party in the DKG. +/// +/// Each party broadcasts their encrypted shares to all other parties. +/// Each recipient can only decrypt the share meant for them using their +/// BFV secret key. #[derive(Derivative, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[derivative(Debug)] pub struct ThresholdShare { - /// The publishers party_id + /// The publisher's party_id pub party_id: u64, - /// The publishers public key share + /// The publisher's TrBFV public key share #[derivative(Debug(format_with = "e3_utils::formatters::hexf"))] pub pk_share: ArcBytes, - /// PVW encrypted sk_sss list with index determining party_id - pub sk_sss: PvwEncrypted, - /// PVW encrypted esi_sss list with index determining party_id - pub esi_sss: Vec>, + /// BFV-encrypted sk_sss - each recipient can decrypt their share + pub sk_sss: BfvEncryptedShares, + /// BFV-encrypted esi_sss - one per ciphertext, each recipient can decrypt their share + pub esi_sss: Vec, } #[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] diff --git a/crates/keyshare/Cargo.toml b/crates/keyshare/Cargo.toml index 278044d111..46da25753f 100644 --- a/crates/keyshare/Cargo.toml +++ b/crates/keyshare/Cargo.toml @@ -20,6 +20,7 @@ e3-multithread = { workspace = true } e3-request = { workspace = true } e3-trbfv = { workspace = true } e3-utils = { workspace = true } +fhe = { workspace = true } fhe-traits = { workspace = true } rand = { workspace = true } rand_chacha = { workspace = true } diff --git a/crates/keyshare/src/encryption_key_collector.rs b/crates/keyshare/src/encryption_key_collector.rs new file mode 100644 index 0000000000..edceb24dcd --- /dev/null +++ b/crates/keyshare/src/encryption_key_collector.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +//! Collector for BFV encryption keys from all parties. +//! +//! Before parties can encrypt their Shamir shares, they need to collect +//! the BFV public keys from all other parties. This actor handles that +//! collection process. + +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Instant, +}; + +use actix::{Actor, Addr, Handler, Message}; +use e3_events::{EncryptionKey, EncryptionKeyCreated}; +use e3_trbfv::PartyId; +use tracing::info; + +use crate::ThresholdKeyshare; + +/// State of the collector +pub enum CollectorState { + /// Currently collecting keys + Collecting, + /// All keys have been collected + Finished, +} + +/// Message sent when all encryption keys have been collected. +/// +/// This contains all parties' BFV public keys, sorted by party_id, +/// ready to be used for encrypting shares. +#[derive(Message)] +#[rtype(result = "()")] +pub struct AllEncryptionKeysCollected { + /// All collected encryption keys, sorted by party_id + pub keys: Vec>, +} + +impl From>> for AllEncryptionKeysCollected { + fn from(value: HashMap>) -> Self { + // Sort by party_id for deterministic ordering + let mut keys: Vec<_> = value.into_values().collect(); + keys.sort_by_key(|k| k.party_id); + AllEncryptionKeysCollected { keys } + } +} + +/// Actor that collects BFV encryption keys from all parties. +/// +/// Once all keys are collected, it sends `AllEncryptionKeysCollected` to the parent +/// `ThresholdKeyshare` actor. +pub struct EncryptionKeyCollector { + /// Set of party IDs we're still waiting for + todo: HashSet, + /// Parent actor to notify when collection is complete + parent: Addr, + /// Current state + state: CollectorState, + /// Collected keys indexed by party_id + keys: HashMap>, +} + +impl EncryptionKeyCollector { + /// Create and start a new collector. + /// + /// # Arguments + /// * `parent` - The ThresholdKeyshare actor to notify when collection is complete + /// * `total` - Total number of parties (keys to collect) + pub fn setup(parent: Addr, total: u64) -> Addr { + let addr = Self { + todo: (0..total).collect(), + parent, + state: CollectorState::Collecting, + keys: HashMap::new(), + } + .start(); + addr + } +} + +impl Actor for EncryptionKeyCollector { + type Context = actix::Context; +} + +impl Handler for EncryptionKeyCollector { + type Result = (); + fn handle(&mut self, msg: EncryptionKeyCreated, _: &mut Self::Context) -> Self::Result { + let start = Instant::now(); + info!("EncryptionKeyCollector: EncryptionKeyCreated received"); + + if let CollectorState::Finished = self.state { + info!("EncryptionKeyCollector is finished, ignoring"); + return; + } + + let pid = msg.key.party_id; + info!("EncryptionKeyCollector: party_id = {}", pid); + + let Some(_) = self.todo.take(&pid) else { + info!( + "Error: {} was not in encryption key collector's ID list", + pid + ); + return; + }; + + info!( + "Inserting encryption key... waiting on: {}", + self.todo.len() + ); + self.keys.insert(pid, msg.key); + + if self.todo.is_empty() { + info!("All encryption keys collected!"); + self.state = CollectorState::Finished; + let event: AllEncryptionKeysCollected = self.keys.clone().into(); + self.parent.do_send(event); + } + + info!( + "Finished processing EncryptionKeyCreated in {:?}", + start.elapsed() + ); + } +} + + diff --git a/crates/keyshare/src/lib.rs b/crates/keyshare/src/lib.rs index 42ffcc1c62..ecf984cc3c 100644 --- a/crates/keyshare/src/lib.rs +++ b/crates/keyshare/src/lib.rs @@ -4,11 +4,13 @@ // without even the implied warranty of MERCHANTABILITY // or FITNESS FOR A PARTICULAR PURPOSE. +mod encryption_key_collector; pub mod ext; mod keyshare; mod repo; mod threshold_keyshare; mod threshold_share_collector; +pub use encryption_key_collector::*; pub use keyshare::*; pub use repo::*; pub use threshold_keyshare::*; diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index d79398f13a..0d716929af 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -10,8 +10,8 @@ use e3_crypto::{Cipher, SensitiveBytes}; use e3_data::Persistable; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CiphertextOutputPublished, ComputeRequest, - ComputeResponse, DecryptionshareCreated, E3id, EnclaveEvent, EnclaveEventData, KeyshareCreated, - PartyId, ThresholdShare, ThresholdShareCreated, + ComputeResponse, DecryptionshareCreated, E3id, EnclaveEvent, EnclaveEventData, EncryptionKey, + EncryptionKeyCreated, KeyshareCreated, PartyId, ThresholdShare, ThresholdShareCreated, }; use e3_fhe::create_crp; use e3_multithread::Multithread; @@ -22,12 +22,14 @@ use e3_trbfv::{ }, gen_esi_sss::{GenEsiSssRequest, GenEsiSssResponse}, gen_pk_share_and_sk_sss::GenPkShareAndSkSssRequest, - shares::{EncryptableVec, Encrypted, PvwEncrypted, ShamirShare, SharedSecret}, + helpers::{deserialize_secret_key, get_share_encryption_params, serialize_secret_key}, + shares::{BfvEncryptedShares, EncryptableVec, Encrypted, ShamirShare, SharedSecret}, TrBFVConfig, TrBFVRequest, TrBFVResponse, }; use e3_utils::{bail, to_ordered_vec, utility_types::ArcBytes}; -use fhe_traits::Serialize; -use rand::SeedableRng; +use fhe::bfv::{PublicKey, SecretKey}; +use fhe_traits::{DeserializeParametrized, Serialize}; +use rand::{rngs::OsRng, SeedableRng}; use rand_chacha::ChaCha20Rng; use std::{ collections::HashMap, @@ -36,6 +38,7 @@ use std::{ }; use tracing::{error, info}; +use crate::encryption_key_collector::{AllEncryptionKeysCollected, EncryptionKeyCollector}; use crate::threshold_share_collector::ThresholdShareCollector; #[derive(Message, Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] @@ -68,11 +71,20 @@ impl From>> for AllThresholdSharesCollected { } } +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct CollectingEncryptionKeysData { + sk_bfv: SensitiveBytes, + pk_bfv: ArcBytes, + ciphernode_selected: CiphernodeSelected, +} + #[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] pub struct GeneratingThresholdShareData { pk_share: Option, sk_sss: Option>, esi_sss: Option>>, + sk_bfv: Option, + pk_bfv: Option, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -80,6 +92,7 @@ pub struct AggregatingDecryptionKey { pk_share: ArcBytes, sk_sss: Encrypted, esi_sss: Vec>, + sk_bfv: SensitiveBytes, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -96,11 +109,12 @@ pub struct Decrypting { es_poly_sum: Vec, } -// TODO: Add GeneratingPvwKey state #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub enum KeyshareState { // Before anything Init, + // Collecting BFV encryption keys from all parties + CollectingEncryptionKeys(CollectingEncryptionKeysData), // Generating TrBFV share material GeneratingThresholdShare(GeneratingThresholdShareData), // Collecting remaining TrBFV shares to aggregate decryption key @@ -123,7 +137,8 @@ impl KeyshareState { true } else { match (self, &new_state) { - (K::Init, K::GeneratingThresholdShare(_)) => true, + (K::Init, K::CollectingEncryptionKeys(_)) => true, + (K::CollectingEncryptionKeys(_), K::GeneratingThresholdShare(_)) => true, (K::GeneratingThresholdShare(_), K::AggregatingDecryptionKey(_)) => true, (K::AggregatingDecryptionKey(_), K::ReadyForDecryption(_)) => true, (K::ReadyForDecryption(_), K::Decrypting(_)) => true, @@ -146,6 +161,7 @@ impl KeyshareState { pub fn variant_name(&self) -> &'static str { match self { Self::Init => "Init", + Self::CollectingEncryptionKeys(_) => "CollectingEncryptionKeys", Self::GeneratingThresholdShare(_) => "GeneratingThresholdShare", Self::AggregatingDecryptionKey(_) => "AggregatingDecryptionKey", Self::ReadyForDecryption(_) => "ReadyForDecryption", @@ -228,6 +244,16 @@ impl ThresholdKeyshareState { } } +impl TryInto for ThresholdKeyshareState { + type Error = anyhow::Error; + fn try_into(self) -> std::result::Result { + match self.state { + KeyshareState::CollectingEncryptionKeys(s) => Ok(s), + _ => Err(anyhow!("Invalid state")), + } + } +} + impl TryInto for ThresholdKeyshareState { type Error = anyhow::Error; fn try_into(self) -> std::result::Result { @@ -279,6 +305,8 @@ pub struct ThresholdKeyshare { bus: BusHandle, cipher: Arc, decryption_key_collector: Option>, + encryption_key_collector: Option>, + collected_encryption_keys: Option>>, multithread: Addr, state: Persistable, } @@ -289,6 +317,8 @@ impl ThresholdKeyshare { bus: params.bus, cipher: params.cipher, decryption_key_collector: None, + encryption_key_collector: None, + collected_encryption_keys: None, multithread: params.multithread, state: params.state, } @@ -318,6 +348,24 @@ impl ThresholdKeyshare { Ok(addr.clone()) } + pub fn ensure_encryption_key_collector( + &mut self, + self_addr: Addr, + ) -> Result> { + let Some(state) = self.state.get() else { + bail!("State not found on threshold keyshare. This should not happen."); + }; + + info!( + "Setting up encryption key collector for addr: {} and {} nodes", + state.address, state.threshold_n + ); + let addr = self + .encryption_key_collector + .get_or_insert_with(|| EncryptionKeyCollector::setup(self_addr, state.threshold_n)); + Ok(addr.clone()) + } + pub fn handle_threshold_share_created( &mut self, msg: ThresholdShareCreated, @@ -330,27 +378,88 @@ impl ThresholdKeyshare { Ok(()) } - /// 1. CiphernodeSelected + pub fn handle_encryption_key_created( + &mut self, + msg: EncryptionKeyCreated, + self_addr: Addr, + ) -> Result<()> { + info!("Received EncryptionKeyCreated forwarding to encryption key collector!"); + let collector = self.ensure_encryption_key_collector(self_addr)?; + collector.do_send(msg); + Ok(()) + } + + /// 1. CiphernodeSelected - Generate BFV keys and start collecting pub fn handle_ciphernode_selected( &mut self, msg: CiphernodeSelected, address: Addr, ) -> Result<()> { - // Ensure the collector is created let _ = self.ensure_collector(address.clone()); - // Initialize State + let _ = self.ensure_encryption_key_collector(address.clone()); + + let params = get_share_encryption_params(); + let mut rng = OsRng; + let sk_bfv = SecretKey::random(¶ms, &mut rng); + let pk_bfv = PublicKey::new(&sk_bfv, &mut rng); + + let sk_bytes = serialize_secret_key(&sk_bfv)?; + let sk_bfv_encrypted = SensitiveBytes::new(sk_bytes, &self.cipher)?; + let pk_bfv_bytes = ArcBytes::from_bytes(&pk_bfv.to_bytes()); + + self.state.try_mutate(|s| { + s.new_state(KeyshareState::CollectingEncryptionKeys( + CollectingEncryptionKeysData { + sk_bfv: sk_bfv_encrypted.clone(), + pk_bfv: pk_bfv_bytes.clone(), + ciphernode_selected: msg.clone(), + }, + )) + })?; + + let state = self.state.get().ok_or(anyhow!("No state"))?; + self.bus.publish(EncryptionKeyCreated { + e3_id: state.e3_id.clone(), + key: Arc::new(EncryptionKey { + party_id: state.party_id, + pk_bfv: pk_bfv_bytes, + }), + external: false, + }); + + Ok(()) + } + + /// 1a. AllEncryptionKeysCollected - All BFV keys received, start share generation + pub fn handle_all_encryption_keys_collected( + &mut self, + msg: AllEncryptionKeysCollected, + address: Addr, + ) -> Result<()> { + info!( + "AllEncryptionKeysCollected - {} keys received", + msg.keys.len() + ); + + self.collected_encryption_keys = Some(msg.keys); + let state = self.state.get().ok_or(anyhow!("No state"))?; + let current: CollectingEncryptionKeysData = state.clone().try_into()?; + self.state.try_mutate(|s| { s.new_state(KeyshareState::GeneratingThresholdShare( GeneratingThresholdShareData { sk_sss: None, pk_share: None, esi_sss: None, + sk_bfv: Some(current.sk_bfv), + pk_bfv: Some(current.pk_bfv), }, )) })?; - address.do_send(GenEsiSss(msg.clone())); - address.do_send(GenPkShareAndSkSss(msg)); + address.do_send(GenEsiSss(current.ciphernode_selected.clone())); + address.do_send(GenPkShareAndSkSss(current.ciphernode_selected)); + Ok(()) } @@ -400,20 +509,25 @@ impl ThresholdKeyshare { let current: GeneratingThresholdShareData = s.clone().try_into()?; let pk_share = current.pk_share; let sk_sss = current.sk_sss; - let next = match (pk_share, sk_sss) { + let sk_bfv = current.sk_bfv; + let pk_bfv = current.pk_bfv; + let next = match (pk_share, sk_sss, &sk_bfv) { // If the other shares are here then transition to aggregation - (Some(pk_share), Some(sk_sss)) => { + (Some(pk_share), Some(sk_sss), Some(sk_bfv_ref)) => { K::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, + sk_bfv: sk_bfv_ref.clone(), }) } // If the other shares are not here yet then dont transition - (None, None) => K::GeneratingThresholdShare(GeneratingThresholdShareData { + (None, None, _) => K::GeneratingThresholdShare(GeneratingThresholdShareData { esi_sss: Some(esi_sss), pk_share: None, sk_sss: None, + sk_bfv, + pk_bfv, }), _ => bail!("Inconsistent state!"), }; @@ -461,6 +575,7 @@ impl ThresholdKeyshare { } /// 3a. GenPkShareAndSkSss + /// 3a. GenPkShareAndSkSss result pub fn handle_gen_pk_share_and_sk_sss_response(&mut self, res: ComputeResponse) -> Result<()> { let ComputeResponse::TrBFV(TrBFVResponse::GenPkShareAndSkSss(output)) = res else { bail!("Error extracting data from compute process") @@ -472,21 +587,30 @@ impl ThresholdKeyshare { info!("try_store_pk_share_and_sk_sss"); let current: GeneratingThresholdShareData = s.clone().try_into()?; let esi_sss = current.esi_sss; - let next = match esi_sss { - // If the esi shares are here then transition to aggregation - Some(esi_sss) => { + let sk_bfv = current.sk_bfv; + let pk_bfv = current.pk_bfv; + let next = match (esi_sss, sk_bfv) { + // If the esi shares and BFV key are here then transition to aggregation + (Some(esi_sss), Some(sk_bfv)) => { KeyshareState::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, + sk_bfv, }) } // If esi shares are not here yet then don't transition - None => KeyshareState::GeneratingThresholdShare(GeneratingThresholdShareData { - esi_sss: None, - pk_share: Some(pk_share), - sk_sss: Some(sk_sss), - }), + (None, sk_bfv) => { + KeyshareState::GeneratingThresholdShare(GeneratingThresholdShareData { + esi_sss: None, + pk_share: Some(pk_share), + sk_sss: Some(sk_sss), + sk_bfv, + pk_bfv, + }) + } + // If we have esi_sss but no sk_bfv, that's an error + (Some(_), None) => bail!("Have esi_sss but no sk_bfv - inconsistent state!"), }; s.new_state(next) })?; @@ -501,7 +625,7 @@ impl ThresholdKeyshare { Ok(()) } - /// 4. SharesGenerated + /// 4. SharesGenerated - Encrypt shares with BFV and publish pub fn handle_shares_generated(&self) -> Result<()> { let Some(ThresholdKeyshareState { state: @@ -519,20 +643,46 @@ impl ThresholdKeyshare { bail!("Invalid state!"); }; - let decrypted = sk_sss.decrypt(&self.cipher)?; - let sk_sss: PvwEncrypted = PvwEncrypted::new(decrypted)?; - let esi_sss = esi_sss + // Get collected BFV public keys from all parties + let encryption_keys = self + .collected_encryption_keys + .as_ref() + .ok_or(anyhow!("No encryption keys collected"))?; + + // Convert to BFV public keys + let params = get_share_encryption_params(); + let recipient_pks: Vec = encryption_keys + .iter() + .map(|k| { + PublicKey::from_bytes(&k.pk_bfv, ¶ms) + .map_err(|e| anyhow!("Failed to deserialize BFV public key: {:?}", e)) + }) + .collect::>()?; + + // Decrypt our shares from local storage + let decrypted_sk_sss: SharedSecret = sk_sss.decrypt(&self.cipher)?; + let decrypted_esi_sss: Vec = esi_sss .into_iter() - .map(|s| PvwEncrypted::new(s.decrypt(&self.cipher)?)) + .map(|s| s.decrypt(&self.cipher)) + .collect::>()?; + + // Encrypt shares for all recipients using BFV + let mut rng = OsRng; + let encrypted_sk_sss = + BfvEncryptedShares::encrypt_all(&decrypted_sk_sss, &recipient_pks, ¶ms, &mut rng)?; + + let encrypted_esi_sss: Vec = decrypted_esi_sss + .iter() + .map(|esi| BfvEncryptedShares::encrypt_all(esi, &recipient_pks, ¶ms, &mut rng)) .collect::>()?; self.bus.publish(ThresholdShareCreated { e3_id, share: Arc::new(ThresholdShare { party_id, - esi_sss, pk_share, - sk_sss, + sk_sss: encrypted_sk_sss, + esi_sss: encrypted_esi_sss, }), external: false, })?; @@ -542,6 +692,7 @@ impl ThresholdKeyshare { /// 5. AllThresholdSharesCollected. This is fired after the ThresholdShareCreated events are /// aggregateed in the decryption_key_collector::ThresholdShareCollector + /// 5. AllThresholdSharesCollected - Decrypt received shares using BFV and aggregate pub fn handle_all_threshold_shares_collected( &self, msg: AllThresholdSharesCollected, @@ -552,37 +703,42 @@ impl ThresholdKeyshare { let party_id = state.party_id as usize; let trbfv_config = state.get_trbfv_config(); - // Shares are in order of party_id - let received_sss: Vec = msg + // Get our BFV secret key from state + let current: AggregatingDecryptionKey = state.clone().try_into()?; + let sk_bytes = current.sk_bfv.access(&cipher)?; + let params = get_share_encryption_params(); + let sk_bfv = deserialize_secret_key(&sk_bytes, ¶ms)?; + let degree = params.degree(); + + // Decrypt our share from each sender using BFV + // Each sender's ThresholdShare contains encrypted shares for all parties + // We extract and decrypt the share meant for us (at index party_id) + let sk_sss_collected: Vec = msg .shares - .clone() - .into_iter() - .map(|ts| ts.sk_sss.clone().pvw_decrypt()) + .iter() + .map(|ts| { + let encrypted = ts + .sk_sss + .clone_share(party_id) + .ok_or(anyhow!("No sk_sss share for party {}", party_id))?; + encrypted.decrypt(&sk_bfv, ¶ms, degree) + }) .collect::>()?; - let received_esi_sss: Vec> = msg + // Similarly decrypt esi_sss for each ciphertext + let esi_sss_collected: Vec> = msg .shares - .into_iter() + .iter() .map(|ts| { ts.esi_sss - .clone() - .into_iter() - .map(|s| s.pvw_decrypt()) - .collect() - }) - .collect::>()?; - - let sk_sss_collected: Vec = received_sss - .into_iter() - .map(|s| s.extract_party_share(party_id)) - .collect::>()?; - - let esi_sss_collected: Vec> = received_esi_sss - .into_iter() - .map(|esi| { - esi.into_iter() - .map(|s| s.extract_party_share(party_id)) - .collect() + .iter() + .map(|esi_shares| { + let encrypted = esi_shares + .clone_share(party_id) + .ok_or(anyhow!("No esi_sss share for party {}", party_id))?; + encrypted.decrypt(&sk_bfv, ¶ms, degree) + }) + .collect::>>() }) .collect::>()?; @@ -750,6 +906,7 @@ impl ThresholdKeyshare { } } +// Will only receive events that are for this specific e3_id // Will only receive events that are for this specific e3_id impl Handler for ThresholdKeyshare { type Result = (); @@ -760,6 +917,9 @@ impl Handler for ThresholdKeyshare { EnclaveEventData::ThresholdShareCreated(data) => { let _ = self.handle_threshold_share_created(data, ctx.address()); } + EnclaveEventData::EncryptionKeyCreated(data) => { + let _ = self.handle_encryption_key_created(data, ctx.address()); + } _ => (), } } @@ -795,6 +955,16 @@ impl Handler for ThresholdKeyshare { } } +impl Handler for ThresholdKeyshare { + type Result = (); + fn handle(&mut self, msg: AllEncryptionKeysCollected, ctx: &mut Self::Context) -> Self::Result { + match self.handle_all_encryption_keys_collected(msg, ctx.address()) { + Err(e) => error!("{e}"), + Ok(_) => (), + } + } +} + impl Handler for ThresholdKeyshare { type Result = ResponseActFuture; fn handle(&mut self, msg: AllThresholdSharesCollected, _: &mut Self::Context) -> Self::Result { diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index a0aac9e06f..40feab390b 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -15,8 +15,8 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CorrelationId, DocumentKind, DocumentMeta, - DocumentReceived, E3RequestComplete, E3id, EType, EnclaveEvent, EnclaveEventData, Event, - PartyId, PublishDocumentRequested, ThresholdShareCreated, + DocumentReceived, E3RequestComplete, E3id, EnclaveErrorType, EnclaveEvent, EnclaveEventData, + EncryptionKeyCreated, Event, PartyId, PublishDocumentRequested, ThresholdShareCreated, }; use e3_utils::retry::{retry_with_backoff, to_retry}; use e3_utils::ArcBytes; @@ -75,6 +75,7 @@ impl DocumentPublisher { match event.get_data() { EnclaveEventData::PublishDocumentRequested(_) => true, EnclaveEventData::ThresholdShareCreated(_) => true, + EnclaveEventData::EncryptionKeyCreated(_) => true, _ => false, } } @@ -380,12 +381,14 @@ pub struct EventConverter { #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] enum ReceivableDocument { ThresholdShareCreated(ThresholdShareCreated), + EncryptionKeyCreated(EncryptionKeyCreated), } impl ReceivableDocument { pub fn get_e3_id(&self) -> &E3id { match self { ReceivableDocument::ThresholdShareCreated(d) => &d.e3_id, + ReceivableDocument::EncryptionKeyCreated(d) => &d.e3_id, } } @@ -405,6 +408,7 @@ impl EventConverter { pub fn setup(bus: &BusHandle) -> Addr { let addr = Self::new(bus).start(); bus.subscribe("ThresholdShareCreated", addr.clone().into()); + bus.subscribe("EncryptionKeyCreated", addr.clone().into()); bus.subscribe("DocumentReceived", addr.clone().into()); addr } @@ -426,19 +430,60 @@ impl EventConverter { .publish(PublishDocumentRequested::new(meta, value))?; Ok(()) } + /// Local node created an encryption key. Send it as a published document + pub fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { + // If this is received from elsewhere + if msg.external { + return Ok(()); + } + let receivable = ReceivableDocument::EncryptionKeyCreated(msg); + let value = ArcBytes::from_bytes(&receivable.to_bytes()?); + let meta = DocumentMeta::new( + receivable.get_e3_id().clone(), + DocumentKind::TrBFV, + vec![], + None, + ); + self.bus.publish(PublishDocumentRequested::new(meta, value)); + Ok(()) + } + /// Local node created an encryption key. Send it as a published document + pub fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { + // If this is received from elsewhere + if msg.external { + return Ok(()); + } + let receivable = ReceivableDocument::EncryptionKeyCreated(msg); + let value = ArcBytes::from_bytes(&receivable.to_bytes()?); + let meta = DocumentMeta::new( + receivable.get_e3_id().clone(), + DocumentKind::TrBFV, + vec![], + None, + ); + self.bus.publish(PublishDocumentRequested::new(meta, value)); + Ok(()) + } /// Received document externally pub fn handle_document_received(&self, msg: DocumentReceived) -> Result<()> { warn!("Converting DocumentReceived..."); let receivable = ReceivableDocument::from_bytes(&msg.value.extract_bytes())?; - let event = match receivable { - ReceivableDocument::ThresholdShareCreated(evt) => ThresholdShareCreated { - external: true, - e3_id: evt.e3_id, - share: evt.share, - }, + match receivable { + ReceivableDocument::ThresholdShareCreated(evt) => { + self.bus.publish(ThresholdShareCreated { + external: true, + e3_id: evt.e3_id, + share: evt.share, + }); + } + ReceivableDocument::EncryptionKeyCreated(evt) => { + self.bus.publish(EncryptionKeyCreated { + external: true, + e3_id: evt.e3_id, + key: evt.key, + }); + } }; - - self.bus.publish(event)?; Ok(()) } } @@ -452,6 +497,7 @@ impl Handler for EventConverter { fn handle(&mut self, msg: EnclaveEvent, ctx: &mut Self::Context) -> Self::Result { match msg.into_data() { EnclaveEventData::ThresholdShareCreated(data) => ctx.notify(data), + EnclaveEventData::EncryptionKeyCreated(data) => ctx.notify(data), EnclaveEventData::DocumentReceived(data) => ctx.notify(data), _ => (), } @@ -468,6 +514,16 @@ impl Handler for EventConverter { } } +impl Handler for EventConverter { + type Result = (); + fn handle(&mut self, msg: EncryptionKeyCreated, _ctx: &mut Self::Context) -> Self::Result { + match self.handle_encryption_key_created(msg) { + Ok(_) => (), + Err(err) => error!("{err}"), + } + } +} + impl Handler for EventConverter { type Result = (); fn handle(&mut self, msg: DocumentReceived, _ctx: &mut Self::Context) -> Self::Result { diff --git a/crates/test-helpers/src/usecase_helpers.rs b/crates/test-helpers/src/usecase_helpers.rs index c5698d7a0e..712e6471d7 100644 --- a/crates/test-helpers/src/usecase_helpers.rs +++ b/crates/test-helpers/src/usecase_helpers.rs @@ -16,18 +16,27 @@ use e3_trbfv::{ gen_pk_share_and_sk_sss::{ gen_pk_share_and_sk_sss, GenPkShareAndSkSssRequest, GenPkShareAndSkSssResponse, }, - shares::{EncryptableVec, PvwEncrypted, PvwEncryptedVecExt, ShamirShare, SharedSecret}, + helpers::get_share_encryption_params, + shares::{BfvEncryptedShares, EncryptableVec, ShamirShare, SharedSecret}, TrBFVConfig, }; use e3_utils::{ArcBytes, SharedRng}; use fhe::{ - bfv::{BfvParameters, PublicKey}, + bfv::{BfvParameters, PublicKey, SecretKey}, mbfv::{AggregateIter, CommonRandomPoly, PublicKeyShare}, }; use fhe_traits::Serialize; +use rand::rngs::OsRng; // The following functions are designed to aid testing our usecases +/// Result of generating shares - includes the shares plus BFV keys for decryption +pub struct GeneratedShares { + pub shares: HashMap, + /// BFV secret keys for each party (for decryption in tests) + pub bfv_secret_keys: Vec, +} + pub fn generate_shares_hash_map( trbfv_config: &TrBFVConfig, esi_per_ct: u64, @@ -35,11 +44,24 @@ pub fn generate_shares_hash_map( crp: &CommonRandomPoly, rng: &SharedRng, cipher: &Cipher, -) -> Result> { - let threshold_n = trbfv_config.num_parties(); +) -> Result { + let threshold_n = trbfv_config.num_parties() as usize; + + // First, generate BFV encryption keys for all parties + let bfv_params = get_share_encryption_params(); + let mut bfv_rng = OsRng; + let mut bfv_secret_keys = Vec::with_capacity(threshold_n); + let mut bfv_public_keys = Vec::with_capacity(threshold_n); + + for _ in 0..threshold_n { + let sk = SecretKey::random(&bfv_params, &mut bfv_rng); + let pk = fhe::bfv::PublicKey::new(&sk, &mut bfv_rng); + bfv_secret_keys.push(sk); + bfv_public_keys.push(pk); + } let mut shares_hash_map = HashMap::new(); - for party_id in 0u64..threshold_n { + for party_id in 0u64..threshold_n as u64 { let GenEsiSssResponse { esi_sss } = gen_esi_sss( &rng, &cipher, @@ -59,24 +81,36 @@ pub fn generate_shares_hash_map( }, )?; - // Simulate actor boundry and SharesGenerated - let sk_sss = PvwEncrypted::new(sk_sss.decrypt(&cipher)?)?; - let esi_sss: Vec> = esi_sss + // Decrypt locally stored secrets + let decrypted_sk_sss: SharedSecret = sk_sss.decrypt(&cipher)?; + let decrypted_esi_sss: Vec = esi_sss .into_iter() - .map(|s| PvwEncrypted::new(s.decrypt(&cipher)?)) + .map(|s| s.decrypt(&cipher)) + .collect::>()?; + + // Encrypt shares for all recipients using BFV + let encrypted_sk_sss = + BfvEncryptedShares::encrypt_all(&decrypted_sk_sss, &bfv_public_keys, &bfv_params, &mut bfv_rng)?; + + let encrypted_esi_sss: Vec = decrypted_esi_sss + .iter() + .map(|esi| BfvEncryptedShares::encrypt_all(esi, &bfv_public_keys, &bfv_params, &mut bfv_rng)) .collect::>()?; shares_hash_map.insert( party_id, ThresholdShare { party_id, - esi_sss, - sk_sss, + esi_sss: encrypted_esi_sss, + sk_sss: encrypted_sk_sss, pk_share, }, ); } - Ok(shares_hash_map) + Ok(GeneratedShares { + shares: shares_hash_map, + bfv_secret_keys, + }) } pub fn get_public_key( @@ -99,36 +133,44 @@ pub fn get_public_key( pub fn get_decryption_keys( shares: Vec, + bfv_secret_keys: &[SecretKey], cipher: &Cipher, trbfv_config: &TrBFVConfig, ) -> Result, SensitiveBytes)>> { - let threshold_n = trbfv_config.num_parties(); - let received_sss = shares - .iter() - .map(|ts| ts.sk_sss.clone().pvw_decrypt()) - .collect::>>()?; - - let received_esi_sss: Vec> = shares - .iter() - .map(|ts| ts.esi_sss.clone().to_vec_decrypted()) - .collect::>>()?; - - // Individualize based on node + let threshold_n = trbfv_config.num_parties() as usize; + let bfv_params = get_share_encryption_params(); + let degree = bfv_params.degree(); + + // Individualize based on node - each party decrypts their share from each sender let mut decryption_keys = HashMap::new(); - for party_id in 0..threshold_n as usize { - let sk_sss_collected = received_sss - .clone() - .into_iter() - .map(|sss| sss.extract_party_share(party_id)) - .collect::>>()?; + for party_id in 0..threshold_n { + let sk_bfv = &bfv_secret_keys[party_id]; - let esi_sss_collected: Vec> = received_esi_sss - .clone() - .into_iter() - .map(|s| { - s.into_iter() - .map(|ss| ss.extract_party_share(party_id)) - .collect() + // Decrypt sk_sss share from each sender using our BFV secret key + let sk_sss_collected: Vec = shares + .iter() + .map(|ts| { + let encrypted = ts + .sk_sss + .clone_share(party_id) + .ok_or_else(|| anyhow::anyhow!("No sk_sss share for party {}", party_id))?; + encrypted.decrypt(sk_bfv, &bfv_params, degree) + }) + .collect::>()?; + + // Similarly decrypt esi_sss + let esi_sss_collected: Vec> = shares + .iter() + .map(|ts| { + ts.esi_sss + .iter() + .map(|esi_shares| { + let encrypted = esi_shares + .clone_share(party_id) + .ok_or_else(|| anyhow::anyhow!("No esi_sss share for party {}", party_id))?; + encrypted.decrypt(sk_bfv, &bfv_params, degree) + }) + .collect::>>() }) .collect::>()?; diff --git a/crates/tests/tests/integration.rs b/crates/tests/tests/integration.rs index 0b13895c1c..53f4c2f561 100644 --- a/crates/tests/tests/integration.rs +++ b/crates/tests/tests/integration.rs @@ -287,6 +287,21 @@ async fn test_trbfv_actor() -> Result<()> { committee_finalized_timer.elapsed(), )); + // First, wait for all EncryptionKeyCreated events (BFV key exchange) + let encryption_keys_timer = Instant::now(); + let expected = vec![ + "EncryptionKeyCreated", + "EncryptionKeyCreated", + "EncryptionKeyCreated", + "EncryptionKeyCreated", + "EncryptionKeyCreated", + ]; + let _ = nodes + .take_history_with_timeout(0, expected.len(), Duration::from_secs(1000)) + .await?; + report.push(("All EncryptionKeyCreated events", encryption_keys_timer.elapsed())); + + // Then wait for all ThresholdShareCreated events let shares_timer = Instant::now(); let expected = vec![ "ThresholdShareCreated", diff --git a/crates/trbfv/Cargo.toml b/crates/trbfv/Cargo.toml index c56f193782..b3a232134c 100644 --- a/crates/trbfv/Cargo.toml +++ b/crates/trbfv/Cargo.toml @@ -24,10 +24,11 @@ serde.workspace = true tracing.workspace = true zeroize.workspace = true +rand.workspace = true + [dev-dependencies] e3-test-helpers.workspace = true e3-fhe.workspace = true e3-events.workspace = true tracing-subscriber.workspace = true tokio.workspace = true -rand.workspace = true diff --git a/crates/trbfv/src/helpers.rs b/crates/trbfv/src/helpers.rs index 7b7f9edf1a..484d5264cd 100644 --- a/crates/trbfv/src/helpers.rs +++ b/crates/trbfv/src/helpers.rs @@ -6,21 +6,49 @@ use crate::shares::ShamirShare; use anyhow::Result; +use e3_bfv_helpers::BfvParamSets; use e3_crypto::{Cipher, SensitiveBytes}; use fhe::mbfv::PublicKeyShare; use fhe::{ - bfv::{self, BfvParameters}, + bfv::{self, BfvParameters, SecretKey}, trbfv::{SmudgingBoundCalculator, SmudgingBoundCalculatorConfig}, }; use fhe_math::rq::Poly; use fhe_traits::{DeserializeWithContext, Serialize}; use num_bigint::BigUint; use petname::Petnames; +use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize}; use std::{ hash::{DefaultHasher, Hash, Hasher}, sync::Arc, }; +/// Simple serialization wrapper for SecretKey coefficients. +#[derive(SerdeSerialize, SerdeDeserialize)] +struct SecretKeyData { + coeffs: Box<[i64]>, +} + +/// Serialize a BFV SecretKey to bytes. +pub fn serialize_secret_key(sk: &SecretKey) -> Result> { + let data = SecretKeyData { + coeffs: sk.coeffs.clone(), + }; + Ok(bincode::serialize(&data)?) +} + +/// Deserialize a BFV SecretKey from bytes. +pub fn deserialize_secret_key(bytes: &[u8], params: &Arc) -> Result { + let data: SecretKeyData = bincode::deserialize(bytes)?; + Ok(SecretKey::new(data.coeffs.to_vec(), params)) +} + +/// TODO: Make this modular +pub fn get_share_encryption_params() -> Arc { + let param_set: e3_bfv_helpers::BfvParamSet = BfvParamSets::Set8192_144115188075855872_2.into(); + param_set.build_arc() +} + pub fn try_poly_from_bytes(bytes: &[u8], params: &BfvParameters) -> Result { Ok(Poly::from_bytes(bytes, params.ctx_at_level(0)?)?) } diff --git a/crates/trbfv/src/shares/bfv_encrypted.rs b/crates/trbfv/src/shares/bfv_encrypted.rs new file mode 100644 index 0000000000..adf383fb09 --- /dev/null +++ b/crates/trbfv/src/shares/bfv_encrypted.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use anyhow::{Context, Result}; +use derivative::Derivative; +use fhe::bfv::{BfvParameters, Ciphertext, Encoding, Plaintext, PublicKey, SecretKey}; +use fhe_traits::{ + DeserializeParametrized, FheDecoder, FheDecrypter, FheEncoder, FheEncrypter, + Serialize as FheSerialize, +}; +use ndarray::Array2; +use rand::{CryptoRng, RngCore}; +use serde::{Deserialize, Serialize}; +use std::ops::Deref; +use std::sync::Arc; + +use super::{ShamirShare, SharedSecret}; + +// Re-export helper functions from helpers module +pub use crate::helpers::{ + deserialize_secret_key, get_share_encryption_params, serialize_secret_key, +}; + +/// A BFV-encrypted Shamir share for secure transmission. +/// +/// Each share is encrypted as multiple BFV ciphertexts (one per modulus level). +/// The recipient can only decrypt using their corresponding BFV secret key. +#[derive(Derivative, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[derivative(Debug)] +pub struct BfvEncryptedShare { + /// BFV ciphertexts, one per modulus level (serialized) + #[derivative(Debug(format_with = "debug_vec_bytes"))] + ciphertexts: Vec>, +} + +/// Debug helper for Vec> +fn debug_vec_bytes(v: &Vec>, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "[{} ciphertexts, total {} bytes]", + v.len(), + v.iter().map(|c| c.len()).sum::() + ) +} + +impl BfvEncryptedShare { + /// Encrypt a Shamir share for a specific recipient. + /// + /// # Arguments + /// * `share` - The Shamir share to encrypt (contains data for all moduli) + /// * `recipient_pk` - The recipient's BFV public key + /// * `params` - BFV parameters for share encryption + /// * `rng` - Random number generator + /// + /// # Returns + /// An encrypted share that can only be decrypted by the recipient + pub fn encrypt( + share: &ShamirShare, + recipient_pk: &PublicKey, + params: &Arc, + rng: &mut R, + ) -> Result { + let data: &Array2 = share.deref(); // Array2 with rows = moduli, cols = coefficients + let num_moduli = data.nrows(); + + let mut ciphertexts = Vec::with_capacity(num_moduli); + + for m in 0..num_moduli { + let row = data.row(m); + let share_vec: Vec = row.to_vec(); + + let pt = Plaintext::try_encode(&share_vec, Encoding::poly(), params) + .context("Failed to encode share as plaintext")?; + + let ct = recipient_pk + .try_encrypt(&pt, rng) + .context("Failed to encrypt share")?; + + ciphertexts.push(ct.to_bytes()); + } + + Ok(Self { ciphertexts }) + } + + /// Decrypt an encrypted share using the recipient's secret key. + /// + /// # Arguments + /// * `sk` - The recipient's BFV secret key + /// * `params` - BFV parameters for share encryption + /// * `degree` - Polynomial degree (for reconstructing the share matrix) + /// + /// # Returns + /// The decrypted Shamir share + pub fn decrypt( + self, + sk: &SecretKey, + params: &Arc, + degree: usize, + ) -> Result { + let num_moduli = self.ciphertexts.len(); + let mut data = Array2::zeros((num_moduli, degree)); + + for (m, ct_bytes) in self.ciphertexts.into_iter().enumerate() { + let ct = Ciphertext::from_bytes(&ct_bytes, params) + .context("Failed to deserialize ciphertext")?; + + let pt = sk + .try_decrypt(&ct) + .context("Failed to decrypt ciphertext")?; + + let decrypted: Vec = Vec::::try_decode(&pt, Encoding::poly()) + .context("Failed to decode plaintext")?; + + for (i, val) in decrypted.into_iter().take(degree).enumerate() { + data[[m, i]] = val; + } + } + + Ok(ShamirShare::new(data)) + } +} + +impl Default for BfvEncryptedShare { + fn default() -> Self { + Self { + ciphertexts: Vec::new(), + } + } +} + +/// A collection of BFV-encrypted shares for all recipients. +/// +/// When a party generates Shamir shares, they encrypt each recipient's share +/// with that recipient's public key. This struct holds all encrypted shares +/// from a single sender. +#[derive(Derivative, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[derivative(Debug)] +pub struct BfvEncryptedShares { + /// Encrypted shares indexed by recipient party_id (0-based) + shares: Vec, +} + +impl BfvEncryptedShares { + /// Encrypt shares for all recipients. + /// + /// # Arguments + /// * `secret` - The SharedSecret containing shares for all parties + /// * `recipient_pks` - Public keys for all recipients, indexed by party_id + /// * `params` - BFV parameters for share encryption + /// * `rng` - Random number generator + pub fn encrypt_all( + secret: &SharedSecret, + recipient_pks: &[PublicKey], + params: &Arc, + rng: &mut R, + ) -> Result { + let num_parties = recipient_pks.len(); + let mut shares = Vec::with_capacity(num_parties); + + for party_id in 0..num_parties { + let share = secret + .extract_party_share(party_id) + .context(format!("Failed to extract share for party {}", party_id))?; + + let encrypted = + BfvEncryptedShare::encrypt(&share, &recipient_pks[party_id], params, rng)?; + + shares.push(encrypted); + } + + Ok(Self { shares }) + } + + /// Get the encrypted share for a specific recipient. + pub fn get_share(&self, party_id: usize) -> Option<&BfvEncryptedShare> { + self.shares.get(party_id) + } + + /// Clone the encrypted share for a specific recipient. + pub fn clone_share(&self, party_id: usize) -> Option { + self.shares.get(party_id).cloned() + } + + /// Number of encrypted shares + pub fn len(&self) -> usize { + self.shares.len() + } + + /// Check if empty + pub fn is_empty(&self) -> bool { + self.shares.is_empty() + } +} + +impl Default for BfvEncryptedShares { + fn default() -> Self { + Self { shares: Vec::new() } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::rngs::OsRng; + + #[test] + fn test_encrypt_decrypt_share() { + let params = get_share_encryption_params(); + let mut rng = OsRng; + + // Generate key pair + let sk = SecretKey::random(¶ms, &mut rng); + let pk = PublicKey::new(&sk, &mut rng); + + // Create a test share (1 modulus row, 512 coefficients) + let degree = params.degree(); + let test_data: Vec = (0..degree as u64).collect(); + let mut data = Array2::zeros((1, degree)); + for (i, val) in test_data.iter().enumerate() { + data[[0, i]] = *val; + } + let share = ShamirShare::new(data.clone()); + + // Encrypt + let encrypted = BfvEncryptedShare::encrypt(&share, &pk, ¶ms, &mut rng) + .expect("Encryption should succeed"); + + // Decrypt + let decrypted = encrypted + .decrypt(&sk, ¶ms, degree) + .expect("Decryption should succeed"); + + // Verify + assert_eq!(share.deref(), decrypted.deref()); + } + + #[test] + fn test_secret_key_serialization() { + let params = get_share_encryption_params(); + let mut rng = OsRng; + + // Generate a secret key + let sk = SecretKey::random(¶ms, &mut rng); + + // Serialize + let bytes = serialize_secret_key(&sk).expect("Serialization should succeed"); + + // Deserialize + let sk_restored = + deserialize_secret_key(&bytes, ¶ms).expect("Deserialization should succeed"); + + // Verify coefficients match + assert_eq!(sk.coeffs, sk_restored.coeffs); + } +} diff --git a/crates/trbfv/src/shares/mod.rs b/crates/trbfv/src/shares/mod.rs index e633ad4085..171cdee7da 100644 --- a/crates/trbfv/src/shares/mod.rs +++ b/crates/trbfv/src/shares/mod.rs @@ -3,9 +3,9 @@ // This file is provided WITHOUT ANY WARRANTY; // without even the implied warranty of MERCHANTABILITY // or FITNESS FOR A PARTICULAR PURPOSE. +pub mod bfv_encrypted; pub mod encrypted; -pub mod pvw; pub mod shares; +pub use bfv_encrypted::*; pub use encrypted::*; -pub use pvw::*; pub use shares::*; diff --git a/crates/trbfv/src/shares/pvw.rs b/crates/trbfv/src/shares/pvw.rs deleted file mode 100644 index 4600204493..0000000000 --- a/crates/trbfv/src/shares/pvw.rs +++ /dev/null @@ -1,66 +0,0 @@ -// SPDX-License-Identifier: LGPL-3.0-only -// -// This file is provided WITHOUT ANY WARRANTY; -// without even the implied warranty of MERCHANTABILITY -// or FITNESS FOR A PARTICULAR PURPOSE. -use anyhow::Result; -use derivative::Derivative; - -use super::DeserializableValue; - -/// Encrypted version of T for secure storage/transmission -#[derive(Derivative, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -#[derivative(Debug)] -// TODO: Currently we are simply serializing the data in preparation for PVW encryption -// expecting to pass in keys to encrypt and decrypt as required at a later point -pub struct PvwEncrypted( - #[derivative(Debug(format_with = "e3_utils::formatters::hexf"))] Vec, - std::marker::PhantomData, -); - -impl PvwEncrypted -where - T: DeserializableValue, -{ - /// Create a new encrypted wrapper from data of type T - pub fn new(data: T) -> Result { - Ok(Self(bincode::serialize(&data)?, std::marker::PhantomData)) - } - - /// Decrypt and deserialize back to type T - pub fn pvw_decrypt(self) -> Result { - let value = self.0; - Ok(bincode::deserialize(&value)?) - } -} - -pub trait PvwEncryptedVecExt { - fn to_vec_decrypted(self) -> Result>; -} - -impl PvwEncryptedVecExt for Vec> -where - T: DeserializableValue, -{ - fn to_vec_decrypted(self) -> Result> { - self.into_iter().map(|s| s.pvw_decrypt()).collect() - } -} - -/// Trait to add decrypt functionality to Vec> -pub trait PvwEncryptableVec { - /// Decrypt all encrypted values in the vector - /// Returns a vector of decrypted values or the first error encountered - fn pvw_encrypt(self) -> Result>>; -} - -impl PvwEncryptableVec for Vec -where - T: DeserializableValue, -{ - fn pvw_encrypt(self) -> Result>> { - self.into_iter() - .map(|encryptable| PvwEncrypted::new(encryptable)) - .collect() - } -} diff --git a/crates/trbfv/tests/integration.rs b/crates/trbfv/tests/integration.rs index b08af1a2af..addc5575d0 100644 --- a/crates/trbfv/tests/integration.rs +++ b/crates/trbfv/tests/integration.rs @@ -83,7 +83,7 @@ async fn test_trbfv_isolation() -> Result<()> { ); // let crp = ArcBytes::from_bytes(crp_raw.to_bytes()); - let shares_hash_map = usecase_helpers::generate_shares_hash_map( + let generated = usecase_helpers::generate_shares_hash_map( &trbfv_config, esi_per_ct as u64, &error_size, @@ -93,9 +93,14 @@ async fn test_trbfv_isolation() -> Result<()> { )?; let pubkey = - usecase_helpers::get_public_key(&shares_hash_map, trbfv_config.params(), &crp_raw)?; - let shares = to_ordered_vec(shares_hash_map); - let decryption_keys = usecase_helpers::get_decryption_keys(shares, &cipher, &trbfv_config)?; + usecase_helpers::get_public_key(&generated.shares, trbfv_config.params(), &crp_raw)?; + let shares = to_ordered_vec(generated.shares); + let decryption_keys = usecase_helpers::get_decryption_keys( + shares, + &generated.bfv_secret_keys, + &cipher, + &trbfv_config, + )?; // Create the inputs let num_votes_per_voter = 3; let num_voters = 30; From 16ef3d09bef6776c7de92c20bf121e92d687654e Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Tue, 16 Dec 2025 18:32:32 +0500 Subject: [PATCH 02/19] chore: formatting --- .../keyshare/src/encryption_key_collector.rs | 2 -- crates/net/src/document_publisher.rs | 26 ++++--------------- crates/test-helpers/src/usecase_helpers.rs | 18 ++++++++----- crates/tests/tests/integration.rs | 5 +++- 4 files changed, 21 insertions(+), 30 deletions(-) diff --git a/crates/keyshare/src/encryption_key_collector.rs b/crates/keyshare/src/encryption_key_collector.rs index edceb24dcd..b7950530de 100644 --- a/crates/keyshare/src/encryption_key_collector.rs +++ b/crates/keyshare/src/encryption_key_collector.rs @@ -129,5 +129,3 @@ impl Handler for EncryptionKeyCollector { ); } } - - diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index 40feab390b..055a75cda3 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -15,7 +15,7 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CorrelationId, DocumentKind, DocumentMeta, - DocumentReceived, E3RequestComplete, E3id, EnclaveErrorType, EnclaveEvent, EnclaveEventData, + DocumentReceived, E3RequestComplete, E3id, EType, EnclaveEvent, EnclaveEventData, EncryptionKeyCreated, Event, PartyId, PublishDocumentRequested, ThresholdShareCreated, }; use e3_utils::retry::{retry_with_backoff, to_retry}; @@ -444,24 +444,8 @@ impl EventConverter { vec![], None, ); - self.bus.publish(PublishDocumentRequested::new(meta, value)); - Ok(()) - } - /// Local node created an encryption key. Send it as a published document - pub fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { - // If this is received from elsewhere - if msg.external { - return Ok(()); - } - let receivable = ReceivableDocument::EncryptionKeyCreated(msg); - let value = ArcBytes::from_bytes(&receivable.to_bytes()?); - let meta = DocumentMeta::new( - receivable.get_e3_id().clone(), - DocumentKind::TrBFV, - vec![], - None, - ); - self.bus.publish(PublishDocumentRequested::new(meta, value)); + self.bus + .publish(PublishDocumentRequested::new(meta, value))?; Ok(()) } /// Received document externally @@ -474,14 +458,14 @@ impl EventConverter { external: true, e3_id: evt.e3_id, share: evt.share, - }); + })?; } ReceivableDocument::EncryptionKeyCreated(evt) => { self.bus.publish(EncryptionKeyCreated { external: true, e3_id: evt.e3_id, key: evt.key, - }); + })?; } }; Ok(()) diff --git a/crates/test-helpers/src/usecase_helpers.rs b/crates/test-helpers/src/usecase_helpers.rs index 712e6471d7..e1d1506df0 100644 --- a/crates/test-helpers/src/usecase_helpers.rs +++ b/crates/test-helpers/src/usecase_helpers.rs @@ -89,12 +89,18 @@ pub fn generate_shares_hash_map( .collect::>()?; // Encrypt shares for all recipients using BFV - let encrypted_sk_sss = - BfvEncryptedShares::encrypt_all(&decrypted_sk_sss, &bfv_public_keys, &bfv_params, &mut bfv_rng)?; + let encrypted_sk_sss = BfvEncryptedShares::encrypt_all( + &decrypted_sk_sss, + &bfv_public_keys, + &bfv_params, + &mut bfv_rng, + )?; let encrypted_esi_sss: Vec = decrypted_esi_sss .iter() - .map(|esi| BfvEncryptedShares::encrypt_all(esi, &bfv_public_keys, &bfv_params, &mut bfv_rng)) + .map(|esi| { + BfvEncryptedShares::encrypt_all(esi, &bfv_public_keys, &bfv_params, &mut bfv_rng) + }) .collect::>()?; shares_hash_map.insert( @@ -165,9 +171,9 @@ pub fn get_decryption_keys( ts.esi_sss .iter() .map(|esi_shares| { - let encrypted = esi_shares - .clone_share(party_id) - .ok_or_else(|| anyhow::anyhow!("No esi_sss share for party {}", party_id))?; + let encrypted = esi_shares.clone_share(party_id).ok_or_else(|| { + anyhow::anyhow!("No esi_sss share for party {}", party_id) + })?; encrypted.decrypt(sk_bfv, &bfv_params, degree) }) .collect::>>() diff --git a/crates/tests/tests/integration.rs b/crates/tests/tests/integration.rs index 53f4c2f561..d1d0033994 100644 --- a/crates/tests/tests/integration.rs +++ b/crates/tests/tests/integration.rs @@ -299,7 +299,10 @@ async fn test_trbfv_actor() -> Result<()> { let _ = nodes .take_history_with_timeout(0, expected.len(), Duration::from_secs(1000)) .await?; - report.push(("All EncryptionKeyCreated events", encryption_keys_timer.elapsed())); + report.push(( + "All EncryptionKeyCreated events", + encryption_keys_timer.elapsed(), + )); // Then wait for all ThresholdShareCreated events let shares_timer = Instant::now(); From ef1b9b595c83351530cdbec295ecad3ed950d428 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Tue, 16 Dec 2025 20:24:27 +0500 Subject: [PATCH 03/19] chore: trigger CI From d4be5d3e0a6f415e74bf38c52587dadc12bb7907 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Wed, 17 Dec 2025 14:37:35 +0500 Subject: [PATCH 04/19] chore: update max document payload size --- crates/net/src/net_interface.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/net/src/net_interface.rs b/crates/net/src/net_interface.rs index 3041c63878..7801b2bed2 100644 --- a/crates/net/src/net_interface.rs +++ b/crates/net/src/net_interface.rs @@ -33,7 +33,7 @@ use tokio::{select, sync::broadcast, sync::mpsc}; use tracing::{debug, error, info, trace, warn}; const PROTOCOL_NAME: StreamProtocol = StreamProtocol::new("/ipfs/kad/1.0.0"); -const MAX_KADEMLIA_PAYLOAD_MB: usize = 10; +const MAX_KADEMLIA_PAYLOAD_MB: usize = 12; const MAX_GOSSIP_MSG_SIZE_KB: usize = 700; use crate::events::{GossipData, NetCommand}; From 4c7a3c07e5ace51dbe4467e62da1274674535484 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 13:24:29 +0500 Subject: [PATCH 05/19] feat: add document chunking --- Cargo.lock | 1 + crates/events/Cargo.toml | 1 + crates/events/src/chunk.rs | 115 ++++++++++ .../src/enclave_event/publish_document/mod.rs | 2 +- .../enclave_event/threshold_share_created.rs | 23 ++ crates/events/src/lib.rs | 2 + crates/keyshare/src/threshold_keyshare.rs | 53 +++-- crates/net/src/chunking/chunkable.rs | 185 +++++++++++++++ crates/net/src/chunking/collector.rs | 201 +++++++++++++++++ crates/net/src/chunking/mod.rs | 14 ++ crates/net/src/document_publisher.rs | 210 ++++++++++++++++-- crates/net/src/lib.rs | 1 + crates/net/src/net_interface.rs | 2 +- crates/trbfv/src/shares/bfv_encrypted.rs | 20 +- .../EnclaveTicketToken.json | 18 +- 15 files changed, 784 insertions(+), 64 deletions(-) create mode 100644 crates/events/src/chunk.rs create mode 100644 crates/net/src/chunking/chunkable.rs create mode 100644 crates/net/src/chunking/collector.rs create mode 100644 crates/net/src/chunking/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 2c9ec89ed9..55aeebe106 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2940,6 +2940,7 @@ dependencies = [ "e3-trbfv", "e3-utils", "futures-util", + "hex", "once_cell", "proptest", "rand 0.8.5", diff --git a/crates/events/Cargo.toml b/crates/events/Cargo.toml index aa58f52e17..0eb2104289 100644 --- a/crates/events/Cargo.toml +++ b/crates/events/Cargo.toml @@ -18,6 +18,7 @@ bs58 = { workspace = true } chrono = { workspace = true } derivative = { workspace = true } futures-util = { workspace = true } +hex = { workspace = true } once_cell = { workspace = true } rand = { workspace = true } serde = { workspace = true } diff --git a/crates/events/src/chunk.rs b/crates/events/src/chunk.rs new file mode 100644 index 0000000000..4576416e36 --- /dev/null +++ b/crates/events/src/chunk.rs @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +/// Unique identifier for a set of chunks, derived from content hash +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ChunkSetId([u8; 32]); + +impl ChunkSetId { + /// Create a ChunkSetId from the original document content + pub fn from_content(content: &[u8]) -> Self { + let mut hasher = Sha256::new(); + hasher.update(content); + let result = hasher.finalize(); + Self(result.into()) + } + + pub fn from_bytes(bytes: [u8; 32]) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } +} + +impl std::fmt::Display for ChunkSetId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "chunk:{}", hex::encode(&self.0[..4])) + } +} + +/// A single chunk of a larger document +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ChunkedDocument { + /// Identifier for the set of chunks this belongs to + pub chunk_id: ChunkSetId, + /// Index of this chunk (0-based) + pub chunk_index: u32, + /// Total number of chunks in the set + pub total_chunks: u32, + /// The actual data + pub data: Vec, +} + +impl ChunkedDocument { + /// Create a new chunk + pub fn new(chunk_id: ChunkSetId, chunk_index: u32, total_chunks: u32, data: Vec) -> Self { + Self { + chunk_id, + chunk_index, + total_chunks, + data, + } + } + + /// Create a single-chunk document (no splitting needed) + pub fn single(data: Vec) -> Self { + let chunk_id = ChunkSetId::from_content(&data); + Self { + chunk_id, + chunk_index: 0, + total_chunks: 1, + data, + } + } + + /// Check if this represents a complete (non-chunked) document + pub fn is_complete_set(&self) -> bool { + self.total_chunks == 1 + } + + /// Get the size of this chunk in bytes + pub fn size(&self) -> usize { + self.data.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_chunk_creation() { + let chunk = ChunkedDocument::single(vec![1, 2, 3, 4]); + assert_eq!(chunk.chunk_index, 0); + assert_eq!(chunk.total_chunks, 1); + assert!(chunk.is_complete_set()); + assert_eq!(chunk.size(), 4); + } + + #[test] + fn test_chunk_set_id() { + let content1 = b"test content 1"; + let content2 = b"test content 2"; + + let id1 = ChunkSetId::from_content(content1); + let id2 = ChunkSetId::from_content(content2); + assert_ne!(id1, id2); + + // Same content should produce same ID + let id3 = ChunkSetId::from_content(content1); + assert_eq!(id1, id3); + + // Test from_bytes round-trip + let bytes = *id1.as_bytes(); + let id4 = ChunkSetId::from_bytes(bytes); + assert_eq!(id1, id4); + } +} diff --git a/crates/events/src/enclave_event/publish_document/mod.rs b/crates/events/src/enclave_event/publish_document/mod.rs index 91b476a885..f0e3e7e9e0 100644 --- a/crates/events/src/enclave_event/publish_document/mod.rs +++ b/crates/events/src/enclave_event/publish_document/mod.rs @@ -11,7 +11,7 @@ use std::fmt::{self, Display}; use actix::Message; use chrono::{serde::ts_seconds, DateTime, Duration, Utc}; use e3_utils::ArcBytes; -use filter::Filter; +pub use filter::Filter; use serde::{Deserialize, Serialize}; use tracing::warn; diff --git a/crates/events/src/enclave_event/threshold_share_created.rs b/crates/events/src/enclave_event/threshold_share_created.rs index 51aac476c5..403f0e8237 100644 --- a/crates/events/src/enclave_event/threshold_share_created.rs +++ b/crates/events/src/enclave_event/threshold_share_created.rs @@ -34,6 +34,29 @@ pub struct ThresholdShare { pub esi_sss: Vec, } +impl ThresholdShare { + /// Extract only the shares meant for a specific party. + pub fn extract_for_party(&self, recipient_party_id: usize) -> Option { + let sk_sss = self.sk_sss.extract_for_party(recipient_party_id)?; + let esi_sss: Option> = self + .esi_sss + .iter() + .map(|shares| shares.extract_for_party(recipient_party_id)) + .collect(); + + esi_sss.map(|esi_sss| Self { + party_id: self.party_id, + pk_share: self.pk_share.clone(), + sk_sss, + esi_sss, + }) + } + + pub fn num_parties(&self) -> usize { + self.sk_sss.len() + } +} + #[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[rtype(result = "()")] pub struct ThresholdShareCreated { diff --git a/crates/events/src/lib.rs b/crates/events/src/lib.rs index eac423f6e7..435575cde9 100644 --- a/crates/events/src/lib.rs +++ b/crates/events/src/lib.rs @@ -5,6 +5,7 @@ // or FITNESS FOR A PARTICULAR PURPOSE. mod bus_handle; +mod chunk; mod correlation_id; mod e3id; mod enclave_event; @@ -19,6 +20,7 @@ mod sequencer; mod traits; pub use bus_handle::*; +pub use chunk::*; pub use correlation_id::*; pub use e3id::*; pub use enclave_event::*; diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 0d716929af..e72f023f2b 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -85,6 +85,7 @@ pub struct GeneratingThresholdShareData { esi_sss: Option>>, sk_bfv: Option, pk_bfv: Option, + collected_encryption_keys: Option>>, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -93,6 +94,7 @@ pub struct AggregatingDecryptionKey { sk_sss: Encrypted, esi_sss: Vec>, sk_bfv: SensitiveBytes, + collected_encryption_keys: Vec>, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -306,7 +308,6 @@ pub struct ThresholdKeyshare { cipher: Arc, decryption_key_collector: Option>, encryption_key_collector: Option>, - collected_encryption_keys: Option>>, multithread: Addr, state: Persistable, } @@ -318,7 +319,6 @@ impl ThresholdKeyshare { cipher: params.cipher, decryption_key_collector: None, encryption_key_collector: None, - collected_encryption_keys: None, multithread: params.multithread, state: params.state, } @@ -417,7 +417,7 @@ impl ThresholdKeyshare { )) })?; - let state = self.state.get().ok_or(anyhow!("No state"))?; + let state = self.state.try_get()?; self.bus.publish(EncryptionKeyCreated { e3_id: state.e3_id.clone(), key: Arc::new(EncryptionKey { @@ -425,7 +425,7 @@ impl ThresholdKeyshare { pk_bfv: pk_bfv_bytes, }), external: false, - }); + })?; Ok(()) } @@ -441,9 +441,7 @@ impl ThresholdKeyshare { msg.keys.len() ); - self.collected_encryption_keys = Some(msg.keys); - let state = self.state.get().ok_or(anyhow!("No state"))?; - let current: CollectingEncryptionKeysData = state.clone().try_into()?; + let current: CollectingEncryptionKeysData = self.state.try_get()?.try_into()?; self.state.try_mutate(|s| { s.new_state(KeyshareState::GeneratingThresholdShare( @@ -453,6 +451,7 @@ impl ThresholdKeyshare { esi_sss: None, sk_bfv: Some(current.sk_bfv), pk_bfv: Some(current.pk_bfv), + collected_encryption_keys: Some(msg.keys), }, )) })?; @@ -511,23 +510,26 @@ impl ThresholdKeyshare { let sk_sss = current.sk_sss; let sk_bfv = current.sk_bfv; let pk_bfv = current.pk_bfv; - let next = match (pk_share, sk_sss, &sk_bfv) { + let collected_encryption_keys = current.collected_encryption_keys; + let next = match (pk_share, sk_sss, &sk_bfv, &collected_encryption_keys) { // If the other shares are here then transition to aggregation - (Some(pk_share), Some(sk_sss), Some(sk_bfv_ref)) => { + (Some(pk_share), Some(sk_sss), Some(sk_bfv_ref), Some(keys)) => { K::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, sk_bfv: sk_bfv_ref.clone(), + collected_encryption_keys: keys.clone(), }) } // If the other shares are not here yet then dont transition - (None, None, _) => K::GeneratingThresholdShare(GeneratingThresholdShareData { + (None, None, _, _) => K::GeneratingThresholdShare(GeneratingThresholdShareData { esi_sss: Some(esi_sss), pk_share: None, sk_sss: None, sk_bfv, pk_bfv, + collected_encryption_keys, }), _ => bail!("Inconsistent state!"), }; @@ -589,28 +591,35 @@ impl ThresholdKeyshare { let esi_sss = current.esi_sss; let sk_bfv = current.sk_bfv; let pk_bfv = current.pk_bfv; - let next = match (esi_sss, sk_bfv) { + let collected_encryption_keys = current.collected_encryption_keys; + let next = match (esi_sss, sk_bfv, &collected_encryption_keys) { // If the esi shares and BFV key are here then transition to aggregation - (Some(esi_sss), Some(sk_bfv)) => { + (Some(esi_sss), Some(sk_bfv), Some(keys)) => { KeyshareState::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, sk_bfv, + collected_encryption_keys: keys.clone(), }) } // If esi shares are not here yet then don't transition - (None, sk_bfv) => { + (None, sk_bfv, _) => { KeyshareState::GeneratingThresholdShare(GeneratingThresholdShareData { esi_sss: None, pk_share: Some(pk_share), sk_sss: Some(sk_sss), sk_bfv, pk_bfv, + collected_encryption_keys, }) } // If we have esi_sss but no sk_bfv, that's an error - (Some(_), None) => bail!("Have esi_sss but no sk_bfv - inconsistent state!"), + (Some(_), None, _) => bail!("Have esi_sss but no sk_bfv - inconsistent state!"), + // If we have shares ready but no encryption keys, that's an error + (Some(_), Some(_), None) => { + bail!("Have shares but no collected encryption keys - inconsistent state!") + } }; s.new_state(next) })?; @@ -633,6 +642,7 @@ impl ThresholdKeyshare { pk_share, sk_sss, esi_sss, + collected_encryption_keys, .. }), party_id, @@ -643,11 +653,8 @@ impl ThresholdKeyshare { bail!("Invalid state!"); }; - // Get collected BFV public keys from all parties - let encryption_keys = self - .collected_encryption_keys - .as_ref() - .ok_or(anyhow!("No encryption keys collected"))?; + // Get collected BFV public keys from all parties (now from persisted state) + let encryption_keys = &collected_encryption_keys; // Convert to BFV public keys let params = get_share_encryption_params(); @@ -699,7 +706,7 @@ impl ThresholdKeyshare { ) -> Result { info!("AllThresholdSharesCollected"); let cipher = self.cipher.clone(); - let state = self.state.get().ok_or(anyhow!("No state found"))?; + let state = self.state.try_get()?; let party_id = state.party_id as usize; let trbfv_config = state.get_trbfv_config(); @@ -781,7 +788,7 @@ impl ThresholdKeyshare { s.new_state(next) })?; - let state = self.state.get().ok_or(anyhow!("No state found"))?; + let state = self.state.try_get()?; let e3_id = state.get_e3_id().clone(); let address = state.get_address().to_owned(); let current: ReadyForDecryption = state.clone().try_into()?; @@ -816,7 +823,7 @@ impl ThresholdKeyshare { })?; let ciphertext_output = msg.ciphertext_output; - let state = self.state.get().ok_or(anyhow!("No state found"))?; + let state = self.state.try_get()?; let decrypting: Decrypting = state.clone().try_into()?; let trbfv_config = state.get_trbfv_config(); let event = ComputeRequest::TrBFV(TrBFVRequest::CalculateDecryptionShare( @@ -839,7 +846,7 @@ impl ThresholdKeyshare { res: ComputeResponse, ) -> Result<()> { let msg: CalculateDecryptionShareResponse = res.try_into()?; - let state = self.state.get().ok_or(anyhow!("No state found"))?; + let state = self.state.try_get()?; let party_id = state.party_id; let node = state.address; let e3_id = state.e3_id; diff --git a/crates/net/src/chunking/chunkable.rs b/crates/net/src/chunking/chunkable.rs new file mode 100644 index 0000000000..7e98adf6b9 --- /dev/null +++ b/crates/net/src/chunking/chunkable.rs @@ -0,0 +1,185 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use anyhow::{bail, Result}; +use e3_events::{ChunkSetId, ChunkedDocument}; +use serde::{de::DeserializeOwned, Serialize}; +use tracing::{debug, info}; + +/// Trait for documents that can be split into chunks for transmission +pub trait Chunkable: Serialize + DeserializeOwned + Clone + Sized { + fn max_chunk_size() -> usize { + 10 * 1024 * 1024 + } + + fn into_chunks(&self) -> Result> { + let bytes = bincode::serialize(self)?; + let max_size = Self::max_chunk_size(); + + debug!( + "Chunking document: {} bytes, max chunk size: {} bytes", + bytes.len(), + max_size + ); + + if bytes.len() <= max_size { + // Small enough, send as single chunk + debug!("Document fits in single chunk"); + return Ok(vec![ChunkedDocument::single(bytes)]); + } + + // Split into multiple chunks + let num_chunks = (bytes.len() + max_size - 1) / max_size; + let chunk_id = ChunkSetId::from_content(&bytes); + + info!( + "Splitting document into {} chunks (chunk_id: {})", + num_chunks, chunk_id + ); + + Ok(bytes + .chunks(max_size) + .enumerate() + .map(|(idx, chunk)| { + ChunkedDocument::new( + chunk_id.clone(), + idx as u32, + num_chunks as u32, + chunk.to_vec(), + ) + }) + .collect()) + } + + /// Reassemble from chunks + fn from_chunks(chunks: Vec) -> Result { + if chunks.is_empty() { + bail!("Cannot reassemble from zero chunks"); + } + + // If single chunk, just deserialize + if chunks.len() == 1 && chunks[0].is_complete_set() { + return Ok(bincode::deserialize(&chunks[0].data)?); + } + + // Validate all chunks are from same set + let chunk_id = &chunks[0].chunk_id; + let total = chunks[0].total_chunks; + + if chunks.len() != total as usize { + bail!("Missing chunks: got {} expected {}", chunks.len(), total); + } + + if !chunks.iter().all(|c| c.chunk_id == *chunk_id) { + bail!("Chunks from different sets"); + } + + // Check for duplicate indices + let mut indices: Vec = chunks.iter().map(|c| c.chunk_index).collect(); + indices.sort_unstable(); + indices.dedup(); + if indices.len() != chunks.len() { + bail!("Duplicate chunk indices detected"); + } + + // Sort by index and concatenate + let mut sorted = chunks; + sorted.sort_by_key(|c| c.chunk_index); + + let bytes: Vec = sorted.into_iter().flat_map(|c| c.data).collect(); + + info!( + "Reassembling document from {} chunks ({} bytes total)", + total, + bytes.len() + ); + + Ok(bincode::deserialize(&bytes)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::{Deserialize, Serialize}; + + #[derive(Clone, Serialize, Deserialize, PartialEq, Debug)] + struct TestDoc { + data: Vec, + } + + impl Chunkable for TestDoc { + fn max_chunk_size() -> usize { + 100 + } + } + + #[test] + fn test_small_document_single_chunk() { + let doc = TestDoc { + data: vec![1, 2, 3, 4, 5], + }; + let chunks = doc.into_chunks().unwrap(); + assert_eq!(chunks.len(), 1); + assert!(chunks[0].is_complete_set()); + + let restored = TestDoc::from_chunks(chunks).unwrap(); + assert_eq!(doc, restored); + } + + #[test] + fn test_large_document_multiple_chunks() { + let doc = TestDoc { + data: vec![42; 500], + }; + let chunks = doc.into_chunks().unwrap(); + assert!(chunks.len() > 1); + + // Verify all chunks have same chunk_id + let first_id = &chunks[0].chunk_id; + assert!(chunks.iter().all(|c| c.chunk_id == *first_id)); + + // Verify chunk indices + for (idx, chunk) in chunks.iter().enumerate() { + assert_eq!(chunk.chunk_index, idx as u32); + assert_eq!(chunk.total_chunks, chunks.len() as u32); + } + + let restored = TestDoc::from_chunks(chunks).unwrap(); + assert_eq!(doc, restored); + } + + #[test] + fn test_missing_chunk_fails() { + let doc = TestDoc { + data: vec![42; 500], + }; + let mut chunks = doc.into_chunks().unwrap(); + chunks.pop(); // Remove last chunk + + let result = TestDoc::from_chunks(chunks); + assert!(result.is_err()); + } + + #[test] + fn test_mixed_chunk_sets_fails() { + let doc1 = TestDoc { + data: vec![42; 500], + }; + let doc2 = TestDoc { + data: vec![99; 500], + }; + + let mut chunks1 = doc1.into_chunks().unwrap(); + let chunks2 = doc2.into_chunks().unwrap(); + + // Mix chunks from different sets + chunks1[0] = chunks2[0].clone(); + + let result = TestDoc::from_chunks(chunks1); + assert!(result.is_err()); + } +} diff --git a/crates/net/src/chunking/collector.rs b/crates/net/src/chunking/collector.rs new file mode 100644 index 0000000000..75b65ffcaf --- /dev/null +++ b/crates/net/src/chunking/collector.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use super::chunkable::Chunkable; +use actix::prelude::*; +use anyhow::Result; +use e3_events::{BusHandle, ChunkSetId, ChunkedDocument, E3id, EventPublisher}; +use std::collections::HashMap; +use std::marker::PhantomData; +use std::time::{Duration, Instant}; +use tracing::{debug, error, info, warn}; + +/// Collects chunks and reassembles them into complete documents +pub struct ChunkCollector { + /// E3 ID this collector is for + e3_id: E3id, + /// Chunks we've received so far, grouped by chunk_id + received_chunks: HashMap>, + /// Expected total chunks per set + expected_totals: HashMap, + /// When we started waiting for each chunk set (for timeout) + start_times: HashMap, + /// Event bus to publish completed documents + bus: BusHandle, + /// Timeout duration + timeout: Duration, + /// Phantom data for type parameter + _phantom: PhantomData, +} + +impl ChunkCollector { + pub fn new(e3_id: E3id, bus: BusHandle, timeout: Duration) -> Self { + Self { + e3_id, + received_chunks: HashMap::new(), + expected_totals: HashMap::new(), + start_times: HashMap::new(), + bus, + timeout, + _phantom: PhantomData, + } + } + + pub fn setup(e3_id: E3id, bus: BusHandle) -> Addr { + Self::new(e3_id, bus, Duration::from_secs(60)).start() + } + + /// Process a received chunk + fn handle_chunk_internal(&mut self, chunk: ChunkedDocument) -> Result> { + let chunk_id = chunk.chunk_id.clone(); + let total = chunk.total_chunks; + + debug!( + "Received chunk {}/{} for chunk_id: {}", + chunk.chunk_index + 1, + total, + chunk_id + ); + + // Track expected total and start time + self.expected_totals + .entry(chunk_id.clone()) + .or_insert(total); + self.start_times + .entry(chunk_id.clone()) + .or_insert_with(Instant::now); + + // Add to received chunks + let chunks = self + .received_chunks + .entry(chunk_id.clone()) + .or_insert_with(Vec::new); + + // Check if we already have this chunk index + if chunks.iter().any(|c| c.chunk_index == chunk.chunk_index) { + debug!("Duplicate chunk {} ignored", chunk.chunk_index); + return Ok(None); + } + + chunks.push(chunk); + + // Check if complete + if chunks.len() == total as usize { + info!( + "All {} chunks received for chunk_id: {}, reassembling...", + total, chunk_id + ); + let chunks = self.received_chunks.remove(&chunk_id).unwrap(); + self.expected_totals.remove(&chunk_id); + self.start_times.remove(&chunk_id); + + let document = T::from_chunks(chunks)?; + return Ok(Some(document)); + } + + debug!( + "Waiting for {}/{} chunks for chunk_id: {}", + chunks.len(), + total, + chunk_id + ); + + Ok(None) + } + + /// Check for timeouts and clean up stale chunk sets + fn check_timeouts(&mut self) { + let now = Instant::now(); + let timed_out: Vec<_> = self + .start_times + .iter() + .filter(|(_, start)| now.duration_since(**start) > self.timeout) + .map(|(id, _)| id.clone()) + .collect(); + + for chunk_id in timed_out { + let received = self + .received_chunks + .get(&chunk_id) + .map(|c| c.len()) + .unwrap_or(0); + let expected = self.expected_totals.get(&chunk_id).copied().unwrap_or(0); + + warn!( + "Chunk set {} timed out (received {}/{} chunks)", + chunk_id, received, expected + ); + + self.received_chunks.remove(&chunk_id); + self.expected_totals.remove(&chunk_id); + self.start_times.remove(&chunk_id); + } + } +} + +impl Actor for ChunkCollector { + type Context = Context; + + fn started(&mut self, ctx: &mut Self::Context) { + info!("ChunkCollector started for E3: {}", self.e3_id); + + // Periodic timeout check every 5 seconds + ctx.run_interval(Duration::from_secs(5), |act, _ctx| { + act.check_timeouts(); + }); + } + + fn stopped(&mut self, _ctx: &mut Self::Context) { + info!("ChunkCollector stopped for E3: {}", self.e3_id); + } +} + +/// Message to send a chunk to the collector +#[derive(Message, Clone, Debug)] +#[rtype(result = "()")] +pub struct ChunkReceived { + pub chunk: ChunkedDocument, + pub _phantom: PhantomData, +} + +impl ChunkReceived { + pub fn new(chunk: ChunkedDocument) -> Self { + Self { + chunk, + _phantom: PhantomData, + } + } +} + +impl Handler> for ChunkCollector +where + T: actix::Message + Send + Into, + T::Result: Send, +{ + type Result = (); + + fn handle(&mut self, msg: ChunkReceived, _ctx: &mut Self::Context) { + match self.handle_chunk_internal(msg.chunk) { + Ok(Some(document)) => { + info!("Document reassembled successfully, publishing to bus"); + // Publish the reassembled document to the event bus + if let Err(e) = self.bus.publish(document) { + error!("Failed to publish reassembled document: {:?}", e); + } + } + Ok(None) => { + // Still waiting for more chunks + } + Err(e) => { + error!("Failed to process chunk: {:?}", e); + } + } + } +} + +// Note: ChunkCollector tests are covered by integration tests with real event types +// that implement Into. Unit tests with mock types would require +// duplicating the handler implementation without the publish constraint. diff --git a/crates/net/src/chunking/mod.rs b/crates/net/src/chunking/mod.rs new file mode 100644 index 0000000000..20b176a526 --- /dev/null +++ b/crates/net/src/chunking/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +mod chunkable; +mod collector; + +pub use chunkable::Chunkable; +pub use collector::{ChunkCollector, ChunkReceived}; +// Re-export chunk types from events crate for convenience +pub use e3_events::{ChunkSetId, ChunkedDocument}; + diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index 055a75cda3..7c3e46d011 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -5,6 +5,7 @@ // or FITNESS FOR A PARTICULAR PURPOSE. use crate::{ + chunking::{ChunkCollector, ChunkReceived, Chunkable, ChunkedDocument}, events::{ call_and_await_response, DocumentPublishedNotification, GossipData, NetCommand, NetEvent, }, @@ -16,7 +17,7 @@ use chrono::{DateTime, Utc}; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CorrelationId, DocumentKind, DocumentMeta, DocumentReceived, E3RequestComplete, E3id, EType, EnclaveEvent, EnclaveEventData, - EncryptionKeyCreated, Event, PartyId, PublishDocumentRequested, ThresholdShareCreated, + EncryptionKeyCreated, Event, Filter, PartyId, PublishDocumentRequested, ThresholdShareCreated, }; use e3_utils::retry::{retry_with_backoff, to_retry}; use e3_utils::ArcBytes; @@ -28,7 +29,13 @@ use std::{ time::{Duration, Instant}, }; use tokio::sync::{broadcast, mpsc}; -use tracing::{debug, error, warn}; +use tracing::{debug, error, info}; + +impl Chunkable for ThresholdShareCreated { + fn max_chunk_size() -> usize { + 10 * 1024 * 1024 + } +} const KADEMLIA_PUT_TIMEOUT: Duration = Duration::from_secs(30); const KADEMLIA_GET_TIMEOUT: Duration = Duration::from_secs(30); @@ -373,22 +380,30 @@ async fn broadcast_document_published_notification( .await } -/// Convert between ThresholdShareCreated and DocumentPublished events +/// Converts between internal events and network documents. +/// +/// Handles party-filtered distribution and chunking: +/// - Outgoing: Splits large ThresholdShares into party-specific chunks +/// - Incoming: Filters documents by party_id and reassembles chunks pub struct EventConverter { bus: BusHandle, + threshold_share_collector: Option>>, + ids: HashMap, } #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] enum ReceivableDocument { ThresholdShareCreated(ThresholdShareCreated), EncryptionKeyCreated(EncryptionKeyCreated), + Chunk(ChunkedDocument), } impl ReceivableDocument { - pub fn get_e3_id(&self) -> &E3id { + pub fn get_e3_id(&self) -> Option<&E3id> { match self { - ReceivableDocument::ThresholdShareCreated(d) => &d.e3_id, - ReceivableDocument::EncryptionKeyCreated(d) => &d.e3_id, + ReceivableDocument::ThresholdShareCreated(d) => Some(&d.e3_id), + ReceivableDocument::EncryptionKeyCreated(d) => Some(&d.e3_id), + ReceivableDocument::Chunk(_) => None, // Chunks don't have E3id directly } } @@ -403,7 +418,19 @@ impl ReceivableDocument { impl EventConverter { pub fn new(bus: &BusHandle) -> Self { - Self { bus: bus.clone() } + Self { + bus: bus.clone(), + threshold_share_collector: None, + ids: HashMap::new(), + } + } + + fn handle_ciphernode_selected(&mut self, event: CiphernodeSelected) -> Result<()> { + let CiphernodeSelected { + e3_id, party_id, .. + } = event; + self.ids.insert(e3_id, party_id); + Ok(()) } pub fn setup(bus: &BusHandle) -> Addr { let addr = Self::new(bus).start(); @@ -412,48 +439,166 @@ impl EventConverter { bus.subscribe("DocumentReceived", addr.clone().into()); addr } - /// Local node created a threshold share. Send it as a published document - pub fn handle_threshold_share_created(&self, msg: ThresholdShareCreated) -> Result<()> { - // If this is received from elsewhere - if msg.external { - return Ok(()); + + /// Initialize chunk collector for a specific E3 + fn ensure_collector(&mut self, e3_id: &E3id) { + if self.threshold_share_collector.is_none() { + debug!("Creating ThresholdShare chunk collector for E3: {}", e3_id); + let collector = + ChunkCollector::::setup(e3_id.clone(), self.bus.clone()); + self.threshold_share_collector = Some(collector); } - let receivable = ReceivableDocument::ThresholdShareCreated(msg); + } + /// Publish a receivable document with party filter + fn publish_filtered( + &self, + receivable: ReceivableDocument, + e3_id: &E3id, + party_id: u64, + ) -> Result<()> { let value = ArcBytes::from_bytes(&receivable.to_bytes()?); let meta = DocumentMeta::new( - receivable.get_e3_id().clone(), + e3_id.clone(), DocumentKind::TrBFV, - vec![], + vec![Filter::Item(party_id)], None, ); self.bus .publish(PublishDocumentRequested::new(meta, value))?; Ok(()) } + + /// Publish chunks for a specific party + fn publish_party_chunks( + &self, + chunks: &[ChunkedDocument], + e3_id: &E3id, + party_id: usize, + ) -> Result<()> { + let party_id_u64 = party_id as u64; + + match chunks { + [single] if single.is_complete_set() => { + debug!( + "Party {} share: single chunk ({} bytes)", + party_id, + single.data.len() + ); + let share = bincode::deserialize(&single.data)?; + self.publish_filtered( + ReceivableDocument::ThresholdShareCreated(share), + e3_id, + party_id_u64, + )?; + } + multiple => { + info!("Party {} share: {} chunks", party_id, multiple.len()); + for (idx, chunk) in multiple.iter().enumerate() { + debug!( + " Chunk {}/{} (chunk_id: {})", + idx + 1, + multiple.len(), + chunk.chunk_id + ); + self.publish_filtered( + ReceivableDocument::Chunk(chunk.clone()), + e3_id, + party_id_u64, + )?; + } + } + } + Ok(()) + } + + /// Local node created a threshold share. Send it as a published document. + /// Uses party-filtered distribution: each party only receives their specific shares. + pub fn handle_threshold_share_created(&self, msg: ThresholdShareCreated) -> Result<()> { + if msg.external { + return Ok(()); + } + + let e3_id = msg.e3_id.clone(); + let num_parties = msg.share.num_parties(); + let total_size_mb = bincode::serialize(&msg)?.len() as f64 / (1024.0 * 1024.0); + + info!( + "Publishing ThresholdShare for E3 {} ({:.2} MB total, {} parties)", + e3_id, total_size_mb, num_parties + ); + + // Publish party-filtered shares + for party_id in 0..num_parties { + let party_share = msg + .share + .extract_for_party(party_id) + .ok_or_else(|| anyhow::anyhow!("Failed to extract share for party {}", party_id))?; + + let party_msg = ThresholdShareCreated { + e3_id: e3_id.clone(), + share: Arc::new(party_share), + external: false, + }; + + let chunks = party_msg.into_chunks()?; + self.publish_party_chunks(&chunks, &e3_id, party_id)?; + } + + Ok(()) + } /// Local node created an encryption key. Send it as a published document pub fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { // If this is received from elsewhere if msg.external { return Ok(()); } + let e3_id = msg.e3_id.clone(); let receivable = ReceivableDocument::EncryptionKeyCreated(msg); let value = ArcBytes::from_bytes(&receivable.to_bytes()?); - let meta = DocumentMeta::new( - receivable.get_e3_id().clone(), - DocumentKind::TrBFV, - vec![], - None, - ); + let meta = DocumentMeta::new(e3_id, DocumentKind::TrBFV, vec![], None); self.bus .publish(PublishDocumentRequested::new(meta, value))?; Ok(()) } /// Received document externally - pub fn handle_document_received(&self, msg: DocumentReceived) -> Result<()> { - warn!("Converting DocumentReceived..."); + /// Check if document passes party filter + fn passes_filter(&self, e3_id: &E3id, meta: &DocumentMeta) -> bool { + match self.ids.get(e3_id) { + Some(party_id) if !meta.matches(party_id) => { + debug!( + "Filtered out: doesn't match party_id {} for E3 {}", + party_id, e3_id + ); + false + } + _ => true, + } + } + + /// Handle a chunk by forwarding to collector + fn handle_chunk(&mut self, chunk: ChunkedDocument, e3_id: &E3id) { + debug!( + "Received chunk {}/{} ({})", + chunk.chunk_index + 1, + chunk.total_chunks, + chunk.chunk_id + ); + self.ensure_collector(e3_id); + if let Some(ref collector) = self.threshold_share_collector { + collector.do_send(ChunkReceived::::new(chunk)); + } + } + + fn handle_document_received(&mut self, msg: DocumentReceived) -> Result<()> { + if !self.passes_filter(&msg.meta.e3_id, &msg.meta) { + return Ok(()); + } + let receivable = ReceivableDocument::from_bytes(&msg.value.extract_bytes())?; + match receivable { ReceivableDocument::ThresholdShareCreated(evt) => { + debug!("Received complete ThresholdShareCreated"); self.bus.publish(ThresholdShareCreated { external: true, e3_id: evt.e3_id, @@ -461,13 +606,17 @@ impl EventConverter { })?; } ReceivableDocument::EncryptionKeyCreated(evt) => { + debug!("Received EncryptionKeyCreated"); self.bus.publish(EncryptionKeyCreated { external: true, e3_id: evt.e3_id, key: evt.key, })?; } - }; + ReceivableDocument::Chunk(chunk) => { + self.handle_chunk(chunk, &msg.meta.e3_id); + } + } Ok(()) } } @@ -483,6 +632,7 @@ impl Handler for EventConverter { EnclaveEventData::ThresholdShareCreated(data) => ctx.notify(data), EnclaveEventData::EncryptionKeyCreated(data) => ctx.notify(data), EnclaveEventData::DocumentReceived(data) => ctx.notify(data), + EnclaveEventData::CiphernodeSelected(data) => ctx.notify(data), _ => (), } } @@ -508,6 +658,18 @@ impl Handler for EventConverter { } } +impl Handler for EventConverter { + type Result = (); + fn handle(&mut self, msg: CiphernodeSelected, _ctx: &mut Self::Context) -> Self::Result { + match self.handle_ciphernode_selected(msg) { + Ok(_) => (), + Err(e) => { + error!("Error handling CiphernodeSelected: {:?}", e); + } + } + } +} + impl Handler for EventConverter { type Result = (); fn handle(&mut self, msg: DocumentReceived, _ctx: &mut Self::Context) -> Self::Result { diff --git a/crates/net/src/lib.rs b/crates/net/src/lib.rs index 0b9868cd08..4e79d2f150 100644 --- a/crates/net/src/lib.rs +++ b/crates/net/src/lib.rs @@ -4,6 +4,7 @@ // without even the implied warranty of MERCHANTABILITY // or FITNESS FOR A PARTICULAR PURPOSE. +pub mod chunking; mod cid; mod dialer; mod document_publisher; diff --git a/crates/net/src/net_interface.rs b/crates/net/src/net_interface.rs index 7801b2bed2..3041c63878 100644 --- a/crates/net/src/net_interface.rs +++ b/crates/net/src/net_interface.rs @@ -33,7 +33,7 @@ use tokio::{select, sync::broadcast, sync::mpsc}; use tracing::{debug, error, info, trace, warn}; const PROTOCOL_NAME: StreamProtocol = StreamProtocol::new("/ipfs/kad/1.0.0"); -const MAX_KADEMLIA_PAYLOAD_MB: usize = 12; +const MAX_KADEMLIA_PAYLOAD_MB: usize = 10; const MAX_GOSSIP_MSG_SIZE_KB: usize = 700; use crate::events::{GossipData, NetCommand}; diff --git a/crates/trbfv/src/shares/bfv_encrypted.rs b/crates/trbfv/src/shares/bfv_encrypted.rs index adf383fb09..53cc78b65d 100644 --- a/crates/trbfv/src/shares/bfv_encrypted.rs +++ b/crates/trbfv/src/shares/bfv_encrypted.rs @@ -6,6 +6,7 @@ use anyhow::{Context, Result}; use derivative::Derivative; +use e3_utils::utility_types::ArcBytes; use fhe::bfv::{BfvParameters, Ciphertext, Encoding, Plaintext, PublicKey, SecretKey}; use fhe_traits::{ DeserializeParametrized, FheDecoder, FheDecrypter, FheEncoder, FheEncrypter, @@ -31,13 +32,13 @@ pub use crate::helpers::{ #[derive(Derivative, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[derivative(Debug)] pub struct BfvEncryptedShare { - /// BFV ciphertexts, one per modulus level (serialized) - #[derivative(Debug(format_with = "debug_vec_bytes"))] - ciphertexts: Vec>, + /// BFV ciphertexts, one per modulus level + #[derivative(Debug(format_with = "debug_vec_arcbytes"))] + ciphertexts: Vec, } -/// Debug helper for Vec> -fn debug_vec_bytes(v: &Vec>, f: &mut std::fmt::Formatter) -> std::fmt::Result { +/// Debug helper for Vec +fn debug_vec_arcbytes(v: &Vec, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!( f, "[{} ciphertexts, total {} bytes]", @@ -79,7 +80,7 @@ impl BfvEncryptedShare { .try_encrypt(&pt, rng) .context("Failed to encrypt share")?; - ciphertexts.push(ct.to_bytes()); + ciphertexts.push(ArcBytes::from_bytes(&ct.to_bytes())); } Ok(Self { ciphertexts }) @@ -184,6 +185,13 @@ impl BfvEncryptedShares { self.shares.get(party_id).cloned() } + /// Extract only the share for a specific party (for bandwidth optimization) + pub fn extract_for_party(&self, party_id: usize) -> Option { + self.shares.get(party_id).map(|share| Self { + shares: vec![share.clone()], + }) + } + /// Number of encrypted shares pub fn len(&self) -> usize { self.shares.len() diff --git a/packages/enclave-contracts/artifacts/contracts/token/EnclaveTicketToken.sol/EnclaveTicketToken.json b/packages/enclave-contracts/artifacts/contracts/token/EnclaveTicketToken.sol/EnclaveTicketToken.json index 17018ad9f1..32e4ef8771 100644 --- a/packages/enclave-contracts/artifacts/contracts/token/EnclaveTicketToken.sol/EnclaveTicketToken.json +++ b/packages/enclave-contracts/artifacts/contracts/token/EnclaveTicketToken.sol/EnclaveTicketToken.json @@ -1148,7 +1148,7 @@ "linkReferences": {}, "deployedLinkReferences": {}, "immutableReferences": { - "3415": [ + "1944": [ { "length": 32, "start": 872 @@ -1174,43 +1174,43 @@ "start": 3711 } ], - "6684": [ + "4931": [ { "length": 32, "start": 3942 } ], - "6686": [ + "4933": [ { "length": 32, "start": 3900 } ], - "6688": [ + "4935": [ { "length": 32, "start": 3858 } ], - "6690": [ + "4937": [ { "length": 32, "start": 4023 } ], - "6692": [ + "4939": [ { "length": 32, "start": 4063 } ], - "6695": [ + "4942": [ { "length": 32, "start": 4727 } ], - "6698": [ + "4945": [ { "length": 32, "start": 4772 @@ -1218,5 +1218,5 @@ ] }, "inputSourceName": "project/contracts/token/EnclaveTicketToken.sol", - "buildInfoId": "solc-0_8_28-572e77328edd43ea3643262ca5f82babf769c493" + "buildInfoId": "solc-0_8_28-bef1dec98cb6cc37e15bd06abee5a353862d85f4" } \ No newline at end of file From 0b2de3e96847ffd4f379bf6e0530f370d8d44576 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 14:25:42 +0500 Subject: [PATCH 06/19] feat: inject params during init --- Cargo.lock | 1 + crates/ciphernode-builder/Cargo.toml | 1 + .../src/ciphernode_builder.rs | 2 + crates/keyshare/src/ext.rs | 5 + crates/keyshare/src/threshold_keyshare.rs | 12 +- crates/net/src/document_publisher.rs | 120 ++++++------------ 6 files changed, 54 insertions(+), 87 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 55aeebe106..8fbf7c3276 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2774,6 +2774,7 @@ dependencies = [ "e3-multithread", "e3-request", "e3-sortition", + "e3-trbfv", "e3-utils", "tracing", ] diff --git a/crates/ciphernode-builder/Cargo.toml b/crates/ciphernode-builder/Cargo.toml index 21a7aa068a..2603e56169 100644 --- a/crates/ciphernode-builder/Cargo.toml +++ b/crates/ciphernode-builder/Cargo.toml @@ -22,5 +22,6 @@ e3-keyshare.workspace = true e3-multithread.workspace = true e3-request.workspace = true e3-sortition.workspace = true +e3-trbfv.workspace = true e3-utils.workspace = true tracing.workspace = true diff --git a/crates/ciphernode-builder/src/ciphernode_builder.rs b/crates/ciphernode-builder/src/ciphernode_builder.rs index 0a88d6793a..54874e5de7 100644 --- a/crates/ciphernode-builder/src/ciphernode_builder.rs +++ b/crates/ciphernode-builder/src/ciphernode_builder.rs @@ -438,12 +438,14 @@ impl CiphernodeBuilder { if let Some(KeyshareKind::Threshold) = self.keyshare { let multithread = self.ensure_multithread(); + let share_encryption_params = e3_trbfv::helpers::get_share_encryption_params(); info!("Setting up ThresholdKeyshareExtension"); e3_builder = e3_builder.with(ThresholdKeyshareExtension::create( &bus, &self.cipher, &multithread, &addr, + share_encryption_params, )) } diff --git a/crates/keyshare/src/ext.rs b/crates/keyshare/src/ext.rs index 39b8db852a..2c0750224e 100644 --- a/crates/keyshare/src/ext.rs +++ b/crates/keyshare/src/ext.rs @@ -117,6 +117,7 @@ pub struct ThresholdKeyshareExtension { cipher: Arc, address: String, multithread: Addr, + share_encryption_params: Arc, } impl ThresholdKeyshareExtension { @@ -125,12 +126,14 @@ impl ThresholdKeyshareExtension { cipher: &Arc, multithread: &Addr, address: &str, + share_encryption_params: Arc, ) -> Box { Box::new(Self { bus: bus.clone(), cipher: cipher.to_owned(), multithread: multithread.clone(), address: address.to_owned(), + share_encryption_params, }) } } @@ -173,6 +176,7 @@ impl E3Extension for ThresholdKeyshareExtension { cipher: self.cipher.clone(), multithread: self.multithread.clone(), state: container, + share_encryption_params: self.share_encryption_params.clone(), }) .start() .into(), @@ -203,6 +207,7 @@ impl E3Extension for ThresholdKeyshareExtension { cipher: self.cipher.clone(), multithread: self.multithread.clone(), state, + share_encryption_params: self.share_encryption_params.clone(), }) .start() .into(); diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index e72f023f2b..85d5b1343c 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -22,11 +22,12 @@ use e3_trbfv::{ }, gen_esi_sss::{GenEsiSssRequest, GenEsiSssResponse}, gen_pk_share_and_sk_sss::GenPkShareAndSkSssRequest, - helpers::{deserialize_secret_key, get_share_encryption_params, serialize_secret_key}, + helpers::{deserialize_secret_key, serialize_secret_key}, shares::{BfvEncryptedShares, EncryptableVec, Encrypted, ShamirShare, SharedSecret}, TrBFVConfig, TrBFVRequest, TrBFVResponse, }; use e3_utils::{bail, to_ordered_vec, utility_types::ArcBytes}; +use fhe::bfv::BfvParameters; use fhe::bfv::{PublicKey, SecretKey}; use fhe_traits::{DeserializeParametrized, Serialize}; use rand::{rngs::OsRng, SeedableRng}; @@ -301,6 +302,7 @@ pub struct ThresholdKeyshareParams { pub cipher: Arc, pub multithread: Addr, pub state: Persistable, + pub share_encryption_params: Arc, } pub struct ThresholdKeyshare { @@ -310,6 +312,7 @@ pub struct ThresholdKeyshare { encryption_key_collector: Option>, multithread: Addr, state: Persistable, + share_encryption_params: Arc, } impl ThresholdKeyshare { @@ -321,6 +324,7 @@ impl ThresholdKeyshare { encryption_key_collector: None, multithread: params.multithread, state: params.state, + share_encryption_params: params.share_encryption_params, } } } @@ -398,7 +402,7 @@ impl ThresholdKeyshare { let _ = self.ensure_collector(address.clone()); let _ = self.ensure_encryption_key_collector(address.clone()); - let params = get_share_encryption_params(); + let params = self.share_encryption_params.clone(); let mut rng = OsRng; let sk_bfv = SecretKey::random(¶ms, &mut rng); let pk_bfv = PublicKey::new(&sk_bfv, &mut rng); @@ -657,7 +661,7 @@ impl ThresholdKeyshare { let encryption_keys = &collected_encryption_keys; // Convert to BFV public keys - let params = get_share_encryption_params(); + let params = self.share_encryption_params.clone(); let recipient_pks: Vec = encryption_keys .iter() .map(|k| { @@ -713,7 +717,7 @@ impl ThresholdKeyshare { // Get our BFV secret key from state let current: AggregatingDecryptionKey = state.clone().try_into()?; let sk_bytes = current.sk_bfv.access(&cipher)?; - let params = get_share_encryption_params(); + let params = self.share_encryption_params.clone(); let sk_bfv = deserialize_secret_key(&sk_bytes, ¶ms)?; let degree = params.degree(); diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index 7c3e46d011..647420d392 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -382,13 +382,15 @@ async fn broadcast_document_published_notification( /// Converts between internal events and network documents. /// -/// Handles party-filtered distribution and chunking: -/// - Outgoing: Splits large ThresholdShares into party-specific chunks -/// - Incoming: Filters documents by party_id and reassembles chunks +/// Responsibilities: +/// - Outgoing: Converts ThresholdShareCreated → party-filtered PublishDocumentRequested +/// - Incoming: Converts DocumentReceived → ThresholdShareCreated/EncryptionKeyCreated +/// - Reassembles chunked documents via ChunkCollector +/// +/// Note: Party filtering is done by DocumentPublisher BEFORE fetching from DHT. pub struct EventConverter { bus: BusHandle, - threshold_share_collector: Option>>, - ids: HashMap, + chunk_collector: Option>>, } #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -399,14 +401,6 @@ enum ReceivableDocument { } impl ReceivableDocument { - pub fn get_e3_id(&self) -> Option<&E3id> { - match self { - ReceivableDocument::ThresholdShareCreated(d) => Some(&d.e3_id), - ReceivableDocument::EncryptionKeyCreated(d) => Some(&d.e3_id), - ReceivableDocument::Chunk(_) => None, // Chunks don't have E3id directly - } - } - pub fn to_bytes(&self) -> Result, bincode::Error> { bincode::serialize(self) } @@ -420,18 +414,10 @@ impl EventConverter { pub fn new(bus: &BusHandle) -> Self { Self { bus: bus.clone(), - threshold_share_collector: None, - ids: HashMap::new(), + chunk_collector: None, } } - fn handle_ciphernode_selected(&mut self, event: CiphernodeSelected) -> Result<()> { - let CiphernodeSelected { - e3_id, party_id, .. - } = event; - self.ids.insert(e3_id, party_id); - Ok(()) - } pub fn setup(bus: &BusHandle) -> Addr { let addr = Self::new(bus).start(); bus.subscribe("ThresholdShareCreated", addr.clone().into()); @@ -440,13 +426,13 @@ impl EventConverter { addr } - /// Initialize chunk collector for a specific E3 - fn ensure_collector(&mut self, e3_id: &E3id) { - if self.threshold_share_collector.is_none() { - debug!("Creating ThresholdShare chunk collector for E3: {}", e3_id); - let collector = - ChunkCollector::::setup(e3_id.clone(), self.bus.clone()); - self.threshold_share_collector = Some(collector); + fn ensure_chunk_collector(&mut self, e3_id: &E3id) { + if self.chunk_collector.is_none() { + debug!("Creating chunk collector for E3: {}", e3_id); + self.chunk_collector = Some(ChunkCollector::::setup( + e3_id.clone(), + self.bus.clone(), + )); } } /// Publish a receivable document with party filter @@ -546,59 +532,29 @@ impl EventConverter { Ok(()) } - /// Local node created an encryption key. Send it as a published document - pub fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { - // If this is received from elsewhere + fn handle_encryption_key_created(&self, msg: EncryptionKeyCreated) -> Result<()> { if msg.external { return Ok(()); } - let e3_id = msg.e3_id.clone(); - let receivable = ReceivableDocument::EncryptionKeyCreated(msg); + let receivable = ReceivableDocument::EncryptionKeyCreated(msg.clone()); let value = ArcBytes::from_bytes(&receivable.to_bytes()?); - let meta = DocumentMeta::new(e3_id, DocumentKind::TrBFV, vec![], None); + let meta = DocumentMeta::new(msg.e3_id, DocumentKind::TrBFV, vec![], None); self.bus .publish(PublishDocumentRequested::new(meta, value))?; Ok(()) } - /// Received document externally - /// Check if document passes party filter - fn passes_filter(&self, e3_id: &E3id, meta: &DocumentMeta) -> bool { - match self.ids.get(e3_id) { - Some(party_id) if !meta.matches(party_id) => { - debug!( - "Filtered out: doesn't match party_id {} for E3 {}", - party_id, e3_id - ); - false - } - _ => true, - } - } - - /// Handle a chunk by forwarding to collector - fn handle_chunk(&mut self, chunk: ChunkedDocument, e3_id: &E3id) { - debug!( - "Received chunk {}/{} ({})", - chunk.chunk_index + 1, - chunk.total_chunks, - chunk.chunk_id - ); - self.ensure_collector(e3_id); - if let Some(ref collector) = self.threshold_share_collector { - collector.do_send(ChunkReceived::::new(chunk)); - } - } + /// Convert received document to internal events. + /// Note: Filtering already happened in DocumentPublisher before DHT fetch. fn handle_document_received(&mut self, msg: DocumentReceived) -> Result<()> { - if !self.passes_filter(&msg.meta.e3_id, &msg.meta) { - return Ok(()); - } - let receivable = ReceivableDocument::from_bytes(&msg.value.extract_bytes())?; match receivable { ReceivableDocument::ThresholdShareCreated(evt) => { - debug!("Received complete ThresholdShareCreated"); + debug!( + "Received ThresholdShareCreated from party {}", + evt.share.party_id + ); self.bus.publish(ThresholdShareCreated { external: true, e3_id: evt.e3_id, @@ -606,7 +562,10 @@ impl EventConverter { })?; } ReceivableDocument::EncryptionKeyCreated(evt) => { - debug!("Received EncryptionKeyCreated"); + debug!( + "Received EncryptionKeyCreated from party {}", + evt.key.party_id + ); self.bus.publish(EncryptionKeyCreated { external: true, e3_id: evt.e3_id, @@ -614,7 +573,15 @@ impl EventConverter { })?; } ReceivableDocument::Chunk(chunk) => { - self.handle_chunk(chunk, &msg.meta.e3_id); + debug!( + "Received chunk {}/{}", + chunk.chunk_index + 1, + chunk.total_chunks + ); + self.ensure_chunk_collector(&msg.meta.e3_id); + if let Some(ref collector) = self.chunk_collector { + collector.do_send(ChunkReceived::::new(chunk)); + } } } Ok(()) @@ -632,7 +599,6 @@ impl Handler for EventConverter { EnclaveEventData::ThresholdShareCreated(data) => ctx.notify(data), EnclaveEventData::EncryptionKeyCreated(data) => ctx.notify(data), EnclaveEventData::DocumentReceived(data) => ctx.notify(data), - EnclaveEventData::CiphernodeSelected(data) => ctx.notify(data), _ => (), } } @@ -658,18 +624,6 @@ impl Handler for EventConverter { } } -impl Handler for EventConverter { - type Result = (); - fn handle(&mut self, msg: CiphernodeSelected, _ctx: &mut Self::Context) -> Self::Result { - match self.handle_ciphernode_selected(msg) { - Ok(_) => (), - Err(e) => { - error!("Error handling CiphernodeSelected: {:?}", e); - } - } - } -} - impl Handler for EventConverter { type Result = (); fn handle(&mut self, msg: DocumentReceived, _ctx: &mut Self::Context) -> Self::Result { From 96acb349faeae1e89eaac2cc86fedd1f218394bf Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 14:27:41 +0500 Subject: [PATCH 07/19] chore: formatting --- crates/net/src/chunking/mod.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/net/src/chunking/mod.rs b/crates/net/src/chunking/mod.rs index 20b176a526..3a09a2692b 100644 --- a/crates/net/src/chunking/mod.rs +++ b/crates/net/src/chunking/mod.rs @@ -9,6 +9,4 @@ mod collector; pub use chunkable::Chunkable; pub use collector::{ChunkCollector, ChunkReceived}; -// Re-export chunk types from events crate for convenience pub use e3_events::{ChunkSetId, ChunkedDocument}; - From c61c80de7fca53f957eb841b78a325a6fe81dae2 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 14:49:53 +0500 Subject: [PATCH 08/19] chore: remove duplicate comments --- crates/keyshare/src/threshold_keyshare.rs | 2 -- crates/net/src/document_publisher.rs | 21 ++++++++------------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 85d5b1343c..7527be5249 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -580,7 +580,6 @@ impl ThresholdKeyshare { Ok(event) } - /// 3a. GenPkShareAndSkSss /// 3a. GenPkShareAndSkSss result pub fn handle_gen_pk_share_and_sk_sss_response(&mut self, res: ComputeResponse) -> Result<()> { let ComputeResponse::TrBFV(TrBFVResponse::GenPkShareAndSkSss(output)) = res else { @@ -917,7 +916,6 @@ impl ThresholdKeyshare { } } -// Will only receive events that are for this specific e3_id // Will only receive events that are for this specific e3_id impl Handler for ThresholdKeyshare { type Result = (); diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index 647420d392..a994e344d6 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -390,7 +390,7 @@ async fn broadcast_document_published_notification( /// Note: Party filtering is done by DocumentPublisher BEFORE fetching from DHT. pub struct EventConverter { bus: BusHandle, - chunk_collector: Option>>, + chunk_collectors: HashMap>>, } #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -414,7 +414,7 @@ impl EventConverter { pub fn new(bus: &BusHandle) -> Self { Self { bus: bus.clone(), - chunk_collector: None, + chunk_collectors: HashMap::new(), } } @@ -426,14 +426,11 @@ impl EventConverter { addr } - fn ensure_chunk_collector(&mut self, e3_id: &E3id) { - if self.chunk_collector.is_none() { + fn ensure_chunk_collector(&mut self, e3_id: &E3id) -> &Addr> { + self.chunk_collectors.entry(e3_id.clone()).or_insert_with(|| { debug!("Creating chunk collector for E3: {}", e3_id); - self.chunk_collector = Some(ChunkCollector::::setup( - e3_id.clone(), - self.bus.clone(), - )); - } + ChunkCollector::::setup(e3_id.clone(), self.bus.clone()) + }) } /// Publish a receivable document with party filter fn publish_filtered( @@ -578,10 +575,8 @@ impl EventConverter { chunk.chunk_index + 1, chunk.total_chunks ); - self.ensure_chunk_collector(&msg.meta.e3_id); - if let Some(ref collector) = self.chunk_collector { - collector.do_send(ChunkReceived::::new(chunk)); - } + let collector = self.ensure_chunk_collector(&msg.meta.e3_id); + collector.do_send(ChunkReceived::::new(chunk)); } } Ok(()) From 7a6506c9f101f03388b8d49b35df3954b61cc97b Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 14:56:51 +0500 Subject: [PATCH 09/19] chore: remove duplicate comments --- crates/net/src/document_publisher.rs | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index a994e344d6..b1c1d664ee 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -426,11 +426,16 @@ impl EventConverter { addr } - fn ensure_chunk_collector(&mut self, e3_id: &E3id) -> &Addr> { - self.chunk_collectors.entry(e3_id.clone()).or_insert_with(|| { - debug!("Creating chunk collector for E3: {}", e3_id); - ChunkCollector::::setup(e3_id.clone(), self.bus.clone()) - }) + fn ensure_chunk_collector( + &mut self, + e3_id: &E3id, + ) -> &Addr> { + self.chunk_collectors + .entry(e3_id.clone()) + .or_insert_with(|| { + debug!("Creating chunk collector for E3: {}", e3_id); + ChunkCollector::::setup(e3_id.clone(), self.bus.clone()) + }) } /// Publish a receivable document with party filter fn publish_filtered( From 47ba25fc2cad9d946e9d1001bb277b19fb7f4cc0 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 15:23:28 +0500 Subject: [PATCH 10/19] chore: add external marking and info comments --- crates/net/src/chunking/chunkable.rs | 5 +++++ crates/net/src/chunking/collector.rs | 9 +++++++-- crates/net/src/document_publisher.rs | 27 ++++++++++++++++++++++++--- 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/crates/net/src/chunking/chunkable.rs b/crates/net/src/chunking/chunkable.rs index 7e98adf6b9..9b59b84ca4 100644 --- a/crates/net/src/chunking/chunkable.rs +++ b/crates/net/src/chunking/chunkable.rs @@ -15,6 +15,11 @@ pub trait Chunkable: Serialize + DeserializeOwned + Clone + Sized { 10 * 1024 * 1024 } + /// Mark this document as received from external source (network). + /// Called after reassembly to prevent re-publication loops. + /// Default implementation does nothing - override for types with external flag. + fn mark_as_external(&mut self) {} + fn into_chunks(&self) -> Result> { let bytes = bincode::serialize(self)?; let max_size = Self::max_chunk_size(); diff --git a/crates/net/src/chunking/collector.rs b/crates/net/src/chunking/collector.rs index 75b65ffcaf..44d3ef4f9b 100644 --- a/crates/net/src/chunking/collector.rs +++ b/crates/net/src/chunking/collector.rs @@ -179,8 +179,12 @@ where fn handle(&mut self, msg: ChunkReceived, _ctx: &mut Self::Context) { match self.handle_chunk_internal(msg.chunk) { - Ok(Some(document)) => { - info!("Document reassembled successfully, publishing to bus"); + Ok(Some(mut document)) => { + info!( + "Document reassembled from chunks, marking as external and publishing to bus" + ); + // Mark as external to prevent re-publication loop + document.mark_as_external(); // Publish the reassembled document to the event bus if let Err(e) = self.bus.publish(document) { error!("Failed to publish reassembled document: {:?}", e); @@ -188,6 +192,7 @@ where } Ok(None) => { // Still waiting for more chunks + debug!("Chunk received, waiting for more chunks"); } Err(e) => { error!("Failed to process chunk: {:?}", e); diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index b1c1d664ee..f9846376a3 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -35,6 +35,10 @@ impl Chunkable for ThresholdShareCreated { fn max_chunk_size() -> usize { 10 * 1024 * 1024 } + + fn mark_as_external(&mut self) { + self.external = true; + } } const KADEMLIA_PUT_TIMEOUT: Duration = Duration::from_secs(30); @@ -423,6 +427,7 @@ impl EventConverter { bus.subscribe("ThresholdShareCreated", addr.clone().into()); bus.subscribe("EncryptionKeyCreated", addr.clone().into()); bus.subscribe("DocumentReceived", addr.clone().into()); + bus.subscribe("E3RequestComplete", addr.clone().into()); addr } @@ -467,8 +472,8 @@ impl EventConverter { match chunks { [single] if single.is_complete_set() => { - debug!( - "Party {} share: single chunk ({} bytes)", + info!( + "Party {} share: single chunk ({} bytes), no chunking needed", party_id, single.data.len() ); @@ -480,7 +485,11 @@ impl EventConverter { )?; } multiple => { - info!("Party {} share: {} chunks", party_id, multiple.len()); + info!( + "Party {} share: using chunking - {} chunks", + party_id, + multiple.len() + ); for (idx, chunk) in multiple.iter().enumerate() { debug!( " Chunk {}/{} (chunk_id: {})", @@ -599,11 +608,23 @@ impl Handler for EventConverter { EnclaveEventData::ThresholdShareCreated(data) => ctx.notify(data), EnclaveEventData::EncryptionKeyCreated(data) => ctx.notify(data), EnclaveEventData::DocumentReceived(data) => ctx.notify(data), + EnclaveEventData::E3RequestComplete(data) => ctx.notify(data), _ => (), } } } +impl Handler for EventConverter { + type Result = (); + fn handle(&mut self, msg: E3RequestComplete, _ctx: &mut Self::Context) -> Self::Result { + debug!( + "EventConverter: E3RequestComplete for {}, cleaning up chunk collector", + msg.e3_id + ); + self.chunk_collectors.remove(&msg.e3_id); + } +} + impl Handler for EventConverter { type Result = (); fn handle(&mut self, msg: ThresholdShareCreated, _ctx: &mut Self::Context) -> Self::Result { From de0d8c82c7ca5ebae3097e2ce399ee6d9c989dd5 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 18 Dec 2025 21:19:16 +0500 Subject: [PATCH 11/19] fix: extract shares --- crates/keyshare/src/threshold_keyshare.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 7527be5249..4b7e7be12e 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -721,16 +721,16 @@ impl ThresholdKeyshare { let degree = params.degree(); // Decrypt our share from each sender using BFV - // Each sender's ThresholdShare contains encrypted shares for all parties - // We extract and decrypt the share meant for us (at index party_id) + // Local share (from self) has all parties' shares, network shares are pre-extracted let sk_sss_collected: Vec = msg .shares .iter() .map(|ts| { + let idx = if ts.sk_sss.len() == 1 { 0 } else { party_id }; let encrypted = ts .sk_sss - .clone_share(party_id) - .ok_or(anyhow!("No sk_sss share for party {}", party_id))?; + .clone_share(idx) + .ok_or(anyhow!("No sk_sss share at index {}", idx))?; encrypted.decrypt(&sk_bfv, ¶ms, degree) }) .collect::>()?; @@ -743,9 +743,10 @@ impl ThresholdKeyshare { ts.esi_sss .iter() .map(|esi_shares| { + let idx = if esi_shares.len() == 1 { 0 } else { party_id }; let encrypted = esi_shares - .clone_share(party_id) - .ok_or(anyhow!("No esi_sss share for party {}", party_id))?; + .clone_share(idx) + .ok_or(anyhow!("No esi_sss share at index {}", idx))?; encrypted.decrypt(&sk_bfv, ¶ms, degree) }) .collect::>>() From c623ce03f1486e2e49d866dd14d0c53c85286cb6 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Fri, 19 Dec 2025 14:46:22 +0500 Subject: [PATCH 12/19] chore: remove option --- crates/keyshare/src/threshold_keyshare.rs | 75 +++++++++-------------- 1 file changed, 29 insertions(+), 46 deletions(-) diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 4b7e7be12e..3a3e9e8a18 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -79,14 +79,14 @@ pub struct CollectingEncryptionKeysData { ciphernode_selected: CiphernodeSelected, } -#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct GeneratingThresholdShareData { pk_share: Option, sk_sss: Option>, esi_sss: Option>>, - sk_bfv: Option, - pk_bfv: Option, - collected_encryption_keys: Option>>, + sk_bfv: SensitiveBytes, + pk_bfv: ArcBytes, + collected_encryption_keys: Vec>, } #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] @@ -453,9 +453,9 @@ impl ThresholdKeyshare { sk_sss: None, pk_share: None, esi_sss: None, - sk_bfv: Some(current.sk_bfv), - pk_bfv: Some(current.pk_bfv), - collected_encryption_keys: Some(msg.keys), + sk_bfv: current.sk_bfv, + pk_bfv: current.pk_bfv, + collected_encryption_keys: msg.keys, }, )) })?; @@ -510,30 +510,25 @@ impl ThresholdKeyshare { info!("try_store_esi_sss"); let current: GeneratingThresholdShareData = s.clone().try_into()?; - let pk_share = current.pk_share; - let sk_sss = current.sk_sss; - let sk_bfv = current.sk_bfv; - let pk_bfv = current.pk_bfv; - let collected_encryption_keys = current.collected_encryption_keys; - let next = match (pk_share, sk_sss, &sk_bfv, &collected_encryption_keys) { + let next = match (current.pk_share, current.sk_sss) { // If the other shares are here then transition to aggregation - (Some(pk_share), Some(sk_sss), Some(sk_bfv_ref), Some(keys)) => { + (Some(pk_share), Some(sk_sss)) => { K::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, - sk_bfv: sk_bfv_ref.clone(), - collected_encryption_keys: keys.clone(), + sk_bfv: current.sk_bfv, + collected_encryption_keys: current.collected_encryption_keys, }) } - // If the other shares are not here yet then dont transition - (None, None, _, _) => K::GeneratingThresholdShare(GeneratingThresholdShareData { + // If the other shares are not here yet then don't transition + (None, None) => K::GeneratingThresholdShare(GeneratingThresholdShareData { esi_sss: Some(esi_sss), pk_share: None, sk_sss: None, - sk_bfv, - pk_bfv, - collected_encryption_keys, + sk_bfv: current.sk_bfv, + pk_bfv: current.pk_bfv, + collected_encryption_keys: current.collected_encryption_keys, }), _ => bail!("Inconsistent state!"), }; @@ -591,38 +586,26 @@ impl ThresholdKeyshare { self.state.try_mutate(|s| { info!("try_store_pk_share_and_sk_sss"); let current: GeneratingThresholdShareData = s.clone().try_into()?; - let esi_sss = current.esi_sss; - let sk_bfv = current.sk_bfv; - let pk_bfv = current.pk_bfv; - let collected_encryption_keys = current.collected_encryption_keys; - let next = match (esi_sss, sk_bfv, &collected_encryption_keys) { - // If the esi shares and BFV key are here then transition to aggregation - (Some(esi_sss), Some(sk_bfv), Some(keys)) => { + let next = match current.esi_sss { + // If the esi shares are here then transition to aggregation + Some(esi_sss) => { KeyshareState::AggregatingDecryptionKey(AggregatingDecryptionKey { esi_sss, pk_share, sk_sss, - sk_bfv, - collected_encryption_keys: keys.clone(), + sk_bfv: current.sk_bfv, + collected_encryption_keys: current.collected_encryption_keys, }) } // If esi shares are not here yet then don't transition - (None, sk_bfv, _) => { - KeyshareState::GeneratingThresholdShare(GeneratingThresholdShareData { - esi_sss: None, - pk_share: Some(pk_share), - sk_sss: Some(sk_sss), - sk_bfv, - pk_bfv, - collected_encryption_keys, - }) - } - // If we have esi_sss but no sk_bfv, that's an error - (Some(_), None, _) => bail!("Have esi_sss but no sk_bfv - inconsistent state!"), - // If we have shares ready but no encryption keys, that's an error - (Some(_), Some(_), None) => { - bail!("Have shares but no collected encryption keys - inconsistent state!") - } + None => KeyshareState::GeneratingThresholdShare(GeneratingThresholdShareData { + esi_sss: None, + pk_share: Some(pk_share), + sk_sss: Some(sk_sss), + sk_bfv: current.sk_bfv, + pk_bfv: current.pk_bfv, + collected_encryption_keys: current.collected_encryption_keys, + }), }; s.new_state(next) })?; From 30b80107ebfcae14652325cd93cf1c552d757c9f Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Fri, 19 Dec 2025 19:38:12 +0500 Subject: [PATCH 13/19] chore: revert chunking --- crates/events/src/chunk.rs | 115 --------------- crates/events/src/lib.rs | 2 - crates/net/src/chunking/chunkable.rs | 190 ------------------------ crates/net/src/chunking/collector.rs | 206 --------------------------- crates/net/src/chunking/mod.rs | 12 -- crates/net/src/document_publisher.rs | 115 ++------------- crates/net/src/lib.rs | 1 - 7 files changed, 10 insertions(+), 631 deletions(-) delete mode 100644 crates/events/src/chunk.rs delete mode 100644 crates/net/src/chunking/chunkable.rs delete mode 100644 crates/net/src/chunking/collector.rs delete mode 100644 crates/net/src/chunking/mod.rs diff --git a/crates/events/src/chunk.rs b/crates/events/src/chunk.rs deleted file mode 100644 index 4576416e36..0000000000 --- a/crates/events/src/chunk.rs +++ /dev/null @@ -1,115 +0,0 @@ -// SPDX-License-Identifier: LGPL-3.0-only -// -// This file is provided WITHOUT ANY WARRANTY; -// without even the implied warranty of MERCHANTABILITY -// or FITNESS FOR A PARTICULAR PURPOSE. - -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; - -/// Unique identifier for a set of chunks, derived from content hash -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct ChunkSetId([u8; 32]); - -impl ChunkSetId { - /// Create a ChunkSetId from the original document content - pub fn from_content(content: &[u8]) -> Self { - let mut hasher = Sha256::new(); - hasher.update(content); - let result = hasher.finalize(); - Self(result.into()) - } - - pub fn from_bytes(bytes: [u8; 32]) -> Self { - Self(bytes) - } - - pub fn as_bytes(&self) -> &[u8; 32] { - &self.0 - } -} - -impl std::fmt::Display for ChunkSetId { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "chunk:{}", hex::encode(&self.0[..4])) - } -} - -/// A single chunk of a larger document -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct ChunkedDocument { - /// Identifier for the set of chunks this belongs to - pub chunk_id: ChunkSetId, - /// Index of this chunk (0-based) - pub chunk_index: u32, - /// Total number of chunks in the set - pub total_chunks: u32, - /// The actual data - pub data: Vec, -} - -impl ChunkedDocument { - /// Create a new chunk - pub fn new(chunk_id: ChunkSetId, chunk_index: u32, total_chunks: u32, data: Vec) -> Self { - Self { - chunk_id, - chunk_index, - total_chunks, - data, - } - } - - /// Create a single-chunk document (no splitting needed) - pub fn single(data: Vec) -> Self { - let chunk_id = ChunkSetId::from_content(&data); - Self { - chunk_id, - chunk_index: 0, - total_chunks: 1, - data, - } - } - - /// Check if this represents a complete (non-chunked) document - pub fn is_complete_set(&self) -> bool { - self.total_chunks == 1 - } - - /// Get the size of this chunk in bytes - pub fn size(&self) -> usize { - self.data.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_chunk_creation() { - let chunk = ChunkedDocument::single(vec![1, 2, 3, 4]); - assert_eq!(chunk.chunk_index, 0); - assert_eq!(chunk.total_chunks, 1); - assert!(chunk.is_complete_set()); - assert_eq!(chunk.size(), 4); - } - - #[test] - fn test_chunk_set_id() { - let content1 = b"test content 1"; - let content2 = b"test content 2"; - - let id1 = ChunkSetId::from_content(content1); - let id2 = ChunkSetId::from_content(content2); - assert_ne!(id1, id2); - - // Same content should produce same ID - let id3 = ChunkSetId::from_content(content1); - assert_eq!(id1, id3); - - // Test from_bytes round-trip - let bytes = *id1.as_bytes(); - let id4 = ChunkSetId::from_bytes(bytes); - assert_eq!(id1, id4); - } -} diff --git a/crates/events/src/lib.rs b/crates/events/src/lib.rs index 435575cde9..eac423f6e7 100644 --- a/crates/events/src/lib.rs +++ b/crates/events/src/lib.rs @@ -5,7 +5,6 @@ // or FITNESS FOR A PARTICULAR PURPOSE. mod bus_handle; -mod chunk; mod correlation_id; mod e3id; mod enclave_event; @@ -20,7 +19,6 @@ mod sequencer; mod traits; pub use bus_handle::*; -pub use chunk::*; pub use correlation_id::*; pub use e3id::*; pub use enclave_event::*; diff --git a/crates/net/src/chunking/chunkable.rs b/crates/net/src/chunking/chunkable.rs deleted file mode 100644 index 9b59b84ca4..0000000000 --- a/crates/net/src/chunking/chunkable.rs +++ /dev/null @@ -1,190 +0,0 @@ -// SPDX-License-Identifier: LGPL-3.0-only -// -// This file is provided WITHOUT ANY WARRANTY; -// without even the implied warranty of MERCHANTABILITY -// or FITNESS FOR A PARTICULAR PURPOSE. - -use anyhow::{bail, Result}; -use e3_events::{ChunkSetId, ChunkedDocument}; -use serde::{de::DeserializeOwned, Serialize}; -use tracing::{debug, info}; - -/// Trait for documents that can be split into chunks for transmission -pub trait Chunkable: Serialize + DeserializeOwned + Clone + Sized { - fn max_chunk_size() -> usize { - 10 * 1024 * 1024 - } - - /// Mark this document as received from external source (network). - /// Called after reassembly to prevent re-publication loops. - /// Default implementation does nothing - override for types with external flag. - fn mark_as_external(&mut self) {} - - fn into_chunks(&self) -> Result> { - let bytes = bincode::serialize(self)?; - let max_size = Self::max_chunk_size(); - - debug!( - "Chunking document: {} bytes, max chunk size: {} bytes", - bytes.len(), - max_size - ); - - if bytes.len() <= max_size { - // Small enough, send as single chunk - debug!("Document fits in single chunk"); - return Ok(vec![ChunkedDocument::single(bytes)]); - } - - // Split into multiple chunks - let num_chunks = (bytes.len() + max_size - 1) / max_size; - let chunk_id = ChunkSetId::from_content(&bytes); - - info!( - "Splitting document into {} chunks (chunk_id: {})", - num_chunks, chunk_id - ); - - Ok(bytes - .chunks(max_size) - .enumerate() - .map(|(idx, chunk)| { - ChunkedDocument::new( - chunk_id.clone(), - idx as u32, - num_chunks as u32, - chunk.to_vec(), - ) - }) - .collect()) - } - - /// Reassemble from chunks - fn from_chunks(chunks: Vec) -> Result { - if chunks.is_empty() { - bail!("Cannot reassemble from zero chunks"); - } - - // If single chunk, just deserialize - if chunks.len() == 1 && chunks[0].is_complete_set() { - return Ok(bincode::deserialize(&chunks[0].data)?); - } - - // Validate all chunks are from same set - let chunk_id = &chunks[0].chunk_id; - let total = chunks[0].total_chunks; - - if chunks.len() != total as usize { - bail!("Missing chunks: got {} expected {}", chunks.len(), total); - } - - if !chunks.iter().all(|c| c.chunk_id == *chunk_id) { - bail!("Chunks from different sets"); - } - - // Check for duplicate indices - let mut indices: Vec = chunks.iter().map(|c| c.chunk_index).collect(); - indices.sort_unstable(); - indices.dedup(); - if indices.len() != chunks.len() { - bail!("Duplicate chunk indices detected"); - } - - // Sort by index and concatenate - let mut sorted = chunks; - sorted.sort_by_key(|c| c.chunk_index); - - let bytes: Vec = sorted.into_iter().flat_map(|c| c.data).collect(); - - info!( - "Reassembling document from {} chunks ({} bytes total)", - total, - bytes.len() - ); - - Ok(bincode::deserialize(&bytes)?) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde::{Deserialize, Serialize}; - - #[derive(Clone, Serialize, Deserialize, PartialEq, Debug)] - struct TestDoc { - data: Vec, - } - - impl Chunkable for TestDoc { - fn max_chunk_size() -> usize { - 100 - } - } - - #[test] - fn test_small_document_single_chunk() { - let doc = TestDoc { - data: vec![1, 2, 3, 4, 5], - }; - let chunks = doc.into_chunks().unwrap(); - assert_eq!(chunks.len(), 1); - assert!(chunks[0].is_complete_set()); - - let restored = TestDoc::from_chunks(chunks).unwrap(); - assert_eq!(doc, restored); - } - - #[test] - fn test_large_document_multiple_chunks() { - let doc = TestDoc { - data: vec![42; 500], - }; - let chunks = doc.into_chunks().unwrap(); - assert!(chunks.len() > 1); - - // Verify all chunks have same chunk_id - let first_id = &chunks[0].chunk_id; - assert!(chunks.iter().all(|c| c.chunk_id == *first_id)); - - // Verify chunk indices - for (idx, chunk) in chunks.iter().enumerate() { - assert_eq!(chunk.chunk_index, idx as u32); - assert_eq!(chunk.total_chunks, chunks.len() as u32); - } - - let restored = TestDoc::from_chunks(chunks).unwrap(); - assert_eq!(doc, restored); - } - - #[test] - fn test_missing_chunk_fails() { - let doc = TestDoc { - data: vec![42; 500], - }; - let mut chunks = doc.into_chunks().unwrap(); - chunks.pop(); // Remove last chunk - - let result = TestDoc::from_chunks(chunks); - assert!(result.is_err()); - } - - #[test] - fn test_mixed_chunk_sets_fails() { - let doc1 = TestDoc { - data: vec![42; 500], - }; - let doc2 = TestDoc { - data: vec![99; 500], - }; - - let mut chunks1 = doc1.into_chunks().unwrap(); - let chunks2 = doc2.into_chunks().unwrap(); - - // Mix chunks from different sets - chunks1[0] = chunks2[0].clone(); - - let result = TestDoc::from_chunks(chunks1); - assert!(result.is_err()); - } -} diff --git a/crates/net/src/chunking/collector.rs b/crates/net/src/chunking/collector.rs deleted file mode 100644 index 44d3ef4f9b..0000000000 --- a/crates/net/src/chunking/collector.rs +++ /dev/null @@ -1,206 +0,0 @@ -// SPDX-License-Identifier: LGPL-3.0-only -// -// This file is provided WITHOUT ANY WARRANTY; -// without even the implied warranty of MERCHANTABILITY -// or FITNESS FOR A PARTICULAR PURPOSE. - -use super::chunkable::Chunkable; -use actix::prelude::*; -use anyhow::Result; -use e3_events::{BusHandle, ChunkSetId, ChunkedDocument, E3id, EventPublisher}; -use std::collections::HashMap; -use std::marker::PhantomData; -use std::time::{Duration, Instant}; -use tracing::{debug, error, info, warn}; - -/// Collects chunks and reassembles them into complete documents -pub struct ChunkCollector { - /// E3 ID this collector is for - e3_id: E3id, - /// Chunks we've received so far, grouped by chunk_id - received_chunks: HashMap>, - /// Expected total chunks per set - expected_totals: HashMap, - /// When we started waiting for each chunk set (for timeout) - start_times: HashMap, - /// Event bus to publish completed documents - bus: BusHandle, - /// Timeout duration - timeout: Duration, - /// Phantom data for type parameter - _phantom: PhantomData, -} - -impl ChunkCollector { - pub fn new(e3_id: E3id, bus: BusHandle, timeout: Duration) -> Self { - Self { - e3_id, - received_chunks: HashMap::new(), - expected_totals: HashMap::new(), - start_times: HashMap::new(), - bus, - timeout, - _phantom: PhantomData, - } - } - - pub fn setup(e3_id: E3id, bus: BusHandle) -> Addr { - Self::new(e3_id, bus, Duration::from_secs(60)).start() - } - - /// Process a received chunk - fn handle_chunk_internal(&mut self, chunk: ChunkedDocument) -> Result> { - let chunk_id = chunk.chunk_id.clone(); - let total = chunk.total_chunks; - - debug!( - "Received chunk {}/{} for chunk_id: {}", - chunk.chunk_index + 1, - total, - chunk_id - ); - - // Track expected total and start time - self.expected_totals - .entry(chunk_id.clone()) - .or_insert(total); - self.start_times - .entry(chunk_id.clone()) - .or_insert_with(Instant::now); - - // Add to received chunks - let chunks = self - .received_chunks - .entry(chunk_id.clone()) - .or_insert_with(Vec::new); - - // Check if we already have this chunk index - if chunks.iter().any(|c| c.chunk_index == chunk.chunk_index) { - debug!("Duplicate chunk {} ignored", chunk.chunk_index); - return Ok(None); - } - - chunks.push(chunk); - - // Check if complete - if chunks.len() == total as usize { - info!( - "All {} chunks received for chunk_id: {}, reassembling...", - total, chunk_id - ); - let chunks = self.received_chunks.remove(&chunk_id).unwrap(); - self.expected_totals.remove(&chunk_id); - self.start_times.remove(&chunk_id); - - let document = T::from_chunks(chunks)?; - return Ok(Some(document)); - } - - debug!( - "Waiting for {}/{} chunks for chunk_id: {}", - chunks.len(), - total, - chunk_id - ); - - Ok(None) - } - - /// Check for timeouts and clean up stale chunk sets - fn check_timeouts(&mut self) { - let now = Instant::now(); - let timed_out: Vec<_> = self - .start_times - .iter() - .filter(|(_, start)| now.duration_since(**start) > self.timeout) - .map(|(id, _)| id.clone()) - .collect(); - - for chunk_id in timed_out { - let received = self - .received_chunks - .get(&chunk_id) - .map(|c| c.len()) - .unwrap_or(0); - let expected = self.expected_totals.get(&chunk_id).copied().unwrap_or(0); - - warn!( - "Chunk set {} timed out (received {}/{} chunks)", - chunk_id, received, expected - ); - - self.received_chunks.remove(&chunk_id); - self.expected_totals.remove(&chunk_id); - self.start_times.remove(&chunk_id); - } - } -} - -impl Actor for ChunkCollector { - type Context = Context; - - fn started(&mut self, ctx: &mut Self::Context) { - info!("ChunkCollector started for E3: {}", self.e3_id); - - // Periodic timeout check every 5 seconds - ctx.run_interval(Duration::from_secs(5), |act, _ctx| { - act.check_timeouts(); - }); - } - - fn stopped(&mut self, _ctx: &mut Self::Context) { - info!("ChunkCollector stopped for E3: {}", self.e3_id); - } -} - -/// Message to send a chunk to the collector -#[derive(Message, Clone, Debug)] -#[rtype(result = "()")] -pub struct ChunkReceived { - pub chunk: ChunkedDocument, - pub _phantom: PhantomData, -} - -impl ChunkReceived { - pub fn new(chunk: ChunkedDocument) -> Self { - Self { - chunk, - _phantom: PhantomData, - } - } -} - -impl Handler> for ChunkCollector -where - T: actix::Message + Send + Into, - T::Result: Send, -{ - type Result = (); - - fn handle(&mut self, msg: ChunkReceived, _ctx: &mut Self::Context) { - match self.handle_chunk_internal(msg.chunk) { - Ok(Some(mut document)) => { - info!( - "Document reassembled from chunks, marking as external and publishing to bus" - ); - // Mark as external to prevent re-publication loop - document.mark_as_external(); - // Publish the reassembled document to the event bus - if let Err(e) = self.bus.publish(document) { - error!("Failed to publish reassembled document: {:?}", e); - } - } - Ok(None) => { - // Still waiting for more chunks - debug!("Chunk received, waiting for more chunks"); - } - Err(e) => { - error!("Failed to process chunk: {:?}", e); - } - } - } -} - -// Note: ChunkCollector tests are covered by integration tests with real event types -// that implement Into. Unit tests with mock types would require -// duplicating the handler implementation without the publish constraint. diff --git a/crates/net/src/chunking/mod.rs b/crates/net/src/chunking/mod.rs deleted file mode 100644 index 3a09a2692b..0000000000 --- a/crates/net/src/chunking/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -// SPDX-License-Identifier: LGPL-3.0-only -// -// This file is provided WITHOUT ANY WARRANTY; -// without even the implied warranty of MERCHANTABILITY -// or FITNESS FOR A PARTICULAR PURPOSE. - -mod chunkable; -mod collector; - -pub use chunkable::Chunkable; -pub use collector::{ChunkCollector, ChunkReceived}; -pub use e3_events::{ChunkSetId, ChunkedDocument}; diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index f9846376a3..2f54e03a4d 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -5,7 +5,6 @@ // or FITNESS FOR A PARTICULAR PURPOSE. use crate::{ - chunking::{ChunkCollector, ChunkReceived, Chunkable, ChunkedDocument}, events::{ call_and_await_response, DocumentPublishedNotification, GossipData, NetCommand, NetEvent, }, @@ -31,16 +30,6 @@ use std::{ use tokio::sync::{broadcast, mpsc}; use tracing::{debug, error, info}; -impl Chunkable for ThresholdShareCreated { - fn max_chunk_size() -> usize { - 10 * 1024 * 1024 - } - - fn mark_as_external(&mut self) { - self.external = true; - } -} - const KADEMLIA_PUT_TIMEOUT: Duration = Duration::from_secs(30); const KADEMLIA_GET_TIMEOUT: Duration = Duration::from_secs(30); const KADEMLIA_BROADCAST_TIMEOUT: Duration = Duration::from_secs(30); @@ -389,19 +378,16 @@ async fn broadcast_document_published_notification( /// Responsibilities: /// - Outgoing: Converts ThresholdShareCreated → party-filtered PublishDocumentRequested /// - Incoming: Converts DocumentReceived → ThresholdShareCreated/EncryptionKeyCreated -/// - Reassembles chunked documents via ChunkCollector /// /// Note: Party filtering is done by DocumentPublisher BEFORE fetching from DHT. pub struct EventConverter { bus: BusHandle, - chunk_collectors: HashMap>>, } #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] enum ReceivableDocument { ThresholdShareCreated(ThresholdShareCreated), EncryptionKeyCreated(EncryptionKeyCreated), - Chunk(ChunkedDocument), } impl ReceivableDocument { @@ -416,10 +402,7 @@ impl ReceivableDocument { impl EventConverter { pub fn new(bus: &BusHandle) -> Self { - Self { - bus: bus.clone(), - chunk_collectors: HashMap::new(), - } + Self { bus: bus.clone() } } pub fn setup(bus: &BusHandle) -> Addr { @@ -427,21 +410,9 @@ impl EventConverter { bus.subscribe("ThresholdShareCreated", addr.clone().into()); bus.subscribe("EncryptionKeyCreated", addr.clone().into()); bus.subscribe("DocumentReceived", addr.clone().into()); - bus.subscribe("E3RequestComplete", addr.clone().into()); addr } - fn ensure_chunk_collector( - &mut self, - e3_id: &E3id, - ) -> &Addr> { - self.chunk_collectors - .entry(e3_id.clone()) - .or_insert_with(|| { - debug!("Creating chunk collector for E3: {}", e3_id); - ChunkCollector::::setup(e3_id.clone(), self.bus.clone()) - }) - } /// Publish a receivable document with party filter fn publish_filtered( &self, @@ -461,53 +432,6 @@ impl EventConverter { Ok(()) } - /// Publish chunks for a specific party - fn publish_party_chunks( - &self, - chunks: &[ChunkedDocument], - e3_id: &E3id, - party_id: usize, - ) -> Result<()> { - let party_id_u64 = party_id as u64; - - match chunks { - [single] if single.is_complete_set() => { - info!( - "Party {} share: single chunk ({} bytes), no chunking needed", - party_id, - single.data.len() - ); - let share = bincode::deserialize(&single.data)?; - self.publish_filtered( - ReceivableDocument::ThresholdShareCreated(share), - e3_id, - party_id_u64, - )?; - } - multiple => { - info!( - "Party {} share: using chunking - {} chunks", - party_id, - multiple.len() - ); - for (idx, chunk) in multiple.iter().enumerate() { - debug!( - " Chunk {}/{} (chunk_id: {})", - idx + 1, - multiple.len(), - chunk.chunk_id - ); - self.publish_filtered( - ReceivableDocument::Chunk(chunk.clone()), - e3_id, - party_id_u64, - )?; - } - } - } - Ok(()) - } - /// Local node created a threshold share. Send it as a published document. /// Uses party-filtered distribution: each party only receives their specific shares. pub fn handle_threshold_share_created(&self, msg: ThresholdShareCreated) -> Result<()> { @@ -517,14 +441,13 @@ impl EventConverter { let e3_id = msg.e3_id.clone(); let num_parties = msg.share.num_parties(); - let total_size_mb = bincode::serialize(&msg)?.len() as f64 / (1024.0 * 1024.0); info!( - "Publishing ThresholdShare for E3 {} ({:.2} MB total, {} parties)", - e3_id, total_size_mb, num_parties + "Publishing ThresholdShare for E3 {} ({} parties)", + e3_id, num_parties ); - // Publish party-filtered shares + // Publish party-filtered shares - each party gets only their share for party_id in 0..num_parties { let party_share = msg .share @@ -537,8 +460,11 @@ impl EventConverter { external: false, }; - let chunks = party_msg.into_chunks()?; - self.publish_party_chunks(&chunks, &e3_id, party_id)?; + self.publish_filtered( + ReceivableDocument::ThresholdShareCreated(party_msg), + &e3_id, + party_id as u64, + )?; } Ok(()) @@ -557,7 +483,7 @@ impl EventConverter { /// Convert received document to internal events. /// Note: Filtering already happened in DocumentPublisher before DHT fetch. - fn handle_document_received(&mut self, msg: DocumentReceived) -> Result<()> { + fn handle_document_received(&self, msg: DocumentReceived) -> Result<()> { let receivable = ReceivableDocument::from_bytes(&msg.value.extract_bytes())?; match receivable { @@ -583,15 +509,6 @@ impl EventConverter { key: evt.key, })?; } - ReceivableDocument::Chunk(chunk) => { - debug!( - "Received chunk {}/{}", - chunk.chunk_index + 1, - chunk.total_chunks - ); - let collector = self.ensure_chunk_collector(&msg.meta.e3_id); - collector.do_send(ChunkReceived::::new(chunk)); - } } Ok(()) } @@ -608,23 +525,11 @@ impl Handler for EventConverter { EnclaveEventData::ThresholdShareCreated(data) => ctx.notify(data), EnclaveEventData::EncryptionKeyCreated(data) => ctx.notify(data), EnclaveEventData::DocumentReceived(data) => ctx.notify(data), - EnclaveEventData::E3RequestComplete(data) => ctx.notify(data), _ => (), } } } -impl Handler for EventConverter { - type Result = (); - fn handle(&mut self, msg: E3RequestComplete, _ctx: &mut Self::Context) -> Self::Result { - debug!( - "EventConverter: E3RequestComplete for {}, cleaning up chunk collector", - msg.e3_id - ); - self.chunk_collectors.remove(&msg.e3_id); - } -} - impl Handler for EventConverter { type Result = (); fn handle(&mut self, msg: ThresholdShareCreated, _ctx: &mut Self::Context) -> Self::Result { diff --git a/crates/net/src/lib.rs b/crates/net/src/lib.rs index 4e79d2f150..0b9868cd08 100644 --- a/crates/net/src/lib.rs +++ b/crates/net/src/lib.rs @@ -4,7 +4,6 @@ // without even the implied warranty of MERCHANTABILITY // or FITNESS FOR A PARTICULAR PURPOSE. -pub mod chunking; mod cid; mod dialer; mod document_publisher; From 0b84ea0b3e179bdfefebb82c1efab33ef63d0daa Mon Sep 17 00:00:00 2001 From: Hamza Khalid <36852564+hmzakhalid@users.noreply.github.com> Date: Mon, 22 Dec 2025 18:22:30 +0500 Subject: [PATCH 14/19] Update crates/events/src/enclave_event/threshold_share_created.rs Co-authored-by: Giacomo --- crates/events/src/enclave_event/threshold_share_created.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/events/src/enclave_event/threshold_share_created.rs b/crates/events/src/enclave_event/threshold_share_created.rs index 403f0e8237..3ba2ff956b 100644 --- a/crates/events/src/enclave_event/threshold_share_created.rs +++ b/crates/events/src/enclave_event/threshold_share_created.rs @@ -30,7 +30,7 @@ pub struct ThresholdShare { pub pk_share: ArcBytes, /// BFV-encrypted sk_sss - each recipient can decrypt their share pub sk_sss: BfvEncryptedShares, - /// BFV-encrypted esi_sss - one per ciphertext, each recipient can decrypt their share + /// BFV-encrypted esi_sss - one per secret key (sk), each recipient can decrypt their share pub esi_sss: Vec, } From 999d5ac6d373b4adc739b1634293aebe9e758d0a Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 25 Dec 2025 00:42:30 +0500 Subject: [PATCH 15/19] feat: move spliting to threshold keyshare --- .../enclave_event/threshold_share_created.rs | 1 + crates/keyshare/src/threshold_keyshare.rs | 53 +++++++++++++++---- crates/net/src/document_publisher.rs | 41 +++++--------- crates/tests/tests/integration.rs | 10 ++-- 4 files changed, 59 insertions(+), 46 deletions(-) diff --git a/crates/events/src/enclave_event/threshold_share_created.rs b/crates/events/src/enclave_event/threshold_share_created.rs index 3ba2ff956b..e9794a0397 100644 --- a/crates/events/src/enclave_event/threshold_share_created.rs +++ b/crates/events/src/enclave_event/threshold_share_created.rs @@ -62,6 +62,7 @@ impl ThresholdShare { pub struct ThresholdShareCreated { pub e3_id: E3id, pub share: Arc, + pub target_party_id: u64, pub external: bool, } diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 3a3e9e8a18..0c2b02f4e2 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -375,7 +375,18 @@ impl ThresholdKeyshare { msg: ThresholdShareCreated, self_addr: Addr, ) -> Result<()> { - info!("Received ThresholdShareCreated forwarding to collector!"); + let state = self.state.try_get()?; + let my_party_id = state.party_id; + + // Filter: only process shares intended for this party + if msg.target_party_id != my_party_id { + return Ok(()); + } + + info!( + "Received ThresholdShareCreated from party {} for us (party {}), forwarding to collector!", + msg.share.party_id, my_party_id + ); let collector = self.ensure_collector(self_addr)?; info!("got collector address!"); collector.do_send(msg); @@ -669,16 +680,36 @@ impl ThresholdKeyshare { .map(|esi| BfvEncryptedShares::encrypt_all(esi, &recipient_pks, ¶ms, &mut rng)) .collect::>()?; - self.bus.publish(ThresholdShareCreated { - e3_id, - share: Arc::new(ThresholdShare { - party_id, - pk_share, - sk_sss: encrypted_sk_sss, - esi_sss: encrypted_esi_sss, - }), - external: false, - })?; + // Create the full share with all parties' encrypted data + let full_share = ThresholdShare { + party_id, + pk_share, + sk_sss: encrypted_sk_sss, + esi_sss: encrypted_esi_sss, + }; + + // Domain-level splitting: publish one ThresholdShareCreated per recipient party + // Each party only receives the share data meant for them + let num_parties = full_share.num_parties(); + info!( + "Publishing ThresholdShare for E3 {} to {} parties", + e3_id, num_parties + ); + + for recipient_party_id in 0..num_parties { + let party_share = full_share + .extract_for_party(recipient_party_id) + .ok_or_else(|| { + anyhow!("Failed to extract share for party {}", recipient_party_id) + })?; + + self.bus.publish(ThresholdShareCreated { + e3_id: e3_id.clone(), + share: Arc::new(party_share), + target_party_id: recipient_party_id as u64, + external: false, + })?; + } Ok(()) } diff --git a/crates/net/src/document_publisher.rs b/crates/net/src/document_publisher.rs index 2f54e03a4d..8f46ce7c4b 100644 --- a/crates/net/src/document_publisher.rs +++ b/crates/net/src/document_publisher.rs @@ -432,40 +432,24 @@ impl EventConverter { Ok(()) } - /// Local node created a threshold share. Send it as a published document. - /// Uses party-filtered distribution: each party only receives their specific shares. + /// Local node created a threshold share (already split per-party by ThresholdKeyshare). + /// Publishes the single-party document with appropriate filter. pub fn handle_threshold_share_created(&self, msg: ThresholdShareCreated) -> Result<()> { if msg.external { return Ok(()); } - - let e3_id = msg.e3_id.clone(); - let num_parties = msg.share.num_parties(); + let target_party_id = msg.target_party_id; info!( - "Publishing ThresholdShare for E3 {} ({} parties)", - e3_id, num_parties + "Publishing ThresholdShare from party {} for target party {} (E3 {})", + msg.share.party_id, target_party_id, msg.e3_id ); - // Publish party-filtered shares - each party gets only their share - for party_id in 0..num_parties { - let party_share = msg - .share - .extract_for_party(party_id) - .ok_or_else(|| anyhow::anyhow!("Failed to extract share for party {}", party_id))?; - - let party_msg = ThresholdShareCreated { - e3_id: e3_id.clone(), - share: Arc::new(party_share), - external: false, - }; - - self.publish_filtered( - ReceivableDocument::ThresholdShareCreated(party_msg), - &e3_id, - party_id as u64, - )?; - } + self.publish_filtered( + ReceivableDocument::ThresholdShareCreated(msg.clone()), + &msg.e3_id, + target_party_id, + )?; Ok(()) } @@ -489,13 +473,14 @@ impl EventConverter { match receivable { ReceivableDocument::ThresholdShareCreated(evt) => { debug!( - "Received ThresholdShareCreated from party {}", - evt.share.party_id + "Received ThresholdShareCreated from party {} for target party {}", + evt.share.party_id, evt.target_party_id ); self.bus.publish(ThresholdShareCreated { external: true, e3_id: evt.e3_id, share: evt.share, + target_party_id: evt.target_party_id, })?; } ReceivableDocument::EncryptionKeyCreated(evt) => { diff --git a/crates/tests/tests/integration.rs b/crates/tests/tests/integration.rs index d1d0033994..4e7b924924 100644 --- a/crates/tests/tests/integration.rs +++ b/crates/tests/tests/integration.rs @@ -305,14 +305,10 @@ async fn test_trbfv_actor() -> Result<()> { )); // Then wait for all ThresholdShareCreated events + // With domain-level splitting, each of the 5 parties publishes 5 events (one per target party) + // Total: 5 parties × 5 targets = 25 events let shares_timer = Instant::now(); - let expected = vec![ - "ThresholdShareCreated", - "ThresholdShareCreated", - "ThresholdShareCreated", - "ThresholdShareCreated", - "ThresholdShareCreated", - ]; + let expected: Vec<&str> = (0..25).map(|_| "ThresholdShareCreated").collect(); let _ = nodes .take_history_with_timeout(0, expected.len(), Duration::from_secs(1000)) .await?; From 9c5aea87894793118d071ecc8c7a53877715df7b Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 25 Dec 2025 00:45:22 +0500 Subject: [PATCH 16/19] fix: make event buffer key a composite --- crates/request/src/context.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/request/src/context.rs b/crates/request/src/context.rs index d45d3b44a0..8a9053ad05 100644 --- a/crates/request/src/context.rs +++ b/crates/request/src/context.rs @@ -80,13 +80,15 @@ impl E3Context { pub fn forward_message(&self, msg: &EnclaveEvent, buffer: &mut EventBuffer) { self.recipients().into_iter().for_each(|(key, recipient)| { + // Use composite key of e3_id:recipient_key to scope buffered events per E3 request + let buffer_key = format!("{}:{}", self.e3_id, key); if let Some(act) = recipient { act.do_send(msg.clone()); - for m in buffer.take(&key) { + for m in buffer.take(&buffer_key) { act.do_send(m); } } else { - buffer.add(&key, msg.clone()); + buffer.add(&buffer_key, msg.clone()); } }); } From c3a4405ad34eb319950adb7abe83cd0a6c08a547 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 25 Dec 2025 02:02:11 +0500 Subject: [PATCH 17/19] feat: add bookend events --- .../encryption_key_collection_failed.rs | 24 ++++ crates/events/src/enclave_event/mod.rs | 10 +- .../threshold_share_collection_failed.rs | 24 ++++ .../keyshare/src/encryption_key_collector.rs | 115 +++++++++++++----- crates/keyshare/src/lib.rs | 2 +- crates/keyshare/src/threshold_keyshare.rs | 66 +++++++++- .../keyshare/src/threshold_share_collector.rs | 115 +++++++++++++++--- 7 files changed, 304 insertions(+), 52 deletions(-) create mode 100644 crates/events/src/enclave_event/encryption_key_collection_failed.rs create mode 100644 crates/events/src/enclave_event/threshold_share_collection_failed.rs diff --git a/crates/events/src/enclave_event/encryption_key_collection_failed.rs b/crates/events/src/enclave_event/encryption_key_collection_failed.rs new file mode 100644 index 0000000000..da330e93fc --- /dev/null +++ b/crates/events/src/enclave_event/encryption_key_collection_failed.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use crate::E3id; +use actix::Message; +use serde::{Deserialize, Serialize}; +use std::fmt::{self, Display}; + +#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[rtype(result = "()")] +pub struct EncryptionKeyCollectionFailed { + pub e3_id: E3id, + pub reason: String, + pub missing_parties: Vec, +} + +impl Display for EncryptionKeyCollectionFailed { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} diff --git a/crates/events/src/enclave_event/mod.rs b/crates/events/src/enclave_event/mod.rs index 9965e80f26..c29cc9a50e 100644 --- a/crates/events/src/enclave_event/mod.rs +++ b/crates/events/src/enclave_event/mod.rs @@ -19,6 +19,7 @@ mod die; mod e3_request_complete; mod e3_requested; mod enclave_error; +mod encryption_key_collection_failed; mod encryption_key_created; mod keyshare_created; mod operator_activation_changed; @@ -28,6 +29,7 @@ mod publickey_aggregated; mod publish_document; mod shutdown; mod test_event; +mod threshold_share_collection_failed; mod threshold_share_created; mod ticket_balance_updated; mod ticket_generated; @@ -48,6 +50,7 @@ pub use die::*; pub use e3_request_complete::*; pub use e3_requested::*; pub use enclave_error::*; +pub use encryption_key_collection_failed::*; pub use encryption_key_created::*; pub use keyshare_created::*; pub use operator_activation_changed::*; @@ -58,6 +61,7 @@ pub use publish_document::*; pub use shutdown::*; use strum::IntoStaticStr; pub use test_event::*; +pub use threshold_share_collection_failed::*; pub use threshold_share_created::*; pub use ticket_balance_updated::*; pub use ticket_generated::*; @@ -115,6 +119,8 @@ pub enum EnclaveEventData { DocumentReceived(DocumentReceived), ThresholdShareCreated(ThresholdShareCreated), EncryptionKeyCreated(EncryptionKeyCreated), + EncryptionKeyCollectionFailed(EncryptionKeyCollectionFailed), + ThresholdShareCollectionFailed(ThresholdShareCollectionFailed), /// This is a test event to use in testing TestEvent(TestEvent), } @@ -327,7 +333,9 @@ impl_into_event_data!( TestEvent, DocumentReceived, ThresholdShareCreated, - EncryptionKeyCreated + EncryptionKeyCreated, + EncryptionKeyCollectionFailed, + ThresholdShareCollectionFailed ); impl TryFrom<&EnclaveEvent> for EnclaveError { diff --git a/crates/events/src/enclave_event/threshold_share_collection_failed.rs b/crates/events/src/enclave_event/threshold_share_collection_failed.rs new file mode 100644 index 0000000000..dd81aac693 --- /dev/null +++ b/crates/events/src/enclave_event/threshold_share_collection_failed.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: LGPL-3.0-only +// +// This file is provided WITHOUT ANY WARRANTY; +// without even the implied warranty of MERCHANTABILITY +// or FITNESS FOR A PARTICULAR PURPOSE. + +use crate::E3id; +use actix::Message; +use serde::{Deserialize, Serialize}; +use std::fmt::{self, Display}; + +#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[rtype(result = "()")] +pub struct ThresholdShareCollectionFailed { + pub e3_id: E3id, + pub reason: String, + pub missing_parties: Vec, +} + +impl Display for ThresholdShareCollectionFailed { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} diff --git a/crates/keyshare/src/encryption_key_collector.rs b/crates/keyshare/src/encryption_key_collector.rs index b7950530de..b9bc2db99b 100644 --- a/crates/keyshare/src/encryption_key_collector.rs +++ b/crates/keyshare/src/encryption_key_collector.rs @@ -4,22 +4,18 @@ // without even the implied warranty of MERCHANTABILITY // or FITNESS FOR A PARTICULAR PURPOSE. -//! Collector for BFV encryption keys from all parties. -//! -//! Before parties can encrypt their Shamir shares, they need to collect -//! the BFV public keys from all other parties. This actor handles that -//! collection process. - use std::{ collections::{HashMap, HashSet}, sync::Arc, - time::Instant, + time::{Duration, Instant}, }; -use actix::{Actor, Addr, Handler, Message}; -use e3_events::{EncryptionKey, EncryptionKeyCreated}; +use actix::{Actor, ActorContext, Addr, AsyncContext, Handler, Message, SpawnHandle}; +use e3_events::{E3id, EncryptionKey, EncryptionKeyCollectionFailed, EncryptionKeyCreated}; use e3_trbfv::PartyId; -use tracing::info; +use tracing::{info, warn}; + +const DEFAULT_COLLECTION_TIMEOUT: Duration = Duration::from_secs(60); use crate::ThresholdKeyshare; @@ -29,6 +25,8 @@ pub enum CollectorState { Collecting, /// All keys have been collected Finished, + /// Collection timed out + TimedOut, } /// Message sent when all encryption keys have been collected. @@ -51,51 +49,69 @@ impl From>> for AllEncryptionKeysCollected { } } +/// Message sent when encryption key collection times out. +#[derive(Message, Clone, Debug)] +#[rtype(result = "()")] +pub struct EncryptionKeyCollectionTimeout; + /// Actor that collects BFV encryption keys from all parties. /// /// Once all keys are collected, it sends `AllEncryptionKeysCollected` to the parent -/// `ThresholdKeyshare` actor. +/// `ThresholdKeyshare` actor. If collection times out, it sends `EncryptionKeyCollectionFailed`. pub struct EncryptionKeyCollector { - /// Set of party IDs we're still waiting for + e3_id: E3id, todo: HashSet, - /// Parent actor to notify when collection is complete parent: Addr, - /// Current state state: CollectorState, - /// Collected keys indexed by party_id keys: HashMap>, + timeout_handle: Option, } impl EncryptionKeyCollector { - /// Create and start a new collector. - /// - /// # Arguments - /// * `parent` - The ThresholdKeyshare actor to notify when collection is complete - /// * `total` - Total number of parties (keys to collect) - pub fn setup(parent: Addr, total: u64) -> Addr { - let addr = Self { + pub fn setup(parent: Addr, total: u64, e3_id: E3id) -> Addr { + let collector = Self { + e3_id, todo: (0..total).collect(), parent, state: CollectorState::Collecting, keys: HashMap::new(), - } - .start(); - addr + timeout_handle: None, + }; + collector.start() } } impl Actor for EncryptionKeyCollector { type Context = actix::Context; + + fn started(&mut self, ctx: &mut Self::Context) { + info!( + e3_id = %self.e3_id, + "EncryptionKeyCollector started, scheduling timeout in {:?}", + DEFAULT_COLLECTION_TIMEOUT + ); + + let handle = ctx.notify_later(EncryptionKeyCollectionTimeout, DEFAULT_COLLECTION_TIMEOUT); + self.timeout_handle = Some(handle); + } } impl Handler for EncryptionKeyCollector { type Result = (); - fn handle(&mut self, msg: EncryptionKeyCreated, _: &mut Self::Context) -> Self::Result { + fn handle(&mut self, msg: EncryptionKeyCreated, ctx: &mut Self::Context) -> Self::Result { let start = Instant::now(); info!("EncryptionKeyCollector: EncryptionKeyCreated received"); - if let CollectorState::Finished = self.state { - info!("EncryptionKeyCollector is finished, ignoring"); + // Ignore if already finished or timed out + if !matches!(self.state, CollectorState::Collecting) { + info!( + "EncryptionKeyCollector is not collecting (state: {:?}), ignoring", + match self.state { + CollectorState::Collecting => "Collecting", + CollectorState::Finished => "Finished", + CollectorState::TimedOut => "TimedOut", + } + ); return; } @@ -119,6 +135,12 @@ impl Handler for EncryptionKeyCollector { if self.todo.is_empty() { info!("All encryption keys collected!"); self.state = CollectorState::Finished; + + // Cancel the timeout since we're done + if let Some(handle) = self.timeout_handle.take() { + ctx.cancel_future(handle); + } + let event: AllEncryptionKeysCollected = self.keys.clone().into(); self.parent.do_send(event); } @@ -129,3 +151,40 @@ impl Handler for EncryptionKeyCollector { ); } } + +impl Handler for EncryptionKeyCollector { + type Result = (); + fn handle( + &mut self, + _: EncryptionKeyCollectionTimeout, + ctx: &mut Self::Context, + ) -> Self::Result { + // Only handle timeout if we're still collecting + if !matches!(self.state, CollectorState::Collecting) { + return; + } + + warn!( + e3_id = %self.e3_id, + missing_parties = ?self.todo, + "Encryption key collection timed out, {} parties missing", + self.todo.len() + ); + + self.state = CollectorState::TimedOut; + + // Notify parent of failure + let missing_parties: Vec = self.todo.iter().copied().collect(); + self.parent.do_send(EncryptionKeyCollectionFailed { + e3_id: self.e3_id.clone(), + reason: format!( + "Timeout waiting for encryption keys from {} parties", + missing_parties.len() + ), + missing_parties, + }); + + // Stop the actor + ctx.stop(); + } +} diff --git a/crates/keyshare/src/lib.rs b/crates/keyshare/src/lib.rs index ecf984cc3c..98f52ab4a5 100644 --- a/crates/keyshare/src/lib.rs +++ b/crates/keyshare/src/lib.rs @@ -10,7 +10,7 @@ mod keyshare; mod repo; mod threshold_keyshare; mod threshold_share_collector; -pub use encryption_key_collector::*; +pub use encryption_key_collector::{AllEncryptionKeysCollected, EncryptionKeyCollector}; pub use keyshare::*; pub use repo::*; pub use threshold_keyshare::*; diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 0c2b02f4e2..6ef70b56e3 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -11,7 +11,8 @@ use e3_data::Persistable; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CiphertextOutputPublished, ComputeRequest, ComputeResponse, DecryptionshareCreated, E3id, EnclaveEvent, EnclaveEventData, EncryptionKey, - EncryptionKeyCreated, KeyshareCreated, PartyId, ThresholdShare, ThresholdShareCreated, + EncryptionKeyCollectionFailed, EncryptionKeyCreated, KeyshareCreated, PartyId, ThresholdShare, + ThresholdShareCollectionFailed, ThresholdShareCreated, }; use e3_fhe::create_crp; use e3_multithread::Multithread; @@ -37,7 +38,7 @@ use std::{ mem, sync::{Arc, Mutex}, }; -use tracing::{error, info}; +use tracing::{error, info, warn}; use crate::encryption_key_collector::{AllEncryptionKeysCollected, EncryptionKeyCollector}; use crate::threshold_share_collector::ThresholdShareCollector; @@ -346,9 +347,11 @@ impl ThresholdKeyshare { "Setting up key collector for addr: {} and {} nodes", state.address, state.threshold_n ); + let e3_id = state.e3_id.clone(); + let threshold_n = state.threshold_n; let addr = self .decryption_key_collector - .get_or_insert_with(|| ThresholdShareCollector::setup(self_addr, state.threshold_n)); + .get_or_insert_with(|| ThresholdShareCollector::setup(self_addr, threshold_n, e3_id)); Ok(addr.clone()) } @@ -364,9 +367,11 @@ impl ThresholdKeyshare { "Setting up encryption key collector for addr: {} and {} nodes", state.address, state.threshold_n ); + let e3_id = state.e3_id.clone(); + let threshold_n = state.threshold_n; let addr = self .encryption_key_collector - .get_or_insert_with(|| EncryptionKeyCollector::setup(self_addr, state.threshold_n)); + .get_or_insert_with(|| EncryptionKeyCollector::setup(self_addr, threshold_n, e3_id)); Ok(addr.clone()) } @@ -1008,3 +1013,56 @@ impl Handler for ThresholdKeyshare { ) } } + +impl Handler for ThresholdKeyshare { + type Result = (); + fn handle( + &mut self, + msg: EncryptionKeyCollectionFailed, + ctx: &mut Self::Context, + ) -> Self::Result { + warn!( + e3_id = %msg.e3_id, + missing_parties = ?msg.missing_parties, + "Encryption key collection failed: {}", + msg.reason + ); + + // Clear the collector reference since it's stopped + self.encryption_key_collector = None; + + // Publish failure event to event bus for sync tracking + if let Err(e) = self.bus.publish(msg) { + error!("Failed to publish EncryptionKeyCollectionFailed: {}", e); + } + + // Stop this actor since we can't proceed without all encryption keys + ctx.stop(); + } +} + +impl Handler for ThresholdKeyshare { + type Result = (); + fn handle( + &mut self, + msg: ThresholdShareCollectionFailed, + ctx: &mut Self::Context, + ) -> Self::Result { + warn!( + e3_id = %msg.e3_id, + missing_parties = ?msg.missing_parties, + "Threshold share collection failed: {}", + msg.reason + ); + + // Clear the collector reference since it's stopped + self.decryption_key_collector = None; + + // Publish failure event to event bus for sync tracking + if let Err(e) = self.bus.publish(msg) { + error!("Failed to publish ThresholdShareCollectionFailed: {}", e); + } + + ctx.stop(); + } +} diff --git a/crates/keyshare/src/threshold_share_collector.rs b/crates/keyshare/src/threshold_share_collector.rs index bd71a1f027..d60834bd4b 100644 --- a/crates/keyshare/src/threshold_share_collector.rs +++ b/crates/keyshare/src/threshold_share_collector.rs @@ -7,71 +7,113 @@ use std::{ collections::{HashMap, HashSet}, sync::Arc, - time::Instant, + time::{Duration, Instant}, }; -use actix::{Actor, Addr, Handler}; -use e3_events::{ThresholdShare, ThresholdShareCreated}; +use actix::{Actor, ActorContext, Addr, AsyncContext, Handler, Message, SpawnHandle}; +use e3_events::{E3id, ThresholdShare, ThresholdShareCollectionFailed, ThresholdShareCreated}; use e3_trbfv::PartyId; -use tracing::info; +use tracing::{info, warn}; use crate::{AllThresholdSharesCollected, ThresholdKeyshare}; -pub enum CollectorState { +const DEFAULT_COLLECTION_TIMEOUT: Duration = Duration::from_secs(120); + +pub(crate) enum CollectorState { Collecting, Finished, + TimedOut, } +/// Message sent when threshold share collection times out. +#[derive(Message, Clone, Debug)] +#[rtype(result = "()")] +pub struct ThresholdShareCollectionTimeout; + pub struct ThresholdShareCollector { + e3_id: E3id, todo: HashSet, parent: Addr, state: CollectorState, shares: HashMap>, + timeout_handle: Option, } impl ThresholdShareCollector { - pub fn setup(parent: Addr, total: u64) -> Addr { - let addr = Self { + pub fn setup( + parent: Addr, + total: u64, + e3_id: E3id, + ) -> Addr { + let collector = Self { + e3_id, todo: (0..total).collect(), parent, state: CollectorState::Collecting, shares: HashMap::new(), - } - .start(); - addr + timeout_handle: None, + }; + collector.start() } } impl Actor for ThresholdShareCollector { type Context = actix::Context; + + fn started(&mut self, ctx: &mut Self::Context) { + info!( + e3_id = %self.e3_id, + "ThresholdShareCollector started, scheduling timeout in {:?}", + DEFAULT_COLLECTION_TIMEOUT + ); + // Schedule timeout + let handle = ctx.notify_later(ThresholdShareCollectionTimeout, DEFAULT_COLLECTION_TIMEOUT); + self.timeout_handle = Some(handle); + } } impl Handler for ThresholdShareCollector { type Result = (); - fn handle(&mut self, msg: ThresholdShareCreated, _: &mut Self::Context) -> Self::Result { + fn handle(&mut self, msg: ThresholdShareCreated, ctx: &mut Self::Context) -> Self::Result { let start = Instant::now(); info!("ThresholdShareCollector: ThresholdShareCreated received by collector"); - if let CollectorState::Finished = self.state { - info!("ThresholdShareCollector is finished so ignoring!"); + + // Ignore if already finished or timed out + if !matches!(self.state, CollectorState::Collecting) { + info!( + "ThresholdShareCollector is not collecting (state: {:?}), ignoring", + match self.state { + CollectorState::Collecting => "Collecting", + CollectorState::Finished => "Finished", + CollectorState::TimedOut => "TimedOut", + } + ); return; - }; + } let pid = msg.share.party_id; info!("ThresholdShareCollector party id: {}", pid); let Some(_) = self.todo.take(&pid) else { info!( - "Error: {} was not in decryption key collectors ID list", + "Error: {} was not in threshold share collector's ID list", pid ); return; }; info!("Inserting... waiting on: {}", self.todo.len()); self.shares.insert(pid, msg.share); - if self.todo.len() == 0 { - info!("We have recieved all the things"); + + if self.todo.is_empty() { + info!("We have received all threshold shares"); self.state = CollectorState::Finished; + + // Cancel the timeout since we're done + if let Some(handle) = self.timeout_handle.take() { + ctx.cancel_future(handle); + } + let event: AllThresholdSharesCollected = self.shares.clone().into(); - self.parent.do_send(event) + self.parent.do_send(event); } info!( "Finished processing ThresholdShareCreated in {:?}", @@ -79,3 +121,40 @@ impl Handler for ThresholdShareCollector { ); } } + +impl Handler for ThresholdShareCollector { + type Result = (); + fn handle( + &mut self, + _: ThresholdShareCollectionTimeout, + ctx: &mut Self::Context, + ) -> Self::Result { + // Only handle timeout if we're still collecting + if !matches!(self.state, CollectorState::Collecting) { + return; + } + + warn!( + e3_id = %self.e3_id, + missing_parties = ?self.todo, + "Threshold share collection timed out, {} parties missing", + self.todo.len() + ); + + self.state = CollectorState::TimedOut; + + // Notify parent of failure + let missing_parties: Vec = self.todo.iter().copied().collect(); + self.parent.do_send(ThresholdShareCollectionFailed { + e3_id: self.e3_id.clone(), + reason: format!( + "Timeout waiting for threshold shares from {} parties", + missing_parties.len() + ), + missing_parties, + }); + + // Stop the actor + ctx.stop(); + } +} From 4aba96e8bfb8de1c2702cf79c0171e851437512c Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 25 Dec 2025 02:09:15 +0500 Subject: [PATCH 18/19] chore: formatting --- crates/keyshare/src/threshold_share_collector.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/crates/keyshare/src/threshold_share_collector.rs b/crates/keyshare/src/threshold_share_collector.rs index d60834bd4b..d0393c34cf 100644 --- a/crates/keyshare/src/threshold_share_collector.rs +++ b/crates/keyshare/src/threshold_share_collector.rs @@ -40,11 +40,7 @@ pub struct ThresholdShareCollector { } impl ThresholdShareCollector { - pub fn setup( - parent: Addr, - total: u64, - e3_id: E3id, - ) -> Addr { + pub fn setup(parent: Addr, total: u64, e3_id: E3id) -> Addr { let collector = Self { e3_id, todo: (0..total).collect(), From 33491098705a6b03014b3662cadd6509b4fba218 Mon Sep 17 00:00:00 2001 From: Hamza Khalid Date: Thu, 25 Dec 2025 02:26:48 +0500 Subject: [PATCH 19/19] fix: add E3 Completed to Threshold keyshare --- crates/keyshare/src/threshold_keyshare.rs | 25 ++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/crates/keyshare/src/threshold_keyshare.rs b/crates/keyshare/src/threshold_keyshare.rs index 6ef70b56e3..e577d81201 100644 --- a/crates/keyshare/src/threshold_keyshare.rs +++ b/crates/keyshare/src/threshold_keyshare.rs @@ -10,9 +10,10 @@ use e3_crypto::{Cipher, SensitiveBytes}; use e3_data::Persistable; use e3_events::{ prelude::*, BusHandle, CiphernodeSelected, CiphertextOutputPublished, ComputeRequest, - ComputeResponse, DecryptionshareCreated, E3id, EnclaveEvent, EnclaveEventData, EncryptionKey, - EncryptionKeyCollectionFailed, EncryptionKeyCreated, KeyshareCreated, PartyId, ThresholdShare, - ThresholdShareCollectionFailed, ThresholdShareCreated, + ComputeResponse, DecryptionshareCreated, Die, E3RequestComplete, E3id, EnclaveEvent, + EnclaveEventData, EncryptionKey, EncryptionKeyCollectionFailed, EncryptionKeyCreated, + KeyshareCreated, PartyId, ThresholdShare, ThresholdShareCollectionFailed, + ThresholdShareCreated, }; use e3_fhe::create_crp; use e3_multithread::Multithread; @@ -949,6 +950,7 @@ impl Handler for ThresholdKeyshare { EnclaveEventData::EncryptionKeyCreated(data) => { let _ = self.handle_encryption_key_created(data, ctx.address()); } + EnclaveEventData::E3RequestComplete(data) => ctx.notify(data), _ => (), } } @@ -1066,3 +1068,20 @@ impl Handler for ThresholdKeyshare { ctx.stop(); } } + +impl Handler for ThresholdKeyshare { + type Result = (); + fn handle(&mut self, _: E3RequestComplete, ctx: &mut Self::Context) -> Self::Result { + self.encryption_key_collector = None; + self.decryption_key_collector = None; + ctx.notify(Die); + } +} + +impl Handler for ThresholdKeyshare { + type Result = (); + fn handle(&mut self, _: Die, ctx: &mut Self::Context) -> Self::Result { + warn!("ThresholdKeyshare is shutting down"); + ctx.stop(); + } +}