Added DataStore struct to act as an injection point for persistence

packages/ciphernode/data/Cargo.toml

@@ -7,3 +7,4 @@ repository = "https://github.com/gnosisguild/enclave/packages/ciphernode"
 
 [dependencies]
 actix = { workspace = true }
+anyhow = { workspace = true }

packages/ciphernode/data/src/data_store.rs

@@ -0,0 +1,49 @@
+use actix::{Addr, Message, Recipient};
+use anyhow::Result;
+
+use crate::InMemDataStore;
+
+#[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
+#[rtype(result = "()")]
+pub struct Insert(pub Vec<u8>, pub Vec<u8>);
+impl Insert {
+    pub fn key(&self) -> Vec<u8> {
+        self.0.clone()
+    }
+
+    pub fn value(&self) -> Vec<u8> {
+        self.1.clone()
+    }
+}
+
+#[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
+#[rtype(result = "Option<Vec<u8>>")]
+pub struct Get(pub Vec<u8>);
+impl Get {
+    pub fn key(&self) -> Vec<u8> {
+        self.0.clone()
+    }
+}
+
+#[derive(Clone)]
+pub struct DataStore(Recipient<Get>, Recipient<Insert>);
+impl DataStore {
+    pub async fn read(&self, msg: Get) -> Result<Option<Vec<u8>>> {
+        Ok(self.0.send(msg).await?)
+    }
+
+    pub fn write(&self, msg: Insert) {
+        self.1.do_send(msg)
+    }
+
+    // use this for testing
+    pub fn from_in_mem(addr: Addr<InMemDataStore>) -> Self {
+        Self(addr.clone().recipient(), addr.clone().recipient())
+    }
+
+    // // use this for production
+    // pub fn from_sled(&data_addr: Addr<SledDb>) -> Self {
+    //   let d = data_addr.clone();
+    //   Self(d.recipient(),d.recipient())
+    // }
+}

packages/ciphernode/data/src/data.rs → packages/ciphernode/data/src/in_mem.rs

@@ -1,29 +1,7 @@
 use actix::{Actor, Context, Handler, Message};
 use std::collections::BTreeMap;
 
-// TODO: replace with sled version
-
-#[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
-#[rtype(result = "()")]
-pub struct Insert(pub Vec<u8>, pub Vec<u8>);
-impl Insert {
-    fn key(&self) -> Vec<u8> {
-        self.0.clone()
-    }
-
-    fn value(&self) -> Vec<u8> {
-        self.1.clone()
-    }
-}
-
-#[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
-#[rtype(result = "Option<Vec<u8>>")]
-pub struct Get(pub Vec<u8>);
-impl Get {
-    fn key(&self) -> Vec<u8> {
-        self.0.clone()
-    }
-}
+use crate::{Get, Insert};
 
 #[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
 #[rtype(result = "Vec<DataOp>")]
@@ -34,17 +12,17 @@ pub enum DataOp {
     Insert(Insert),
 }
 
-pub struct Data {
+pub struct InMemDataStore {
     db: BTreeMap<Vec<u8>, Vec<u8>>,
     log: Vec<DataOp>,
     capture: bool,
 }
 
-impl Actor for Data {
+impl Actor for InMemDataStore {
     type Context = Context<Self>;
 }
 
-impl Data {
+impl InMemDataStore {
     pub fn new(capture: bool) -> Self {
         Self {
             db: BTreeMap::new(),
@@ -54,7 +32,7 @@ impl Data {
     }
 }
 
-impl Handler<Insert> for Data {
+impl Handler<Insert> for InMemDataStore {
     type Result = ();
     fn handle(&mut self, event: Insert, _: &mut Self::Context) {
         // insert data into sled
@@ -66,15 +44,15 @@ impl Handler<Insert> for Data {
     }
 }
 
-impl Handler<Get> for Data {
+impl Handler<Get> for InMemDataStore {
     type Result = Option<Vec<u8>>;
     fn handle(&mut self, event: Get, _: &mut Self::Context) -> Option<Vec<u8>> {
         let key = event.key();
         self.db.get(&key).map(|r| r.clone())
     }
 }
 
-impl Handler<GetLog> for Data {
+impl Handler<GetLog> for InMemDataStore {
     type Result = Vec<DataOp>;
     fn handle(&mut self, _: GetLog, _: &mut Self::Context) -> Vec<DataOp> {
         self.log.clone()

packages/ciphernode/data/src/lib.rs

@@ -1,2 +1,4 @@
-mod data;
-pub use data::*;
+mod in_mem;
+mod data_store;
+pub use in_mem::*;
+pub use data_store::*;

packages/ciphernode/enclave_node/src/ciphernode.rs

@@ -1,7 +1,7 @@
 use actix::{Actor, Addr, Context};
 use alloy::primitives::Address;
 use anyhow::Result;
-use data::Data;
+use data::{DataStore, InMemDataStore};
 use enclave_core::EventBus;
 use evm::{CiphernodeRegistrySol, EnclaveSolReader};
 use logger::SimpleLogger;
@@ -21,7 +21,7 @@ use crate::app_config::AppConfig;
 pub struct MainCiphernode {
     addr: Address,
     bus: Addr<EventBus>,
-    data: Addr<Data>,
+    data: DataStore,
     sortition: Addr<Sortition>,
     selector: Addr<CiphernodeSelector>,
     e3_manager: Addr<E3RequestRouter>,
@@ -32,7 +32,7 @@ impl MainCiphernode {
     pub fn new(
         addr: Address,
         bus: Addr<EventBus>,
-        data: Addr<Data>,
+        data: DataStore,
         sortition: Addr<Sortition>,
         selector: Addr<CiphernodeSelector>,
         p2p: Addr<P2p>,
@@ -57,7 +57,8 @@ impl MainCiphernode {
             rand_chacha::ChaCha20Rng::from_rng(OsRng).expect("Failed to create RNG"),
         ));
         let bus = EventBus::new(true).start();
-        let data = Data::new(true).start(); // TODO: Use a sled backed Data Actor
+        // TODO: switch to Sled actor
+        let data = DataStore::from_in_mem(InMemDataStore::new(true).start());
         let sortition = Sortition::attach(bus.clone());
         let selector =
             CiphernodeSelector::attach(bus.clone(), sortition.clone(), &address.to_string());

packages/ciphernode/keyshare/src/keyshare.rs

@@ -1,6 +1,6 @@
 use actix::prelude::*;
 use anyhow::{anyhow, Context, Result};
-use data::{Data, Get, Insert};
+use data::{DataStore, Get, Insert};
 use enclave_core::{
     CiphernodeSelected, CiphertextOutputPublished, DecryptionshareCreated, Die, EnclaveErrorType,
     EnclaveEvent, EventBus, FromError, KeyshareCreated,
@@ -10,7 +10,7 @@ use std::sync::Arc;
 
 pub struct Keyshare {
     fhe: Arc<Fhe>,
-    data: Addr<Data>,
+    data: DataStore,
     bus: Addr<EventBus>,
     address: String,
 }
@@ -20,7 +20,7 @@ impl Actor for Keyshare {
 }
 
 impl Keyshare {
-    pub fn new(bus: Addr<EventBus>, data: Addr<Data>, fhe: Arc<Fhe>, address: &str) -> Self {
+    pub fn new(bus: Addr<EventBus>, data: DataStore, fhe: Arc<Fhe>, address: &str) -> Self {
         Self {
             bus,
             fhe,
@@ -64,11 +64,11 @@ impl Handler<CiphernodeSelected> for Keyshare {
         // best practice would be as you boot up a node you enter in a configured password from
         // which we derive a kdf which gets used to generate this key
         self.data
-            .do_send(Insert(format!("{}/sk", e3_id).into(), sk));
+            .write(Insert(format!("{}/sk", e3_id).into(), sk));
 
         // save public key against e3_id/pk
         self.data
-            .do_send(Insert(format!("{}/pk", e3_id).into(), pubkey.clone()));
+            .write(Insert(format!("{}/pk", e3_id).into(), pubkey.clone()));
 
         // broadcast the KeyshareCreated message
         let event = EnclaveEvent::from(KeyshareCreated {
@@ -109,7 +109,7 @@ impl Handler<Die> for Keyshare {
 
 async fn on_decryption_requested(
     fhe: Arc<Fhe>,
-    data: Addr<Data>,
+    data: DataStore,
     bus: Addr<EventBus>,
     event: CiphertextOutputPublished,
     address: String,
@@ -120,7 +120,7 @@ async fn on_decryption_requested(
     } = event;
 
     // get secret key by id from data
-    let Some(unsafe_secret) = data.send(Get(format!("{}/sk", e3_id).into())).await? else {
+    let Some(unsafe_secret) = data.read(Get(format!("{}/sk", e3_id).into())).await? else {
         return Err(anyhow::anyhow!("Secret key not stored for {}", e3_id));
     };
 

packages/ciphernode/router/src/hooks.rs

@@ -1,7 +1,7 @@
 use crate::EventHook;
 use actix::{Actor, Addr};
 use aggregator::{PlaintextAggregator, PublicKeyAggregator};
-use data::Data;
+use data::DataStore;
 use enclave_core::{E3Requested, EnclaveEvent, EventBus};
 use fhe::{Fhe, SharedRng};
 use keyshare::Keyshare;
@@ -28,7 +28,7 @@ impl LazyFhe {
 
 pub struct LazyKeyshare;
 impl LazyKeyshare {
-    pub fn create(bus: Addr<EventBus>, data: Addr<Data>, address: &str) -> EventHook {
+    pub fn create(bus: Addr<EventBus>, data: DataStore, address: &str) -> EventHook {
         let address = address.to_string();
         Box::new(move |ctx, evt| {
             // Save Ciphernode on CiphernodeSelected

packages/ciphernode/tests/tests/test_aggregation_and_decryption.rs

@@ -1,4 +1,4 @@
-use data::Data;
+use data::{DataStore, InMemDataStore};
 use enclave_core::{
     CiphernodeAdded, CiphernodeSelected, CiphertextOutputPublished, DecryptionshareCreated,
     E3RequestComplete, E3Requested, E3id, EnclaveEvent, EventBus, GetHistory, KeyshareCreated,
@@ -31,7 +31,8 @@ use tokio::{sync::mpsc::channel, time::sleep};
 // Simulating a local node
 async fn setup_local_ciphernode(bus: Addr<EventBus>, rng: SharedRng, logging: bool, addr: &str) {
     // create data actor for saving data
-    let data = Data::new(logging).start(); // TODO: Use a sled backed Data Actor
+    let data_actor = InMemDataStore::new(logging).start(); // TODO: Use a sled backed Data Actor
+    let store = DataStore::from_in_mem(data_actor);
 
     // create ciphernode actor for managing ciphernode flow
     let sortition = Sortition::attach(bus.clone());
@@ -47,7 +48,7 @@ async fn setup_local_ciphernode(bus: Addr<EventBus>, rng: SharedRng, logging: bo
             bus.clone(),
             sortition.clone(),
         ))
-        .add_hook(LazyKeyshare::create(bus.clone(), data.clone(), addr))
+        .add_hook(LazyKeyshare::create(bus.clone(), store.clone(), addr))
         .build();
 
     SimpleLogger::attach(addr, bus.clone());