From 5b7d7950f992e113eca37418e1029f276590e0ec Mon Sep 17 00:00:00 2001
From: Joey Kraut <joey@renegade.fi>
Date: Wed, 17 Apr 2024 17:17:05 -0700
Subject: [PATCH] offline-phase: lowgear: Implement input auth methods for
 lowgear

---
 offline-phase/src/lib.rs         | 11 +++++----
 offline-phase/src/lowgear/mod.rs |  1 +
 offline-phase/src/structs.rs     | 40 ++++++++++++++++++++++++++++----
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/offline-phase/src/lib.rs b/offline-phase/src/lib.rs
index f6a37d5..049f548 100644
--- a/offline-phase/src/lib.rs
+++ b/offline-phase/src/lib.rs
@@ -189,9 +189,10 @@ pub(crate) mod test_helpers {
         let mut lowgear2 = LowGear::new(net2);
 
         // Setup the lowgear instances
-        let params = &lowgear1.params;
-        let keypair1 = BGVKeypair::gen(params);
-        let keypair2 = BGVKeypair::gen(params);
+        let params1 = &lowgear1.params;
+        let params2 = &lowgear2.params;
+        let keypair1 = BGVKeypair::gen(params1);
+        let keypair2 = BGVKeypair::gen(params2);
 
         let mac_share1 = Scalar::random(&mut rng);
         let mac_share2 = Scalar::random(&mut rng);
@@ -204,9 +205,9 @@ pub(crate) mod test_helpers {
 
         // Set the exchanged values
         lowgear1.other_pk = Some(keypair2.public_key());
-        lowgear1.other_mac_enc = Some(encrypt_all(mac_share2, &keypair2.public_key(), params));
+        lowgear1.other_mac_enc = Some(encrypt_all(mac_share2, &keypair2.public_key(), params1));
         lowgear2.other_pk = Some(keypair1.public_key());
-        lowgear2.other_mac_enc = Some(encrypt_all(mac_share1, &keypair1.public_key(), params));
+        lowgear2.other_mac_enc = Some(encrypt_all(mac_share1, &keypair1.public_key(), params2));
 
         (lowgear1, lowgear2)
     }
diff --git a/offline-phase/src/lowgear/mod.rs b/offline-phase/src/lowgear/mod.rs
index 5c25c78..462b794 100644
--- a/offline-phase/src/lowgear/mod.rs
+++ b/offline-phase/src/lowgear/mod.rs
@@ -115,6 +115,7 @@ impl<C: CurveGroup, N: MpcNetwork<C> + Unpin> LowGear<C, N> {
             self.inverse_tuples.clone(),
             self.shared_bits.clone(),
             self.shared_randomness.clone(),
+            self.input_masks.clone(),
             self.triples.clone(),
         ))
     }
diff --git a/offline-phase/src/structs.rs b/offline-phase/src/structs.rs
index 5c967be..00cd0a6 100644
--- a/offline-phase/src/structs.rs
+++ b/offline-phase/src/structs.rs
@@ -40,6 +40,8 @@ pub struct LowGearPrep<C: CurveGroup> {
     pub bits: ValueMacBatch<C>,
     /// The shared random values
     pub shared_randomness: ValueMacBatch<C>,
+    /// The input masks
+    pub input_masks: InputMasks<C>,
     /// The shared Beaver triplets
     pub triplets: (ValueMacBatch<C>, ValueMacBatch<C>, ValueMacBatch<C>),
 }
@@ -51,9 +53,10 @@ impl<C: CurveGroup> LowGearPrep<C> {
         inverse_pairs: (ValueMacBatch<C>, ValueMacBatch<C>),
         bits: ValueMacBatch<C>,
         shared_randomness: ValueMacBatch<C>,
+        input_masks: InputMasks<C>,
         triplets: (ValueMacBatch<C>, ValueMacBatch<C>, ValueMacBatch<C>),
     ) -> Self {
-        Self { params, inverse_pairs, bits, shared_randomness, triplets }
+        Self { params, inverse_pairs, bits, shared_randomness, input_masks, triplets }
     }
 
     /// Create an empty `LowGearPrep`
@@ -63,6 +66,7 @@ impl<C: CurveGroup> LowGearPrep<C> {
             inverse_pairs: (ValueMacBatch::new(vec![]), ValueMacBatch::new(vec![])),
             bits: ValueMacBatch::new(vec![]),
             shared_randomness: ValueMacBatch::new(vec![]),
+            input_masks: InputMasks::default(),
             triplets: (
                 ValueMacBatch::new(vec![]),
                 ValueMacBatch::new(vec![]),
@@ -94,6 +98,30 @@ impl<C: CurveGroup> LowGearPrep<C> {
 }
 
 impl<C: CurveGroup> PreprocessingPhase<C> for LowGearPrep<C> {
+    fn get_mac_key_share(&self) -> Scalar<C> {
+        self.params.mac_key_share
+    }
+
+    fn next_local_input_mask(&mut self) -> (Scalar<C>, ScalarShare<C>) {
+        self.input_masks.get_local_mask()
+    }
+
+    fn next_local_input_mask_batch(
+        &mut self,
+        num_values: usize,
+    ) -> (Vec<Scalar<C>>, Vec<ScalarShare<C>>) {
+        let (masks, mask_shares) = self.input_masks.get_local_mask_batch(num_values);
+        (masks, mask_shares.into_inner())
+    }
+
+    fn next_counterparty_input_mask(&mut self) -> ScalarShare<C> {
+        self.input_masks.get_counterparty_mask()
+    }
+
+    fn next_counterparty_input_mask_batch(&mut self, num_values: usize) -> Vec<ScalarShare<C>> {
+        self.input_masks.get_counterparty_mask_batch(num_values).into_inner()
+    }
+
     fn next_shared_bit(&mut self) -> ScalarShare<C> {
         self.bits.split_off(1).into_inner()[0]
     }
@@ -380,9 +408,13 @@ mod test {
     #[tokio::test]
     async fn test_lowgear_offline_phase() {
         // Setup the mock offline phase
+        const N: usize = 100;
         let (prep1, prep2) = mock_lowgear_with_triples(
-            100, // num_triples
-            |mut lowgear| async move { lowgear.get_offline_result().unwrap() },
+            N, // num_triples
+            |mut lowgear| async move {
+                lowgear.generate_input_masks(N).await.unwrap();
+                lowgear.get_offline_result().unwrap()
+            },
         )
         .await;
 
@@ -398,7 +430,7 @@ mod test {
                 let b_shared = fabric.share_scalar(b, PARTY1);
 
                 let c = a_shared * b_shared;
-                c.open().await
+                c.open_authenticated().await.unwrap()
             },
             prep1,
             prep2,