Multithreaded Decryption

ipa-core/src/cli/crypto/hybrid_decrypt.rs

@@ -14,7 +14,10 @@ use crate::{
         U128Conversions,
     },
     hpke::{KeyRegistry, PrivateKeyOnly},
-    report::hybrid::{EncryptedHybridReport, HybridReport},
+    report::{
+        hybrid::{EncryptedHybridReport, HybridReport},
+        hybrid_info::HybridInfo,
+    },
     test_fixture::Reconstruct,
 };
 
@@ -104,7 +107,7 @@ impl HybridDecryptArgs {
         for (dec_report1, (dec_report2, dec_report3)) in
             decrypted_reports1.zip(decrypted_reports2.zip(decrypted_reports3))
         {
-            match (dec_report1, dec_report2, dec_report3) {
+            match (dec_report1.0, dec_report2.0, dec_report3.0) {
                 (
                     HybridReport::Impression(impression_report1),
                     HybridReport::Impression(impression_report2),
@@ -125,7 +128,7 @@ impl HybridDecryptArgs {
                     ]
                     .reconstruct()
                     .as_u128();
-                    let key_id = impression_report1.info.key_id;
+                    let key_id = dec_report1.1.impression.key_id;
 
                     writeln!(writer, "i,{match_key},{breakdown_key},{key_id}")?;
                 }
@@ -149,11 +152,13 @@ impl HybridDecryptArgs {
                     ]
                     .reconstruct()
                     .as_u128();
-                    let key_id = conversion_report1.info.key_id;
-                    let conversion_site_domain = conversion_report1.info.conversion_site_domain;
-                    let timestamp = conversion_report1.info.timestamp;
-                    let epsilon = conversion_report1.info.epsilon;
-                    let sensitivity = conversion_report1.info.sensitivity;
+
+                    let key_id = dec_report1.1.conversion.key_id;
+                    let conversion_site_domain = dec_report1.1.conversion.conversion_site_domain;
+                    let timestamp = dec_report1.1.conversion.timestamp;
+                    let epsilon = dec_report1.1.conversion.epsilon;
+                    let sensitivity = dec_report1.1.conversion.sensitivity;
+
                     writeln!(writer, "c,{match_key},{value},{key_id},{conversion_site_domain},{timestamp},{epsilon},{sensitivity}")?;
                 }
                 _ => {
@@ -172,17 +177,18 @@ struct DecryptedHybridReports {
 }
 
 impl Iterator for DecryptedHybridReports {
-    type Item = HybridReport<BA8, BA3>;
+    type Item = (HybridReport<BA8, BA3>, HybridInfo);
 
     fn next(&mut self) -> Option<Self::Item> {
         let mut line = String::new();
         if self.reader.read_line(&mut line).unwrap() > 0 {
             let encrypted_report_bytes = hex::decode(line.trim()).unwrap();
             let enc_report =
                 EncryptedHybridReport::from_bytes(encrypted_report_bytes.into()).unwrap();
-            let dec_report: HybridReport<BA8, BA3> =
-                enc_report.decrypt(&self.key_registry).unwrap();
-            Some(dec_report)
+            let (dec_report, info) = enc_report
+                .decrypt_and_return_info(&self.key_registry)
+                .unwrap();
+            Some((dec_report, info))
         } else {
             None
         }

ipa-core/src/cli/crypto/hybrid_encrypt.rs

@@ -21,14 +21,22 @@ use crate::{
     config::{KeyRegistries, NetworkConfig},
     error::BoxError,
     hpke::{KeyRegistry, PublicKeyOnly},
-    report::hybrid::{HybridReport, DEFAULT_KEY_ID},
+    report::{
+        hybrid::{HybridReport, DEFAULT_KEY_ID},
+        hybrid_info::HybridInfo,
+    },
     secret_sharing::IntoShares,
     test_fixture::hybrid::TestHybridRecord,
 };
 
-/// Encryptor takes 3 arguments: `report_id`, helper that the shares must be encrypted towards
-/// and the actual share ([`HybridReport`]) to encrypt.
-type EncryptorInput = (usize, usize, HybridReport<BreakdownKey, TriggerValue>);
+/// Encryptor takes 4 arguments: `report_id`, helper that the shares must be encrypted towards,
+/// AAD info, and the actual share ([`HybridReport`]) to encrypt.
+type EncryptorInput = (
+    usize,
+    usize,
+    HybridReport<BreakdownKey, TriggerValue>,
+    HybridInfo,
+);
 /// Encryptor sends report id and encrypted bytes down to file worker to write those bytes
 /// down
 type EncryptorOutput = (usize, Vec<u8>);
@@ -91,7 +99,8 @@ impl HybridEncryptArgs {
 
         let mut worker_pool = ReportWriter::new(key_registries, &self.output_dir);
         for (report_id, record) in input.iter::<TestHybridRecord>().enumerate() {
-            worker_pool.submit(report_id, record.share())?;
+            let info = record.create_hybrid_info();
+            worker_pool.submit(report_id, record.share(), &info)?;
         }
 
         worker_pool.join()?;
@@ -130,11 +139,12 @@ impl EncryptorPool {
                         std::thread::Builder::new()
                             .name(format!("encryptor-{i}"))
                             .spawn(move || {
-                                for (i, helper_id, report) in rx {
+                                for (i, helper_id, report, info) in rx {
                                     let key_registry = &key_registries[helper_id];
                                     let output = report.encrypt(
                                         DEFAULT_KEY_ID,
                                         key_registry,
+                                        &info,
                                         &mut thread_rng(),
                                     )?;
                                     file_writer[helper_id].send((i, output))?;
@@ -206,9 +216,12 @@ impl ReportWriter {
         &mut self,
         report_id: usize,
         shares: [HybridReport<BreakdownKey, TriggerValue>; 3],
+        info: &HybridInfo,
     ) -> UnitResult {
         for (i, share) in shares.into_iter().enumerate() {
-            self.encryptor_pool.encrypt_share((report_id, i, share))?;
+            //  todo: maybe a smart pointer to avoid cloning
+            self.encryptor_pool
+                .encrypt_share((report_id, i, share, info.clone()))?;
         }
 
         Ok(())

ipa-core/src/query/runner/hybrid.rs

@@ -46,6 +46,7 @@ use crate::{
         replicated::semi_honest::AdditiveShare as Replicated, BitDecomposed, TransposeFrom,
         Vectorizable,
     },
+    seq_join::seq_join,
     sharding::{ShardConfiguration, Sharded},
 };
 
@@ -130,10 +131,12 @@ where
                 }))
             })
             .try_flatten()
-            .take(sz);
+            .take(sz)
+            .map(|v| async move { v });
+
         let (decrypted_reports, resharded_tags) = reshard_aad(
             ctx.narrow(&HybridStep::ReshardByTag),
-            stream,
+            seq_join(ctx.active_work(), stream),
             |ctx, _, tag| tag.shard_picker(ctx.shard_count()),
         )
         .await?;
@@ -242,8 +245,15 @@ mod tests {
         let shares: [Vec<HybridReport<BA8, BA3>>; 3] = records.iter().cloned().share();
         for (buf, shares) in zip(&mut buffers, shares) {
             for (i, share) in shares.into_iter().enumerate() {
+                let info = records[i].create_hybrid_info();
                 share
-                    .delimited_encrypt_to(key_id, key_registry.as_ref(), &mut rng, &mut buf[i % s])
+                    .delimited_encrypt_to(
+                        key_id,
+                        key_registry.as_ref(),
+                        &info,
+                        &mut rng,
+                        &mut buf[i % s],
+                    )
                     .unwrap();
             }
         }