Repository: cloudflare/daphne
Pull request #43 (2022-07-13 19:53:16)

Add ReportStore and HelperStateStore structs
This PR implements ReportStore and HelperStateStore structs that are necessary for storing report shares and helper state.

nakatsuka-y: Currently only the basic functionality is implemented.
Test code will be added in the next couple of commits.

Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 27ade3b..478f4fa 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -313,7 +313,7 @@ pub struct DapLeaderUncommitted {
 }
 
 /// The Helper's state during the aggregation flow.
-#[derive(Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct DapHelperState {
     pub(crate) seq: Vec<(VdafState, Nonce)>,
 }
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index b8e3625..706ea1f 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -444,6 +444,18 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     &early_rejects,
                 )?;
 
+                // Check that helper state with task_id and agg_job_id does not exist.
+                if self
+                    .get_helper_state(&agg_init_req.task_id, &agg_init_req.agg_job_id)
+                    .await?
+                    .is_some()
+                {
+                    // TODO spec: Consider an explicit abort for this case.
+                    return Err(DapAbort::BadRequest(
+                        "unexpected message for aggregation job (already exists)".into(),
+                    ));
+                }
+
                 let agg_resp = match transition {
                     DapHelperTransition::Continue(state, agg_resp) => {
                         self.put_helper_state(


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 2224261..9fac475 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -4,39 +4,114 @@
 use crate::{
     auth::BearerToken,
     constants::{MEDIA_TYPE_AGG_INIT_REQ, MEDIA_TYPE_AGG_SHARE_REQ, MEDIA_TYPE_COLLECT_REQ},
-    hpke::{HpkeDecrypter, HpkeReceiverConfig},
+    hpke::HpkeReceiverConfig,
     messages::{
         AggregateInitializeReq, AggregateResp, AggregateShareReq, CollectReq, HpkeCiphertext, Id,
         Interval, Nonce, Report, ReportShare, TransitionFailure, TransitionVar,
     },
     roles::{DapAggregator, DapHelper, DapLeader},
-    testing::{MockAggregator, HPKE_RECEIVER_CONFIG_LIST, LEADER_BEARER_TOKEN},
-    DapAbort, DapError, DapMeasurement, DapRequest, Prio3Config, VdafConfig,
+    testing::{
+        BucketInfo, MockAggregator, ReportStore, HPKE_RECEIVER_CONFIG_LIST, LEADER_BEARER_TOKEN,
+    },
+    DapAbort, DapMeasurement, DapRequest, Prio3Config, VdafConfig,
 };
 use assert_matches::assert_matches;
 use prio::codec::{Decode, Encode};
-use rand::prelude::*;
-use std::vec;
+use rand::{thread_rng, Rng};
+use std::{ops::DerefMut, vec};
 
-#[tokio::test]
-async fn http_post_aggregate_unauthorized_request() {
+fn gen_test_agg_init_req(
+    agg: &MockAggregator,
+    report_shares: Vec<ReportShare>,
+) -> DapRequest<BearerToken> {
     let mut rng = thread_rng();
-    let helper = MockAggregator::new();
-    let task_id = helper.nominal_task_id();
-    let task_config = helper.get_task_config_for(task_id).unwrap();
+    let task_id = agg.nominal_task_id();
+    let task_config = agg.get_task_config_for(task_id).unwrap();
 
-    let mut req = DapRequest {
+    DapRequest {
         media_type: Some(MEDIA_TYPE_AGG_INIT_REQ),
         payload: AggregateInitializeReq {
             task_id: task_id.clone(),
             agg_job_id: Id(rng.gen()),
             agg_param: Vec::default(),
-            report_shares: Vec::default(),
+            report_shares,
         }
         .get_encoded(),
         url: task_config.helper_url.join("/aggregate").unwrap(),
         sender_auth: None,
-    };
+    }
+}
+
+fn gen_test_agg_share_req(
+    agg: &MockAggregator,
+    report_count: u64,
+    checksum: [u8; 32],
+) -> DapRequest<BearerToken> {
+    let task_id = agg.nominal_task_id();
+    let task_config = agg.get_task_config_for(task_id).unwrap();
+
+    DapRequest {
+        media_type: Some(MEDIA_TYPE_AGG_SHARE_REQ),
+        payload: AggregateShareReq {
+            task_id: task_id.clone(),
+            batch_interval: Interval::default(),
+            agg_param: Vec::default(),
+            report_count,
+            checksum,
+        }
+        .get_encoded(),
+        url: task_config.helper_url.join("/aggregate_share").unwrap(),
+        sender_auth: None,
+    }
+}
+
+fn gen_test_collect_req(agg: &MockAggregator) -> DapRequest<BearerToken> {
+    let task_id = agg.nominal_task_id();
+    let task_config = agg.get_task_config_for(task_id).unwrap();
+
+    DapRequest {
+        media_type: Some(MEDIA_TYPE_COLLECT_REQ),
+        payload: CollectReq {
+            task_id: task_id.clone(),
+            batch_interval: Interval::default(),
+            agg_param: Vec::default(),
+        }
+        .get_encoded(),
+        url: task_config.leader_url.join("/collect").unwrap(),
+        sender_auth: None,
+    }
+}
+
+fn gen_test_report(task_id: &Id) -> Report {
+    // Construct HPKE receiver config List.
+    let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
+        serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
+            .expect("failed to parse hpke_receiver_config_list");
+
+    // Construct HPKE config list.
+    let mut hpke_config_list = Vec::with_capacity(hpke_receiver_config_list.len());
+    for receiver_config in hpke_receiver_config_list {
+        hpke_config_list.push(receiver_config.config);
+    }
+
+    // Construct report.
+    let vdaf_config: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Count);
+    let report = vdaf_config
+        .produce_report(
+            &hpke_config_list,
+            1637361337,
+            task_id,
+            DapMeasurement::U64(1),
+        )
+        .unwrap();
+
+    report
+}
+
+#[tokio::test]
+async fn http_post_aggregate_unauthorized_request() {
+    let helper = MockAggregator::new();
+    let mut req = gen_test_agg_init_req(&helper, Vec::default());
 
     // Expect failure due to missing bearer token.
     assert_matches!(
@@ -55,22 +130,7 @@ async fn http_post_aggregate_unauthorized_request() {
 #[tokio::test]
 async fn http_post_aggregate_share_unauthorized_request() {
     let helper = MockAggregator::new();
-    let task_id = helper.nominal_task_id();
-    let task_config = helper.get_task_config_for(task_id).unwrap();
-
-    let mut req = DapRequest {
-        media_type: Some(MEDIA_TYPE_AGG_SHARE_REQ),
-        payload: AggregateShareReq {
-            task_id: task_id.clone(),
-            batch_interval: Interval::default(),
-            agg_param: Vec::default(),
-            report_count: 0,
-            checksum: [0; 32],
-        }
-        .get_encoded(),
-        url: task_config.helper_url.join("/aggregate_share").unwrap(),
-        sender_auth: None,
-    };
+    let mut req = gen_test_agg_share_req(&helper, 0, [0; 32]);
 
     // Expect failure due to missing bearer token.
     assert_matches!(
@@ -89,20 +149,7 @@ async fn http_post_aggregate_share_unauthorized_request() {
 #[tokio::test]
 async fn http_post_collect_unauthorized_request() {
     let leader = MockAggregator::new();
-    let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
-
-    let mut req = DapRequest {
-        media_type: Some(MEDIA_TYPE_COLLECT_REQ),
-        payload: CollectReq {
-            task_id: task_id.clone(),
-            batch_interval: Interval::default(),
-            agg_param: Vec::default(),
-        }
-        .get_encoded(),
-        url: task_config.leader_url.join("/collect").unwrap(),
-        sender_auth: None,
-    };
+    let mut req = gen_test_collect_req(&leader);
 
     // Expect failure due to missing bearer token.
     assert_matches!(
@@ -119,34 +166,23 @@ async fn http_post_collect_unauthorized_request() {
 }
 
 #[tokio::test]
-async fn http_post_aggregate_invalid_ciphertext() {
+async fn http_post_aggregate_failure_hpke_decrypt_error() {
     let helper = MockAggregator::new();
-    let task_id = helper.nominal_task_id();
-    let task_config = helper.get_task_config_for(task_id).unwrap();
-
-    let req = DapRequest {
-        media_type: Some(MEDIA_TYPE_AGG_INIT_REQ),
-        payload: AggregateInitializeReq {
-            task_id: task_id.clone(),
-            agg_job_id: Id([1; 32]),
-            agg_param: b"this is an aggregation parameter".to_vec(),
-            report_shares: vec![ReportShare {
-                nonce: Nonce {
-                    time: 1637361337,
-                    rand: [1; 16],
-                },
-                extensions: Vec::default(),
-                encrypted_input_share: HpkeCiphertext {
-                    config_id: 23,
-                    enc: b"invalid encapsulated key".to_vec(),
-                    payload: b"invalid ciphertext".to_vec(),
-                },
-            }],
-        }
-        .get_encoded(),
-        url: task_config.helper_url.join("/aggregate").unwrap(),
-        sender_auth: Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string())),
-    };
+    let report_shares = vec![ReportShare {
+        nonce: Nonce {
+            time: 1637361337,
+            rand: [1; 16],
+        },
+        extensions: Vec::default(),
+        encrypted_input_share: HpkeCiphertext {
+            config_id: 23,
+            enc: b"invalid encapsulated key".to_vec(),
+            payload: b"invalid ciphertext".to_vec(),
+        },
+    }];
+    let sender_auth = Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string()));
+    let mut req = gen_test_agg_init_req(&helper, report_shares);
+    req.sender_auth = sender_auth;
 
     // Get AggregateResp and then extract the transition data from inside.
     let agg_resp =
@@ -162,51 +198,21 @@ async fn http_post_aggregate_invalid_ciphertext() {
 }
 
 #[tokio::test]
-async fn http_post_aggregate_valid_ciphertext() {
+async fn http_post_aggregate_transition_continue() {
     let helper = MockAggregator::new();
     let task_id = helper.nominal_task_id();
-    let task_config = helper.get_task_config_for(task_id).unwrap();
 
-    // Construct HPKE receiver config List.
-    let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
-        serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
-            .expect("failed to parse hpke_receiver_config_list");
-
-    // Construct HPKE config list.
-    let mut hpke_config_list = Vec::with_capacity(hpke_receiver_config_list.len());
-    for receiver_config in hpke_receiver_config_list {
-        hpke_config_list.push(receiver_config.config);
-    }
+    let report = gen_test_report(task_id);
+    let report_shares = vec![ReportShare {
+        nonce: report.nonce,
+        extensions: report.extensions,
+        // 1st share is for Leader and the rest is for Helpers (note that there is only 1 helper).
+        encrypted_input_share: report.encrypted_input_shares[1].clone(),
+    }];
+    let sender_auth = Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string()));
 
-    // Construct report.
-    let vdaf_config: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Count);
-    let report = vdaf_config
-        .produce_report(
-            &hpke_config_list,
-            1637361337,
-            task_id,
-            DapMeasurement::U64(1),
-        )
-        .unwrap();
-
-    // Construct DapRequest.
-    let req = DapRequest {
-        media_type: Some(MEDIA_TYPE_AGG_INIT_REQ),
-        payload: AggregateInitializeReq {
-            task_id: task_id.clone(),
-            agg_job_id: Id([1; 32]),
-            agg_param: b"this is an aggregation parameter".to_vec(),
-            report_shares: vec![ReportShare {
-                nonce: report.nonce,
-                extensions: report.extensions,
-                // 1st share is for Leader and the rest is for Helpers (note that there is only 1 helper).
-                encrypted_input_share: report.encrypted_input_shares[1].clone(),
-            }],
-        }
-        .get_encoded(),
-        url: task_config.helper_url.join("/aggregate").unwrap(),
-        sender_auth: Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string())),
-    };
+    let mut req = gen_test_agg_init_req(&helper, report_shares);
+    req.sender_auth = sender_auth;
 
     // Get AggregateResp and then extract the transition data from inside.
     let agg_resp =
@@ -218,71 +224,78 @@ async fn http_post_aggregate_valid_ciphertext() {
     assert_matches!(transition.var, TransitionVar::Continued(_));
 }
 
-#[test]
-fn hpke_decrypter() {
-    // Construct mock aggregator.
-    let aggregator = MockAggregator::new();
-    let task_id = aggregator.nominal_task_id();
+#[tokio::test]
+async fn http_post_aggregate_failure_batch_collected() {
+    let helper = MockAggregator::new();
+    let task_id = helper.nominal_task_id();
+    let task_config = helper.get_task_config_for(task_id).unwrap();
+
+    let report = gen_test_report(task_id);
+    let report_shares = vec![ReportShare {
+        nonce: report.nonce.clone(),
+        extensions: report.extensions,
+        // 1st share is for Leader and the rest is for Helpers (note that there is only 1 helper).
+        encrypted_input_share: report.encrypted_input_shares[1].clone(),
+    }];
+    let sender_auth = Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string()));
 
-    // Initialize variables for mock report.
-    let info = b"info string";
-    let aad = b"associated data";
-    let plaintext = b"plaintext";
-    let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
-        serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
-            .expect("failed to parse hpke_receiver_config_list");
-    let config = &hpke_receiver_config_list[0].config;
-    let (enc, ciphertext) = config.encrypt(info, aad, plaintext).unwrap();
+    let mut req = gen_test_agg_init_req(&helper, report_shares);
+    req.sender_auth = sender_auth;
 
-    // Construct mock report.
-    let report = Report {
-        task_id: Id([23; 32]),
-        nonce: Nonce {
-            time: 1637364244,
-            rand: [2; 16],
-        },
-        extensions: Vec::default(),
-        encrypted_input_shares: vec![HpkeCiphertext {
-            config_id: 23,
-            enc: enc,
-            payload: ciphertext,
-        }],
-    };
-
-    // Expect false due to non-existing config ID.
-    assert_eq!(aggregator.can_hpke_decrypt(&task_id, 0), false);
-
-    // Expect true due to existing config ID.
-    assert_eq!(
-        aggregator.can_hpke_decrypt(&task_id, report.encrypted_input_shares[0].config_id),
-        true
-    );
+    // This is to ensure that the lock is released before we call http_post_aggregate.
+    {
+        let mut report_store_mutex_guard = helper.report_store.lock().expect("lock() failed");
+        let report_store = report_store_mutex_guard.deref_mut();
 
-    // Expect decryption to fail.
+        // Create a new instance of report store associated to the bucket_info.
+        let bucket_info = BucketInfo::new(task_config, task_id, &report.nonce);
+        report_store.insert(bucket_info.clone(), ReportStore::new());
+
+        // Intentionally mark report store as `collected`.
+        report_store
+            .get_mut(&bucket_info)
+            .expect("report_store not found")
+            .process_mark_collected();
+    }
+
+    // Get AggregateResp and then extract the transition data from inside.
+    let agg_resp =
+        AggregateResp::get_decoded(&helper.http_post_aggregate(&req).await.unwrap().payload)
+            .unwrap();
+    let transition = &agg_resp.transitions[0];
+
+    // Expect failure due to report store marked as collected.
     assert_matches!(
-        aggregator.hpke_decrypt(
-            &report.task_id,
-            info,
-            aad,
-            &HpkeCiphertext {
-                config_id: 0,
-                enc: vec![],
-                payload: b"ciphertext".to_vec(),
-            }
-        ),
-        Err(DapError::Transition(TransitionFailure::HpkeUnknownConfigId))
+        transition.var,
+        TransitionVar::Failed(TransitionFailure::BatchCollected)
     );
+}
+
+#[tokio::test]
+async fn http_post_aggregate_abort_helper_state_overwritten() {
+    let helper = MockAggregator::new();
+    let task_id = helper.nominal_task_id();
+
+    let report = gen_test_report(task_id);
+    let report_shares = vec![ReportShare {
+        nonce: report.nonce.clone(),
+        extensions: report.extensions,
+        // 1st share is for Leader and the rest is for Helpers (note that there is only 1 helper).
+        encrypted_input_share: report.encrypted_input_shares[1].clone(),
+    }];
+    let sender_auth = Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string()));
+
+    let mut req = gen_test_agg_init_req(&helper, report_shares);
+    req.sender_auth = sender_auth;
+
+    // Send aggregate request.
+    let _ = helper.http_post_aggregate(&req).await;
+
+    // Send another aggregate request.
+    let res = helper.http_post_aggregate(&req).await;
 
-    // Expect decryption to succeed.
-    assert_eq!(
-        aggregator
-            .hpke_decrypt(
-                &report.task_id,
-                info,
-                aad,
-                &report.encrypted_input_shares[0]
-            )
-            .unwrap(),
-        plaintext
+    // Expect failure due to overwriting existing helper state.
+    assert_matches!(res.unwrap_err(), DapAbort::BadRequest(e) =>
+        assert_eq!(e, "unexpected message for aggregation job (already exists)")
     );
 }
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 144fddc..c642843 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -15,7 +15,12 @@ use crate::{
     DapResponse, DapTaskConfig,
 };
 use async_trait::async_trait;
-use std::collections::HashMap;
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::{HashMap, HashSet},
+    ops::DerefMut,
+    sync::{Arc, Mutex},
+};
 use url::Url;
 
 pub const TASK_LIST: &str = r#"{
@@ -61,27 +66,32 @@ pub const COLLECTOR_BEARER_TOKEN: &str = "syfRfvcvNFF5MJk4Y-B7xjRIqD_iNzhaaEB9mY
 pub(crate) struct MockAggregator {
     tasks: HashMap<Id, DapTaskConfig>,
     hpke_receiver_config_list: Vec<HpkeReceiverConfig>,
+    pub(crate) report_store: Arc<Mutex<HashMap<BucketInfo, ReportStore>>>,
+    helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapHelperState>>>,
 }
 
 #[allow(dead_code)]
 impl MockAggregator {
-    pub fn new() -> Self {
-        // Construct task list
+    pub(crate) fn new() -> Self {
         let tasks: HashMap<Id, DapTaskConfig> =
             serde_json::from_str(TASK_LIST).expect("failed to parse task list");
 
-        // Construct HPKE receiver config List
         let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
             serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
                 .expect("failed to parse hpke_receiver_config_list");
 
+        let report_store = Arc::new(Mutex::new(HashMap::new()));
+        let helper_state_store = Arc::new(Mutex::new(HashMap::new()));
+
         Self {
             tasks,
             hpke_receiver_config_list,
+            report_store,
+            helper_state_store,
         }
     }
 
-    pub fn get_hpke_receiver_config_for(&self, hpke_config_id: u8) -> Option<&HpkeReceiverConfig> {
+    fn get_hpke_receiver_config_for(&self, hpke_config_id: u8) -> Option<&HpkeReceiverConfig> {
         for hpke_receiver_config in self.hpke_receiver_config_list.iter() {
             if hpke_config_id == hpke_receiver_config.config.id {
                 return Some(hpke_receiver_config);
@@ -91,7 +101,7 @@ impl MockAggregator {
     }
 
     /// Task to use for nominal tests.
-    pub fn nominal_task_id(&self) -> &Id {
+    pub(crate) fn nominal_task_id(&self) -> &Id {
         // Just use the first key in the hash map.
         self.tasks.keys().next().as_ref().unwrap()
     }
@@ -195,32 +205,103 @@ impl DapAggregator<BearerToken> for MockAggregator {
 impl DapHelper<BearerToken> for MockAggregator {
     async fn mark_aggregated(
         &self,
-        _task_id: &Id,
-        _report_shares: &[ReportShare],
+        task_id: &Id,
+        report_shares: &[ReportShare],
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError> {
-        // Return empty HashMap (for now).
-        // TODO(nakatsuka-y) Implement correct functionality.
-        let early_fails: HashMap<Nonce, TransitionFailure> = HashMap::new();
-        return Ok(early_fails);
+        let mut early_fails: HashMap<Nonce, TransitionFailure> = HashMap::new();
+
+        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 report_share in report_shares.iter() {
+            let bucket_info = BucketInfo::new(
+                self.tasks
+                    .get(task_id)
+                    .ok_or_else(|| DapError::fatal("no task found"))?,
+                task_id,
+                &report_share.nonce,
+            );
+
+            if !report_store.contains_key(&bucket_info) {
+                report_store.insert(bucket_info.clone(), ReportStore::new());
+            }
+
+            match report_store
+                .get_mut(&bucket_info)
+                .unwrap()
+                .process_put_processed(report_share.nonce.clone())
+            {
+                Ok(()) => (),
+                Err(failure_reason) => {
+                    early_fails.insert(report_share.nonce.clone(), failure_reason);
+                }
+            }
+        }
+
+        Ok(early_fails)
     }
 
     async fn put_helper_state(
         &self,
-        _task_id: &Id,
-        _agg_job_id: &Id,
-        _helper_state: &DapHelperState,
+        task_id: &Id,
+        agg_job_id: &Id,
+        helper_state: &DapHelperState,
     ) -> Result<(), DapError> {
-        // Return empty Ok (for now).
-        // TODO(nakatsuka-y) Implement correct functionality.
+        let helper_state_info = HelperStateInfo {
+            task_id: task_id.clone(),
+            agg_job_id: agg_job_id.clone(),
+        };
+
+        let mut helper_state_store_mutex_guard = self
+            .helper_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+
+        let helper_state_store = helper_state_store_mutex_guard.deref_mut();
+
+        if helper_state_store.contains_key(&helper_state_info) {
+            return Err(DapError::Fatal(
+                "overwriting existing helper state".to_string(),
+            ));
+        }
+
+        // NOTE: This code is only correct for VDAFs with exactly one round of preparation.
+        // For VDAFs with more rounds, the helper state blob will need to be updated here.
+        helper_state_store.insert(helper_state_info, helper_state.clone());
+
         Ok(())
     }
 
     async fn get_helper_state(
         &self,
-        _task_id: &Id,
-        _agg_job_id: &Id,
+        task_id: &Id,
+        agg_job_id: &Id,
     ) -> Result<Option<DapHelperState>, DapError> {
-        unreachable!("not implemented");
+        let helper_state_info = HelperStateInfo {
+            task_id: task_id.clone(),
+            agg_job_id: agg_job_id.clone(),
+        };
+
+        let mut helper_state_store_mutex_guard = self
+            .helper_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+
+        let helper_state_store = helper_state_store_mutex_guard.deref_mut();
+
+        // NOTE: This code is only correct for VDAFs with exactly one round of preparation.
+        // For VDAFs with more rounds, the helper state blob will need to be updated here.
+        if helper_state_store.contains_key(&helper_state_info) {
+            let helper_state = helper_state_store.remove(&helper_state_info);
+
+            return Ok(helper_state);
+        }
+
+        Ok(None)
     }
 }
 
@@ -272,3 +353,91 @@ impl DapLeader<BearerToken> for MockAggregator {
         unreachable!("not implemented");
     }
 }
+
+/// Information associated to a certain helper state for a given task ID and aggregate job ID.
+#[derive(Clone, Eq, Hash, PartialEq, Deserialize, Serialize)]
+pub(crate) struct HelperStateInfo {
+    task_id: Id,
+    agg_job_id: Id,
+}
+
+/// Information associated to a certain report for a given task and nonce to decide which bucket it would be put into.
+#[derive(Clone, Eq, Hash, PartialEq, Deserialize, Serialize)]
+pub(crate) struct BucketInfo {
+    task_id: Id,
+    window: u64,
+}
+
+impl BucketInfo {
+    pub(crate) fn new(task_config: &DapTaskConfig, task_id: &Id, nonce: &Nonce) -> Self {
+        let window = nonce.time - (nonce.time % task_config.min_batch_duration);
+
+        Self {
+            task_id: task_id.clone(),
+            window,
+        }
+    }
+}
+
+// TODO(nakatsuka-y) remove dead_code once all functions are used
+#[allow(dead_code)]
+pub(crate) struct ReportStore {
+    pending: Vec<Report>,
+    processed: HashSet<Nonce>,
+    collected: bool,
+}
+
+// TODO(nakatsuka-y) remove dead_code once all functions are used
+#[allow(dead_code)]
+impl ReportStore {
+    pub(crate) fn new() -> Self {
+        let pending: Vec<Report> = Vec::new();
+        let processed: HashSet<Nonce> = HashSet::new();
+        let collected = false;
+
+        Self {
+            pending,
+            processed,
+            collected,
+        }
+    }
+
+    fn checked_process(&mut self, nonce: &Nonce) -> Result<(), TransitionFailure> {
+        let observed = self.processed.contains(nonce);
+        if observed && !self.collected {
+            return Err(TransitionFailure::ReportReplayed);
+        } else if !observed && self.collected {
+            return Err(TransitionFailure::BatchCollected);
+        }
+        self.processed.insert(nonce.clone());
+        Ok(())
+    }
+
+    pub(crate) fn process_delete_all(&mut self) {
+        self.pending.clear();
+        self.processed.clear();
+        self.collected = false;
+    }
+
+    pub(crate) fn process_get_pending(&mut self, reports_requested: u64) -> Vec<Report> {
+        let reports_drained =
+            std::cmp::min(reports_requested.try_into().unwrap(), self.pending.len());
+        let reports: Vec<Report> = self.pending.drain(..reports_drained).collect();
+        reports
+    }
+
+    pub(crate) fn process_put_pending(&mut self, report: Report) -> Result<(), TransitionFailure> {
+        self.checked_process(&report.nonce)?;
+        self.pending.push(report);
+        Ok(())
+    }
+
+    pub(crate) fn process_put_processed(&mut self, nonce: Nonce) -> Result<(), TransitionFailure> {
+        self.checked_process(&nonce)?;
+        Ok(())
+    }
+
+    pub(crate) fn process_mark_collected(&mut self) {
+        self.collected = true;
+    }
+}