From 794c9ae119a93d5f83fa8f4f40c1db6b4c6a7d2b Mon Sep 17 00:00:00 2001
From: Richa Jain <richajain.dce@gmail.com>
Date: Tue, 10 Sep 2024 21:06:52 +0800
Subject: [PATCH] In case decryption fails, do not error out

---
 ipa-core/src/cli/crypto.rs | 113 +++++++++++++++++++------------------
 1 file changed, 59 insertions(+), 54 deletions(-)

diff --git a/ipa-core/src/cli/crypto.rs b/ipa-core/src/cli/crypto.rs
index 69acd1e11..2f77238d0 100644
--- a/ipa-core/src/cli/crypto.rs
+++ b/ipa-core/src/cli/crypto.rs
@@ -18,7 +18,7 @@ use crate::{
         U128Conversions,
     },
     hpke::{KeyRegistry, PrivateKeyOnly},
-    report::{EncryptedOprfReport, EventType, OprfReport, DEFAULT_KEY_ID},
+    report::{EncryptedOprfReport, EventType, InvalidReportError, OprfReport, DEFAULT_KEY_ID},
     secret_sharing::IntoShares,
     test_fixture::{ipa::TestRawDataRecord, Reconstruct},
 };
@@ -146,7 +146,7 @@ impl DecryptedReports {
 }
 
 impl Iterator for DecryptedReports {
-    type Item = OprfReport<BA8, BA3, BA20>;
+    type Item = Result<OprfReport<BA8, BA3, BA20>, InvalidReportError>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let mut line = String::new();
@@ -154,8 +154,7 @@ impl Iterator for DecryptedReports {
             let encrypted_report_bytes = hex::decode(line.trim()).unwrap();
             let enc_report =
                 EncryptedOprfReport::from_bytes(encrypted_report_bytes.as_slice()).unwrap();
-            let dec_report: OprfReport<BA8, BA3, BA20> =
-                enc_report.decrypt(&self.key_registry).unwrap();
+            let dec_report = enc_report.decrypt(&self.key_registry);
             Some(dec_report)
         } else {
             None
@@ -182,57 +181,63 @@ pub async fn decrypt_and_reconstruct(args: DecryptArgs) -> Result<(), BoxError>
             .open(args.output_file)?,
     );
 
-    for (dec_report1, (dec_report2, dec_report3)) in
-        decrypted_reports1.zip(decrypted_reports2.zip(decrypted_reports3))
+    for (idx, (dec_report1, (dec_report2, dec_report3))) in decrypted_reports1
+        .zip(decrypted_reports2.zip(decrypted_reports3))
+        .enumerate()
     {
-        let timestamp = [
-            dec_report1.timestamp,
-            dec_report2.timestamp,
-            dec_report3.timestamp,
-        ]
-        .reconstruct()
-        .as_u128();
-
-        let match_key = [
-            dec_report1.match_key,
-            dec_report2.match_key,
-            dec_report3.match_key,
-        ]
-        .reconstruct()
-        .as_u128();
-
-        // these aren't reconstucted, so we explictly make sure
-        // they are consistent across all three files, then set
-        // it to the first one (without loss of generality)
-        assert_eq!(dec_report1.event_type, dec_report2.event_type);
-        assert_eq!(dec_report2.event_type, dec_report3.event_type);
-        let is_trigger_report = dec_report1.event_type == EventType::Trigger;
-
-        let breakdown_key = [
-            dec_report1.breakdown_key,
-            dec_report2.breakdown_key,
-            dec_report3.breakdown_key,
-        ]
-        .reconstruct()
-        .as_u128();
-
-        let trigger_value = [
-            dec_report1.trigger_value,
-            dec_report2.trigger_value,
-            dec_report3.trigger_value,
-        ]
-        .reconstruct()
-        .as_u128();
-
-        writeln!(
-            writer,
-            "{},{},{},{},{}",
-            timestamp,
-            match_key,
-            u8::from(is_trigger_report),
-            breakdown_key,
-            trigger_value,
-        )?;
+        match (dec_report1, dec_report2, dec_report3) {
+            (Ok(dec_report1), Ok(dec_report2), Ok(dec_report3)) => {
+                let timestamp = [
+                    dec_report1.timestamp,
+                    dec_report2.timestamp,
+                    dec_report3.timestamp,
+                ]
+                .reconstruct()
+                .as_u128();
+
+                let match_key = [
+                    dec_report1.match_key,
+                    dec_report2.match_key,
+                    dec_report3.match_key,
+                ]
+                .reconstruct()
+                .as_u128();
+
+                // these aren't reconstucted, so we explictly make sure
+                // they are consistent across all three files, then set
+                // it to the first one (without loss of generality)
+                assert_eq!(dec_report1.event_type, dec_report2.event_type);
+                assert_eq!(dec_report2.event_type, dec_report3.event_type);
+                let is_trigger_report = dec_report1.event_type == EventType::Trigger;
+
+                let breakdown_key = [
+                    dec_report1.breakdown_key,
+                    dec_report2.breakdown_key,
+                    dec_report3.breakdown_key,
+                ]
+                .reconstruct()
+                .as_u128();
+
+                let trigger_value = [
+                    dec_report1.trigger_value,
+                    dec_report2.trigger_value,
+                    dec_report3.trigger_value,
+                ]
+                .reconstruct()
+                .as_u128();
+
+                writeln!(
+                    writer,
+                    "{},{},{},{},{}",
+                    timestamp,
+                    match_key,
+                    u8::from(is_trigger_report),
+                    breakdown_key,
+                    trigger_value,
+                )?;
+            }
+            _ => println!("Decryption failed for record no {idx}"),
+        }
     }
 
     Ok(())