Merge pull request private-attribution#1331 from andyleiserson/padding

ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs

@@ -52,6 +52,7 @@ use crate::{
 pub async fn breakdown_reveal_aggregation<C, BK, TV, HV, const B: usize>(
     ctx: C,
     attributed_values: Vec<SecretSharedAttributionOutputs<BK, TV>>,
+    padding_params: &PaddingParameters,
 ) -> Result<BitDecomposed<Replicated<Boolean, B>>, Error>
 where
     C: Context,
@@ -63,13 +64,12 @@ where
     BitDecomposed<Replicated<Boolean, B>>:
         for<'a> TransposeFrom<&'a [Replicated<TV>; B], Error = Infallible>,
 {
-    let dp_padding_params = PaddingParameters::default();
     // Apply DP padding for Breakdown Reveal Aggregation
     let attributed_values_padded =
         apply_dp_padding::<_, AttributionOutputs<Replicated<BK>, Replicated<TV>>, B>(
             ctx.narrow(&AggregationStep::PaddingDp),
             attributed_values,
-            dp_padding_params,
+            padding_params,
         )
         .await?;
 
@@ -205,6 +205,7 @@ pub mod tests {
         },
         protocol::ipa_prf::{
             aggregation::breakdown_reveal::breakdown_reveal_aggregation,
+            oprf_padding::PaddingParameters,
             prf_sharding::{AttributionOutputsTestInput, SecretSharedAttributionOutputs},
         },
         secret_sharing::{
@@ -257,12 +258,16 @@ pub mod tests {
                         })
                         .collect();
                     let r: Vec<Replicated<BA8>> =
-                        breakdown_reveal_aggregation::<_, BA5, BA3, BA8, 32>(ctx, aos)
-                            .map_ok(|d: BitDecomposed<Replicated<Boolean, 32>>| {
-                                Vec::transposed_from(&d).unwrap()
-                            })
-                            .await
-                            .unwrap();
+                        breakdown_reveal_aggregation::<_, BA5, BA3, BA8, 32>(
+                            ctx,
+                            aos,
+                            &PaddingParameters::relaxed(),
+                        )
+                        .map_ok(|d: BitDecomposed<Replicated<Boolean, 32>>| {
+                            Vec::transposed_from(&d).unwrap()
+                        })
+                        .await
+                        .unwrap();
                     r
                 })
                 .await

ipa-core/src/protocol/ipa_prf/mod.rs

@@ -269,7 +269,7 @@ where
     let padded_input_rows = apply_dp_padding::<_, OPRFIPAInputRow<BK, TV, TS>, B>(
         ctx.narrow(&Step::PaddingDp),
         input_rows,
-        dp_padding_params,
+        &dp_padding_params,
     )
     .await?;
 
@@ -297,6 +297,7 @@ where
         prfd_inputs,
         attribution_window_seconds,
         &row_count_histogram,
+        &dp_padding_params,
     )
     .await?;
 
@@ -449,7 +450,14 @@ pub mod tests {
             ]; // trigger value of 2 attributes to earlier source row with breakdown 1 and trigger
                // value of 5 attributes to source row with breakdown 2.
             let dp_params = DpMechanism::NoDp;
-            let padding_params = PaddingParameters::relaxed();
+            let padding_params = if cfg!(feature = "shuttle") {
+                // To reduce runtime. There is also a hard upper limit in the shuttle
+                // config (`max_steps`), that may need to be increased to support larger
+                // runs.
+                PaddingParameters::no_padding()
+            } else {
+                PaddingParameters::relaxed()
+            };
 
             let mut result: Vec<_> = world
                 .semi_honest(records.into_iter(), |ctx, input_rows| async move {
@@ -489,7 +497,7 @@ pub mod tests {
             ]; // trigger value of 2 attributes to earlier source row with breakdown 1 and trigger
                // value of 5 attributes to source row with breakdown 2.
             let dp_params = DpMechanism::NoDp;
-            let padding_params = PaddingParameters::relaxed();
+            let padding_params = PaddingParameters::no_padding();
 
             let mut result: Vec<_> = world
                 .malicious(records.into_iter(), |ctx, input_rows| async move {

ipa-core/src/protocol/ipa_prf/oprf_padding/mod.rs

@@ -278,7 +278,7 @@ where
 pub async fn apply_dp_padding<C, T, const B: usize>(
     ctx: C,
     mut input: Vec<T>,
-    padding_params: PaddingParameters,
+    padding_params: &PaddingParameters,
 ) -> Result<Vec<T>, Error>
 where
     C: Context,
@@ -291,7 +291,7 @@ where
         ctx.narrow(&PaddingDpStep::PaddingDpPass1),
         input,
         Role::H3,
-        &padding_params,
+        padding_params,
     )
     .await?;
 
@@ -300,7 +300,7 @@ where
         ctx.narrow(&PaddingDpStep::PaddingDpPass2),
         input,
         Role::H2,
-        &padding_params,
+        padding_params,
     )
     .await?;
 
@@ -309,7 +309,7 @@ where
         ctx.narrow(&PaddingDpStep::PaddingDpPass3),
         input,
         Role::H1,
-        &padding_params,
+        padding_params,
     )
     .await?;
 

ipa-core/src/protocol/ipa_prf/prf_sharding/mod.rs

@@ -39,6 +39,7 @@ use crate::{
                 comparison_and_subtraction_sequential::{compare_gt, integer_sub},
                 expand_shared_array_in_place,
             },
+            oprf_padding::PaddingParameters,
             prf_sharding::step::{
                 AttributionPerRowStep as PerRowStep, AttributionStep as Step,
                 AttributionWindowStep as WindowStep,
@@ -469,6 +470,7 @@ pub async fn attribute_cap_aggregate<
     input_rows: Vec<PrfShardedIpaInputRow<BK, TV, TS>>,
     attribution_window_seconds: Option<NonZeroU32>,
     histogram: &[usize],
+    padding_parameters: &PaddingParameters,
 ) -> Result<BitDecomposed<Replicated<Boolean, B>>, Error>
 where
     C: UpgradableContext + 'ctx,
@@ -544,9 +546,12 @@ where
         aggregate_values_proof_chunk(B, usize::try_from(TV::BITS).unwrap()),
     );
     let user_contributions = flattened_user_results.try_collect::<Vec<_>>().await?;
-    let result =
-        breakdown_reveal_aggregation::<_, _, _, HV, B>(validator.context(), user_contributions)
-            .await;
+    let result = breakdown_reveal_aggregation::<_, _, _, HV, B>(
+        validator.context(),
+        user_contributions,
+        padding_parameters,
+    )
+    .await;
     validator.validate().await?;
     result
 }
@@ -891,7 +896,9 @@ pub mod tests {
             Field, U128Conversions,
         },
         helpers::repeat_n,
-        protocol::ipa_prf::prf_sharding::attribute_cap_aggregate,
+        protocol::ipa_prf::{
+            oprf_padding::PaddingParameters, prf_sharding::attribute_cap_aggregate,
+        },
         rand::Rng,
         secret_sharing::{
             replicated::semi_honest::AdditiveShare as Replicated, IntoShares, SharedValue,
@@ -1077,7 +1084,11 @@ pub mod tests {
                 .malicious(records.into_iter(), |ctx, input_rows| async move {
                     Vec::transposed_from(
                         &attribute_cap_aggregate::<_, BA5, BA3, BA16, BA20, 5, 32>(
-                            ctx, input_rows, None, &histogram,
+                            ctx,
+                            input_rows,
+                            None,
+                            &histogram,
+                            &PaddingParameters::relaxed(),
                         )
                         .await
                         .unwrap(),
@@ -1138,6 +1149,7 @@ pub mod tests {
                             input_rows,
                             NonZeroU32::new(ATTRIBUTION_WINDOW_SECONDS),
                             &histogram,
+                            &PaddingParameters::relaxed(),
                         )
                         .await
                         .unwrap(),
@@ -1175,6 +1187,7 @@ pub mod tests {
                         input_rows,
                         None,
                         histogram_ref,
+                        &PaddingParameters::relaxed(),
                     )
                     .await
                     .unwrap()
@@ -1261,7 +1274,13 @@ pub mod tests {
                             BA20,
                             { SaturatingSumType::BITS as usize },
                             256,
-                        >(ctx, input_rows, None, &HISTOGRAM)
+                        >(
+                            ctx,
+                            input_rows,
+                            None,
+                            &HISTOGRAM,
+                            &PaddingParameters::relaxed(),
+                        )
                         .await
                         .unwrap(),
                     )