Architecture fixes (#114)

* Extract FHE to a library

* Ciphernode -> Keyshare

* Remove serializers

* Fix decryption and aggregation hacks

packages/ciphernode/core/src/e3_request.rs

@@ -1,17 +1,17 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use actix::{Actor, Addr, Context, Handler, Recipient};
 
 use crate::{
-    Ciphernode, CommitteeMeta, E3id, EnclaveEvent, EventBus, Fhe, PlaintextAggregator,
+    Keyshare, CommitteeMeta, E3id, EnclaveEvent, EventBus, Fhe, PlaintextAggregator,
     PublicKeyAggregator, Subscribe,
 };
 
 #[derive(Default)]
 // TODO: Set this up with a Typestate pattern
 pub struct E3RequestContext {
-    pub ciphernode: Option<Addr<Ciphernode>>,
-    pub fhe: Option<Addr<Fhe>>,
+    pub keyshare: Option<Addr<Keyshare>>,
+    pub fhe: Option<Arc<Fhe>>,
     pub plaintext: Option<Addr<PlaintextAggregator>>,
     pub publickey: Option<Addr<PublicKeyAggregator>>,
     pub meta: Option<CommitteeMeta>,
@@ -38,8 +38,8 @@ impl E3RequestContext {
     fn recipients(&self) -> Vec<(String, Option<Recipient<EnclaveEvent>>)> {
         vec![
             (
-                "ciphernode".to_owned(),
-                self.ciphernode.clone().map(|addr| addr.into()),
+                "keyshare".to_owned(),
+                self.keyshare.clone().map(|addr| addr.into()),
             ),
             (
                 "plaintext".to_owned(),
@@ -66,12 +66,6 @@ impl E3RequestContext {
     }
 }
 
-struct E3RequestBuffers {
-    ciphernode: Vec<EnclaveEvent>,
-    publickey: Vec<EnclaveEvent>,
-    plaintext: Vec<EnclaveEvent>,
-}
-
 pub type ActorFactory = Box<dyn FnMut(&mut E3RequestContext, EnclaveEvent)>;
 
 // TODO: setup typestate pattern so that we have to place factories within correct order of

packages/ciphernode/core/src/events.rs

@@ -345,15 +345,13 @@ impl EnclaveEvent {
 #[cfg(test)]
 mod tests {
     use super::EnclaveEvent;
-    use crate::{
-        events::extract_enclave_event_name, serializers::PublicKeyShareSerializer, E3id,
-        KeyshareCreated,
-    };
+    use crate::{events::extract_enclave_event_name, E3id, KeyshareCreated};
     use alloy_primitives::address;
     use fhe::{
         bfv::{BfvParametersBuilder, SecretKey},
         mbfv::{CommonRandomPoly, PublicKeyShare},
     };
+    use fhe_traits::Serialize;
     use rand::SeedableRng;
     use rand_chacha::ChaCha20Rng;
     use std::error::Error;
@@ -384,7 +382,7 @@ mod tests {
         let crp = CommonRandomPoly::new(&params, &mut rng)?;
         let sk_share = { SecretKey::random(&params, &mut rng) };
         let pk_share = { PublicKeyShare::new(&sk_share, crp.clone(), &mut rng)? };
-        let pubkey = PublicKeyShareSerializer::to_bytes(pk_share, params.clone(), crp.clone())?;
+        let pubkey = pk_share.to_bytes();
         let kse = EnclaveEvent::from(KeyshareCreated {
             e3_id: E3id::from(1001),
             pubkey,

packages/ciphernode/core/src/fhe.rs

@@ -1,95 +1,44 @@
-use crate::{
-    ordered_set::OrderedSet,
-    serializers::{
-        CiphertextSerializer, DecryptionShareSerializer, PublicKeySerializer,
-        PublicKeyShareSerializer, SecretKeySerializer,
-    },
-    ActorFactory, E3Requested, EnclaveEvent,
-};
-use actix::{Actor, Context, Handler, Message};
+use crate::{ordered_set::OrderedSet, ActorFactory, E3Requested, EnclaveEvent};
 use anyhow::*;
 use fhe::{
-    bfv::{BfvParameters, BfvParametersBuilder, Encoding, Plaintext, PublicKey, SecretKey},
+    bfv::{
+        BfvParameters, BfvParametersBuilder, Ciphertext, Encoding, Plaintext, PublicKey, SecretKey,
+    },
     mbfv::{AggregateIter, CommonRandomPoly, DecryptionShare, PublicKeyShare},
 };
-use fhe_traits::{FheDecoder, Serialize};
-use rand::SeedableRng;
-use rand_chacha::{ChaCha20Rng, ChaCha8Rng};
-use std::{
-    hash::Hash,
-    sync::{Arc, Mutex},
-};
+use fhe_traits::{DeserializeParametrized, FheDecoder, Serialize};
+use rand_chacha::ChaCha20Rng;
+use std::sync::{Arc, Mutex};
 
-#[derive(Message, Clone, Debug, PartialEq, Eq, Hash)]
-#[rtype(result = "Result<(Vec<u8>, Vec<u8>)>")]
-pub struct GenerateKeyshare {
-    // responder_pk: Vec<u8>, // TODO: use this to encrypt the secret data
-}
-
-#[derive(Message, Clone, Debug, PartialEq, Eq)]
-#[rtype(result = "Result<(Vec<u8>)>")]
 pub struct GetAggregatePublicKey {
     pub keyshares: OrderedSet<Vec<u8>>,
 }
 
-#[derive(Message, Clone, Debug, PartialEq, Eq)]
-#[rtype(result = "Result<(Vec<u8>)>")]
 pub struct GetAggregatePlaintext {
     pub decryptions: OrderedSet<Vec<u8>>,
+    pub ciphertext_output: Vec<u8>,
 }
 
-#[derive(Message, Clone, Debug, PartialEq, Eq)]
-#[rtype(result = "Result<(Vec<u8>)>")]
 pub struct DecryptCiphertext {
     pub unsafe_secret: Vec<u8>,
     pub ciphertext: Vec<u8>,
 }
 
 pub type SharedRng = Arc<Mutex<ChaCha20Rng>>;
 
-/// Fhe library adaptor. All FHE computations should happen through this actor.
+/// Fhe library adaptor.
+#[derive(Clone)]
 pub struct Fhe {
     params: Arc<BfvParameters>,
     crp: CommonRandomPoly,
     rng: SharedRng,
 }
 
-impl Actor for Fhe {
-    type Context = Context<Self>;
-}
-
 impl Fhe {
     pub fn new(params: Arc<BfvParameters>, crp: CommonRandomPoly, rng: SharedRng) -> Self {
         Self { params, crp, rng }
     }
 
-    // Deprecated
-    pub fn try_default() -> Result<Self> {
-        // TODO: The production bootstrapping of this will involve receiving a crp bytes and param
-        // input form the event
-        let moduli = &vec![0x3FFFFFFF000001];
-        let degree = 2048usize;
-        let plaintext_modulus = 1032193u64;
-        let rng = Arc::new(Mutex::new(ChaCha20Rng::from_entropy()));
-        let crp = CommonRandomPoly::new(
-            &BfvParametersBuilder::new()
-                .set_degree(degree)
-                .set_plaintext_modulus(plaintext_modulus)
-                .set_moduli(&moduli)
-                .build_arc()?,
-            &mut *rng.lock().unwrap(),
-        )?
-        .to_bytes();
-
-        Ok(Fhe::from_raw_params(
-            moduli,
-            degree,
-            plaintext_modulus,
-            &crp,
-            rng,
-        )?)
-    }
-
     pub fn from_raw_params(
         moduli: &[u64],
         degree: usize,
@@ -109,77 +58,56 @@ impl Fhe {
             rng,
         ))
     }
-}
 
-impl Handler<GenerateKeyshare> for Fhe {
-    type Result = Result<(Vec<u8>, Vec<u8>)>;
-    fn handle(&mut self, _event: GenerateKeyshare, _: &mut Self::Context) -> Self::Result {
+    pub fn generate_keyshare(&self) -> Result<(Vec<u8>, Vec<u8>)> {
         let sk_share = { SecretKey::random(&self.params, &mut *self.rng.lock().unwrap()) };
         let pk_share =
             { PublicKeyShare::new(&sk_share, self.crp.clone(), &mut *self.rng.lock().unwrap())? };
         Ok((
-            SecretKeySerializer::to_bytes(sk_share, self.params.clone())?,
-            PublicKeyShareSerializer::to_bytes(pk_share, self.params.clone(), self.crp.clone())?,
+            SecretKeySerializer::to_bytes(sk_share)?,
+            pk_share.to_bytes(),
         ))
     }
-}
 
-impl Handler<DecryptCiphertext> for Fhe {
-    type Result = Result<Vec<u8>>;
-    fn handle(&mut self, msg: DecryptCiphertext, _: &mut Self::Context) -> Self::Result {
+    pub fn decrypt_ciphertext(&self, msg: DecryptCiphertext) -> Result<Vec<u8>> {
         let DecryptCiphertext {
-            unsafe_secret, // TODO: fix security issues with sending secrets between actors
+            unsafe_secret,
             ciphertext,
         } = msg;
 
-        let secret_key = SecretKeySerializer::from_bytes(&unsafe_secret)?;
-        let ct = Arc::new(CiphertextSerializer::from_bytes(&ciphertext)?);
-        let inner = DecryptionShare::new(&secret_key, &ct, &mut *self.rng.lock().unwrap()).unwrap();
-
-        Ok(DecryptionShareSerializer::to_bytes(
-            inner,
-            self.params.clone(),
-            ct.clone(),
-        )?)
+        let secret_key = SecretKeySerializer::from_bytes(&unsafe_secret, self.params.clone())?;
+        let ct = Arc::new(
+            Ciphertext::from_bytes(&ciphertext, &self.params)
+                .context("Error deserializing ciphertext")?,
+        );
+        let decryption_share =
+            DecryptionShare::new(&secret_key, &ct, &mut *self.rng.lock().unwrap()).unwrap();
+        Ok(decryption_share.to_bytes())
     }
-}
 
-impl Handler<GetAggregatePublicKey> for Fhe {
-    type Result = Result<Vec<u8>>;
-
-    fn handle(&mut self, msg: GetAggregatePublicKey, _: &mut Self::Context) -> Self::Result {
+    pub fn get_aggregate_public_key(&self, msg: GetAggregatePublicKey) -> Result<Vec<u8>> {
         let public_key: PublicKey = msg
             .keyshares
             .iter()
-            .map(|k| PublicKeyShareSerializer::from_bytes(k))
-            .collect::<Result<Vec<_>>>()?
-            .into_iter()
+            .map(|k| PublicKeyShare::deserialize(k, &self.params, self.crp.clone()))
             .aggregate()?;
 
         Ok(public_key.to_bytes())
     }
-}
 
-impl Handler<GetAggregatePlaintext> for Fhe {
-    type Result = Result<Vec<u8>>;
-    fn handle(&mut self, msg: GetAggregatePlaintext, _: &mut Self::Context) -> Self::Result {
+    pub fn get_aggregate_plaintext(&self, msg: GetAggregatePlaintext) -> Result<Vec<u8>> {
+        let arc_ct = Arc::new(Ciphertext::from_bytes(
+            &msg.ciphertext_output,
+            &self.params,
+        )?);
+
         let plaintext: Plaintext = msg
             .decryptions
             .iter()
-            .map(|k| DecryptionShareSerializer::from_bytes(k))
-            .collect::<Result<Vec<_>>>()?
-            .into_iter()
+            .map(|k| DecryptionShare::deserialize(k, &self.params, arc_ct.clone()))
             .aggregate()?;
 
-        // XXX: how do we know what the expected output of the plaintext is in order to decrypt
-        // here for serialization?
-        // This would be dependent on the computation that is running.
-        // For now assuming testcase of Vec<u64> and currently represents a "HARDCODED" program
-        // output format of Vec<u64>
-        // This could be determined based on the "program" config events
-
         let decoded = Vec::<u64>::try_decode(&plaintext, Encoding::poly())?;
-        let decoded = &decoded[0..2]; // TODO: this will be computation dependent
         Ok(bincode::serialize(&decoded)?)
     }
 }
@@ -201,11 +129,45 @@ impl FheFactory {
                 ..
             } = data;
 
-            ctx.fhe = Some(
+            ctx.fhe = Some(Arc::new(
                 Fhe::from_raw_params(&moduli, degree, plaintext_modulus, &crp, rng.clone())
-                    .unwrap()
-                    .start(),
-            );
+                    .unwrap(),
+            ));
+        })
+    }
+}
+
+struct SecretKeySerializer {
+    pub inner: SecretKey,
+}
+
+impl SecretKeySerializer {
+    pub fn to_bytes(inner: SecretKey) -> Result<Vec<u8>> {
+        let value = Self { inner };
+        Ok(value.unsafe_serialize()?)
+    }
+
+    pub fn from_bytes(bytes: &[u8], params: Arc<BfvParameters>) -> Result<SecretKey> {
+        Ok(Self::deserialize(bytes, params)?.inner)
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SecretKeyData {
+    coeffs: Box<[i64]>,
+}
+
+impl SecretKeySerializer {
+    pub fn unsafe_serialize(&self) -> Result<Vec<u8>> {
+        Ok(bincode::serialize(&SecretKeyData {
+            coeffs: self.inner.coeffs.clone(),
+        })?)
+    }
+
+    pub fn deserialize(bytes: &[u8], params: Arc<BfvParameters>) -> Result<SecretKeySerializer> {
+        let SecretKeyData { coeffs } = bincode::deserialize(&bytes)?;
+        Ok(Self {
+            inner: SecretKey::new(coeffs.to_vec(), &params),
         })
     }
 }