Add malicious security support to oneshot bench

private-attribution · Nov 6, 2024 · a496811 · a496811
1 parent 52c58eb
commit a496811
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 32 deletions.
diff --git a/ipa-core/benches/oneshot/ipa.rs b/ipa-core/benches/oneshot/ipa.rs
@@ -48,7 +48,7 @@ struct Args {
     #[arg(short = 'c', long, default_value = "8")]
     per_user_cap: u32,
     /// The number of breakdown keys.
-    #[arg(short = 'b', long, default_value = "16")]
+    #[arg(short = 'b', long, default_value = "32")]
     breakdown_keys: u32,
     /// The maximum trigger value.
     #[arg(short = 't', long, default_value = "5")]
@@ -75,7 +75,7 @@ struct Args {
     active_work: Option<NonZeroUsize>,
     /// Desired security model for IPA protocol
     #[arg(short = 'm', long, value_enum, default_value_t=IpaSecurityModel::Malicious)]
-    mode: IpaSecurityModel,
+    security_model: IpaSecurityModel,
     /// Needed for benches.
     #[arg(long, hide = true)]
     bench: bool,
@@ -150,10 +150,10 @@ async fn run(args: Args) -> Result<(), Error> {
     tracing::trace!("Preparation complete in {:?}", _prep_time.elapsed());
 
     let _protocol_time = Instant::now();
-    test_oprf_ipa::<BenchField>(&world, raw_data, &expected_results, args.config()).await;
+    test_oprf_ipa::<BenchField>(&world, raw_data, &expected_results, args.config(), args.security_model).await;
     tracing::info!(
         "{m:?} IPA for {q} records took {t:?}",
-        m = args.mode,
+        m = args.security_model,
         q = args.query_size,
         t = _protocol_time.elapsed()
     );

diff --git a/ipa-core/src/test_fixture/ipa.rs b/ipa-core/src/test_fixture/ipa.rs
@@ -20,7 +20,7 @@ use crate::{
     },
 };
 
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 #[cfg_attr(feature = "clap", derive(clap::ValueEnum))]
 pub enum IpaSecurityModel {
     SemiHonest,
@@ -186,6 +186,7 @@ pub async fn test_oprf_ipa<F>(
     records: Vec<TestRawDataRecord>,
     expected_results: &[u32],
     config: IpaQueryConfig,
+    security_model: IpaSecurityModel,
 ) where
     F: PrimeField + ExtendableField + IntoShares<semi_honest::AdditiveShare<F>>,
     rand::distributions::Standard: rand::distributions::Distribution<F>,
@@ -209,47 +210,65 @@ pub async fn test_oprf_ipa<F>(
         },
     };
     let padding_params = PaddingParameters::default();
-    let result: Vec<_> = if config.per_user_credit_cap == 256 {
-        // Note that many parameters are different in this case, not just the credit cap.
-        // This config is needed for collect_steps coverage.
+
+    let result: Vec<_> = if security_model == IpaSecurityModel::SemiHonest && matches!(config, IpaQueryConfig {
+        per_user_credit_cap: 8,
+        max_breakdown_key: 32,
+        ..
+    }) {
         world.semi_honest(
             records.into_iter(),
-            |ctx, input_rows: Vec<OPRFIPAInputRow<BA5, BA8, BA20>>| async move {
-                oprf_ipa::<_, BA5, BA8, BA32, BA20, 8, 32>(ctx, input_rows, aws, dp_params, padding_params)
+            |ctx, input_rows: Vec<OPRFIPAInputRow<BA5, BA3, BA20>>| async move {
+                oprf_ipa::<_, BA5, BA3, BA32, BA20, 3, 32>(ctx, input_rows, aws, dp_params, padding_params)
                     .await
                     .unwrap()
             },
         )
-    } else {
-        // In these configurations, the credit cap is the only parameter that changes.
+    } else if security_model == IpaSecurityModel::SemiHonest && matches!(config, IpaQueryConfig {
+        per_user_credit_cap: 8,
+        max_breakdown_key: 256,
+        ..
+    }) {
         world.semi_honest(
             records.into_iter(),
             |ctx, input_rows: Vec<OPRFIPAInputRow<BA8, BA3, BA20>>| async move {
-
-                match config.per_user_credit_cap {
-                    8 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 3, 256>(ctx, input_rows, aws, dp_params, padding_params)
-                    .await
-                    .unwrap(),
-                    16 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 4, 256>(ctx, input_rows, aws, dp_params, padding_params)
-                    .await
-                    .unwrap(),
-                    32 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 5, 256>(ctx, input_rows, aws, dp_params, padding_params)
+                oprf_ipa::<_, BA8, BA3, BA32, BA20, 3, 256>(ctx, input_rows, aws, dp_params, padding_params)
                     .await
-                    .unwrap(),
-                    64 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 6, 256>(ctx, input_rows, aws, dp_params, padding_params)
+                    .unwrap()
+            },
+        )
+    } else if security_model == IpaSecurityModel::Malicious && matches!(config, IpaQueryConfig {
+        per_user_credit_cap: 8,
+        max_breakdown_key: 32,
+        ..
+    }) {
+        world.malicious(
+            records.into_iter(),
+            |ctx, input_rows: Vec<OPRFIPAInputRow<BA5, BA3, BA20>>| async move {
+                oprf_ipa::<_, BA5, BA3, BA32, BA20, 3, 32>(ctx, input_rows, aws, dp_params, padding_params)
                     .await
-                    .unwrap(),
-                    128 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 7, 256>(ctx, input_rows, aws, dp_params, padding_params)
+                    .unwrap()
+            },
+        )
+    } else if security_model == IpaSecurityModel::Malicious && matches!(config, IpaQueryConfig {
+        per_user_credit_cap: 8,
+        max_breakdown_key: 256,
+        ..
+    }) {
+        world.malicious(
+            records.into_iter(),
+            |ctx, input_rows: Vec<OPRFIPAInputRow<BA8, BA3, BA20>>| async move {
+                oprf_ipa::<_, BA8, BA3, BA32, BA20, 3, 256>(ctx, input_rows, aws, dp_params, padding_params)
                     .await
-                    .unwrap(),
-                    _ =>
-                    panic!(
-                        "Invalid value specified for per-user cap: {:?}. Must be one of 8, 16, 32, 64, or 128.",
-                        config.per_user_credit_cap
-                    ),
-                }
+                    .unwrap()
             },
         )
+    } else {
+        panic!(
+            "Unsupported configuration: per_user_credit_cap = {:?}, max_breakdown_key = {:?}.",
+            config.per_user_credit_cap,
+            config.max_breakdown_key,
+        )
     }
     .await
     .reconstruct();