Add test for hydration

packages/ciphernode/core/src/events.rs

@@ -120,10 +120,13 @@ pub enum EnclaveEvent {
         id: EventId,
         data: E3RequestComplete,
     },
-    // CommitteeSelected,
-    // OutputDecrypted,
-    // CiphernodeRegistered,
-    // CiphernodeDeregistered,
+    Shutdown {
+        id: EventId,
+        data: Shutdown,
+    }, // CommitteeSelected,
+       // OutputDecrypted,
+       // CiphernodeRegistered,
+       // CiphernodeDeregistered,
 }
 
 impl EnclaveEvent {
@@ -147,6 +150,7 @@ impl EnclaveEvent {
             EnclaveEvent::CiphernodeAdded { .. } => true,
             EnclaveEvent::CiphernodeRemoved { .. } => true,
             EnclaveEvent::E3RequestComplete { .. } => true,
+            EnclaveEvent::Shutdown { .. } => true,
             _ => false,
         }
     }
@@ -166,6 +170,7 @@ impl From<EnclaveEvent> for EventId {
             EnclaveEvent::CiphernodeRemoved { id, .. } => id,
             EnclaveEvent::EnclaveError { id, .. } => id,
             EnclaveEvent::E3RequestComplete { id, .. } => id,
+            EnclaveEvent::Shutdown { id, .. } => id,
         }
     }
 }
@@ -196,6 +201,7 @@ impl EnclaveEvent {
             EnclaveEvent::CiphernodeRemoved { data, .. } => format!("{}", data),
             EnclaveEvent::E3RequestComplete { data, .. } => format!("{}", data),
             EnclaveEvent::EnclaveError { data, .. } => format!("{:?}", data),
+            EnclaveEvent::Shutdown { data, .. } => format!("{:?}", data),
             // _ => "<omitted>".to_string(),
         }
     }
@@ -206,6 +212,8 @@ pub trait FromError {
     fn from_error(err_type: EnclaveErrorType, error: Self::Error) -> Self;
 }
 
+
+// TODO: These From traits should be handled by a macro
 impl From<KeyshareCreated> for EnclaveEvent {
     fn from(data: KeyshareCreated) -> Self {
         EnclaveEvent::KeyshareCreated {
@@ -305,6 +313,15 @@ impl From<EnclaveError> for EnclaveEvent {
     }
 }
 
+impl From<Shutdown> for EnclaveEvent {
+    fn from(data: Shutdown) -> Self {
+        EnclaveEvent::Shutdown {
+            id: EventId::from(data.clone()),
+            data: data.clone(),
+        }
+    }
+}
+
 impl FromError for EnclaveEvent {
     type Error = anyhow::Error;
     fn from_error(err_type: EnclaveErrorType, error: Self::Error) -> Self {
@@ -505,6 +522,17 @@ impl Display for Die {
     }
 }
 
+
+/// Represents a shutdown event triggered by SIG TERM
+#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[rtype(result = "()")]
+pub struct Shutdown;
+impl Display for Shutdown {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Shutdown",)
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct Seed(pub [u8; 32]);
 impl From<Seed> for u64 {

packages/ciphernode/data/src/data_store.rs

@@ -64,7 +64,7 @@ impl DataStore {
     }
 
     /// Construct a data store from an InMemStore actor
-    pub fn from_in_mem(addr: Addr<InMemStore>) -> Self {
+    pub fn from_in_mem(addr: &Addr<InMemStore>) -> Self {
         Self {
             get: addr.clone().recipient(),
             insert: addr.clone().recipient(),
@@ -88,7 +88,7 @@ impl DataStore {
     /// #[actix_rt::main]
     /// async fn main() -> Result<()>{  
     ///   let addr = InMemStore::new(false).start();
-    ///   let store = DataStore::from_in_mem(addr);
+    ///   let store = DataStore::from_in_mem(&addr);
     ///   assert_eq!(store.base("//foo")
     ///     .scope("bar")
     ///     .scope("/baz")

packages/ciphernode/enclave_node/src/aggregator.rs

@@ -55,7 +55,7 @@ impl MainAggregator {
             rand_chacha::ChaCha20Rng::from_rng(OsRng).expect("Failed to create RNG"),
         ));
 
-        let store = DataStore::from_in_mem(InMemStore::new(true).start());
+        let store = DataStore::from_in_mem(&InMemStore::new(true).start());
         let repositories = store.repositories();
         let sortition = Sortition::attach(bus.clone(), repositories.sortition());
         let signer = pull_eth_signer_from_env("PRIVATE_KEY").await?;

packages/ciphernode/enclave_node/src/ciphernode.rs

@@ -61,7 +61,7 @@ impl MainCiphernode {
         ));
         let bus = EventBus::new(true).start();
         // TODO: switch to Sled actor
-        let store = DataStore::from_in_mem(InMemStore::new(true).start());
+        let store = DataStore::from_in_mem(&InMemStore::new(true).start());
         let repositories = store.repositories();
 
         let sortition = Sortition::attach(bus.clone(), repositories.sortition());

packages/ciphernode/keyshare/Cargo.toml

@@ -11,3 +11,4 @@ actix = { workspace = true }
 anyhow = { workspace = true }
 serde = { workspace = true }
 async-trait = { workspace = true }
+tracing = { workspace = true }

packages/ciphernode/keyshare/src/keyshare.rs

@@ -9,6 +9,7 @@ use enclave_core::{
 use fhe::{DecryptCiphertext, Fhe};
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
+use tracing::warn;
 
 pub struct Keyshare {
     fhe: Arc<Fhe>,
@@ -84,6 +85,7 @@ impl Handler<EnclaveEvent> for Keyshare {
             EnclaveEvent::CiphernodeSelected { data, .. } => ctx.notify(data),
             EnclaveEvent::CiphertextOutputPublished { data, .. } => ctx.notify(data),
             EnclaveEvent::E3RequestComplete { .. } => ctx.notify(Die),
+            EnclaveEvent::Shutdown { .. } => ctx.notify(Die),
             _ => (),
         }
     }
@@ -162,6 +164,7 @@ impl Handler<CiphertextOutputPublished> for Keyshare {
 impl Handler<Die> for Keyshare {
     type Result = ();
     fn handle(&mut self, _: Die, ctx: &mut Self::Context) -> Self::Result {
+        warn!("Keyshare is shutting down");
         ctx.stop()
     }
 }

packages/ciphernode/router/src/ciphernode_selector.rs

@@ -1,7 +1,7 @@
 /// CiphernodeSelector is an actor that determines if a ciphernode is part of a committee and if so
 /// forwards a CiphernodeSelected event to the event bus
 use actix::prelude::*;
-use enclave_core::{CiphernodeSelected, EnclaveEvent, EventBus, Subscribe};
+use enclave_core::{CiphernodeSelected, E3Requested, EnclaveEvent, EventBus, Shutdown, Subscribe};
 use sortition::{GetHasNode, Sortition};
 use tracing::info;
 
@@ -28,24 +28,32 @@ impl CiphernodeSelector {
         let addr = CiphernodeSelector::new(bus.clone(), sortition, address).start();
 
         bus.do_send(Subscribe::new("E3Requested", addr.clone().recipient()));
+        bus.do_send(Subscribe::new("Shutdown", addr.clone().recipient()));
 
         addr
     }
 }
 
 impl Handler<EnclaveEvent> for CiphernodeSelector {
+    type Result = ();
+    fn handle(&mut self, msg: EnclaveEvent, ctx: &mut Self::Context) -> Self::Result {
+        match msg {
+            EnclaveEvent::E3Requested { data, .. } => ctx.notify(data),
+            EnclaveEvent::Shutdown { data, .. } => ctx.notify(data),
+            _ => (),
+        }
+    }
+}
+
+impl Handler<E3Requested> for CiphernodeSelector {
     type Result = ResponseFuture<()>;
 
-    fn handle(&mut self, event: EnclaveEvent, _ctx: &mut Self::Context) -> Self::Result {
+    fn handle(&mut self, data: E3Requested, _ctx: &mut Self::Context) -> Self::Result {
         let address = self.address.clone();
         let sortition = self.sortition.clone();
         let bus = self.bus.clone();
 
         Box::pin(async move {
-            let EnclaveEvent::E3Requested { data, .. } = event else {
-                return;
-            };
-
             let seed = data.seed;
             let size = data.threshold_m;
 
@@ -70,3 +78,11 @@ impl Handler<EnclaveEvent> for CiphernodeSelector {
         })
     }
 }
+
+impl Handler<Shutdown> for CiphernodeSelector {
+    type Result = ();
+    fn handle(&mut self, _msg: Shutdown, ctx: &mut Self::Context) -> Self::Result {
+        tracing::info!("Killing CiphernodeSelector");
+        ctx.stop();
+    }
+}

packages/ciphernode/router/src/context.rs

@@ -82,6 +82,14 @@ impl E3RequestContext {
         });
     }
 
+    pub fn forward_message_now(&self, msg: &EnclaveEvent) {
+        self.recipients().into_iter().for_each(|(_, recipient)| {
+            if let Some(act) = recipient {
+                act.do_send(msg.clone());
+            }
+        });
+    }
+
     /// Accept a DataStore ID and a Keystore actor address
     pub fn set_keyshare(&mut self, value: Addr<Keyshare>) {
         self.keyshare = Some(value);

packages/ciphernode/router/src/e3_request_router.rs

@@ -3,6 +3,7 @@ use crate::E3RequestContext;
 use crate::E3RequestContextParams;
 use crate::E3RequestContextSnapshot;
 use crate::RepositoriesFactory;
+use actix::AsyncContext;
 use actix::{Actor, Addr, Context, Handler};
 use anyhow::*;
 use async_trait::async_trait;
@@ -12,6 +13,7 @@ use data::FromSnapshotWithParams;
 use data::Repository;
 use data::Snapshot;
 use enclave_core::E3RequestComplete;
+use enclave_core::Shutdown;
 use enclave_core::{E3id, EnclaveEvent, EventBus, Subscribe};
 use serde::Deserialize;
 use serde::Serialize;
@@ -105,7 +107,13 @@ impl Actor for E3RequestRouter {
 
 impl Handler<EnclaveEvent> for E3RequestRouter {
     type Result = ();
-    fn handle(&mut self, msg: EnclaveEvent, _: &mut Self::Context) -> Self::Result {
+    fn handle(&mut self, msg: EnclaveEvent, ctx: &mut Self::Context) -> Self::Result {
+        // If we are shuttomg down then bail on anything else
+        if let EnclaveEvent::Shutdown { data, .. } = msg {
+            ctx.notify(data);
+            return;
+        }
+
         let Some(e3_id) = msg.get_e3_id() else {
             return;
         };
@@ -155,6 +163,16 @@ impl Handler<EnclaveEvent> for E3RequestRouter {
     }
 }
 
+impl Handler<Shutdown> for E3RequestRouter {
+    type Result = ();
+    fn handle(&mut self, msg: Shutdown, _ctx: &mut Self::Context) -> Self::Result {
+        let shutdown_evt = EnclaveEvent::from(msg);
+        for (_, ctx) in self.contexts.iter() {
+            ctx.forward_message_now(&shutdown_evt)
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct E3RequestRouterSnapshot {
     contexts: Vec<E3id>,