Repository: cloudflare/daphne
Pull request #194 (2023-01-11 17:32:53)

Separate the early metadata check from the worker implementation.
Add unit testing of the check now that it is easier to test.


Patch:
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 3fb81bc..5945d69 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -992,3 +992,31 @@ where
 
     Ok(())
 }
+
+/// Check for transition failures due to:
+///
+/// * the report having already been processed
+/// * the report having already been collected
+/// * the report not being within time bounds
+///
+/// Returns `Some(TransitionFailure)` if there is a problem,
+/// or `None` if no transition failure occurred.
+pub fn early_metadata_check(
+    metadata: &ReportMetadata,
+    processed: bool,
+    collected: bool,
+    min_time: u64,
+    max_time: u64,
+) -> Option<TransitionFailure> {
+    if processed {
+        Some(TransitionFailure::ReportReplayed)
+    } else if collected {
+        Some(TransitionFailure::BatchCollected)
+    } else if metadata.time < min_time {
+        Some(TransitionFailure::ReportDropped)
+    } else if metadata.time > max_time {
+        Some(TransitionFailure::ReportTooEarly)
+    } else {
+        None
+    }
+}
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 94fb7db..1dc0305 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -49,7 +49,7 @@ use daphne::{
         BatchSelector, CollectReq, CollectResp, HpkeCiphertext, Id, PartialBatchSelector, Report,
         ReportId, ReportMetadata, TransitionFailure,
     },
-    roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
+    roles::{early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     taskprov::{bad_request, get_taskprov_task_config},
     DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
     DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig, DapVersion,
@@ -561,14 +561,10 @@ where
         {
             for metadata in span.get(bucket).unwrap() {
                 let processed = reports_processed.contains(&metadata.id);
-                if processed && !collected {
-                    early_fails.insert(metadata.id.clone(), TransitionFailure::ReportReplayed);
-                } else if !processed && collected {
-                    early_fails.insert(metadata.id.clone(), TransitionFailure::BatchCollected);
-                } else if metadata.time < min_time {
-                    early_fails.insert(metadata.id.clone(), TransitionFailure::ReportDropped);
-                } else if metadata.time > max_time {
-                    early_fails.insert(metadata.id.clone(), TransitionFailure::ReportTooEarly);
+                if let Some(failure) =
+                    early_metadata_check(metadata, processed, collected, min_time, max_time)
+                {
+                    early_fails.insert(metadata.id.clone(), failure);
                 }
             }
         }


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 6fba57f..46fbe33 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -12,11 +12,12 @@ use crate::{
     messages::{
         taskprov, AggregateContinueReq, AggregateInitializeReq, AggregateResp, AggregateShareReq,
         AggregateShareResp, BatchSelector, CollectReq, CollectResp, Extension, HpkeKemId, Id,
-        Interval, PartialBatchSelector, Query, Report, ReportShare, Time, Transition,
-        TransitionFailure, TransitionVar,
+        Interval, PartialBatchSelector, Query, Report, ReportId, ReportMetadata, ReportShare, Time,
+        Transition, TransitionFailure, TransitionVar,
     },
-    roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
+    roles::{early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     taskprov::TaskprovVersion,
+    test_version, test_versions,
     testing::{AggStore, DapBatchBucketOwned, MockAggregator, MockAggregatorReportSelector},
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateShare, DapCollectJob, DapGlobalConfig, DapLeaderTransition,
@@ -1593,3 +1594,78 @@ async fn e2e_taskprov(version: DapVersion) {
 }
 
 async_test_version! { e2e_taskprov, Draft02 }
+
+fn early_metadata_checks(version: DapVersion) {
+    let t = Test::new(version);
+    let mut rng = thread_rng();
+    let metadata = ReportMetadata {
+        id: ReportId(rng.gen()),
+        time: t.now,
+        extensions: vec![],
+    };
+    // We declare these so the first call to early_metadata_check() is more readable.
+    let processed = false;
+    let collected = false;
+    // A current not processed, not collected report is OK!
+    assert_matches!(
+        early_metadata_check(&metadata, processed, collected, t.now - 100, t.now + 100),
+        None
+    );
+    // A current processed, not collected report is TransitionFailure::ReportReplayed.
+    assert_matches!(
+        early_metadata_check(&metadata, true, false, t.now - 100, t.now + 100),
+        Some(TransitionFailure::ReportReplayed)
+    );
+    // A current not processed but collected report is TransitionFailure::BatchCollected.
+    assert_matches!(
+        early_metadata_check(&metadata, false, true, t.now - 100, t.now + 100),
+        Some(TransitionFailure::BatchCollected)
+    );
+    // A current processed and collected report is TransitionFailure::ReportReplayed.
+    assert_matches!(
+        early_metadata_check(&metadata, true, true, t.now - 100, t.now + 100),
+        Some(TransitionFailure::ReportReplayed)
+    );
+    // A not collected and not processed report at the future boundary is OK.
+    let metadata = ReportMetadata {
+        id: ReportId(rng.gen()),
+        time: t.now + 100,
+        extensions: vec![],
+    };
+    assert_matches!(
+        early_metadata_check(&metadata, false, false, t.now - 100, t.now + 100),
+        None
+    );
+    // A not collected and not processed report too far in the future is TransitionFailure::ReportTooEarly.
+    let metadata = ReportMetadata {
+        id: ReportId(rng.gen()),
+        time: t.now + 101,
+        extensions: vec![],
+    };
+    assert_matches!(
+        early_metadata_check(&metadata, false, false, t.now - 100, t.now + 100),
+        Some(TransitionFailure::ReportTooEarly)
+    );
+    // A not collected and not processed report at the past boundary is OK.
+    let metadata = ReportMetadata {
+        id: ReportId(rng.gen()),
+        time: t.now - 100,
+        extensions: vec![],
+    };
+    assert_matches!(
+        early_metadata_check(&metadata, false, false, t.now - 100, t.now + 100),
+        None
+    );
+    // A not collected and not processed report too far in the past is TransitionFailure::ReportDropped.
+    let metadata = ReportMetadata {
+        id: ReportId(rng.gen()),
+        time: t.now - 101,
+        extensions: vec![],
+    };
+    assert_matches!(
+        early_metadata_check(&metadata, false, false, t.now - 100, t.now + 100),
+        Some(TransitionFailure::ReportDropped)
+    );
+}
+
+test_versions! { early_metadata_checks }