Repository: cloudflare/daphne
Pull request #523 (2024-03-05 23:28:01)

Simplify `DapReportInitializer`
Long ago, an implementation of this trait was responsible for some of the early report validation checks, including checking for replays and batch overlaps. These days it is used solely for initializing preparation for a set of reports.

However, `MockAggregator` currently does a bit more, in particular what `DapAggregator::try_put_agg_share_span()` already does. After removing this code, we can simplify the `DapReportInitializer` API so that it does only what it's supposed to do.


Patch:
diff --git a/daphne/src/protocol/aggregator.rs b/daphne/src/protocol/aggregator.rs
index d8cd9ed..65eabe1 100644
--- a/daphne/src/protocol/aggregator.rs
+++ b/daphne/src/protocol/aggregator.rs
@@ -436,14 +436,7 @@ impl DapTaskConfig {
             }
         }
         let initialized_reports = initializer
-            .initialize_reports(
-                true,
-                task_id,
-                self,
-                part_batch_sel,
-                agg_param,
-                consumed_reports,
-            )
+            .initialize_reports(true, self, agg_param, consumed_reports)
             .await?;
 
         assert_eq!(initialized_reports.len(), helper_shares.len());
@@ -569,14 +562,7 @@ impl DapTaskConfig {
                 .map_err(|e| DapAbort::from_codec_error(e, *task_id))?;
 
         let initialized_reports = initializer
-            .initialize_reports(
-                false,
-                task_id,
-                self,
-                &agg_job_init_req.part_batch_sel,
-                &agg_param,
-                consumed_reports,
-            )
+            .initialize_reports(false, self, &agg_param, consumed_reports)
             .await?;
 
         Ok(initialized_reports)
diff --git a/daphne/src/roles/aggregator.rs b/daphne/src/roles/aggregator.rs
index ea8c66e..3024d65 100644
--- a/daphne/src/roles/aggregator.rs
+++ b/daphne/src/roles/aggregator.rs
@@ -11,9 +11,7 @@ use crate::{
     constants::DapMediaType,
     error::DapAbort,
     hpke::{HpkeConfig, HpkeDecrypter},
-    messages::{
-        BatchId, BatchSelector, HpkeConfigList, PartialBatchSelector, ReportId, TaskId, Time,
-    },
+    messages::{BatchId, BatchSelector, HpkeConfigList, ReportId, TaskId, Time},
     metrics::{DaphneMetrics, DaphneRequestType},
     protocol::aggregator::{EarlyReportStateConsumed, EarlyReportStateInitialized},
     DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapError, DapGlobalConfig,
@@ -26,12 +24,10 @@ use crate::{
 pub trait DapReportInitializer {
     /// Initialize a sequence of reports that are in the "consumed" state by initializing VDAF
     /// preparation.
-    async fn initialize_reports<'req>(
+    async fn initialize_reports(
         &self,
         is_leader: bool,
-        task_id: &TaskId,
         task_config: &DapTaskConfig,
-        part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
         consumed_reports: Vec<EarlyReportStateConsumed>,
     ) -> Result<Vec<EarlyReportStateInitialized>, DapError>;
diff --git a/daphne/src/roles/mod.rs b/daphne/src/roles/mod.rs
index 7788ea4..fc2f93e 100644
--- a/daphne/src/roles/mod.rs
+++ b/daphne/src/roles/mod.rs
@@ -184,6 +184,13 @@ mod test {
     use std::{collections::HashMap, sync::Arc, time::SystemTime, vec};
     use url::Url;
 
+    fn num_agg_job_reqs_for_version(version: DapVersion) -> usize {
+        match version {
+            DapVersion::Draft02 => 2,
+            DapVersion::Draft09 | DapVersion::Latest => 1,
+        }
+    }
+
     fn empty_report_extensions_for_version(version: DapVersion) -> Option<Vec<Extension>> {
         match version {
             DapVersion::Draft02 => Some(Vec::new()),
@@ -1197,9 +1204,8 @@ mod test {
         let task_config = t.helper.unchecked_get_task_config(task_id).await;
 
         let report = t.gen_test_report(task_id).await;
-        let (_, req) = t
-            .gen_test_agg_job_init_req(task_id, version, DapAggregationParam::Empty, vec![report])
-            .await;
+        let req = t.gen_test_upload_req(report.clone(), task_id).await;
+        leader::handle_upload_req(&*t.leader, &req).await.unwrap();
 
         // Add mock data to the aggreagte store backend. This is done in its own scope so that the lock
         // is released before running the test. Otherwise the test will deadlock.
@@ -1222,27 +1228,21 @@ mod test {
             );
         }
 
-        // Get AggregationJobResp and then extract the transition data from inside.
-        let agg_job_resp = AggregationJobResp::get_decoded(
-            &helper::handle_agg_job_req(&*t.helper, &req)
-                .await
-                .unwrap()
-                .payload,
-        )
-        .unwrap();
-        let transition = &agg_job_resp.transitions[0];
+        let query = task_config.query_for_current_batch_window(t.now);
+        let req = t.gen_test_coll_job_req(query, task_id).await;
+        leader::handle_coll_job_req(&*t.leader, &req).await.unwrap();
 
-        assert_eq!(t.helper.audit_log.invocations(), 1);
+        leader::process(&*t.leader, "leader.com", 100)
+            .await
+            .unwrap();
 
-        // Expect failure due to report store marked as collected.
-        assert_matches!(
-            transition.var,
-            TransitionVar::Failed(TransitionFailure::BatchCollected)
-        );
+        assert_metrics_include!(t.leader_registry, {
+            r#"report_counter{env="test_leader",host="leader.com",status="rejected_batch_collected"}"#: 1,
+        });
 
         assert_metrics_include!(t.helper_registry, {
             r#"report_counter{env="test_helper",host="helper.org",status="rejected_batch_collected"}"#: 1,
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: 1,
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version),
             r#"aggregation_job_counter{env="test_helper",host="helper.org",status="started"}"#: 1,
         });
     }
@@ -1826,13 +1826,8 @@ mod test {
             .await
             .unwrap();
 
-        let agg_job_req_count = match version {
-            DapVersion::Draft02 => 2,
-            DapVersion::Draft09 | DapVersion::Latest => 1,
-        };
-
         assert_metrics_include!(t.helper_registry, {
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: agg_job_req_count,
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version),
             r#"inbound_request_counter{env="test_helper",host="helper.org",type="collect"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="aggregated"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="collected"}"#: 1,
@@ -1875,13 +1870,8 @@ mod test {
             .await
             .unwrap();
 
-        let agg_job_req_count = match version {
-            DapVersion::Draft02 => 2,
-            DapVersion::Draft09 | DapVersion::Latest => 1,
-        };
-
         assert_metrics_include!(t.helper_registry, {
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: agg_job_req_count,
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version),
             r#"inbound_request_counter{env="test_helper",host="helper.org",type="collect"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="aggregated"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="collected"}"#: 1,
@@ -1981,13 +1971,8 @@ mod test {
             .await
             .unwrap();
 
-        let agg_job_req_count = match version {
-            DapVersion::Draft02 => 2,
-            DapVersion::Draft09 | DapVersion::Latest => 1,
-        };
-
         assert_metrics_include!(t.helper_registry, {
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: agg_job_req_count,
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version),
             r#"inbound_request_counter{env="test_helper",host="helper.org",type="collect"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="aggregated"}"#: 1,
             r#"report_counter{env="test_helper",host="helper.org",status="collected"}"#: 1,
@@ -2154,13 +2139,8 @@ mod test {
             .await
             .unwrap();
 
-        let agg_job_req_count = match version {
-            DapVersion::Draft02 => 2,
-            DapVersion::Draft09 | DapVersion::Latest => 1,
-        } * 2;
-
         assert_metrics_include!(t.helper_registry, {
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: agg_job_req_count,
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version) * 2,
             r#"inbound_request_counter{env="test_helper",host="helper.org",type="collect"}"#: 2,
             r#"report_counter{env="test_helper",host="helper.org",status="aggregated"}"#: 2,
             r#"report_counter{env="test_helper",host="helper.org",status="collected"}"#: 2,
diff --git a/daphne_server/src/roles/aggregator.rs b/daphne_server/src/roles/aggregator.rs
index 2cfba6d..6d5865a 100644
--- a/daphne_server/src/roles/aggregator.rs
+++ b/daphne_server/src/roles/aggregator.rs
@@ -10,10 +10,7 @@ use daphne::{
     error::DapAbort,
     fatal_error,
     hpke::{HpkeConfig, HpkeDecrypter},
-    messages::{
-        BatchId, BatchSelector, HpkeCiphertext, PartialBatchSelector, TaskId, Time,
-        TransitionFailure,
-    },
+    messages::{BatchId, BatchSelector, HpkeCiphertext, TaskId, Time, TransitionFailure},
     metrics::DaphneMetrics,
     roles::{aggregator::MergeAggShareError, DapAggregator, DapReportInitializer},
     DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapBatchBucket, DapError,
@@ -357,12 +354,10 @@ impl DapAggregator<DaphneAuth> for crate::App {
 
 #[async_trait]
 impl DapReportInitializer for crate::App {
-    async fn initialize_reports<'req>(
+    async fn initialize_reports(
         &self,
         is_leader: bool,
-        _task_id: &TaskId,
         task_config: &DapTaskConfig,
-        _part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
         consumed_reports: Vec<EarlyReportStateConsumed>,
     ) -> Result<Vec<EarlyReportStateInitialized>, DapError> {


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 617ae9c..783827f 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -15,9 +15,7 @@ use crate::{
         Report, ReportId, TaskId, Time, TransitionFailure,
     },
     metrics::{prometheus::DaphnePromMetrics, DaphneMetrics},
-    protocol::aggregator::{
-        EarlyReportState, EarlyReportStateConsumed, EarlyReportStateInitialized,
-    },
+    protocol::aggregator::{EarlyReportStateConsumed, EarlyReportStateInitialized},
     roles::{
         aggregator::MergeAggShareError,
         helper,
@@ -70,50 +68,38 @@ pub struct AggregationJobTest {
     pub(crate) helper_registry: prometheus::Registry,
     pub(crate) leader_metrics: DaphnePromMetrics,
     pub(crate) helper_metrics: DaphnePromMetrics,
-    pub(crate) leader_reports_processed: Arc<Mutex<HashSet<ReportId>>>,
-    pub(crate) helper_reports_processed: Arc<Mutex<HashSet<ReportId>>>,
 }
 
-// NOTE(cjpatton) This implementation of the report initializer is not feature complete. Since
-// [`AggrregationJobTest`], is only used to test the aggregation flow, features that are not
-// directly relevant to the tests aren't implemented.
+fn initialize_reports(
+    is_leader: bool,
+    task_config: &DapTaskConfig,
+    agg_param: &DapAggregationParam,
+    consumed_reports: Vec<EarlyReportStateConsumed>,
+) -> Result<Vec<EarlyReportStateInitialized>, DapError> {
+    consumed_reports
+        .into_iter()
+        .map(|consumed| {
+            EarlyReportStateInitialized::initialize(
+                is_leader,
+                &task_config.vdaf_verify_key,
+                &task_config.vdaf,
+                agg_param,
+                consumed,
+            )
+        })
+        .collect()
+}
+
 #[async_trait]
 impl DapReportInitializer for AggregationJobTest {
-    async fn initialize_reports<'req>(
+    async fn initialize_reports(
         &self,
         is_leader: bool,
-        _task_id: &TaskId,
         task_config: &DapTaskConfig,
-        _part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
         consumed_reports: Vec<EarlyReportStateConsumed>,
     ) -> Result<Vec<EarlyReportStateInitialized>, DapError> {
-        let mut reports_processed = if is_leader {
-            self.leader_reports_processed.lock().unwrap()
-        } else {
-            self.helper_reports_processed.lock().unwrap()
-        };
-
-        Ok(consumed_reports
-            .into_iter()
-            .map(|consumed| {
-                if reports_processed.contains(&consumed.metadata().id) {
-                    Ok(
-                        consumed
-                            .into_initialized_rejected_due_to(TransitionFailure::ReportReplayed),
-                    )
-                } else {
-                    reports_processed.insert(consumed.metadata().id);
-                    EarlyReportStateInitialized::initialize(
-                        is_leader,
-                        &task_config.vdaf_verify_key,
-                        &task_config.vdaf,
-                        agg_param,
-                        consumed,
-                    )
-                }
-            })
-            .collect::<Result<Vec<_>, _>>()?)
+        initialize_reports(is_leader, task_config, agg_param, consumed_reports)
     }
 }
 
@@ -179,8 +165,6 @@ impl AggregationJobTest {
             helper_registry,
             leader_metrics,
             helper_metrics,
-            leader_reports_processed: Default::default(),
-            helper_reports_processed: Default::default(),
         }
     }
 
@@ -701,24 +685,6 @@ impl MockAggregator {
         self.peer.is_some()
     }
 
-    /// Conducts checks on a received report to see whether:
-    /// 1) the report falls into a batch that has been already collected, or
-    fn check_report_has_been_collected(
-        &self,
-        task_id: &TaskId,
-        bucket: &DapBatchBucket,
-    ) -> Option<TransitionFailure> {
-        // Check AggStateStore to see whether the report is part of a batch that has already
-        // been collected.
-        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
-        let agg_store = guard.entry(*task_id).or_default();
-        if matches!(agg_store.get(bucket), Some(inner_agg_store) if inner_agg_store.collected) {
-            return Some(TransitionFailure::BatchCollected);
-        }
-
-        None
-    }
-
     fn get_hpke_receiver_config_for(&self, hpke_config_id: u8) -> Option<&HpkeReceiverConfig> {
         self.hpke_receiver_config_list
             .iter()
@@ -830,48 +796,14 @@ impl DapAuthorizedSender<BearerToken> for MockAggregator {
 
 #[async_trait]
 impl DapReportInitializer for MockAggregator {
-    async fn initialize_reports<'req>(
+    async fn initialize_reports(
         &self,
         is_leader: bool,
-        task_id: &TaskId,
         task_config: &DapTaskConfig,
-        part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
         consumed_reports: Vec<EarlyReportStateConsumed>,
     ) -> Result<Vec<EarlyReportStateInitialized>, DapError> {
-        let span = task_config.batch_span_for_meta(
-            part_batch_sel,
-            consumed_reports.iter().filter(|report| report.is_ready()),
-        )?;
-
-        let mut early_fails = HashMap::new();
-        for (bucket, ((), report_ids_and_time)) in span.iter() {
-            for (id, _) in report_ids_and_time {
-                // Check whether Report has been collected or replayed.
-                if let Some(transition_failure) =
-                    self.check_report_has_been_collected(task_id, bucket)
-                {
-                    early_fails.insert(*id, transition_failure);
-                };
-            }
-        }
-
-        Ok(consumed_reports
-            .into_iter()
-            .map(|consumed| {
-                if let Some(failure) = early_fails.get(&consumed.metadata().id) {
-                    Ok(consumed.into_initialized_rejected_due_to(*failure))
-                } else {
-                    EarlyReportStateInitialized::initialize(
-                        is_leader,
-                        &task_config.vdaf_verify_key,
-                        &task_config.vdaf,
-                        agg_param,
-                        consumed,
-                    )
-                }
-            })
-            .collect::<Result<Vec<_>, _>>()?)
+        initialize_reports(is_leader, task_config, agg_param, consumed_reports)
     }
 }