Repository: cloudflare/daphne
Pull request #57 (2022-07-22 00:10:40)

Implement `get_agg_share()` and test `http_post_aggregate_share()`
This PR implements necessary methods to test the aggregate share request/response procedure between leader/helper.

It also add a unit test that tests when a leader obtains an empty aggregate share.
The e2e unit test add the procedure where the Leader picks a collect job, produces its encrypted aggregate share, and lastly prepares an AggregateShareReq and sends it to the Helper to obtain the helper's encrypted aggregate share.

This partially implements #20, in particular the beginning of the Collection phase.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 478f4fa..90d8f57 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -366,7 +366,7 @@ pub struct DapOutputShare {
 }
 
 /// An aggregate share computed by combining a set of output shares.
-#[derive(Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct DapAggregateShare {
     pub(crate) report_count: u64,
     pub(crate) checksum: [u8; 32],
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 295a428..ca923d0 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -55,7 +55,7 @@ pub(crate) enum VdafMessage {
     Prio3ShareField128(Prio3PrepareShare<Field128, 16>),
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub(crate) enum VdafAggregateShare {
     Field64(prio::vdaf::AggregateShare<Field64>),
     Field128(prio::vdaf::AggregateShare<Field128>),


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 25749dd..8e39abd 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -9,9 +9,9 @@ use crate::{
     },
     hpke::HpkeReceiverConfig,
     messages::{
-        AggregateContinueReq, AggregateInitializeReq, AggregateResp, AggregateShareReq, CollectReq,
-        HpkeCiphertext, Id, Interval, Nonce, Report, ReportShare, Transition, TransitionFailure,
-        TransitionVar,
+        AggregateContinueReq, AggregateInitializeReq, AggregateResp, AggregateShareReq,
+        AggregateShareResp, CollectReq, HpkeCiphertext, Id, Interval, Nonce, Report, ReportShare,
+        Transition, TransitionFailure, TransitionVar,
     },
     roles::{DapAggregator, DapHelper, DapLeader},
     testing::{
@@ -577,4 +577,48 @@ async fn e2e() {
         .handle_final_agg_resp(leader_uncommitted, agg_resp)
         .unwrap();
     leader.put_out_shares(task_id, out_shares).await.unwrap();
+
+    // TODO(nakatsuka-y) Call http_post_collect to post a collect job.
+    //                   Assume for now that the leader has already picked a collect job.
+    let collect_req = CollectReq {
+        task_id: task_id.clone(),
+        batch_interval: task_config.current_batch_window(now),
+        agg_param: Vec::default(),
+    };
+
+    // Leader: Get Leader's encrypted aggregate share.
+    let leader_agg_share = leader
+        .get_agg_share(&collect_req.task_id, &collect_req.batch_interval)
+        .await
+        .unwrap();
+
+    let _leader_enc_agg_share = task_config
+        .vdaf
+        .produce_leader_encrypted_agg_share(
+            &task_config.collector_hpke_config,
+            &collect_req.task_id,
+            &collect_req.batch_interval,
+            &leader_agg_share,
+        )
+        .unwrap();
+
+    // Leader: Prepare AggregateShareReq.
+    let agg_share_req = AggregateShareReq {
+        task_id: collect_req.task_id.clone(),
+        batch_interval: collect_req.batch_interval.clone(),
+        agg_param: collect_req.agg_param.clone(),
+        report_count: leader_agg_share.report_count,
+        checksum: leader_agg_share.checksum,
+    };
+
+    // Leader: Send AggregateShareReq to Helper and receive AggregateShareResp.
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_AGG_SHARE_REQ),
+        payload: agg_share_req.get_encoded(),
+        url: task_config.helper_url.join("/aggregate_share").unwrap(),
+        sender_auth: Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string())),
+    };
+    let res = leader.http_post_aggregate_share(&req).await.unwrap();
+    let agg_share_resp = AggregateShareResp::get_decoded(&res.payload).unwrap();
+    let _helper_enc_agg_share = agg_share_resp.encrypted_agg_share;
 }
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 546bc46..50d6042 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -30,7 +30,7 @@ pub const TASK_LIST: &str = r#"{
         "helper_url": "http://helper.com:8788",
         "collector_hpke_config": "f40020000100010020a761d90c8c76d3d76349a3794a439a1572ab1fb8f13531d69744c92ea7757d7f",
         "min_batch_duration": 3600,
-        "min_batch_size": 100,
+        "min_batch_size": 1,
         "vdaf": {
             "prio3": "count"
         },
@@ -221,18 +221,51 @@ impl DapAggregator<BearerToken> for MockAggregator {
 
     async fn get_agg_share(
         &self,
-        _task_id: &Id,
-        _batch_interval: &Interval,
+        task_id: &Id,
+        batch_interval: &Interval,
     ) -> Result<DapAggregateShare, DapError> {
-        unreachable!("not implemented");
+        // Lock agg_store.
+        let mut agg_store_mutex_guard = self
+            .agg_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let agg_store = agg_store_mutex_guard.deref_mut();
+
+        // Fetch aggregate shares.
+        let mut agg_share = DapAggregateShare::default();
+        for (inner_bucket_info, agg_share_delta) in agg_store.iter() {
+            if task_id == &inner_bucket_info.task_id
+                && batch_interval.start <= inner_bucket_info.window
+                && batch_interval.end() > inner_bucket_info.window
+            {
+                agg_share.merge(agg_share_delta.clone())?;
+            }
+        }
+
+        Ok(agg_share)
     }
 
     async fn mark_collected(
         &self,
-        _task_id: &Id,
-        _batch_interval: &Interval,
+        task_id: &Id,
+        batch_interval: &Interval,
     ) -> Result<(), DapError> {
-        unreachable!("not implemented");
+        let mut report_store_mutex_guard = self
+            .report_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let report_store = report_store_mutex_guard.deref_mut();
+
+        for (inner_bucket_info, store) in report_store.iter_mut() {
+            if task_id == &inner_bucket_info.task_id
+                && batch_interval.start <= inner_bucket_info.window
+                && batch_interval.end() > inner_bucket_info.window
+            {
+                store.collected = true;
+            }
+        }
+
+        Ok(())
     }
 }
 
@@ -411,17 +444,12 @@ impl DapLeader<BearerToken> for MockAggregator {
         let report_store = report_store_mutex_guard.deref_mut();
 
         // Fetch reports.
-        let mut bucket_info = BucketInfo {
-            task_id: task_id.clone(),
-            window: batch_interval.start,
-        };
-        // TODO(nakatsuka-y) Should the batch window be inclusive?
-        while bucket_info.window < batch_interval.end() {
-            let num_reports_remaining = agg_rate - reports.len();
-
-            // If there are reports in this bucket, append into reports.
-            // If not, move on to the next bucket.
-            if let Some(ref mut store) = report_store.get_mut(&bucket_info) {
+        for (inner_bucket_info, store) in report_store.iter_mut() {
+            if task_id == &inner_bucket_info.task_id
+                && batch_interval.start <= inner_bucket_info.window
+                && batch_interval.end() > inner_bucket_info.window
+            {
+                let num_reports_remaining = agg_rate - reports.len();
                 let num_reports_drained = std::cmp::min(num_reports_remaining, store.pending.len());
                 let mut reports_drained: Vec<Report> =
                     store.pending.drain(..num_reports_drained).collect();
@@ -437,8 +465,6 @@ impl DapLeader<BearerToken> for MockAggregator {
                     break;
                 }
             }
-
-            bucket_info.window += task_config.min_batch_duration;
         }
 
         Ok(reports)
@@ -507,7 +533,7 @@ impl BucketInfo {
 pub(crate) struct ReportStore {
     pub(crate) pending: Vec<Report>,
     processed: HashSet<Nonce>,
-    collected: bool,
+    pub(crate) collected: bool,
 }
 
 // TODO(nakatsuka-y) remove dead_code once all functions are used