Python formatting (#1238)

* add test reading in encrypted file. compiles but test fails

* tests passing

* run formatter on encrypt_report.py

in-market-test/v2/hpke/encrypt_report.py

@@ -1,27 +1,27 @@
 import argparse
-from enum import Enum
 import math
 import os
 import secrets
-from typing import Any
+from enum import Enum
 
 import pyhpke
 from cryptography.hazmat.primitives.asymmetric import x25519
 
+
 class EventType(Enum):
     SOURCE = 0
     TRIGGER = 1
 
     def to_bytes(self):
         return self.value.to_bytes(1, "little")
 
+
 class ShareType(Enum):
     MATCH_KEY = 0
     TIMESTAMP = 1
     BREAKDOWN = 2
     TRIGGER_VALUE = 3
 
-
     def bit_count(self) -> int:
         match self:
             case ShareType.MATCH_KEY:
@@ -35,13 +35,10 @@ def bit_count(self) -> int:
 
         raise Exception("Invalid share type")
 
-
     def byte_count(self) -> int:
         return math.ceil(self.bit_count() / 8)
 
 
-
-
 class IPAReportInfo:
     DOMAIN = "private-attribution"
 
@@ -169,8 +166,8 @@ def generate_random_share(share_type: ShareType) -> bytes:
     @classmethod
     def create_shares(
         cls,
-        value: bytes, 
-        share_type: ShareType
+        value: bytes,
+        share_type: ShareType,
     ) -> tuple["IPAShare", "IPAShare", "IPAShare"]:
         first_share = IPAShare.generate_random_share(share_type)
         second_share = IPAShare.generate_random_share(share_type)
@@ -183,7 +180,6 @@ def create_shares(
         )
 
 
-
 class IPAReport:
     __slots__ = [
         "mk_encap_key_ciphertext",
@@ -213,13 +209,19 @@ def encrypt(self, public_key_string: str) -> bytes:
             + self.breakdown_key.to_bytes()
             + self.trigger_value.to_bytes()
         )
-        encrypted = encrypt_share(share_data, self.info.event_type, self.info.site_domain, public_key_string, self.info.helper_domain)
+        encrypted = encrypt_share(
+            share_data,
+            self.info.event_type,
+            self.info.site_domain,
+            public_key_string,
+            self.info.helper_domain,
+        )
         return (
             self.mk_encap_key_ciphertext
             + encrypted.encrypted_to_bytes()
             + encrypted.ipa_report_info_to_bytes()
         )
-    
+
 
 def generate_report_per_helper(
     mk_share: IPAShare,
@@ -232,11 +234,11 @@ def generate_report_per_helper(
     helper_domain: str,
 ) -> bytes:
     mk_encap_key_ciphertext = encrypt_share(
-        mk_share.to_bytes(), 
-        event_type=event_type, 
-        site_domain=site_domain, 
-        public_key_string=pub_key, 
-        helper_domain=helper_domain
+        mk_share.to_bytes(),
+        event_type=event_type,
+        site_domain=site_domain,
+        public_key_string=pub_key,
+        helper_domain=helper_domain,
     ).encrypted_to_bytes()
 
     return IPAReport(
@@ -271,54 +273,64 @@ def encrypt_to_file(
         for line_num, line in enumerate(f_in):
             # File format: <timestamp>,<match_key>,<event_type>,<breakdown_key>,<trigger_value>
             values = line.split(",")
-            assert len(values) >= 5, f"Corrupted file: line {line_num} has less than 5 values"
+            assert (
+                len(values) >= 5
+            ), f"Corrupted file: line {line_num} has less than 5 values"
             timestamp = int(values[0].strip())
             match_key = int(values[1].strip())
             event_type = EventType(int(values[2].strip()))
             breakdown_key = int(values[3].strip())
             trigger_value = int(values[4].strip())
 
-            mk_share = IPAShare.create_shares(match_key.to_bytes(8, "little"), ShareType.MATCH_KEY)
-            ts_share = IPAShare.create_shares(timestamp.to_bytes(3, "little"), ShareType.TIMESTAMP)
-            bk_share = IPAShare.create_shares(breakdown_key.to_bytes(1, "little"), ShareType.BREAKDOWN)
-            tv_share = IPAShare.create_shares(trigger_value.to_bytes(1, "little"), ShareType.TRIGGER_VALUE)
+            mk_share = IPAShare.create_shares(
+                match_key.to_bytes(8, "little"), ShareType.MATCH_KEY
+            )
+            ts_share = IPAShare.create_shares(
+                timestamp.to_bytes(3, "little"), ShareType.TIMESTAMP
+            )
+            bk_share = IPAShare.create_shares(
+                breakdown_key.to_bytes(1, "little"), ShareType.BREAKDOWN
+            )
+            tv_share = IPAShare.create_shares(
+                trigger_value.to_bytes(1, "little"), ShareType.TRIGGER_VALUE
+            )
 
             encrypted_reports_1.append(
                 generate_report_per_helper(
-                    mk_share=mk_share[0], 
-                    ts_share=ts_share[0], 
-                    bk_share=bk_share[0], 
-                    tv_share=tv_share[0], 
-                    site_domain=site_domain, 
+                    mk_share=mk_share[0],
+                    ts_share=ts_share[0],
+                    bk_share=bk_share[0],
+                    tv_share=tv_share[0],
+                    site_domain=site_domain,
                     event_type=event_type,
-                    pub_key=pub_key, 
-                    helper_domain=helper_domain
+                    pub_key=pub_key,
+                    helper_domain=helper_domain,
                 )
             )
 
             encrypted_reports_2.append(
                 generate_report_per_helper(
-                    mk_share=mk_share[1], 
-                    ts_share=ts_share[1], 
-                    bk_share=bk_share[1], 
-                    tv_share=tv_share[1], 
-                    site_domain=site_domain, 
+                    mk_share=mk_share[1],
+                    ts_share=ts_share[1],
+                    bk_share=bk_share[1],
+                    tv_share=tv_share[1],
+                    site_domain=site_domain,
                     event_type=event_type,
-                    pub_key=pub_key, 
-                    helper_domain=helper_domain
+                    pub_key=pub_key,
+                    helper_domain=helper_domain,
                 )
             )
 
             encrypted_reports_3.append(
                 generate_report_per_helper(
-                    mk_share=mk_share[2], 
-                    ts_share=ts_share[2], 
-                    bk_share=bk_share[2], 
-                    tv_share=tv_share[2], 
-                    site_domain=site_domain, 
+                    mk_share=mk_share[2],
+                    ts_share=ts_share[2],
+                    bk_share=bk_share[2],
+                    tv_share=tv_share[2],
+                    site_domain=site_domain,
                     event_type=event_type,
-                    pub_key=pub_key, 
-                    helper_domain=helper_domain
+                    pub_key=pub_key,
+                    helper_domain=helper_domain,
                 )
             )
 
@@ -335,7 +347,7 @@ def encrypt_to_file(
                 f_out.write(i.hex() + "\n")
                 if verbose:
                     print(i.hex())
-        
+
         file_out_3 = os.path.join(dir_out, file_out_prefix + "_h3")
         with open(file_out_3, "w") as f_out:
             for i in encrypted_reports_3:

ipa-core/Cargo.toml

@@ -197,7 +197,7 @@ bench = false
 
 [[bin]]
 name = "crypto_util"
-required-features = ["cli", "test-fixture", "web-app"]
+required-features = ["cli", "test-fixture", "web-app", "in-memory-infra"]
 bench = false
 
 [[bench]]

ipa-core/src/cli/crypto.rs

@@ -244,9 +244,12 @@ mod tests {
         fs::File,
         io::{BufRead, BufReader, Write},
         path::Path,
+        sync::Arc,
     };
 
+    use bytes::BufMut;
     use clap::Parser;
+    use hpke::Deserializable;
     use rand::thread_rng;
     use tempfile::{tempdir, NamedTempFile};
 
@@ -255,7 +258,17 @@ mod tests {
             crypto::{decrypt_and_reconstruct, encrypt, DecryptArgs, EncryptArgs},
             CsvSerializer,
         },
-        test_fixture::{EventGenerator, EventGeneratorConfig},
+        ff::{boolean_array::BA16, U128Conversions},
+        helpers::{
+            query::{IpaQueryConfig, QuerySize},
+            BodyStream,
+        },
+        hpke::{IpaPrivateKey, KeyRegistry, PrivateKeyOnly},
+        query::OprfIpaQuery,
+        test_fixture::{
+            ipa::TestRawDataRecord, join3v, EventGenerator, EventGeneratorConfig, Reconstruct,
+            TestWorld,
+        },
     };
 
     fn are_files_equal(file1: &Path, file2: &Path) {
@@ -524,4 +537,134 @@ public_key = "cfdbaaff16b30aa8a4ab07eaad2cdd80458208a1317aefbb807e46dce596617e"
 
         are_files_equal(input_file.path(), &decrypt_output);
     }
+
+    #[tokio::test]
+    async fn encrypt_and_execute_query() {
+        const EXPECTED: &[u128] = &[0, 8, 5];
+
+        let records: Vec<TestRawDataRecord> = vec![
+            TestRawDataRecord {
+                timestamp: 0,
+                user_id: 12345,
+                is_trigger_report: false,
+                breakdown_key: 2,
+                trigger_value: 0,
+            },
+            TestRawDataRecord {
+                timestamp: 4,
+                user_id: 68362,
+                is_trigger_report: false,
+                breakdown_key: 1,
+                trigger_value: 0,
+            },
+            TestRawDataRecord {
+                timestamp: 10,
+                user_id: 12345,
+                is_trigger_report: true,
+                breakdown_key: 0,
+                trigger_value: 5,
+            },
+            TestRawDataRecord {
+                timestamp: 12,
+                user_id: 68362,
+                is_trigger_report: true,
+                breakdown_key: 0,
+                trigger_value: 2,
+            },
+            TestRawDataRecord {
+                timestamp: 20,
+                user_id: 68362,
+                is_trigger_report: false,
+                breakdown_key: 1,
+                trigger_value: 0,
+            },
+            TestRawDataRecord {
+                timestamp: 30,
+                user_id: 68362,
+                is_trigger_report: true,
+                breakdown_key: 1,
+                trigger_value: 7,
+            },
+        ];
+        let query_size = QuerySize::try_from(records.len()).unwrap();
+        let mut input_file = NamedTempFile::new().unwrap();
+
+        for event in records {
+            let _ = event.to_csv(input_file.as_file_mut());
+            writeln!(input_file.as_file()).unwrap();
+        }
+        input_file.as_file_mut().flush().unwrap();
+
+        let output_dir = tempdir().unwrap();
+        let network_file = write_network_file();
+        let encrypt_args =
+            build_encrypt_args(input_file.path(), output_dir.path(), network_file.path());
+        let _ = encrypt(&encrypt_args);
+
+        let enc1 = output_dir.path().join("helper1.enc");
+        let enc2 = output_dir.path().join("helper2.enc");
+        let enc3 = output_dir.path().join("helper3.enc");
+
+        let mut buffers: [_; 3] = std::array::from_fn(|_| Vec::new());
+        for (i, path) in [enc1, enc2, enc3].iter().enumerate() {
+            let file = File::open(path).unwrap();
+            let reader = BufReader::new(file);
+            for line in reader.lines() {
+                let line = line.unwrap();
+                let encrypted_report_bytes = hex::decode(line.trim()).unwrap();
+                println!("{}", encrypted_report_bytes.len());
+                buffers[i].put_u16_le(encrypted_report_bytes.len().try_into().unwrap());
+                buffers[i].put_slice(encrypted_report_bytes.as_slice());
+            }
+        }
+
+        let world = TestWorld::default();
+        let contexts = world.contexts();
+
+        let mk_private_keys = vec![
+            hex::decode("53d58e022981f2edbf55fec1b45dbabd08a3442cb7b7c598839de5d7a5888bff")
+                .expect("manually provided for test"),
+            hex::decode("3a0a993a3cfc7e8d381addac586f37de50c2a14b1a6356d71e94ca2afaeb2569")
+                .expect("manually provided for test"),
+            hex::decode("1fb5c5274bf85fbe6c7935684ef05499f6cfb89ac21640c28330135cc0e8a0f7")
+                .expect("manually provided for test"),
+        ];
+
+        #[allow(clippy::large_futures)]
+        let results = join3v(buffers.into_iter().zip(contexts).zip(mk_private_keys).map(
+            |((buffer, ctx), mk_private_key)| {
+                let query_config = IpaQueryConfig {
+                    num_multi_bits: 3,
+                    per_user_credit_cap: 8,
+                    attribution_window_seconds: None,
+                    max_breakdown_key: 3,
+                    with_dp: 0,
+                    epsilon: 1.0,
+                    plaintext_match_keys: false,
+                };
+                let input = BodyStream::from(buffer);
+
+                let private_registry =
+                    Arc::new(KeyRegistry::<PrivateKeyOnly>::from_keys([PrivateKeyOnly(
+                        IpaPrivateKey::from_bytes(&mk_private_key)
+                            .expect("manually constructed for test"),
+                    )]));
+
+                OprfIpaQuery::<BA16, KeyRegistry<PrivateKeyOnly>>::new(
+                    query_config,
+                    private_registry,
+                )
+                .execute(ctx, query_size, input)
+            },
+        ))
+        .await;
+
+        assert_eq!(
+            results.reconstruct()[0..3]
+                .iter()
+                .map(U128Conversions::as_u128)
+                .collect::<Vec<u128>>(),
+            EXPECTED
+        );
+    }
 }

ipa-core/src/cli/mod.rs

@@ -1,6 +1,11 @@
 #[cfg(feature = "web-app")]
 mod clientconf;
-#[cfg(all(feature = "test-fixture", feature = "web-app", feature = "cli"))]
+#[cfg(all(
+    feature = "test-fixture",
+    feature = "web-app",
+    feature = "cli",
+    feature = "in-memory-infra"
+))]
 pub mod crypto;
 mod csv;
 mod ipa_output;

ipa-core/src/query/mod.rs

@@ -10,4 +10,5 @@ pub use processor::{
     NewQueryError, PrepareQueryError, Processor as QueryProcessor, QueryCompletionError,
     QueryInputError, QueryStatusError,
 };
+pub use runner::OprfIpaQuery;
 pub use state::QueryStatus;