Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Prepare for decryption #24

Merged
merged 2 commits into from
Aug 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -85,3 +85,18 @@ async fn on_computation_requested(

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

// TODO: Complete this

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
Loading