Repository: cloudflare/daphne
Pull request #331 (2023-05-30 21:35:42)

Add support for Prio3SumVec
This PR should solve https://github.com/cloudflare/daphne/issues/267 by adding support for Prio3SumVec. I added a test vector to check that the shares are prepared properly.

@cjpatton I left a couple of comments with `TODO(tholop)` with my questions, I'm still exploring the codebase so I might be misunderstanding some things

tholop: Thanks @cjpatton, I pushed fixes for your comments and I'll open a new PR soon for the `Histogram` and `Sum` (not `Count` indeed) types.

Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index f1d3be2..240d5b0 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -483,6 +483,7 @@ impl AsRef<DapTaskConfig> for DapTaskConfig {
 pub enum DapMeasurement {
     U64(u64),
     U32Vec(Vec<u32>),
+    U128Vec(Vec<u128>),
 }
 
 /// The aggregate result computed by the Collector.
@@ -730,6 +731,10 @@ pub enum Prio3Config {
     /// The sum of 64-bit, unsigned integers. Each measurement is an integer in range `[0,
     /// 2^bits)`.
     Sum { bits: usize },
+
+    /// The element-wise sum of vectors. Each vector has `len` elements.
+    /// Each element is a 64-bit unsigned integer in range `[0,2^bits)`.
+    SumVec { bits: usize, len: usize },
 }
 
 /// DAP sender role.
diff --git a/daphne/src/vdaf/prio3.rs b/daphne/src/vdaf/prio3.rs
index 7279cfa..ef03061 100644
--- a/daphne/src/vdaf/prio3.rs
+++ b/daphne/src/vdaf/prio3.rs
@@ -60,6 +60,10 @@ pub(crate) fn prio3_shard(
             let vdaf = Prio3::new_sum(2, *bits)?;
             Ok(shard!(vdaf, &(measurement as u128), nonce))
         }
+        (Prio3Config::SumVec { bits, len }, DapMeasurement::U128Vec(measurement)) => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            Ok(shard!(vdaf, &measurement, nonce))
+        }
         _ => panic!("prio3_shard: unexpected VDAF config"),
     }
 }
@@ -147,6 +151,21 @@ pub(crate) fn prio3_prepare_init(
                 VdafMessage::Prio3ShareField128(share),
             ))
         }
+        Prio3Config::SumVec { bits, len } => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            let (state, share) = prep_init!(
+                vdaf,
+                verify_key,
+                agg_id,
+                nonce,
+                public_share_data,
+                input_share_data
+            );
+            Ok((
+                VdafState::Prio3Field128(state),
+                VdafMessage::Prio3ShareField128(share),
+            ))
+        }
     }
 }
 
@@ -213,6 +232,16 @@ pub(crate) fn prio3_leader_prepare_finish(
             let agg_share = VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?);
             (agg_share, outbound)
         }
+        (
+            Prio3Config::SumVec { bits, len },
+            VdafState::Prio3Field128(state),
+            VdafMessage::Prio3ShareField128(share),
+        ) => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            let (out_share, outbound) = leader_prep_fin!(vdaf, state, share, helper_share_data);
+            let agg_share = VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?);
+            (agg_share, outbound)
+        }
         _ => panic!("prio3_leader_prepare_finish: {ERR_FIELD_TYPE}"),
     };
 
@@ -262,6 +291,11 @@ pub(crate) fn prio3_helper_prepare_finish(
             let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
             VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?)
         }
+        (Prio3Config::SumVec { bits, len }, VdafState::Prio3Field128(state)) => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
+            VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?)
+        }
         _ => panic!("prio3_helper_prepare_finish: {ERR_FIELD_TYPE}"),
     };
 
@@ -313,6 +347,12 @@ pub(crate) fn prio3_decode_prepare_state(
                 Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
             ))
         }
+        Prio3Config::SumVec { bits, len } => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            Ok(VdafState::Prio3Field128(
+                Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
+            ))
+        }
     }
 }
 
@@ -362,5 +402,10 @@ pub(crate) fn prio3_unshard<M: IntoIterator<Item = Vec<u8>>>(
             let agg_res = unshard!(vdaf, num_measurements, agg_shares)?;
             Ok(DapAggregateResult::U128(agg_res))
         }
+        Prio3Config::SumVec { bits, len } => {
+            let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
+            let agg_res = unshard!(vdaf, num_measurements, agg_shares)?;
+            Ok(DapAggregateResult::U128Vec(agg_res))
+        }
     }
 }
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index 19fbe24..e75ce99 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -831,16 +831,21 @@ impl<'srv> DaphneWorker<'srv> {
             .ok_or_else(|| int_err("task ID is not valid URL-safe base64"))?;
 
         // VDAF config.
-        let vdaf = match (cmd.vdaf.typ.as_ref(), cmd.vdaf.bits) {
-            ("Prio3Count", None) => VdafConfig::Prio3(Prio3Config::Count),
-            ("Prio3Sum", Some(bits)) => {
+        let vdaf = match (cmd.vdaf.typ.as_ref(), cmd.vdaf.bits, cmd.vdaf.length) {
+            ("Prio3Count", None, None) => VdafConfig::Prio3(Prio3Config::Count),
+            ("Prio3Sum", Some(bits), None) => {
                 let bits = bits.parse().map_err(int_err)?;
                 VdafConfig::Prio3(Prio3Config::Sum { bits })
             }
+            ("Prio3SumVec", Some(bits), Some(length)) => {
+                let bits = bits.parse().map_err(int_err)?;
+                let len = length.parse().map_err(int_err)?;
+                VdafConfig::Prio3(Prio3Config::SumVec { bits, len })
+            }
             _ => return Err(int_err("command failed: unrecognized VDAF")),
         };
 
-        // VDAF verificaiton key.
+        // VDAF verification key.
         let vdaf_verify_key_data = decode_base64url_vec(cmd.vdaf_verify_key.as_bytes())
             .ok_or_else(|| int_err("VDAF verify key is not valid URL-safe base64"))?;
         let vdaf_verify_key = vdaf
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 7536c32..9d25a01 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -698,6 +698,8 @@ pub(crate) struct InternalTestVdaf {
     typ: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     bits: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    length: Option<String>,
 }
 
 #[derive(Deserialize)]


Test patch:
diff --git a/daphne/src/vdaf/prio3_test.rs b/daphne/src/vdaf/prio3_test.rs
index 20bf784..97998b6 100644
--- a/daphne/src/vdaf/prio3_test.rs
+++ b/daphne/src/vdaf/prio3_test.rs
@@ -46,6 +46,23 @@ fn prepare_histogram() {
     .unwrap();
 }
 
+#[test]
+fn prepare_sum_vec() {
+    test_prepare(
+        &Prio3Config::SumVec { bits: 23, len: 1 },
+        DapMeasurement::U128Vec(vec![(1 << 23) - 1]),
+        DapAggregateResult::U128Vec(vec![(1 << 23) - 1]),
+    )
+    .unwrap();
+
+    test_prepare(
+        &Prio3Config::SumVec { bits: 23, len: 3 },
+        DapMeasurement::U128Vec(vec![1, 0, 42]),
+        DapAggregateResult::U128Vec(vec![1, 0, 42]),
+    )
+    .unwrap();
+}
+
 fn test_prepare(
     config: &Prio3Config,
     measurement: DapMeasurement,