Refactor to kill context in a single place

gnosisguild · Oct 9, 2024 · dbaaa26 · dbaaa26
1 parent 948149e
commit dbaaa26
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 79 deletions.
diff --git 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<EnclaveEvent> 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
@@ -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,103 +30,84 @@ pub struct LazyKeyshare;
 impl LazyKeyshare {
     pub fn create(bus: Addr<EventBus>, data: Addr<Data>, 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())
         })
     }
 }
 
 pub struct LazyPlaintextAggregator;
 impl LazyPlaintextAggregator {
     pub fn create(bus: Addr<EventBus>, sortition: Addr<Sortition>) -> 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(),
+            );
         })
     }
 }
 
 pub struct LazyPublicKeyAggregator;
 impl LazyPublicKeyAggregator {
     pub fn create(bus: Addr<EventBus>, sortition: Addr<Sortition>) -> 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(),
+            );
         })
     }
 }