From 80743d1ee9e2ec7de63ede3c2fb458e6baf36a34 Mon Sep 17 00:00:00 2001 From: Alex Koshelev Date: Thu, 26 Dec 2024 12:47:01 -0500 Subject: [PATCH 1/2] Use sharded shuffle in Hybrid It turns out, we've never used sharded shuffle in Hybrid, running local shuffles (per shard) instead. This change fixes that. Also, drop `Shuffle` implementation for sharded contexts to prevent this from happening again. --- ipa-core/src/protocol/hybrid/breakdown_reveal.rs | 6 +++--- ipa-core/src/protocol/hybrid/mod.rs | 6 +++--- ipa-core/src/protocol/hybrid/oprf.rs | 8 +++++--- ipa-core/src/protocol/ipa_prf/shuffle/mod.rs | 8 ++++---- ipa-core/src/query/runner/hybrid.rs | 4 ++-- 5 files changed, 17 insertions(+), 15 deletions(-) diff --git a/ipa-core/src/protocol/hybrid/breakdown_reveal.rs b/ipa-core/src/protocol/hybrid/breakdown_reveal.rs index d94bf001c..9af2b7ede 100644 --- a/ipa-core/src/protocol/hybrid/breakdown_reveal.rs +++ b/ipa-core/src/protocol/hybrid/breakdown_reveal.rs @@ -20,7 +20,7 @@ use crate::{ AGGREGATE_DEPTH, }, oprf_padding::{apply_dp_padding, PaddingParameters}, - shuffle::Shuffle, + shuffle::ShardedShuffle, }, BooleanProtocols, RecordId, }, @@ -66,7 +66,7 @@ pub async fn breakdown_reveal_aggregation( padding_params: &PaddingParameters, ) -> Result>, Error> where - C: UpgradableContext + Shuffle + ShardedContext, + C: UpgradableContext + ShardedShuffle + ShardedContext, Boolean: FieldSimd, Replicated: BooleanProtocols, B>, BK: BooleanArray + U128Conversions, @@ -94,7 +94,7 @@ where let attributions = ctx .narrow(&Step::Shuffle) - .shuffle(attributed_values_padded) + .sharded_shuffle(attributed_values_padded) .instrument(info_span!("shuffle_attribution_outputs")) .await?; diff --git a/ipa-core/src/protocol/hybrid/mod.rs b/ipa-core/src/protocol/hybrid/mod.rs index 5217365b6..b6de81fc6 100644 --- a/ipa-core/src/protocol/hybrid/mod.rs +++ b/ipa-core/src/protocol/hybrid/mod.rs @@ -33,7 +33,7 @@ use crate::{ ipa_prf::{ oprf_padding::{apply_dp_padding, PaddingParameters}, prf_eval::PrfSharing, - shuffle::Shuffle, + shuffle::ShardedShuffle, }, prss::FromPrss, BooleanProtocols, @@ -79,7 +79,7 @@ pub async fn hybrid_protocol<'ctx, C, BK, V, HV, const SS_BITS: usize, const B: where C: UpgradableContext + 'ctx - + Shuffle + + ShardedShuffle + ShardedContext + FinalizerContext>, BK: BreakdownKey, @@ -121,7 +121,7 @@ where let shuffled_input_rows = ctx .narrow(&Step::InputShuffle) - .shuffle(padded_input_rows) + .sharded_shuffle(padded_input_rows) .instrument(info_span!("shuffle_inputs")) .await?; diff --git a/ipa-core/src/protocol/hybrid/oprf.rs b/ipa-core/src/protocol/hybrid/oprf.rs index 29d13d755..8bed507ea 100644 --- a/ipa-core/src/protocol/hybrid/oprf.rs +++ b/ipa-core/src/protocol/hybrid/oprf.rs @@ -181,9 +181,11 @@ where // reshard reports based on OPRF values. This ensures at the end of this function // reports with the same value end up on the same shard. - reshard_try_stream(ctx, report_stream, |ctx, _, report| { - report.match_key % ctx.shard_count() - }) + reshard_try_stream( + ctx.narrow(&HybridStep::ReshardByPrf), + report_stream, + |ctx, _, report| report.match_key % ctx.shard_count(), + ) .await } diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs b/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs index 0ee65d169..189b3e75b 100644 --- a/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs +++ b/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs @@ -9,7 +9,7 @@ use crate::{ context::{Context, MaliciousContext, SemiHonestContext}, ipa_prf::shuffle::sharded::ShuffleContext, }, - sharding::{ShardBinding, Sharded}, + sharding::{Sharded}, }; mod base; @@ -21,6 +21,7 @@ use base::shuffle_protocol as base_shuffle; use malicious::{malicious_sharded_shuffle, malicious_shuffle}; use sharded::shuffle as sharded_shuffle; pub use sharded::{MaliciousShuffleable, Shuffleable}; +use crate::sharding::NotSharded; /// This struct stores some intermediate messages during the shuffle. /// In a maliciously secure shuffle, @@ -63,7 +64,7 @@ pub trait Shuffle: Context { S: MaliciousShuffleable; } -impl Shuffle for SemiHonestContext<'_, T> { +impl Shuffle for SemiHonestContext<'_, NotSharded> { fn shuffle(self, shares: Vec) -> impl Future, Error>> + Send where S: MaliciousShuffleable, @@ -73,7 +74,7 @@ impl Shuffle for SemiHonestContext<'_, T> { } } -impl Shuffle for MaliciousContext<'_, T> { +impl Shuffle for MaliciousContext<'_, NotSharded> { fn shuffle(self, shares: Vec) -> impl Future, Error>> + Send where S: MaliciousShuffleable, @@ -84,7 +85,6 @@ impl Shuffle for MaliciousContext<'_, T> { /// Trait used by protocols to invoke either semi-honest or malicious sharded shuffle, /// depending on the type of context being used. -#[allow(dead_code)] pub trait ShardedShuffle: ShuffleContext { fn sharded_shuffle( self, diff --git a/ipa-core/src/query/runner/hybrid.rs b/ipa-core/src/query/runner/hybrid.rs index e41fbc152..e3b4f6311 100644 --- a/ipa-core/src/query/runner/hybrid.rs +++ b/ipa-core/src/query/runner/hybrid.rs @@ -35,7 +35,7 @@ use crate::{ oprf::{CONV_CHUNK, PRF_CHUNK}, step::HybridStep, }, - ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::Shuffle}, + ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::ShardedShuffle}, prss::{Endpoint, FromPrss}, step::ProtocolStep::Hybrid, Gate, @@ -73,7 +73,7 @@ impl Query { impl Query where C: UpgradableContext - + Shuffle + + ShardedShuffle + ShardedContext + FinalizerContext>, HV: BooleanArray + U128Conversions, From 251d1f3984f11348b8392166928df166c11529a5 Mon Sep 17 00:00:00 2001 From: Alex Koshelev Date: Thu, 26 Dec 2024 16:00:41 -0500 Subject: [PATCH 2/2] Fix compact gate steps for sharded shuffle --- .../src/protocol/hybrid/breakdown_reveal.rs | 6 ++---- ipa-core/src/protocol/hybrid/step.rs | 19 +++++++++++++++++-- .../src/protocol/ipa_prf/shuffle/malicious.rs | 8 ++++---- ipa-core/src/protocol/ipa_prf/shuffle/mod.rs | 3 ++- ipa-core/src/protocol/ipa_prf/shuffle/step.rs | 5 +++++ 5 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ipa-core/src/protocol/hybrid/breakdown_reveal.rs b/ipa-core/src/protocol/hybrid/breakdown_reveal.rs index 9af2b7ede..d851a05ac 100644 --- a/ipa-core/src/protocol/hybrid/breakdown_reveal.rs +++ b/ipa-core/src/protocol/hybrid/breakdown_reveal.rs @@ -14,11 +14,9 @@ use crate::{ dzkp_validator::DZKPValidator, Context, DZKPUpgraded, MaliciousProtocolSteps, ShardedContext, UpgradableContext, }, + hybrid::step::AggregationStep as Step, ipa_prf::{ - aggregation::{ - aggregate_values, aggregate_values_proof_chunk, step::AggregationStep as Step, - AGGREGATE_DEPTH, - }, + aggregation::{aggregate_values, aggregate_values_proof_chunk, AGGREGATE_DEPTH}, oprf_padding::{apply_dp_padding, PaddingParameters}, shuffle::ShardedShuffle, }, diff --git a/ipa-core/src/protocol/hybrid/step.rs b/ipa-core/src/protocol/hybrid/step.rs index 2a98488fc..7f4de200a 100644 --- a/ipa-core/src/protocol/hybrid/step.rs +++ b/ipa-core/src/protocol/hybrid/step.rs @@ -5,7 +5,7 @@ pub(crate) enum HybridStep { ReshardByTag, #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="report_padding_dp")] PaddingDp, - #[step(child = crate::protocol::ipa_prf::shuffle::step::OPRFShuffleStep)] + #[step(child = crate::protocol::ipa_prf::shuffle::step::ShardedShuffleStep)] InputShuffle, #[step(child = crate::protocol::ipa_prf::boolean_ops::step::Fp25519ConversionStep)] ConvertFp25519, @@ -19,7 +19,7 @@ pub(crate) enum HybridStep { GroupBySum, #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)] GroupBySumValidate, - #[step(child = crate::protocol::ipa_prf::aggregation::step::AggregationStep)] + #[step(child = AggregationStep)] Aggregate, #[step(child = FinalizeSteps)] Finalize, @@ -40,3 +40,18 @@ pub(crate) enum FinalizeSteps { #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)] Validate, } + +#[derive(CompactStep)] +pub(crate) enum AggregationStep { + #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="padding_dp")] + PaddingDp, + #[step(child = crate::protocol::ipa_prf::shuffle::step::ShardedShuffleStep)] + Shuffle, + Reveal, + #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)] + RevealValidate, // only partly used -- see code + #[step(count = 4, child = crate::protocol::ipa_prf::aggregation::step::AggregateChunkStep, name = "chunks")] + Aggregate(usize), + #[step(count = 4, child = crate::protocol::context::step::DzkpValidationProtocolStep)] + AggregateValidate(usize), +} diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs b/ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs index b0745f568..f73e1e973 100644 --- a/ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs +++ b/ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs @@ -30,7 +30,7 @@ use crate::{ h1_shuffle_for_shard, h2_shuffle_for_shard, h3_shuffle_for_shard, MaliciousShuffleable, ShuffleShare, Shuffleable, }, - step::{OPRFShuffleStep, VerifyShuffleStep}, + step::{OPRFShuffleStep, ShardedShuffleStep, VerifyShuffleStep}, IntermediateShuffleMessages, }, prss::SharedRandomness, @@ -179,11 +179,11 @@ where // prepare keys let amount_of_keys: usize = (usize::try_from(S::Share::BITS).unwrap() + 31) / 32; - let keys = setup_keys(ctx.narrow(&OPRFShuffleStep::SetupKeys), amount_of_keys).await?; + let keys = setup_keys(ctx.narrow(&ShardedShuffleStep::SetupKeys), amount_of_keys).await?; // compute and append tags to rows let shares_and_tags: Vec> = - compute_and_add_tags(ctx.narrow(&OPRFShuffleStep::GenerateTags), &keys, shares).await?; + compute_and_add_tags(ctx.narrow(&ShardedShuffleStep::GenerateTags), &keys, shares).await?; let (shuffled_shares, messages) = match ctx.role() { Role::H1 => h1_shuffle_for_shard(ctx.clone(), shares_and_tags).await, @@ -193,7 +193,7 @@ where // verify the shuffle verify_shuffle::<_, S>( - ctx.narrow(&OPRFShuffleStep::VerifyShuffle), + ctx.narrow(&ShardedShuffleStep::VerifyShuffle), &keys, &shuffled_shares, messages, diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs b/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs index 189b3e75b..b9d7c40b3 100644 --- a/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs +++ b/ipa-core/src/protocol/ipa_prf/shuffle/mod.rs @@ -9,7 +9,7 @@ use crate::{ context::{Context, MaliciousContext, SemiHonestContext}, ipa_prf::shuffle::sharded::ShuffleContext, }, - sharding::{Sharded}, + sharding::Sharded, }; mod base; @@ -21,6 +21,7 @@ use base::shuffle_protocol as base_shuffle; use malicious::{malicious_sharded_shuffle, malicious_shuffle}; use sharded::shuffle as sharded_shuffle; pub use sharded::{MaliciousShuffleable, Shuffleable}; + use crate::sharding::NotSharded; /// This struct stores some intermediate messages during the shuffle. diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/step.rs b/ipa-core/src/protocol/ipa_prf/shuffle/step.rs index 6a4ff2050..6ca43b5ad 100644 --- a/ipa-core/src/protocol/ipa_prf/shuffle/step.rs +++ b/ipa-core/src/protocol/ipa_prf/shuffle/step.rs @@ -29,6 +29,9 @@ pub(crate) enum VerifyShuffleStep { #[derive(CompactStep)] pub(crate) enum ShardedShuffleStep { + SetupKeys, + #[step(child = crate::protocol::boolean::step::EightBitStep)] + GenerateTags, /// Depending on the helper position inside the MPC ring, generate Ã, B̃ or both. PseudoRandomTable, /// Permute the input according to the PRSS shared between H1 and H2. @@ -46,6 +49,8 @@ pub(crate) enum ShardedShuffleStep { TransferXY, /// H2 and H3 interaction - Exchange `C_1` and `C_2`. TransferC, + #[step(child = crate::protocol::ipa_prf::shuffle::step::VerifyShuffleStep)] + VerifyShuffle, } #[derive(CompactStep)]