WIP: Parallelize bottlenecks

tcerqueira · tcerqueira · commit e028124241c9 · 2024-10-11T19:07:41.000+01:00
diff --git a/mix-node/src/crypto.rs b/mix-node/src/crypto.rs
@@ -1,5 +1,4 @@
 use anyhow::Context;
-use bitvec::vec::BitVec;
 use elastic_elgamal::{
     group::Ristretto,
     sharing::{ActiveParticipant, PublicKeySet},
@@ -11,6 +10,7 @@ use std::sync::LazyLock;
 use thiserror::Error;
 
 pub type Ciphertext = elastic_elgamal::Ciphertext<Ristretto>;
+pub type Bits = bitvec::vec::BitVec;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DecryptionShare {
@@ -52,6 +52,26 @@ pub fn remix(
     Ok(())
 }
 
+pub fn encrypt(pub_key: &PublicKey<Ristretto>, bits: &Bits) -> Vec<Ciphertext> {
+    bits.as_raw_slice()
+        .into_par_iter()
+        .enumerate()
+        .flat_map(|(chunk_idx, &chunk)| {
+            let chunk_size = 8 * std::mem::size_of_val(&chunk);
+            let start_bit = chunk_idx * chunk_size;
+            let end_bit = std::cmp::min(start_bit + chunk_size, bits.len());
+
+            (0..end_bit - start_bit)
+                .into_par_iter()
+                .map(move |bit_offset| {
+                    let mut rng = rand::thread_rng();
+                    let bit = (chunk >> bit_offset) & 1;
+                    pub_key.encrypt(bit as u64, &mut rng)
+                })
+        })
+        .collect::<Vec<_>>()
+}
+
 pub fn decryption_share_for(
     active_participant: &ActiveParticipant<Ristretto>,
     ciphertext: &[Ciphertext],
@@ -66,47 +86,46 @@ pub fn decryption_share_for(
     DecryptionShare::new(active_participant.index(), share)
 }
 
-// PERF: parallelize and maybe async this
+// PERF: parallelize this
 pub fn decrypt_shares(
     key_set: &PublicKeySet<Ristretto>,
     enc: &[Ciphertext],
     shares: &[DecryptionShare],
-) -> anyhow::Result<BitVec> {
+) -> anyhow::Result<Bits> {
     if shares.iter().any(|s| s.share.len() != enc.len()) {
-        return Err(anyhow::anyhow!(
-            "mismatch of lengths between encrypted ciphertext a decryption shares"
-        ));
+        anyhow::bail!("mismatch of lengths between encrypted ciphertext a decryption shares");
     }
     // Transpose vectors
-    let rows = shares.len();
-    let cols = enc.len();
-    let transposed = (0..cols).map(|col| {
-        (0..rows)
-            .map(|row| (shares[row].index, shares[row].share[col]))
-            .collect::<Vec<_>>()
+    let transposed = (0..enc.len()).into_par_iter().map(|ct_idx| {
+        shares
+            .into_par_iter()
+            .map(move |s| (s.index, s.share[ct_idx]))
     });
 
     transposed
         .zip(enc)
         .map(|(shares, enc)| {
-            let dec_iter = shares.into_iter().filter_map(|(i, (share, proof))| {
-                let share = CandidateDecryption::from_bytes(&share.to_bytes())?;
-                let verification = key_set.verify_share(share, *enc, i, &proof).ok()?;
-                Some((i, verification))
-            });
+            let dec_iter: Vec<_> = shares
+                .filter_map(|(i, (share, proof))| {
+                    let share = CandidateDecryption::from_bytes(&share.to_bytes())?;
+                    let verification = key_set.verify_share(share, *enc, i, &proof).ok()?;
+                    Some((i, verification))
+                })
+                .collect();
             let combined = key_set
                 .params()
-                .combine_shares(dec_iter)
+                .combine_shares(dec_iter.into_iter())
                 .context("failed to combine shares")?;
             Ok(combined
                 .decrypt(*enc, &LOOKUP_TABLE)
                 .context("decrypted values out of range of lookup table")?
                 == 1u64)
         })
-        .collect::<anyhow::Result<BitVec>>()
+        .collect::<anyhow::Result<Vec<_>>>()
+        .map(Bits::from_iter)
 }
 
-pub fn hamming_distance(x_code: BitVec, y_code: BitVec) -> usize {
+pub fn hamming_distance(x_code: Bits, y_code: Bits) -> usize {
     // Q: What if x and y are different sizes?
     (x_code ^ y_code).count_ones()
 }
diff --git a/mix-node/src/rest/routes.rs b/mix-node/src/rest/routes.rs
@@ -1,13 +1,11 @@
 use super::error::Error;
 use crate::{
-    crypto::{self, Ciphertext, CryptoError, DecryptionShare},
+    crypto::{self, Bits, Ciphertext, CryptoError, DecryptionShare},
     rokio, AppState, EncryptedCodes,
 };
 use anyhow::Context;
 use axum::{extract::State, response::Json};
-use bitvec::vec::BitVec;
 use elastic_elgamal::{group::Ristretto, sharing::PublicKeySet};
-use rayon::prelude::*;
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
@@ -45,21 +43,14 @@ pub async fn public_key_set(State(state): State<Arc<AppState>>) -> Json<PublicKe
     Json(state.pub_key_set().clone())
 }
 
-#[tracing::instrument(skip(state, plaintext))]
+#[tracing::instrument(skip(state, bits))]
 pub async fn encrypt(
     State(state): State<Arc<AppState>>,
-    Json(plaintext): Json<BitVec>,
+    Json(bits): Json<Bits>,
 ) -> Json<Vec<Ciphertext>> {
     let ciphertexts = rokio::spawn(move || {
         let pub_key = state.crypto.active_participant.key_set().shared_key();
-        plaintext
-            .into_iter()
-            .par_bridge()
-            .map(|msg| {
-                let mut rng = rand::thread_rng();
-                pub_key.encrypt(msg as u64, &mut rng)
-            })
-            .collect::<Vec<_>>()
+        crypto::encrypt(pub_key, &bits)
     })
     .await;
 
@@ -127,13 +118,38 @@ pub async fn hamming_distance(
     let (x_code, y_code) = (Arc::new(x_code), Arc::new(y_code));
 
     // Decrypt
-    let (x_shares, y_shares) = tokio::join!(
-        request_all_shares(&x_code, &state),
-        request_all_shares(&y_code, &state)
-    );
-    // PERF: parallelize and rokio
-    let x_decrypt = crypto::decrypt_shares(state.pub_key_set(), &x_code, &x_shares)?;
-    let y_decrypt = crypto::decrypt_shares(state.pub_key_set(), &y_code, &y_shares)?;
+    let (x_shares, y_shares) = {
+        let (x_state, y_state) = (Arc::clone(&state), Arc::clone(&state));
+        let (x_inner_code, y_inner_code) = (Arc::clone(&x_code), Arc::clone(&y_code));
+
+        let (mut x_shares, mut y_shares, x_self_share, y_self_share) = tokio::join!(
+            request_all_shares(&x_code, &state),
+            request_all_shares(&y_code, &state),
+            rokio::spawn(move || {
+                crypto::decryption_share_for(&x_state.crypto.active_participant, &x_inner_code)
+            }),
+            rokio::spawn(move || {
+                crypto::decryption_share_for(&y_state.crypto.active_participant, &y_inner_code)
+            })
+        );
+        x_shares.push(x_self_share);
+        y_shares.push(y_self_share);
+
+        (x_shares, y_shares)
+    };
+
+    let x_decrypt = {
+        let state = Arc::clone(&state);
+        rokio::spawn(move || crypto::decrypt_shares(state.pub_key_set(), &x_code, &x_shares))
+    };
+    let y_decrypt = {
+        let state = Arc::clone(&state);
+        rokio::spawn(move || crypto::decrypt_shares(state.pub_key_set(), &y_code, &y_shares))
+    };
+    let (x_decrypt, y_decrypt) = {
+        let (x_decrypt, y_decrypt) = tokio::join!(x_decrypt, y_decrypt);
+        (x_decrypt?, y_decrypt?)
+    };
 
     // Hamming
     let hamming_distance = crypto::hamming_distance(x_decrypt, y_decrypt);
diff --git a/mix-node/tests/common/mod.rs b/mix-node/tests/common/mod.rs
@@ -1,12 +1,16 @@
+#![allow(dead_code)]
+
 use bitvec::prelude::*;
 use elastic_elgamal::{group::Ristretto, DiscreteLogTable, Keypair, PublicKey, SecretKey};
-use mix_node::{crypto::Ciphertext, EncryptedCodes};
+use mix_node::{
+    crypto::{self, Bits, Ciphertext},
+    EncryptedCodes,
+};
 use rand::{CryptoRng, Rng};
 use std::iter;
 
 pub const N_BITS: usize = mix_node::N_BITS / 2;
 
-#[allow(unused)]
 pub fn set_up_payload() -> (EncryptedCodes, Keypair<Ristretto>) {
     let mut rng = rand::thread_rng();
     let new_iris_code = BitVec::<_, Lsb0>::from_slice(&rng.gen::<[u8; N_BITS / 8]>());
@@ -20,11 +24,10 @@ pub fn set_up_payload() -> (EncryptedCodes, Keypair<Ristretto>) {
 
     // Encrypt
     let receiver = Keypair::generate(&mut rng);
-    let dec_key = receiver.secret().clone();
     let enc_key = receiver.public().clone();
 
-    let enc_new_user: Vec<_> = encrypt_bits(&new_user[..], &enc_key, &mut rng).collect();
-    let enc_archived_user: Vec<_> = encrypt_bits(&archived_user[..], &enc_key, &mut rng).collect();
+    let enc_new_user: Vec<_> = crypto::encrypt(receiver.public(), &new_user);
+    let enc_archived_user: Vec<_> = crypto::encrypt(receiver.public(), &archived_user);
 
     (
         EncryptedCodes {
@@ -36,7 +39,11 @@ pub fn set_up_payload() -> (EncryptedCodes, Keypair<Ristretto>) {
     )
 }
 
-#[allow(unused)]
+pub fn set_up_iris_code(size: usize) -> Bits {
+    let mut rng = rand::thread_rng();
+    (0..size).map(|_| rng.gen::<bool>()).collect::<Bits>()
+}
+
 pub fn encode_bits<T: BitStore, O: BitOrder>(
     bits: &BitSlice<T, O>,
 ) -> impl Iterator<Item = bool> + '_ {
@@ -49,7 +56,6 @@ pub fn encode_bits<T: BitStore, O: BitOrder>(
     })
 }
 
-#[allow(unused)]
 pub fn decode_bits<T: BitStore, O: BitOrder>(
     bits: &BitSlice<T, O>,
 ) -> impl Iterator<Item = bool> + '_ {
@@ -63,7 +69,6 @@ pub fn decode_bits<T: BitStore, O: BitOrder>(
     })
 }
 
-#[allow(unused)]
 pub fn encrypt_bits<'a, T: BitStore, O: BitOrder>(
     bits: &'a BitSlice<T, O>,
     ek: &'a PublicKey<Ristretto>,
@@ -72,7 +77,6 @@ pub fn encrypt_bits<'a, T: BitStore, O: BitOrder>(
     bits.iter().map(|bit| ek.encrypt(*bit as u32, rng))
 }
 
-#[allow(unused)]
 pub fn decrypt_bits<'a>(
     ct: &'a [Ciphertext],
     pk: &'a SecretKey<Ristretto>,
diff --git a/mix-node/tests/web_integration.rs b/mix-node/tests/web_integration.rs
@@ -5,7 +5,7 @@ use elastic_elgamal::{group::Ristretto, sharing::PublicKeySet, Ciphertext};
 use format as f;
 use mix_node::{
     config::get_configuration,
-    crypto::{self, DecryptionShare},
+    crypto::{self, Bits, DecryptionShare},
     rest::routes::HammingResponse,
     test_helpers::{self, TestApp},
     EncryptedCodes,
@@ -41,9 +41,8 @@ async fn test_mix_node() -> anyhow::Result<()> {
         .await?;
 
     // Decrypt
-    let dec_new_user: BitVec<u8, Lsb0> =
-        common::decrypt_bits(&enc_new_user, receiver.secret()).collect();
-    let dec_archived_user: BitVec<u8, Lsb0> =
+    let dec_new_user: Bits = common::decrypt_bits(&enc_new_user, receiver.secret()).collect();
+    let dec_archived_user: Bits =
         common::decrypt_bits(&enc_archived_user, receiver.secret()).collect();
 
     // Assert result
@@ -63,7 +62,9 @@ async fn test_mix_node_bad_request() -> anyhow::Result<()> {
     config.application.auth_token = None;
     let TestApp { port, .. } = test_helpers::create_app(config).await;
 
-    let (mut codes, _dec_key) = common::set_up_payload();
+    // let code = common::set_up_iris_code(mix_node::N_BITS);
+
+    let (mut codes, _receiver) = common::set_up_payload();
     // Remove elements to cause a size mismatch
     codes.x_code.pop();
     codes.x_code.pop();
@@ -107,9 +108,9 @@ async fn test_mix_node_authorized() -> anyhow::Result<()> {
     config.application.auth_token = Some(Secret::new(auth_token.to_string()));
     let TestApp { port, .. } = test_helpers::create_app(config).await;
 
-    let (codes, _dec_key) = common::set_up_payload();
+    let (codes, _receiver) = common::set_up_payload();
 
-    // Bad request + Serialization
+    // Auth
     let client = reqwest::Client::new();
     let response = client
         .post(format!("http://localhost:{port}/remix"))
@@ -136,7 +137,6 @@ async fn test_mix_node_public_key() -> anyhow::Result<()> {
     assert_eq!(response.status(), StatusCode::OK);
 
     let _body: PublicKeySet<Ristretto> = response.json().await?;
-
     Ok(())
 }
 
@@ -145,7 +145,7 @@ async fn test_mix_node_encrypt() -> anyhow::Result<()> {
     let config = get_configuration()?;
     let TestApp { port, .. } = test_helpers::create_app(config).await;
 
-    let payload = bitvec![0, 1, 0, 1, 1, 0, 0, 1];
+    let payload = common::set_up_iris_code(mix_node::N_BITS);
 
     let client = reqwest::Client::new();
     let response = client
@@ -173,14 +173,17 @@ async fn test_network_decrypt_shares() -> anyhow::Result<()> {
         .await?;
     assert_eq!(response.status(), StatusCode::OK);
 
-    // Encrypt client side
-    let mut rng = rand::thread_rng();
-    let payload = bitvec![0, 1, 0, 1, 1, 0, 0, 1];
+    // Encrypt server side
+    let payload: Bits = bitvec![0, 1, 0, 1, 1, 0, 0, 1];
+
     let pub_key: PublicKeySet<Ristretto> = response.json().await?;
-    let encrypted: Vec<_> = payload
-        .iter()
-        .map(|pt| pub_key.shared_key().encrypt(*pt as u64, &mut rng))
-        .collect();
+    let encrypted: Vec<_> = client
+        .post(f!("http://localhost:{}/encrypt", nodes[0].port))
+        .json(&payload)
+        .send()
+        .await?
+        .json()
+        .await?;
 
     // Decrypt
     let mut shares = vec![];
@@ -207,7 +210,6 @@ async fn test_network_hamming_distance() -> anyhow::Result<()> {
     let [TestApp { port, .. }, ..] = test_helpers::create_network(3, 2).await[..] else {
         return Err(anyhow::anyhow!("there's not at least one node"));
     };
-    let mut rng = rand::thread_rng();
 
     // Request public key
     let client = reqwest::Client::new();
@@ -220,10 +222,10 @@ async fn test_network_hamming_distance() -> anyhow::Result<()> {
 
     // Codes are the same, expected hamming distance = 0
     let payload = {
-        let code: Vec<_> = [0, 1, 0, 1, 1, 0, 0, 1u64]
-            .into_iter()
-            .map(|m| pub_key.shared_key().encrypt(m, &mut rng))
-            .collect();
+        let code = crypto::encrypt(
+            pub_key.shared_key(),
+            &common::set_up_iris_code(mix_node::N_BITS),
+        );
         EncryptedCodes {
             x_code: code.clone(),
             y_code: code,
