Repository: cloudflare/daphne
Pull request #204 (2023-01-20 22:59:36)

Have both Leader and Helper check for task expiration
Move the expiry check into `VdafConfig::consume_report_share()`, thereby ensuring that both Leader and Helper performn this check. This also cleans up some of the control flow in the `roles` module.


Patch:
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 4797fc7..1855542 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -421,12 +421,11 @@ where
             .vdaf
             .produce_agg_init_req(
                 self,
-                &task_config.vdaf_verify_key,
                 task_id,
+                task_config,
                 &agg_job_id,
                 part_batch_sel,
                 reports,
-                task_config.version,
             )
             .await?;
         let (state, agg_init_req) = match transition {
@@ -718,12 +717,7 @@ where
 
                 let transition = task_config
                     .vdaf
-                    .handle_agg_init_req(
-                        self,
-                        &task_config.vdaf_verify_key,
-                        &agg_init_req,
-                        task_config.version,
-                    )
+                    .handle_agg_init_req(self, task_config, &agg_init_req)
                     .await?;
 
                 // Check that helper state with task_id and agg_job_id does not exist.
@@ -740,18 +734,12 @@ where
                     DapHelperTransition::Continue(mut state, mut agg_resp) => {
                         let mut i = 0;
                         while i < state.seq.len() {
-                            let (_vdaf_state, time, report_id) = &state.seq[i];
-
-                            let early_result =
-                                early_rejects.get(&agg_resp.transitions[i].report_id);
-                            let result = if time >= &task_config.expiration {
-                                Some(&TransitionFailure::TaskExpired)
-                            } else {
-                                early_result
-                            };
+                            let (_vdaf_state, _time, report_id) = &state.seq[i];
 
                             // TODO Emit metrics for failure reasons.
-                            if let Some(failure) = result {
+                            if let Some(failure) =
+                                early_rejects.get(&agg_resp.transitions[i].report_id)
+                            {
                                 // Mark reports that were rejected early as TransitionVar::Failed.
                                 agg_resp.transitions[i].var = TransitionVar::Failed(*failure);
 
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 13de1d2..4687ac6 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -24,7 +24,7 @@ use crate::{
     },
     DapAbort, DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
     DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement, DapOutputShare,
-    DapVersion, VdafConfig,
+    DapTaskConfig, DapVersion, VdafConfig,
 };
 use prio::{
     codec::{CodecError, Decode, Encode, ParameterizedEncode},
@@ -288,8 +288,8 @@ impl VdafConfig {
         )
     }
 
-    /// Consume an encrypted input share sent in a report by the Client and return the Prepare
-    /// step. This is run by an Aggregator.
+    /// Consume a report share sent by the Client and return the initial Prepare step. This is run
+    /// by each Aggregator.
     ///
     /// # Inputs
     ///
@@ -309,14 +309,17 @@ impl VdafConfig {
         &self,
         decrypter: &impl HpkeDecrypter<'_>,
         is_leader: bool,
-        verify_key: &VdafVerifyKey,
         task_id: &Id,
+        task_config: &DapTaskConfig,
         metadata: &ReportMetadata,
         public_share: &[u8],
         encrypted_input_share: &HpkeCiphertext,
-        version: DapVersion,
     ) -> Result<(VdafState, VdafMessage), DapError> {
-        let input_share_text = match version {
+        if metadata.time >= task_config.expiration {
+            return Err(DapError::Transition(TransitionFailure::TaskExpired));
+        }
+
+        let input_share_text = match task_config.version {
             DapVersion::Draft02 => CTX_INPUT_SHARE_DRAFT02,
             DapVersion::Draft03 => CTX_INPUT_SHARE_DRAFT03,
             _ => return Err(unimplemented_version()),
@@ -334,7 +337,7 @@ impl VdafConfig {
 
         let mut aad = Vec::with_capacity(58);
         task_id.encode(&mut aad);
-        metadata.encode_with_param(&version, &mut aad);
+        metadata.encode_with_param(&task_config.version, &mut aad);
         // TODO spec: Consider folding the public share into a field called "header".
         encode_u32_bytes(&mut aad, public_share);
 
@@ -344,7 +347,7 @@ impl VdafConfig {
         // For Draft02, the encoded input share is the VDAF-specific payload, but for Draft03 and
         // later it is a serialized PlaintextInputShare.  For simplicity in later code, we wrap the Draft02
         // payload into a PlaintextInputShare.
-        let input_share = match version {
+        let input_share = match task_config.version {
             DapVersion::Draft02 => PlaintextInputShare {
                 extensions: vec![],
                 payload: encoded_input_share,
@@ -353,7 +356,7 @@ impl VdafConfig {
         };
 
         let agg_id = usize::from(!is_leader);
-        match (self, verify_key) {
+        match (self, &task_config.vdaf_verify_key) {
             (Self::Prio3(ref prio3_config), VdafVerifyKey::Prio3(ref verify_key)) => {
                 Ok(prio3_prepare_init(
                     prio3_config,
@@ -398,12 +401,11 @@ impl VdafConfig {
     pub(crate) async fn produce_agg_init_req(
         &self,
         decrypter: &impl HpkeDecrypter<'_>,
-        verify_key: &VdafVerifyKey,
         task_id: &Id,
+        task_config: &DapTaskConfig,
         agg_job_id: &Id,
         part_batch_sel: &PartialBatchSelector,
         reports: Vec<Report>,
-        version: DapVersion,
     ) -> Result<DapLeaderTransition<AggregateInitializeReq>, DapAbort> {
         let mut processed = HashSet::with_capacity(reports.len());
         let mut states = Vec::with_capacity(reports.len());
@@ -432,12 +434,11 @@ impl VdafConfig {
                 .consume_report_share(
                     decrypter,
                     true, // is_leader
-                    verify_key,
                     task_id,
+                    task_config,
                     &report.metadata,
                     &report.public_share,
                     &leader_share,
-                    version,
                 )
                 .await
             {
@@ -508,9 +509,8 @@ impl VdafConfig {
     pub(crate) async fn handle_agg_init_req(
         &self,
         decrypter: &impl HpkeDecrypter<'_>,
-        verify_key: &VdafVerifyKey,
+        task_config: &DapTaskConfig,
         agg_init_req: &AggregateInitializeReq,
-        version: DapVersion,
     ) -> Result<DapHelperTransition<AggregateResp>, DapAbort> {
         let num_reports = agg_init_req.report_shares.len();
         let mut processed = HashSet::with_capacity(num_reports);
@@ -526,12 +526,11 @@ impl VdafConfig {
                 .consume_report_share(
                     decrypter,
                     false, // is_leader
-                    verify_key,
                     &agg_init_req.task_id,
+                    task_config,
                     &report_share.metadata,
                     &report_share.public_share,
                     &report_share.encrypted_input_share,
-                    version,
                 )
                 .await
             {


Test patch:
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 48d09b3..2b2317d 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -7,12 +7,12 @@ use crate::{
     messages::{
         AggregateContinueReq, AggregateInitializeReq, AggregateResp, BatchSelector, HpkeAeadId,
         HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Interval, PartialBatchSelector,
-        Report, ReportId, Transition, TransitionFailure, TransitionVar,
+        Report, ReportId, Time, Transition, TransitionFailure, TransitionVar,
     },
     test_version, test_versions, DapAbort, DapAggregateResult, DapAggregateShare, DapError,
     DapHelperState, DapHelperTransition, DapLeaderState, DapLeaderTransition, DapLeaderUncommitted,
-    DapMeasurement, DapOutputShare, DapVersion, Prio3Config, VdafAggregateShare, VdafConfig,
-    VdafMessage, VdafState, VdafVerifyKey,
+    DapMeasurement, DapOutputShare, DapQueryConfig, DapTaskConfig, DapVersion, Prio3Config,
+    VdafAggregateShare, VdafConfig, VdafMessage, VdafState,
 };
 use assert_matches::assert_matches;
 use hpke_rs::HpkePublicKey;
@@ -23,6 +23,7 @@ use prio::vdaf::{
 };
 use rand::prelude::*;
 use std::{collections::HashMap, fmt::Debug, time::SystemTime};
+use url::Url;
 
 impl<M: Debug> DapLeaderTransition<M> {
     pub(crate) fn unwrap_continue(self) -> (DapLeaderState, M) {
@@ -70,6 +71,7 @@ const TEST_VDAF: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Count);
 async fn roundtrip_report(version: DapVersion) {
     let t = Test::new(TEST_VDAF, version);
     let report = t
+        .task_config
         .vdaf
         .produce_report(
             &t.client_hpke_config_list,
@@ -84,12 +86,11 @@ async fn roundtrip_report(version: DapVersion) {
         .consume_report_share(
             &t.leader_hpke_receiver_config,
             true, // is_leader
-            &t.vdaf_verify_key,
             &t.task_id,
+            &t.task_config,
             &report.metadata,
             &report.public_share,
             &report.encrypted_input_shares[0],
-            version,
         )
         .await
         .unwrap();
@@ -98,12 +99,11 @@ async fn roundtrip_report(version: DapVersion) {
         .consume_report_share(
             &t.helper_hpke_receiver_config,
             false, // is_leader
-            &t.vdaf_verify_key,
             &t.task_id,
+            &t.task_config,
             &report.metadata,
             &report.public_share,
             &report.encrypted_input_shares[1],
-            version,
         )
         .await
         .unwrap();
@@ -161,7 +161,7 @@ fn roundtrip_report_unsupported_hpke_suite(version: DapVersion) {
         },
     ];
 
-    let res = t.vdaf.produce_report(
+    let res = t.task_config.vdaf.produce_report(
         &unsupported_hpke_config_list,
         t.now,
         &t.task_id,
@@ -577,23 +577,20 @@ async fn helper_state_serialization(version: DapVersion) {
 
 async_test_versions! { helper_state_serialization }
 
-pub(crate) struct Test<'a> {
-    now: u64,
-    vdaf: &'a VdafConfig,
+pub(crate) struct Test {
+    now: Time,
     task_id: Id,
     agg_job_id: Id,
-    vdaf_verify_key: VdafVerifyKey,
+    task_config: DapTaskConfig,
     leader_hpke_receiver_config: HpkeReceiverConfig,
     helper_hpke_receiver_config: HpkeReceiverConfig,
     early_rejects: HashMap<ReportId, TransitionFailure>,
     client_hpke_config_list: Vec<HpkeConfig>,
-    collector_hpke_config: HpkeConfig,
     collector_hpke_receiver_config: HpkeReceiverConfig,
-    version: DapVersion,
 }
 
-impl<'a> Test<'a> {
-    pub(crate) fn new(vdaf: &'a VdafConfig, version: DapVersion) -> Test {
+impl Test {
+    pub(crate) fn new(vdaf: &VdafConfig, version: DapVersion) -> Test {
         let mut rng = thread_rng();
         let now = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
@@ -614,17 +611,25 @@ impl<'a> Test<'a> {
 
         Test {
             now,
-            vdaf,
             task_id,
             agg_job_id,
-            vdaf_verify_key,
             leader_hpke_receiver_config,
             helper_hpke_receiver_config,
             early_rejects: HashMap::default(),
             client_hpke_config_list: vec![leader_hpke_config, helper_hpke_config],
-            collector_hpke_config,
             collector_hpke_receiver_config,
-            version,
+            task_config: DapTaskConfig {
+                version,
+                leader_url: Url::parse("http://dummy.url").unwrap(),
+                helper_url: Url::parse("http://dummy.url").unwrap(),
+                time_precision: 500,
+                expiration: now + 500,
+                min_batch_size: 10,
+                query: DapQueryConfig::TimeInterval,
+                vdaf: vdaf.clone(),
+                vdaf_verify_key,
+                collector_hpke_config,
+            },
         }
     }
 
@@ -633,13 +638,14 @@ impl<'a> Test<'a> {
 
         for measurement in measurements.into_iter() {
             reports.push(
-                self.vdaf
+                self.task_config
+                    .vdaf
                     .produce_report(
                         &self.client_hpke_config_list,
                         self.now,
                         &self.task_id,
                         measurement,
-                        self.version,
+                        self.task_config.version,
                     )
                     .unwrap(),
             );
@@ -651,15 +657,15 @@ impl<'a> Test<'a> {
         &self,
         reports: Vec<Report>,
     ) -> DapLeaderTransition<AggregateInitializeReq> {
-        self.vdaf
+        self.task_config
+            .vdaf
             .produce_agg_init_req(
                 &self.leader_hpke_receiver_config,
-                &self.vdaf_verify_key,
                 &self.task_id,
+                &self.task_config,
                 &self.agg_job_id,
                 &PartialBatchSelector::TimeInterval,
                 reports,
-                self.version,
             )
             .await
             .unwrap()
@@ -670,12 +676,12 @@ impl<'a> Test<'a> {
         agg_init_req: AggregateInitializeReq,
     ) -> DapHelperTransition<AggregateResp> {
         let agg_resp = self
+            .task_config
             .vdaf
             .handle_agg_init_req(
                 &self.helper_hpke_receiver_config,
-                &self.vdaf_verify_key,
+                &self.task_config,
                 &agg_init_req,
-                self.version,
             )
             .await
             .unwrap();
@@ -696,7 +702,8 @@ impl<'a> Test<'a> {
         leader_state: DapLeaderState,
         agg_resp: AggregateResp,
     ) -> DapLeaderTransition<AggregateContinueReq> {
-        self.vdaf
+        self.task_config
+            .vdaf
             .handle_agg_resp(&self.task_id, &self.agg_job_id, leader_state, agg_resp)
             .unwrap()
     }
@@ -706,7 +713,8 @@ impl<'a> Test<'a> {
         leader_state: DapLeaderState,
         agg_resp: AggregateResp,
     ) -> DapAbort {
-        self.vdaf
+        self.task_config
+            .vdaf
             .handle_agg_resp(&self.task_id, &self.agg_job_id, leader_state, agg_resp)
             .err()
             .expect("handle_agg_resp() succeeded; expected failure")
@@ -717,7 +725,8 @@ impl<'a> Test<'a> {
         helper_state: DapHelperState,
         agg_cont_req: &AggregateContinueReq,
     ) -> DapHelperTransition<AggregateResp> {
-        self.vdaf
+        self.task_config
+            .vdaf
             .handle_agg_cont_req(helper_state, agg_cont_req)
             .unwrap()
     }
@@ -727,7 +736,8 @@ impl<'a> Test<'a> {
         helper_state: DapHelperState,
         agg_cont_req: &AggregateContinueReq,
     ) -> DapAbort {
-        self.vdaf
+        self.task_config
+            .vdaf
             .handle_agg_cont_req(helper_state, agg_cont_req)
             .err()
             .expect("handle_agg_cont_req() succeeded; expected failure")
@@ -738,7 +748,8 @@ impl<'a> Test<'a> {
         leader_uncommitted: DapLeaderUncommitted,
         agg_resp: AggregateResp,
     ) -> Vec<DapOutputShare> {
-        self.vdaf
+        self.task_config
+            .vdaf
             .handle_final_agg_resp(leader_uncommitted, agg_resp)
             .unwrap()
     }
@@ -748,13 +759,14 @@ impl<'a> Test<'a> {
         batch_selector: &BatchSelector,
         agg_share: &DapAggregateShare,
     ) -> HpkeCiphertext {
-        self.vdaf
+        self.task_config
+            .vdaf
             .produce_leader_encrypted_agg_share(
-                &self.collector_hpke_config,
+                &self.task_config.collector_hpke_config,
                 &self.task_id,
                 batch_selector,
                 agg_share,
-                self.version,
+                self.task_config.version,
             )
             .unwrap()
     }
@@ -764,13 +776,14 @@ impl<'a> Test<'a> {
         batch_selector: &BatchSelector,
         agg_share: &DapAggregateShare,
     ) -> HpkeCiphertext {
-        self.vdaf
+        self.task_config
+            .vdaf
             .produce_helper_encrypted_agg_share(
-                &self.collector_hpke_config,
+                &self.task_config.collector_hpke_config,
                 &self.task_id,
                 batch_selector,
                 agg_share,
-                self.version,
+                self.task_config.version,
             )
             .unwrap()
     }
@@ -781,14 +794,15 @@ impl<'a> Test<'a> {
         report_count: u64,
         enc_agg_shares: Vec<HpkeCiphertext>,
     ) -> DapAggregateResult {
-        self.vdaf
+        self.task_config
+            .vdaf
             .consume_encrypted_agg_shares(
                 &self.collector_hpke_receiver_config,
                 &self.task_id,
                 batch_selector,
                 report_count,
                 enc_agg_shares,
-                self.version,
+                self.task_config.version,
             )
             .await
             .unwrap()
@@ -812,9 +826,9 @@ impl<'a> Test<'a> {
         let (leader_state, agg_init) = self.produce_agg_init_req(reports).await.unwrap_continue();
         let (helper_state, agg_resp) = self.handle_agg_init_req(agg_init).await.unwrap_continue();
         let got = DapHelperState::get_decoded(
-            &self.vdaf,
+            &self.task_config.vdaf,
             &helper_state
-                .get_encoded(&self.vdaf)
+                .get_encoded(&self.task_config.vdaf)
                 .expect("failed to encode helper state"),
         )
         .expect("failed to decode helper state");