Add malicious security support to oneshot bench

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,17 @@ 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()
     );

ipa-core/src/test_fixture/ipa.rs

@@ -20,7 +20,7 @@
     },
 };
 
-#[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 @@
     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,107 @@
         },
     };
     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)
-                    .await
-                    .unwrap()
+            |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)
-                    .await
-                    .unwrap(),
-                    64 => oprf_ipa::<_, BA8, BA3, BA32, BA20, 6, 256>(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)
-                    .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
-                    ),
-                }
+                oprf_ipa::<_, BA8, BA3, BA32, BA20, 3, 256>(
+                    ctx,
+                    input_rows,
+                    aws,
+                    dp_params,
+                    padding_params,
+                )
+                .await
+                .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()
+            },
+        )
+    } 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()
             },
         )
+    } else {
+        panic!(
+            "Unsupported configuration: per_user_credit_cap = {:?}, max_breakdown_key = {:?}.",
+            config.per_user_credit_cap, config.max_breakdown_key,
+        )
     }
     .await
     .reconstruct();

scripts/coverage-ci

@@ -21,7 +21,7 @@ done
 # integration tests run without relaxed dp, except for these
 cargo test --release --test "ipa_with_relaxed_dp" --no-default-features --features "cli web-app real-world-infra test-fixture compact-gate relaxed-dp"
 
-cargo test --bench oneshot_ipa --no-default-features --features "enable-benches compact-gate" -- -n 62 -c 16
+cargo test --bench oneshot_ipa --no-default-features --features "enable-benches compact-gate" -- -n 62 -c 8
 cargo test --bench criterion_arithmetic --no-default-features --features "enable-benches compact-gate"
 
 # compact gate + in-memory-infra