Skip to content

Commit

Permalink
Prepare for decryption
Browse files Browse the repository at this point in the history
  • Loading branch information
ryardley committed Aug 29, 2024
1 parent 64c5059 commit 8141bfb
Show file tree
Hide file tree
Showing 6 changed files with 156 additions and 21 deletions.
17 changes: 16 additions & 1 deletion packages/ciphernode/core/src/ciphernode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::{
eventbus::EventBus,
events::{ComputationRequested, EnclaveEvent, KeyshareCreated},
fhe::{Fhe, GenerateKeyshare},
Subscribe,
DecryptionRequested, Subscribe,
};
use actix::prelude::*;
use anyhow::Result;
Expand Down Expand Up @@ -83,5 +83,20 @@ async fn on_computation_requested(
let event = EnclaveEvent::from(KeyshareCreated { pubkey, e3_id });
bus.do_send(event);

Ok(())
}

async fn on_decryption_requested(
fhe: Addr<Fhe>,
data: Addr<Data>,
bus: Addr<EventBus>,
event: DecryptionRequested,
) -> Result<()> {
let DecryptionRequested { e3_id, ciphertext } = event;

// get secret key by id from data



Ok(())
}
3 changes: 2 additions & 1 deletion packages/ciphernode/core/src/committee.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,8 @@ impl Handler<EnclaveEvent> for CommitteeManager {

key.do_send(Die);
self.keys.remove(&data.e3_id);
} // _ => (),
},
_ => (),
}
}
}
12 changes: 12 additions & 0 deletions packages/ciphernode/core/src/eventbus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ impl Subscribe {
#[rtype(result = "Vec<EnclaveEvent>")]
pub struct GetHistory;

#[derive(Message)]
#[rtype(result = "()")]
pub struct ResetHistory;


/// Central EventBus for each node. Actors publish events to this bus by sending it EnclaveEvents.
/// All events sent to this bus are assumed to be published over the network via pubsub.
Expand Down Expand Up @@ -73,6 +77,14 @@ impl Handler<GetHistory> for EventBus {
self.history.clone()
}
}
impl Handler<ResetHistory> for EventBus {
type Result = ();

fn handle(&mut self, _: ResetHistory, _: &mut Context<Self>) {
self.history.clear()
}
}


impl Handler<EnclaveEvent> for EventBus {
type Result = ();
Expand Down
29 changes: 28 additions & 1 deletion packages/ciphernode/core/src/events.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
use crate::fhe::{WrappedPublicKey, WrappedPublicKeyShare};
use crate::{
fhe::{WrappedPublicKey, WrappedPublicKeyShare},
WrappedCiphertext,
};
use actix::Message;
use bincode;
use serde::{Deserialize, Serialize};
Expand Down Expand Up @@ -64,6 +67,10 @@ pub enum EnclaveEvent {
id: EventId,
data: PublicKeyAggregated,
},
DecryptionRequested {
id: EventId,
data: DecryptionRequested
}
// CommitteeSelected,
// OutputDecrypted,
// CiphernodeRegistered,
Expand All @@ -90,6 +97,7 @@ impl From<EnclaveEvent> for EventId {
EnclaveEvent::KeyshareCreated { id, .. } => id,
EnclaveEvent::ComputationRequested { id, .. } => id,
EnclaveEvent::PublicKeyAggregated { id, .. } => id,
EnclaveEvent::DecryptionRequested { id, .. } => id,
}
}
}
Expand Down Expand Up @@ -121,6 +129,18 @@ impl From<PublicKeyAggregated> for EnclaveEvent {
}
}


impl From<DecryptionRequested> for EnclaveEvent {
fn from(data: DecryptionRequested) -> Self {
EnclaveEvent::DecryptionRequested {
id: EventId::from(data.clone()),
data: data.clone(),
}
}
}



impl fmt::Display for EnclaveEvent {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&format!("{}({})", self.event_type(), self.get_id()))
Expand Down Expand Up @@ -154,6 +174,13 @@ pub struct ComputationRequested {
// availability_duration: ??, // TODO:
}

#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[rtype(result = "()")]
pub struct DecryptionRequested {
pub e3_id: E3id,
pub ciphertext: WrappedCiphertext,
}

fn extract_enclave_event_name(s: &str) -> &str {
let bytes = s.as_bytes();
for (i, &item) in bytes.iter().enumerate() {
Expand Down
59 changes: 56 additions & 3 deletions packages/ciphernode/core/src/fhe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use crate::ordered_set::OrderedSet;
use actix::{Actor, Context, Handler, Message};
use anyhow::*;
use fhe::{
bfv::{BfvParameters, BfvParametersBuilder, PublicKey, SecretKey},
bfv::{BfvParameters, BfvParametersBuilder, Ciphertext, PublicKey, SecretKey},
mbfv::{AggregateIter, CommonRandomPoly, PublicKeyShare},
};
use fhe_traits::{Deserialize, DeserializeParametrized, Serialize};
Expand Down Expand Up @@ -123,7 +123,6 @@ impl serde::Serialize for WrappedPublicKeyShare {
}
}


/// Wrapped PublicKey. This is wrapped to provide an inflection point
/// as we use this library elsewhere we only implement traits as we need them
/// and avoid exposing underlying structures from fhe.rs
Expand Down Expand Up @@ -214,6 +213,60 @@ impl WrappedSecretKey {
}
}

/// Wrapped Ciphertext. This is wrapped to provide an inflection point
/// as we use this library elsewhere we only implement traits as we need them
/// and avoid exposing underlying structures from fhe.rs
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WrappedCiphertext {
inner: Ciphertext,
params: Arc<BfvParameters>,
}

impl WrappedCiphertext {
pub fn from_fhe_rs(inner: Ciphertext, params: Arc<BfvParameters>) -> Self {
Self { inner, params }
}
}

impl Hash for WrappedCiphertext {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.inner.to_bytes().hash(state)
}
}

/// Deserialize from serde to WrappedPublicKey
impl<'de> serde::Deserialize<'de> for WrappedCiphertext {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
// Intermediate struct of bytes for deserialization
#[derive(serde::Deserialize)]
struct DeserializedBytes {
par: Vec<u8>,
bytes: Vec<u8>,
}
let DeserializedBytes { par, bytes } = DeserializedBytes::deserialize(deserializer)?;
let params = Arc::new(BfvParameters::try_deserialize(&par).unwrap());
let inner = Ciphertext::from_bytes(&bytes, &params).map_err(serde::de::Error::custom)?;
std::result::Result::Ok(WrappedCiphertext::from_fhe_rs(inner, params))
}
}
impl serde::Serialize for WrappedCiphertext {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
use serde::ser::SerializeStruct;
let bytes = self.inner.to_bytes();
let par_bytes = self.params.to_bytes();
// Intermediate struct of bytes
let mut state = serializer.serialize_struct("Ciphertext", 2)?;
state.serialize_field("par_bytes", &par_bytes)?;
state.serialize_field("bytes", &bytes)?;
state.end()
}
}
/// Fhe library adaptor. All FHE computations should happen through this actor.
pub struct Fhe {
params: Arc<BfvParameters>,
Expand Down Expand Up @@ -274,7 +327,7 @@ impl Handler<GetAggregatePublicKey> for Fhe {
}
}

fn serialize_box_i64(boxed: Box<[i64]>) -> Vec<u8> {
pub fn serialize_box_i64(boxed: Box<[i64]>) -> Vec<u8> {
let vec = boxed.into_vec();
let mut bytes = Vec::with_capacity(vec.len() * mem::size_of::<i64>());
for &num in &vec {
Expand Down
57 changes: 42 additions & 15 deletions packages/ciphernode/core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,31 +10,31 @@ mod enclave_contract;
mod eventbus;
mod events;
mod fhe;
mod logger;
mod ordered_set;
mod p2p;
mod logger;

// TODO: this is too permissive
pub use data::*;
pub use actix::prelude::*;
pub use ciphernode::*;
pub use committee::*;
pub use committee_key::*;
pub use data::*;
pub use eventbus::*;
pub use events::*;
pub use fhe::*;
pub use p2p::*;
pub use actix::prelude::*;
pub use logger::*;
pub use p2p::*;

pub use data::*;
pub use actix::prelude::*;
pub use ciphernode::*;
pub use committee::*;
pub use committee_key::*;
pub use data::*;
pub use eventbus::*;
pub use events::*;
pub use fhe::*;
pub use p2p::*;
pub use actix::prelude::*;

// pub struct Core {
// pub name: String,
Expand All @@ -53,6 +53,7 @@ pub use actix::prelude::*;
// }
// }

// TODO: move these out to a test folder
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
Expand All @@ -65,13 +66,15 @@ mod tests {
events::{ComputationRequested, E3id, EnclaveEvent, KeyshareCreated, PublicKeyAggregated},
fhe::{Fhe, WrappedPublicKey, WrappedPublicKeyShare},
p2p::P2p,
DecryptionRequested, ResetHistory, WrappedCiphertext,
};
use actix::prelude::*;
use anyhow::*;
use fhe::{
bfv::{BfvParameters, BfvParametersBuilder, PublicKey, SecretKey},
bfv::{BfvParameters, BfvParametersBuilder, Encoding, Plaintext, PublicKey, SecretKey},
mbfv::{AggregateIter, CommonRandomPoly, PublicKeyShare},
};
use fhe_traits::{FheEncoder, FheEncrypter};
use rand::SeedableRng;
use rand_chacha::ChaCha20Rng;
use tokio::sync::Mutex;
Expand Down Expand Up @@ -144,24 +147,29 @@ mod tests {
plaintext_modulus: u64,
rng1: ChaCha20Rng,
rng2: ChaCha20Rng,
) -> Result<Addr<Fhe>> {
) -> Result<(Addr<Fhe>, Arc<BfvParameters>, CommonRandomPoly)> {
let (params, crp) = setup_bfv_params(&moduli, degree, plaintext_modulus, rng1)?;
Ok(Fhe::new(params, crp, rng2)?.start())
Ok((
Fhe::new(params.clone(), crp.clone(), rng2)?.start(),
params,
crp,
))
}

#[actix::test]
async fn test_public_key_aggregation() -> Result<()> {
async fn test_public_key_aggregation_and_decryption() -> Result<()> {
// Setup EventBus
let bus = EventBus::new(true).start();

// Setup global FHE actor
let fhe = setup_global_fhe_actor(
let (fhe, ..) = setup_global_fhe_actor(
&vec![0x3FFFFFFF000001],
2048,
1032193,
ChaCha20Rng::seed_from_u64(42),
ChaCha20Rng::seed_from_u64(42),
)?;

setup_local_ciphernode(bus.clone(), fhe.clone(), true);
setup_local_ciphernode(bus.clone(), fhe.clone(), true);
setup_local_ciphernode(bus.clone(), fhe.clone(), true);
Expand Down Expand Up @@ -196,7 +204,7 @@ mod tests {
let (p2, rng) = generate_pk_share(params.clone(), crp.clone(), rng)?;
let (p3, _) = generate_pk_share(params.clone(), crp.clone(), rng)?;

let aggregated: PublicKey = vec![p1.clone(), p2.clone(), p3.clone()]
let pubkey: PublicKey = vec![p1.clone(), p2.clone(), p3.clone()]
.iter()
.map(|k| k.clone_inner())
.aggregate()?;
Expand Down Expand Up @@ -224,12 +232,27 @@ mod tests {
e3_id: e3_id.clone()
}),
EnclaveEvent::from(PublicKeyAggregated {
pubkey: WrappedPublicKey::from_fhe_rs(aggregated, params),
pubkey: WrappedPublicKey::from_fhe_rs(pubkey.clone(), params.clone()),
e3_id: e3_id.clone()
})
]
);

// Aggregate decryption
bus.send(ResetHistory).await?;

let yes = 12376213u64;
let no = 873827u64;

let pt = Plaintext::try_encode(&vec![yes, no], Encoding::poly(), &params)?;

let ciphertext = pubkey.try_encrypt(&pt, &mut ChaCha20Rng::seed_from_u64(42))?;

bus.do_send(EnclaveEvent::from(DecryptionRequested {
ciphertext: WrappedCiphertext::from_fhe_rs(ciphertext, params),
e3_id: e3_id.clone(),
}));

Ok(())
}

Expand Down Expand Up @@ -268,7 +291,7 @@ mod tests {

bus.do_send(evt_1.clone());
bus.do_send(evt_2.clone());

sleep(Duration::from_millis(1)).await; // need to push to next tick

// check the history of the event bus
Expand All @@ -280,7 +303,11 @@ mod tests {
"P2p did not transmit events to the network"
);

assert_eq!(history, vec![evt_1, evt_2], "P2p must not retransmit forwarded event to event bus");
assert_eq!(
history,
vec![evt_1, evt_2],
"P2p must not retransmit forwarded event to event bus"
);

Ok(())
}
Expand Down

0 comments on commit 8141bfb

Please sign in to comment.