feat: trbfv integration test [skip-line-limit]

crates/aggregator/src/threshold_plaintext_aggregator.rs

@@ -22,7 +22,7 @@ use e3_trbfv::{
     TrBFVConfig, TrBFVRequest,
 };
 use e3_utils::utility_types::ArcBytes;
-use tracing::{error, info, trace};
+use tracing::{debug, error, info, trace};
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct Collecting {
@@ -247,6 +247,12 @@ impl Handler<DecryptionshareCreated> for ThresholdPlaintextAggregator {
     type Result = ResponseActFuture<Self, Result<()>>;
 
     fn handle(&mut self, event: DecryptionshareCreated, _: &mut Self::Context) -> Self::Result {
+        let Some(ThresholdPlaintextAggregatorState::Collecting(Collecting { .. })) =
+            self.state.get()
+        else {
+            debug!(state=?self.state, "Aggregator has been closed for collecting so ignoring this event.");
+            return Box::pin(fut::ready(Ok(())));
+        };
         info!(event=?event, "Processing DecryptionShareCreated...");
         let address = event.node.clone();
         let party_id = event.party_id;

crates/bfv-helpers/Cargo.toml

@@ -15,8 +15,9 @@ rand.workspace = true
 anyhow.workspace = true
 fhe-util = { git = "https://github.com/gnosisguild/fhe.rs" }
 fhe-math = { git = "https://github.com/gnosisguild/fhe.rs" }
-greco = { package = "zkfhe-greco", git = "https://github.com/gnosisguild/zkfhe-generator"}
+greco = { package = "zkfhe-greco", git = "https://github.com/gnosisguild/zkfhe-generator" }
 num-bigint = { workspace = true }
+thiserror = { workspace = true }
 
 [dev-dependencies]
 hex.workspace = true

crates/bfv-helpers/src/client.rs

@@ -35,12 +35,12 @@ pub fn bfv_encrypt<T>(
     public_key: Vec<u8>,
     degree: usize,
     plaintext_modulus: u64,
-    moduli: [u64; 1],
+    moduli: &[u64],
 ) -> Result<Vec<u8>>
 where
     Plaintext: for<'a> FheEncoder<&'a T, Error = FheError>,
 {
-    let params = build_bfv_params_arc(degree, plaintext_modulus, &moduli, None);
+    let params = build_bfv_params_arc(degree, plaintext_modulus, moduli, None);
 
     let pk = PublicKey::from_bytes(&public_key, &params)
         .map_err(|e| anyhow!("Error deserializing public key:{e}"))?;
@@ -86,7 +86,7 @@ pub fn bfv_verifiable_encrypt<T>(
     public_key: Vec<u8>,
     degree: usize,
     plaintext_modulus: u64,
-    moduli: [u64; 1],
+    moduli: Vec<u64>,
 ) -> Result<VerifiableEncryptionResult>
 where
     Plaintext: for<'a> FheEncoder<&'a T, Error = FheError>,
@@ -146,7 +146,7 @@ mod tests {
 
         let num = [1u64];
         let encrypted_data =
-            bfv_encrypt(num, pk.to_bytes(), degree, plaintext_modulus, moduli).unwrap();
+            bfv_encrypt(num, pk.to_bytes(), degree, plaintext_modulus, &moduli).unwrap();
 
         let ct = Ciphertext::from_bytes(&encrypted_data, &params).unwrap();
         let pt = sk.try_decrypt(&ct).unwrap();
@@ -175,7 +175,7 @@ mod tests {
             pk.to_bytes(),
             degree,
             plaintext_modulus,
-            moduli,
+            &moduli,
         )
         .unwrap();
 
@@ -202,8 +202,14 @@ mod tests {
         let pk = PublicKey::new(&sk, &mut rng);
 
         let num = [1u64];
-        let encrypted_data =
-            bfv_verifiable_encrypt(num, pk.to_bytes(), degree, plaintext_modulus, moduli).unwrap();
+        let encrypted_data = bfv_verifiable_encrypt(
+            num,
+            pk.to_bytes(),
+            degree,
+            plaintext_modulus,
+            moduli.to_vec(),
+        )
+        .unwrap();
 
         let ct = Ciphertext::from_bytes(&encrypted_data.encrypted_data, &params).unwrap();
         let pt = sk.try_decrypt(&ct).unwrap();
@@ -232,7 +238,7 @@ mod tests {
             pk.to_bytes(),
             degree,
             plaintext_modulus,
-            moduli,
+            moduli.to_vec(),
         )
         .unwrap();
 

crates/bfv-helpers/src/lib.rs

@@ -9,8 +9,22 @@ mod util;
 
 use alloy_dyn_abi::{DynSolType, DynSolValue};
 use alloy_primitives::U256;
-use fhe::bfv::{BfvParameters, BfvParametersBuilder};
+use fhe::bfv::{BfvParameters, BfvParametersBuilder, Encoding, Plaintext};
+use fhe_traits::FheDecoder;
 use std::sync::Arc;
+use thiserror::Error as ThisError;
+
+#[derive(ThisError, Debug)]
+pub enum Error {
+    #[error("Plaintext decoding failed")]
+    PlaintextDecodeFailed,
+    // TODO: add errors from client.rs
+    #[error("Input was not encoded correctly")]
+    BadEncoding,
+}
+
+/// Result that returns a type T or a BfvHelpersError
+type Result<T> = std::result::Result<T, Error>;
 
 /// Predefined BFV parameters for common use cases
 pub mod params {
@@ -337,6 +351,35 @@ pub fn decode_bfv_params_arc(bytes: &[u8]) -> Arc<BfvParameters> {
     Arc::new(decode_bfv_params(bytes))
 }
 
+/// Decode Plaintext to a Vec<u64>
+pub fn decode_plaintext_to_vec_u64(value: &Plaintext) -> Result<Vec<u64>> {
+    let decoded = Vec::<u64>::try_decode(&value, Encoding::poly())
+        .map_err(|_| Error::PlaintextDecodeFailed)?;
+
+    Ok(decoded)
+}
+
+/// Convert from a Vec<u64> to Vec<u8>
+pub fn encode_vec_u64_to_bytes(value: &[u64]) -> Vec<u8> {
+    let mut bytes = Vec::new();
+    for num in &value.to_vec() {
+        bytes.extend_from_slice(&num.to_le_bytes());
+    }
+    bytes
+}
+
+/// Decode bytes to Vec<u64>
+pub fn decode_bytes_to_vec_u64(bytes: &[u8]) -> Result<Vec<u64>> {
+    if bytes.len() % 8 != 0 {
+        return Err(Error::BadEncoding);
+    }
+
+    Ok(bytes
+        .chunks_exact(8)
+        .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap()))
+        .collect())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/cli/src/cli.rs

@@ -61,9 +61,9 @@ pub struct Cli {
     #[arg(long = "otel", global = true)]
     pub otel: Option<ValidUrl>,
 
-    // TODO: expose this as a feature once we have this hooked up to the bin integration test
-    #[arg(long, hide = true)]
-    pub experimental_trbfv: Option<String>,
+    /// Enable the experimental TrBFV threshold feature.
+    #[arg(long, global = true)]
+    pub experimental_trbfv: bool,
 }
 
 impl Cli {
@@ -158,7 +158,9 @@ impl Cli {
         }
 
         match self.command {
-            Commands::Start { peers } => start::execute(config, peers).await?,
+            Commands::Start { peers } => {
+                start::execute(config, peers, self.experimental_trbfv).await?
+            }
             Commands::Init { .. } => {
                 bail!("Cannot run `enclave init` when a configuration exists.");
             }
@@ -180,6 +182,7 @@ impl Cli {
                     self.verbose,
                     self.config,
                     self.otel.clone().map(Into::into),
+                    self.experimental_trbfv,
                 )
                 .await?
             }

crates/cli/src/nodes.rs

@@ -78,15 +78,33 @@ pub async fn execute(
     verbose: u8,
     config_string: Option<String>,
     otel: Option<String>,
+    experimental_trbfv: bool,
 ) -> Result<()> {
     match command {
         NodeCommands::Up { detach, exclude } => {
-            nodes_up::execute(config, detach, exclude, verbose, config_string, otel).await?
+            nodes_up::execute(
+                config,
+                detach,
+                exclude,
+                verbose,
+                config_string,
+                otel,
+                experimental_trbfv,
+            )
+            .await?
         }
         NodeCommands::Down => nodes_down::execute().await?,
         NodeCommands::Ps => nodes_ps::execute().await?,
         NodeCommands::Daemon { exclude } => {
-            nodes_daemon::execute(config, exclude, verbose, config_string, otel).await?
+            nodes_daemon::execute(
+                config,
+                exclude,
+                verbose,
+                config_string,
+                otel,
+                experimental_trbfv,
+            )
+            .await?
         }
         NodeCommands::Start { id } => nodes_start::execute(&id).await?,
         NodeCommands::Status { id } => nodes_status::execute(&id).await?,

crates/cli/src/nodes_daemon.rs

@@ -14,6 +14,15 @@ pub async fn execute(
     verbose: u8,
     config_string: Option<String>,
     otel: Option<String>,
+    experimental_trbfv: bool,
 ) -> Result<()> {
-    daemon::execute(config, exclude, verbose, config_string, otel).await
+    daemon::execute(
+        config,
+        exclude,
+        verbose,
+        config_string,
+        otel,
+        experimental_trbfv,
+    )
+    .await
 }

crates/cli/src/nodes_up.rs

@@ -15,6 +15,16 @@ pub async fn execute(
     verbose: u8,
     config_string: Option<String>,
     otel: Option<String>,
+    experimental_trbfv: bool,
 ) -> Result<()> {
-    up::execute(config, detach, exclude, verbose, config_string, otel).await
+    up::execute(
+        config,
+        detach,
+        exclude,
+        verbose,
+        config_string,
+        otel,
+        experimental_trbfv,
+    )
+    .await
 }

crates/cli/src/start.rs

@@ -11,7 +11,11 @@ use e3_entrypoint::helpers::listen_for_shutdown;
 use tracing::{info, instrument};
 
 #[instrument(skip_all)]
-pub async fn execute(mut config: AppConfig, peers: Vec<String>) -> Result<()> {
+pub async fn execute(
+    mut config: AppConfig,
+    peers: Vec<String>,
+    experimental_trbfv: bool,
+) -> Result<()> {
     owo();
 
     let Some(address) = config.address() else {
@@ -31,12 +35,15 @@ pub async fn execute(mut config: AppConfig, peers: Vec<String>) -> Result<()> {
                 &config,
                 pubkey_write_path,
                 plaintext_write_path,
+                experimental_trbfv,
             )
             .await?
         }
 
         // Launch in ciphernode configuration
-        NodeRole::Ciphernode => e3_entrypoint::start::start::execute(&config, address).await?,
+        NodeRole::Ciphernode => {
+            e3_entrypoint::start::start::execute(&config, address, experimental_trbfv).await?
+        }
     };
 
     info!(

crates/entrypoint/src/nodes/daemon.rs

@@ -53,6 +53,7 @@ impl LaunchCommand {
         verbose: u8,
         maybe_config_string: &Option<String>,
         maybe_otel: &Option<String>,
+        experimental_trbfv: bool,
     ) -> Result<CommandParams> {
         let enclave_bin = env::current_exe()?.display().to_string();
         let mut args = vec![];
@@ -80,6 +81,10 @@ impl LaunchCommand {
             args.push(peer.to_string());
         }
 
+        if experimental_trbfv {
+            args.push("--experimental-trbfv".to_string());
+        }
+
         Ok((enclave_bin, args))
     }
 }
@@ -91,6 +96,7 @@ fn extract_commands(
     verbose: u8,
     maybe_config_string: Option<String>,
     maybe_otel: Option<String>,
+    experimental_trbfv: bool,
 ) -> Result<CommandMap> {
     let mut exclude_list = exclude.clone();
 
@@ -111,7 +117,12 @@ fn extract_commands(
 
     let mut cmds = HashMap::new();
     for item in filtered.iter() {
-        let params = item.to_params(verbose, &maybe_config_string, &maybe_otel)?;
+        let params = item.to_params(
+            verbose,
+            &maybe_config_string,
+            &maybe_otel,
+            experimental_trbfv,
+        )?;
         cmds.insert(item.name.clone(), params);
     }
 
@@ -125,6 +136,7 @@ pub async fn execute(
     verbose: u8,
     maybe_config_string: Option<String>,
     maybe_otel: Option<String>,
+    experimental_trbfv: bool,
 ) -> Result<()> {
     let command_map = extract_commands(
         config.nodes(),
@@ -133,6 +145,7 @@ pub async fn execute(
         verbose,
         maybe_config_string,
         maybe_otel,
+        experimental_trbfv,
     )?;
 
     let process_manager = Arc::new(Mutex::new(ProcessManager::from(command_map)));