Merge pull request #1283 from akoshelev/hybrid-query-type

Hybrid and sharded suffle query types

ipa-core/src/bin/test_mpc.rs

@@ -85,6 +85,11 @@ enum TestAction {
  /// All helpers add their shares locally and set the resulting share to be the
  /// sum. No communication is required to run the circuit.
  AddInPrimeField,
+ /// A test protocol for sharded MPCs. The goal here is to use
+ /// both shard-to-shard and helper-to-helper communication channels.
+ /// This is exactly what shuffle does and that's why it is picked
+ /// for this purpose.
+ ShardedShuffle,
 }
 
 #[tokio::main]
@@ -102,6 +107,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
  match args.action {
  TestAction::Multiply => multiply(&args, &clients).await,
  TestAction::AddInPrimeField => add(&args, &clients).await,
+ TestAction::ShardedShuffle => sharded_shuffle(&args, &clients).await,
  };
 
  Ok(())
@@ -159,3 +165,7 @@ async fn add(args: &Args, helper_clients: &[MpcHelperClient; 3]) {
  FieldType::Fp32BitPrime => add_in_field::<Fp32BitPrime>(args, helper_clients).await,
  };
 }
+
+async fn sharded_shuffle(_args: &Args, _helper_clients: &[MpcHelperClient; 3]) {
+ unimplemented!()
+}

ipa-core/src/helpers/transport/query/hybrid.rs

@@ -0,0 +1,32 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq)]
+#[cfg_attr(feature = "clap", derive(clap::Args))]
+pub struct HybridQueryParams {
+ #[cfg_attr(feature = "clap", arg(long, default_value = "8"))]
+ pub per_user_credit_cap: u32,
+ #[cfg_attr(feature = "clap", arg(long, default_value = "5"))]
+ pub max_breakdown_key: u32,
+ #[cfg_attr(feature = "clap", arg(short = 'd', long, default_value = "1"))]
+ pub with_dp: u32,
+ #[cfg_attr(feature = "clap", arg(short = 'e', long, default_value = "5.0"))]
+ pub epsilon: f64,
+ #[cfg_attr(feature = "clap", arg(long))]
+ #[serde(default)]
+ pub plaintext_match_keys: bool,
+}
+
+#[cfg(test)]
+impl Eq for HybridQueryParams {}
+
+impl Default for HybridQueryParams {
+ fn default() -> Self {
+ Self {
+ per_user_credit_cap: 8,
+ max_breakdown_key: 20,
+ with_dp: 1,
+ epsilon: 0.10,
+ plaintext_match_keys: false,
+ }
+ }
+}

ipa-core/src/helpers/transport/query/mod.rs

@@ -1,8 +1,11 @@
+mod hybrid;
+
 use std::{
  fmt::{Debug, Display, Formatter},
  num::NonZeroU32,
 };
 
+pub use hybrid::HybridQueryParams;
 use serde::{Deserialize, Deserializer, Serialize};
 
 use crate::{
@@ -198,16 +201,21 @@ pub enum QueryType {
  TestMultiply,
  #[cfg(any(test, feature = "test-fixture", feature = "cli"))]
  TestAddInPrimeField,
+ #[cfg(any(test, feature = "test-fixture", feature = "cli"))]
+ TestShardedShuffle,
  SemiHonestOprfIpa(IpaQueryConfig),
  MaliciousOprfIpa(IpaQueryConfig),
+ SemiHonestHybrid(HybridQueryParams),
 }
 
 impl QueryType {
  /// TODO: strum
  pub const TEST_MULTIPLY_STR: &'static str = "test-multiply";
  pub const TEST_ADD_STR: &'static str = "test-add";
+ pub const TEST_SHARDED_SHUFFLE_STR: &'static str = "test-sharded-shuffle";
  pub const SEMI_HONEST_OPRF_IPA_STR: &'static str = "semi-honest-oprf-ipa";
  pub const MALICIOUS_OPRF_IPA_STR: &'static str = "malicious-oprf-ipa";
+ pub const SEMI_HONEST_HYBRID_STR: &'static str = "semi-honest-hybrid";
 }
 
 /// TODO: should this `AsRef` impl (used for `Substep`) take into account config of IPA?
@@ -218,8 +226,11 @@ impl AsRef<str> for QueryType {
  QueryType::TestMultiply => Self::TEST_MULTIPLY_STR,
  #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
  QueryType::TestAddInPrimeField => Self::TEST_ADD_STR,
+ #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
+ QueryType::TestShardedShuffle => Self::TEST_SHARDED_SHUFFLE_STR,
  QueryType::SemiHonestOprfIpa(_) => Self::SEMI_HONEST_OPRF_IPA_STR,
  QueryType::MaliciousOprfIpa(_) => Self::MALICIOUS_OPRF_IPA_STR,
+ QueryType::SemiHonestHybrid(_) => Self::SEMI_HONEST_HYBRID_STR,
  }
  }
 }

ipa-core/src/net/http_serde.rs

@@ -152,6 +152,8 @@ pub mod query {
  match self.query_type {
  #[cfg(any(test, feature = "test-fixture", feature = "cli"))]
  QueryType::TestMultiply | QueryType::TestAddInPrimeField => Ok(()),
+ #[cfg(any(test, feature = "test-fixture", feature = "cli"))]
+ QueryType::TestShardedShuffle => Ok(()),
  QueryType::SemiHonestOprfIpa(config) | QueryType::MaliciousOprfIpa(config) => {
  write!(
  f,
@@ -170,6 +172,22 @@ pub mod query {
  write!(f, "&attribution_window_seconds={}", window.get())?;
  }
 
+ Ok(())
+ }
+ QueryType::SemiHonestHybrid(config) => {
+ write!(
+ f,
+ "&per_user_credit_cap={}&max_breakdown_key={}&with_dp={}&epsilon={}",
+ config.per_user_credit_cap,
+ config.max_breakdown_key,
+ config.with_dp,
+ config.epsilon,
+ )?;
+
+ if config.plaintext_match_keys {
+ write!(f, "&plaintext_match_keys=true")?;
+ }
+
  Ok(())
  }
  }

ipa-core/src/query/executor.rs

@@ -44,7 +44,7 @@ use crate::{
  Gate,
  },
  query::{
- runner::{OprfIpaQuery, QueryResult},
+ runner::{HybridQuery, OprfIpaQuery, QueryResult},
  state::RunningQuery,
  },
  sync::Arc,
@@ -94,6 +94,13 @@ pub fn execute<R: PrivateKeyRegistry>(
  Box::pin(execute_test_multiply::<Fp32BitPrime>(prss, gateway, input))
  })
  }
+ #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
+ (QueryType::TestShardedShuffle, _) => do_query(
+ config,
+ gateway,
+ input,
+ |_prss, _gateway, _config, _input| unimplemented!(),
+ ),
  #[cfg(any(test, feature = "weak-field"))]
  (QueryType::TestAddInPrimeField, FieldType::Fp31) => {
  do_query(config, gateway, input, |prss, gateway, _config, input| {
@@ -138,6 +145,19 @@ pub fn execute<R: PrivateKeyRegistry>(
  )
  },
  ),
+ (QueryType::SemiHonestHybrid(query_params), _) => do_query(
+ config,
+ gateway,
+ input,
+ move |prss, gateway, config, input| {
+ let ctx = SemiHonestContext::new(prss, gateway);
+ Box::pin(
+ HybridQuery::<_, BA32, R>::new(query_params, key_registry)
+ .execute(ctx, config.size, input)
+ .then(|res| ready(res.map(|out| Box::new(out) as Box<dyn Result>))),
+ )
+ },
+ ),
  }
 }
 

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

@@ -0,0 +1,37 @@
+use std::{marker::PhantomData, sync::Arc};
+
+use crate::{
+ error::Error,
+ helpers::{
+ query::{HybridQueryParams, QuerySize},
+ BodyStream,
+ },
+ hpke::PrivateKeyRegistry,
+ secret_sharing::{replicated::semi_honest::AdditiveShare as ReplicatedShare, SharedValue},
+};
+
+pub struct Query<C, HV, R: PrivateKeyRegistry> {
+ _config: HybridQueryParams,
+ _key_registry: Arc<R>,
+ phantom_data: PhantomData<(C, HV)>,
+}
+
+impl<C, HV: SharedValue, R: PrivateKeyRegistry> Query<C, HV, R> {
+ pub fn new(query_params: HybridQueryParams, key_registry: Arc<R>) -> Self {
+ Self {
+ _config: query_params,
+ _key_registry: key_registry,
+ phantom_data: PhantomData,
+ }
+ }
+
+ #[tracing::instrument("hybrid_query", skip_all, fields(sz=%query_size))]
+ pub async fn execute(
+ self,
+ _ctx: C,
+ query_size: QuerySize,
+ _input_stream: BodyStream,
+ ) -> Result<Vec<ReplicatedShare<HV>>, Error> {
+ unimplemented!()
+ }
+}

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

@@ -1,11 +1,13 @@
 #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
 mod add_in_prime_field;
+mod hybrid;
 mod oprf_ipa;
 #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
 mod test_multiply;
 
 #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
 pub(super) use add_in_prime_field::execute as test_add_in_prime_field;
+pub use hybrid::Query as HybridQuery;
 #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
 pub(super) use test_multiply::execute_test_multiply;