Skip to content

Commit

Permalink
add CLI command to run hybrid query
Browse files Browse the repository at this point in the history
  • Loading branch information
eriktaubeneck committed Dec 4, 2024
1 parent dca5be7 commit ac2c6e2
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 1 deletion.
64 changes: 63 additions & 1 deletion ipa-core/src/bin/report_collector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@ use ipa_core::{
},
config::{KeyRegistries, NetworkConfig},
ff::{boolean_array::BA32, FieldType},
helpers::query::{DpMechanism, IpaQueryConfig, QueryConfig, QuerySize, QueryType},
helpers::query::{
DpMechanism, HybridQueryParams, IpaQueryConfig, QueryConfig, QuerySize, QueryType,
},
net::{Helper, IpaHttpClient},
report::{EncryptedOprfReportStreams, DEFAULT_KEY_ID},
test_fixture::{
Expand Down Expand Up @@ -132,6 +134,13 @@ enum ReportCollectorCommand {
#[clap(flatten)]
ipa_query_config: IpaQueryConfig,
},
MaliciousHybrid {
#[clap(flatten)]
encrypted_inputs: EncryptedInputs,

#[clap(flatten)]
hybrid_query_config: HybridQueryParams,
},
}

#[derive(Debug, clap::Args)]
Expand Down Expand Up @@ -227,6 +236,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
)
.await?
}
ReportCollectorCommand::MaliciousHybrid {
ref encrypted_inputs,
hybrid_query_config,
} => hybrid(&args, hybrid_query_config, &clients, encrypted_inputs).await?,
};

Ok(())
Expand Down Expand Up @@ -329,6 +342,55 @@ fn write_ipa_output_file(
Ok(())
}

async fn hybrid(
args: &Args,
hybrid_query_config: HybridQueryParams,
helper_clients: &[IpaHttpClient<Helper>; 3],
encrypted_inputs: &EncryptedInputs,
) -> Result<(), Box<dyn Error>> {
let query_type = QueryType::MaliciousHybrid(hybrid_query_config);

let files = [
&encrypted_inputs.enc_input_file1,
&encrypted_inputs.enc_input_file2,
&encrypted_inputs.enc_input_file3,
];

// despite the name, this is generic enough to work with hybrid
let encrypted_report_streams = EncryptedOprfReportStreams::from(files);

let query_config = QueryConfig {
size: QuerySize::try_from(encrypted_report_streams.query_size).unwrap(),
field_type: FieldType::Fp32BitPrime,
query_type,
};

let query_id = helper_clients[0]
.create_query(query_config)
.await
.expect("Unable to create query!");

tracing::info!("Starting query for OPRF");
// the value for histogram values (BA32) must be kept in sync with the server-side
// implementation, otherwise a runtime reconstruct error will be generated.
// see ipa-core/src/query/executor.rs
let actual = run_query_and_validate::<BA32>(
encrypted_report_streams.streams,
encrypted_report_streams.query_size,
helper_clients,
query_id,
hybrid_query_config.into(),
)
.await;

if let Some(ref path) = args.output_file {
write_ipa_output_file(path, &actual)?;
} else {
println!("{}", serde_json::to_string_pretty(&actual)?);
}
Ok(())
}

async fn ipa(
args: &Args,
security_model: IpaSecurityModel,
Expand Down
19 changes: 19 additions & 0 deletions ipa-core/src/helpers/transport/query/hybrid.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use serde::{Deserialize, Serialize};

use super::IpaQueryConfig;

#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq)]
#[cfg_attr(feature = "clap", derive(clap::Args))]
pub struct HybridQueryParams {
Expand Down Expand Up @@ -27,3 +29,20 @@ impl Default for HybridQueryParams {
}
}
}

// `IpaQueryConfig` is a super set of this, and to avoid an almost entire duplication of
// `run_query_and_validate`, we just convert this over. the unused fields (`per_user_credit_cap`
// and `attribution_window_seconds`) aren't used in that function.
// Once we deprecate OprfIpa, we can simply swap out IpaQueryConfig with HybridQueryParams
#[cfg(all(feature = "web-app", feature = "cli"))]
impl From<HybridQueryParams> for IpaQueryConfig {
fn from(params: HybridQueryParams) -> Self {
Self {
max_breakdown_key: params.max_breakdown_key,
with_dp: params.with_dp,
epsilon: params.epsilon,
plaintext_match_keys: params.plaintext_match_keys,
..Default::default()
}
}
}

0 comments on commit ac2c6e2

Please sign in to comment.