diff --git a/packages/ciphernode/router/src/e3_request_router.rs b/packages/ciphernode/router/src/e3_request_router.rs index 1669da76..835d2ec7 100644 --- a/packages/ciphernode/router/src/e3_request_router.rs +++ b/packages/ciphernode/router/src/e3_request_router.rs @@ -3,6 +3,7 @@ use crate::CommitteeMetaFactory; use super::CommitteeMeta; use aggregator::PlaintextAggregator; use aggregator::PublicKeyAggregator; +use enclave_core::Die; use enclave_core::E3RequestComplete; use enclave_core::{E3id, EnclaveEvent, EventBus, Subscribe}; use fhe::Fhe; @@ -71,6 +72,20 @@ impl E3RequestContext { } }); } + + pub fn cleanup(&mut self) { + if let Some(keyshare) = self.keyshare.take() { + keyshare.do_send(Die); + } + if let Some(plaintext) = self.plaintext.take() { + plaintext.do_send(Die); + } + if let Some(publickey) = self.publickey.take() { + publickey.do_send(Die); + } + self.fhe = None; + self.meta = None; + } } /// Format of the hook that needs to be passed to E3RequestRouter @@ -128,6 +143,8 @@ impl Handler for E3RequestRouter { // Send to bus so all other actors can react to a request being complete. self.bus.do_send(event); + context.cleanup(); + self.contexts.remove(&e3_id); } } } @@ -148,7 +165,7 @@ impl E3RequestRouterBuilder { contexts: HashMap::new(), hooks: self.hooks, buffer: EventBuffer::default(), - bus: self.bus.clone() + bus: self.bus.clone(), }; let addr = e3r.start(); diff --git a/packages/ciphernode/router/src/hooks.rs b/packages/ciphernode/router/src/hooks.rs index 77ffc072..9e50c45a 100644 --- a/packages/ciphernode/router/src/hooks.rs +++ b/packages/ciphernode/router/src/hooks.rs @@ -2,7 +2,7 @@ use crate::EventHook; use actix::{Actor, Addr}; use aggregator::{PlaintextAggregator, PublicKeyAggregator}; use data::Data; -use enclave_core::{Die, E3Requested, EnclaveEvent, EventBus}; +use enclave_core::{E3Requested, EnclaveEvent, EventBus}; use fhe::{Fhe, SharedRng}; use keyshare::Keyshare; use sortition::Sortition; @@ -30,23 +30,18 @@ pub struct LazyKeyshare; impl LazyKeyshare { pub fn create(bus: Addr, data: Addr, address: &str) -> EventHook { let address = address.to_string(); - Box::new(move |ctx, evt| match evt { - EnclaveEvent::CiphernodeSelected { .. } => { - let Some(ref fhe) = ctx.fhe else { - return; - }; + Box::new(move |ctx, evt| { + // Save Ciphernode on CiphernodeSelected + let EnclaveEvent::CiphernodeSelected { .. } = evt else { + return; + }; - ctx.keyshare = - Some(Keyshare::new(bus.clone(), data.clone(), fhe.clone(), &address).start()) - } - EnclaveEvent::E3RequestComplete { .. } => { - let Some(actor) = ctx.keyshare.take() else { - return; - }; + let Some(ref fhe) = ctx.fhe else { + return; + }; - actor.do_send(Die); - } - _ => (), + ctx.keyshare = + Some(Keyshare::new(bus.clone(), data.clone(), fhe.clone(), &address).start()) }) } } @@ -54,38 +49,31 @@ impl LazyKeyshare { pub struct LazyPlaintextAggregator; impl LazyPlaintextAggregator { pub fn create(bus: Addr, sortition: Addr) -> EventHook { - Box::new(move |ctx, evt| match evt { - EnclaveEvent::CiphertextOutputPublished { data, .. } => { - let Some(ref fhe) = ctx.fhe else { - return; - }; - let Some(ref meta) = ctx.meta else { - return; - }; - - // Save plaintext aggregator - ctx.plaintext = Some( - PlaintextAggregator::new( - fhe.clone(), - bus.clone(), - sortition.clone(), - data.e3_id, - meta.threshold_m, - meta.seed, - data.ciphertext_output, - meta.src_chain_id, - ) - .start(), - ); - } - EnclaveEvent::E3RequestComplete { .. } => { - let Some(actor) = ctx.plaintext.take() else { - return; - }; + Box::new(move |ctx, evt| { + // Save plaintext aggregator + let EnclaveEvent::CiphertextOutputPublished { data, .. } = evt else { + return; + }; + let Some(ref fhe) = ctx.fhe else { + return; + }; + let Some(ref meta) = ctx.meta else { + return; + }; - actor.do_send(Die); - } - _ => (), + ctx.plaintext = Some( + PlaintextAggregator::new( + fhe.clone(), + bus.clone(), + sortition.clone(), + data.e3_id, + meta.threshold_m, + meta.seed, + data.ciphertext_output, + meta.src_chain_id, + ) + .start(), + ); }) } } @@ -93,40 +81,33 @@ impl LazyPlaintextAggregator { pub struct LazyPublicKeyAggregator; impl LazyPublicKeyAggregator { pub fn create(bus: Addr, sortition: Addr) -> EventHook { - Box::new(move |ctx, evt| match evt { + Box::new(move |ctx, evt| { // Saving the publickey aggregator with deps on E3Requested - EnclaveEvent::E3Requested { data, .. } => { - let Some(ref fhe) = ctx.fhe else { - println!("fhe was not on ctx"); - return; - }; - let Some(ref meta) = ctx.meta else { - println!("meta was not on ctx"); - return; - }; - - ctx.publickey = Some( - PublicKeyAggregator::new( - fhe.clone(), - bus.clone(), - sortition.clone(), - data.e3_id, - meta.threshold_m, - meta.seed, - meta.src_chain_id, - ) - .start(), - ); - } - EnclaveEvent::E3RequestComplete { .. } => { - let Some(actor) = ctx.publickey.take() else { - return; - }; + let EnclaveEvent::E3Requested { data, .. } = evt else { + return; + }; - actor.do_send(Die); - } + let Some(ref fhe) = ctx.fhe else { + println!("fhe was not on ctx"); + return; + }; + let Some(ref meta) = ctx.meta else { + println!("meta was not on ctx"); + return; + }; - _ => (), + ctx.publickey = Some( + PublicKeyAggregator::new( + fhe.clone(), + bus.clone(), + sortition.clone(), + data.e3_id, + meta.threshold_m, + meta.seed, + meta.src_chain_id, + ) + .start(), + ); }) } }