From ac2c6e28b3483c04ea5a173916d0c5d185189a92 Mon Sep 17 00:00:00 2001
From: Erik Taubeneck <erik.taubeneck@gmail.com>
Date: Tue, 3 Dec 2024 16:13:51 -0800
Subject: [PATCH] add CLI command to run hybrid query

---
 ipa-core/src/bin/report_collector.rs          | 64 ++++++++++++++++++-
 .../src/helpers/transport/query/hybrid.rs     | 19 ++++++
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/ipa-core/src/bin/report_collector.rs b/ipa-core/src/bin/report_collector.rs
index 126004fcc..18779b545 100644
--- a/ipa-core/src/bin/report_collector.rs
+++ b/ipa-core/src/bin/report_collector.rs
@@ -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::{
@@ -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)]
@@ -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(())
@@ -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,
diff --git a/ipa-core/src/helpers/transport/query/hybrid.rs b/ipa-core/src/helpers/transport/query/hybrid.rs
index 63c80160e..1ea2a8732 100644
--- a/ipa-core/src/helpers/transport/query/hybrid.rs
+++ b/ipa-core/src/helpers/transport/query/hybrid.rs
@@ -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 {
@@ -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()
+        }
+    }
+}