Repository: cloudflare/daphne
Pull request #24 (2022-06-24 02:06:45)

Clean up
Based on #23 (merge that first).

Miscellaneous things:

* Fix minor issues with APIs
* Rename `prio3_prepare_start()`
* Triage remaining TODOs
* Update READMEs


Patch:
diff --git a/README.md b/README.md
index 27340d4..405862a 100644
--- a/README.md
+++ b/README.md
@@ -1,30 +1,59 @@
 # Daphne
 
-## Integration tests
+Daphne is a Rust implementation of the Distributed Aggregation Protocol
+([DAP])(https://github.com/ietf-wg-ppm/draft-ietf-ppm-dap) standard. DAP is
+under active development in the PPM working group of the IETF.
 
-[miniflare](https://miniflare.dev/) is used to test integration of Daphne-Worker
-with Daphne. These tests can be run from Docker:
+This software is intended to support experimental DAP deployments and is not yet
+suitable for use in production. Daphne will evolve along with the DAP draft:
+Backwards compatibility with previous drafts won't be guaranteed until the draft
+itself begins to stabilize. API-breaking changes between releases should also be
+expected.
+
+This repository contains two crates:
+
+* `daphne` (aka "Daphne") -- Implementation of the core DAP protocol logic for
+  Clients, Aggregators, and Collectors. This crate does not provide the
+  complete, end-to-end functionality of any party. Instead, it defines traits
+  for the functionalities that a concrete instantantiation of the protocol is
+  required to implement. We call these functionalities "roles".
+
+* `daphne_worker` (aka "Daphne-Worker") -- Implements a backend for the
+  Aggregator roles based on [Cloudflare
+  Workers](https://workers.cloudflare.com/). This crate also implements the
+  various HTTP endpoints defined in the DAP spec.
+
+## Testing
+
+> NOTE These instructions will be changed as part of a planned refactor of
+> Daphne-Worker and the tests. See issue
+> [#1](https://github.com/cloudflare/daphne/issues/1) for details.
+
+The `daphne` crate relies on unit tests. The `daphne_worker` crate contains a
+number of integration and end-to-end tests that use
+[miniflare](https://miniflare.dev/) to mock the Cloudflare Workers platform.
+These can via docker-compose:
 
 ```
 docker-compose up --build --abort-on-container-exit --exit-code-from test
 ```
 
-## Development environment
+However for development it's often more convenient to run miniflare directly.
+Prerequisites:
 
-While iterating on the code it's often helpful to bypass docker-compose so that
-you don't have to wait for Docker containers to rebuild. Here are instructions
-for getting the Leader and Helper running without Docker.
+* miniflare>=2.5.1
 
-The Worker is built on [workers-rs](https://github.com/cloudflare/workers-rs).
-You'll need wrangler >= 1.19.5 in order to build it. We mock the Workers
-platform locally using [miniflare](https://github.com/cloudflare/miniflare). To
-get it to work you'll also need to upgrade node to the latest version.
+    ```
+    nvm use 18.4.0 && npm install -g miniflare
+    ```
 
-```
-nvm use 18.4.0 && npm install -g miniflare@2.5.1
-```
+* wrangler>=1.19.12
+
+    ```
+    npm install -g @cloudflare/wrangler
+    ```
 
-To run the Leader: from the `daphne_worker` directory, do
+To run the Leader, first navigate to the `daphne_worker` directory. Then run
 
 ```
 miniflare -p 8787 --env=tests/backend/leader.env --binding=DAP_ENV=dev
@@ -42,7 +71,8 @@ You can now `curl` the Leader or the Helper, or run the end-to-end tests via
 DAP_ENV=dev cargo test --features=test_e2e -- --test-threads 1
 ```
 
-To run [Janus](https://github.com/divviup/janus) interop tests, do
+The `daphne_worker` crate also has integration tests for
+[Janus](https://github.com/divviup/janus). To run these, do
 
 ```
 DAP_ENV=dev cargo test --features=test_janus -- --test-threads 1
diff --git a/daphne/src/hpke.rs b/daphne/src/hpke.rs
index 02697cd..6144f72 100644
--- a/daphne/src/hpke.rs
+++ b/daphne/src/hpke.rs
@@ -99,7 +99,7 @@ impl HpkeSecretKey {
 
 impl HpkeDecrypter for HpkeSecretKey {
     fn get_hpke_config_for(&self, _task_id: &Id) -> Option<&HpkeConfig> {
-        // TODO(MVP) Have HpkeSecretKey subsume the config.
+        // TODO(issue#12) Have HpkeSecretKey subsume the config.
         unreachable!("not implemented");
     }
 
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 7a8ec00..710a542 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -1,15 +1,22 @@
 // Copyright (c) 2022 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-//! Daphne is an implementation of the [DAP](https://datatracker.ietf.org/doc/draft-ietf-ppm-dap/)
-//! protocol. Daphne includes supports for the following [VDAFs](https://datatracker.ietf.org/doc/draft-irtf-cfrg-vdaf/):
-//!  - Prio3Aes128Count
-//!  - Prio3Aes128Histogram
-//!  - Prio3Aes128Sum
+//! This crate implements the core protocol logic for the Distributed Aggregation Protocol
+//! ([DAP](https://github.com/ietf-wg-ppm/draft-ietf-ppm-dap)) standard under development in the
+//! PPM working group of the IETF. See [`VdafConfig`] for a listing of supported
+//! [VDAFs](https://github.com/cfrg/draft-irtf-cfrg-vdaf).
 //!
-//! Additional Prio3 variants can be added easily as need.
+//! Daphne does not provide the complete, end-to-end functionality of any party in the protocol.
+//! Instead, it defines traits for the functionalities that a concrete instantiation of the
+//! protocol is required to implement. For example, the `daphne_worker` crate implements a backend
+//! for the DAP Leader and Helper. See the [`crate::roles`](roles) module for details.
 //!
-//! Daphne is not compatible with DAP tasks whose batch lifetime is longer than one aggregation.
+//! Daphne is not yet feature complete. Known issues include:
+//!
+//! * Daphne is not compatible with DAP tasks whose maximum batch lifetime is longer than one.
+//!
+//! * Daphne does handle aborts precisely as specified. In particular, some fields in the "Problem
+//! Details" document are missing.
 
 use crate::{
     messages::{CollectResp, HpkeConfig, Interval, Nonce, TransitionFailure},
@@ -519,4 +526,4 @@ mod messages_test;
 pub mod roles;
 #[cfg(test)]
 mod roles_test;
-mod vdaf;
+pub mod vdaf;
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 3e767de..355269c 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -7,6 +7,7 @@ use crate::{
     constants::{
         MEDIA_TYPE_AGG_CONT_REQ, MEDIA_TYPE_AGG_CONT_RESP, MEDIA_TYPE_AGG_INIT_REQ,
         MEDIA_TYPE_AGG_INIT_RESP, MEDIA_TYPE_AGG_SHARE_REQ, MEDIA_TYPE_AGG_SHARE_RESP,
+        MEDIA_TYPE_HPKE_CONFIG,
     },
     hpke::HpkeDecrypter,
     messages::{
@@ -81,9 +82,8 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
     async fn mark_collected(&self, task_id: &Id, batch_interval: &Interval)
         -> Result<(), DapError>;
 
-    /// Handle HTTP GET to `/hpke_config?task_id=<task_id>`. Returns the encoded HPKE config to put
-    /// in the body of the response.
-    async fn http_get_hpke_config(&self, req: &DapRequest<S>) -> Result<Vec<u8>, DapAbort> {
+    /// Handle HTTP GET to `/hpke_config?task_id=<task_id>`.
+    async fn http_get_hpke_config(&self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
         // Parse the task ID from the query string, ensuring that it is the only query parameter.
         let mut id = None;
         for (k, v) in req.url.query_pairs() {
@@ -105,7 +105,10 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
             let hpke_config = self
                 .get_hpke_config_for(task_id)
                 .ok_or(DapAbort::UnrecognizedTask)?;
-            Ok(hpke_config.get_encoded())
+            Ok(DapResponse {
+                media_type: Some(MEDIA_TYPE_HPKE_CONFIG),
+                payload: hpke_config.get_encoded(),
+            })
         } else {
             Err(DapAbort::BadRequest("missing query parameter".into()))
         }
@@ -141,12 +144,8 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     /// Data type used to guide selection of a set of reports for aggregation.
     type ReportSelector;
 
-    /// Store a sequence of reports to use later on. Each input is the encoded report sent in the
-    /// body of an HTTP request.
-    //
-    // TODO(MVP) Just put one report with this interface.
-    async fn put_reports<I: IntoIterator<Item = Report>>(&self, reports: I)
-        -> Result<(), DapError>;
+    /// Store a report for use later on.
+    async fn put_report(&self, report: &Report) -> Result<(), DapError>;
 
     /// Fetch a sequence of reports to aggregate. The reports returned are removed from persistent
     /// storage.
@@ -205,7 +204,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
             return Err(DapAbort::UnrecognizedHpkeConfig);
         }
 
-        Ok(self.put_reports([report]).await?)
+        Ok(self.put_report(&report).await?)
     }
 
     /// Handle HTTP POST to `/collect`. The input is a [`CollectReq`](crate::messages::CollectReq).
@@ -239,7 +238,12 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         Ok(self.init_collect_job(&collect_req).await?)
     }
 
-    /// Run an aggregation job if there are pending reports and process all pending collect jobs.
+    /// Run an aggregation job if there are pending reports, then process all pending collect jobs.
+    ///
+    /// This method is geared primarily towards testing. It also demonstrates how to properly
+    /// synchronize collect and aggregation jobs. If used in a large DAP deployment, it is likely
+    /// create a bottleneck. Such deployments can improve throughput by running many aggregation
+    /// jobs in parallel.
     async fn process(
         &self,
         task_id: &Id,
@@ -319,7 +323,10 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
             telem.reports_aggregated += out_shares_count;
         }
 
-        // Process pending collect jobs.
+        // Process pending collect jobs. We wait until all aggregation jobs are finished before
+        // proceeding to this step. This is to prevent a race condition involving an aggregate
+        // share computed during a collect job and any output shares computed during an aggregation
+        // job.
         let pending = self.get_pending_collect_jobs(task_id).await?;
         for (collect_id, collect_req) in pending.into_iter() {
             let leader_agg_share = self
@@ -422,7 +429,7 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     .get_task_config_for(&agg_init_req.task_id)
                     .ok_or(DapAbort::UnrecognizedTask)?;
 
-                let early_fails = self
+                let early_rejects = self
                     .mark_aggregated(&agg_init_req.task_id, &agg_init_req.report_shares)
                     .await?;
 
@@ -430,7 +437,7 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     self,
                     &task_config.vdaf_verify_key,
                     &agg_init_req,
-                    |nonce| early_fails.get(nonce).copied(),
+                    &early_rejects,
                 )?;
 
                 let agg_resp = match transition {
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 76be214..01f142b 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -13,7 +13,7 @@ use crate::{
     },
     vdaf::prio3::{
         prio3_encode_prepare_message, prio3_helper_prepare_finish, prio3_leader_prepare_finish,
-        prio3_prepare_start, prio3_shard, prio3_unshard, Prio3Error,
+        prio3_prepare_init, prio3_shard, prio3_unshard, Prio3Error,
     },
     DapAbort, DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
     DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement, DapOutputShare,
@@ -26,7 +26,10 @@ use prio::{
 };
 use rand::prelude::*;
 use serde::{Deserialize, Serialize};
-use std::{collections::HashSet, convert::TryInto};
+use std::{
+    collections::{HashMap, HashSet},
+    convert::TryInto,
+};
 
 const CTX_INPUT_SHARE: &[u8] = b"dap-01 input share";
 const CTX_AGG_SHARE: &[u8] = b"dap-01 aggregate share";
@@ -106,8 +109,6 @@ impl VdafConfig {
     /// * `task_id` is the DAP task for which this report is being generated.
     ///
     /// * `measurement` is the measurement.
-    //
-    // TODO We may need to change this API to accommodate a public parameter for the VDAF.
     pub fn produce_report(
         &self,
         hpke_config_list: &[HpkeConfig],
@@ -181,9 +182,6 @@ impl VdafConfig {
     /// * `nonce` is the report nonce.
     ///
     /// * `encrypted_input_share` is the encrypted input share.
-    //
-    // TODO spec: Note in the spec that we MUST ignore extensions we don't recognize, as they may
-    // be intended for the other aggregator. (Does this make sense?)
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn consume_report_share<D: HpkeDecrypter>(
         &self,
@@ -219,7 +217,7 @@ impl VdafConfig {
         let agg_id = if is_leader { 0 } else { 1 };
         match (self, verify_key) {
             (Self::Prio3(ref prio3_config), VdafVerifyKey::Prio3(ref verify_key)) => {
-                Ok(prio3_prepare_start(
+                Ok(prio3_prepare_init(
                     prio3_config,
                     verify_key,
                     agg_id,
@@ -339,19 +337,18 @@ impl VdafConfig {
     ///
     /// * `agg_init_req` is the request sent by the Leader.
     ///
-    /// * `early_tran_fail_for` is a callback used to determine if a report should be rejected based on
-    /// the sequence of aggregate and collect requests made for the task so far. Its input is the
-    /// nonce of report share in `agg_init_req` and its return value is an optional transition failure.
-    pub(crate) fn handle_agg_init_req<D, F>(
+    /// * `early_rejects` is a tableindicating the set of reports in `agg_init_req` that the Helper
+    /// knows in advance it must reject. Each key is the report's nonce and the corresponding value
+    /// is the transition failure the Helper is to transmit.
+    pub(crate) fn handle_agg_init_req<D>(
         &self,
         decrypter: &D,
         verify_key: &VdafVerifyKey,
         agg_init_req: &AggregateInitializeReq,
-        early_tran_fail_for: F,
+        early_rejects: &HashMap<Nonce, TransitionFailure>,
     ) -> Result<DapHelperTransition<AggregateResp>, DapAbort>
     where
         D: HpkeDecrypter,
-        F: Fn(&Nonce) -> Option<TransitionFailure>,
     {
         let num_reports = agg_init_req.report_shares.len();
         let mut processed = HashSet::with_capacity(num_reports);
@@ -363,10 +360,10 @@ impl VdafConfig {
             }
             processed.insert(report_share.nonce.clone());
 
-            if let Some(failure) = early_tran_fail_for(&report_share.nonce) {
+            if let Some(failure) = early_rejects.get(&report_share.nonce) {
                 transitions.push(Transition {
                     nonce: report_share.nonce.clone(),
-                    var: TransitionVar::Failed(failure),
+                    var: TransitionVar::Failed(*failure),
                 });
                 continue;
             }
diff --git a/daphne/src/vdaf/prio3.rs b/daphne/src/vdaf/prio3.rs
index 7e5440c..5a3b4d2 100644
--- a/daphne/src/vdaf/prio3.rs
+++ b/daphne/src/vdaf/prio3.rs
@@ -83,7 +83,7 @@ macro_rules! prep_start {
 }
 
 /// Consume an input share and return the corresponding VDAF step and message.
-pub(crate) fn prio3_prepare_start(
+pub(crate) fn prio3_prepare_init(
     config: &Prio3Config,
     verify_key: &[u8; 16],
     agg_id: usize,
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index 549dc71..2ec0f1c 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -50,7 +50,7 @@ pub(crate) struct DaphneWorkerConfig<D> {
     /// configured as the DAP Leader, i.e., if `DAP_AGGREGATOR_ROLE == "leader"`.
     pub(crate) collector_bearer_tokens: Option<HashMap<Id, BearerToken>>,
 
-    // TODO(MVP) Make `bucket_Key` and `bucket_count` unique per task.
+    // TODO(issue#12) Make `bucket_key` and `bucket_count` unique per task.
     bucket_key: Seed<16>,
     bucket_count: u64,
 }
@@ -78,7 +78,6 @@ impl<D> DaphneWorkerConfig<D> {
             ctx.var("DAP_HPKE_CONFIG_LIST")?.to_string().as_ref(),
         )
         .map_err(|e| Error::RustError(format!("Failed to parse DAP_HPKE_CONFIG_LIST: {}", e)))?;
-        // TODO(MVP) Encode as JSON objects rather than hex blobs.
         let mut hpke_config_list = Vec::with_capacity(hpke_config_list_hex.len());
         for hex in hpke_config_list_hex {
             let bytes = hex::decode(hex).map_err(int_err)?;
@@ -192,7 +191,6 @@ impl<D> DaphneWorkerConfig<D> {
             Seed::get_decoded(&hex::decode(json_bucket_key).map_err(int_err)?).map_err(int_err)?;
 
         let hpke_config_list_hex: Vec<String> = serde_json::from_str(json_hpke_config_list)?;
-        // TODO(MVP) Encode as JSON objects rather than hex blobs.
         let mut hpke_config_list = Vec::with_capacity(hpke_config_list_hex.len());
         for hex in hpke_config_list_hex {
             let bytes = hex::decode(hex).map_err(int_err)?;
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 3d29bfe..a9578a2 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -228,27 +228,19 @@ impl<D> DapAggregator<BearerToken> for DaphneWorkerConfig<D> {
 impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
     type ReportSelector = InternalAggregateInfo;
 
-    async fn put_reports<I: IntoIterator<Item = Report>>(
-        &self,
-        reports: I,
-    ) -> std::result::Result<(), DapError> {
+    async fn put_report(&self, report: &Report) -> std::result::Result<(), DapError> {
         let durable_namespace = self.durable_object("DAP_REPORT_STORE")?;
-        for report in reports.into_iter() {
-            let task_config = self
-                .get_task_config_for(&report.task_id)
-                .ok_or_else(|| DapError::fatal(INT_ERR_UNRECOGNIZED_TASK))?;
-            let durable_name =
-                self.durable_report_store_name(task_config, &report.task_id, &report.nonce);
-            let stub = durable_namespace.id_from_name(&durable_name)?.get_stub()?;
-            let mut resp = durable_post!(stub, DURABLE_REPORT_STORE_PUT_PENDING, &report).await?;
-            match resp.json().await? {
-                ReportStoreResult::Ok => (),
-                // TODO(MVP) Don't quit early if we get a transition failure. This will become a
-                // problem if we ever upload multiple reports
-                ReportStoreResult::Err(t) => return Err(DapError::Transition(t)),
-            };
+        let task_config = self
+            .get_task_config_for(&report.task_id)
+            .ok_or_else(|| DapError::fatal(INT_ERR_UNRECOGNIZED_TASK))?;
+        let durable_name =
+            self.durable_report_store_name(task_config, &report.task_id, &report.nonce);
+        let stub = durable_namespace.id_from_name(&durable_name)?.get_stub()?;
+        let mut resp = durable_post!(stub, DURABLE_REPORT_STORE_PUT_PENDING, &report).await?;
+        match resp.json().await? {
+            ReportStoreResult::Ok => Ok(()),
+            ReportStoreResult::Err(t) => return Err(DapError::Transition(t)),
         }
-        Ok(())
     }
 
     async fn get_reports(
diff --git a/daphne_worker/src/durable/leader_state_store.rs b/daphne_worker/src/durable/leader_state_store.rs
index 81114e9..7f1f18f 100644
--- a/daphne_worker/src/durable/leader_state_store.rs
+++ b/daphne_worker/src/durable/leader_state_store.rs
@@ -76,7 +76,7 @@ impl DurableObject for LeaderStateStore {
                 // TODO Consider disallowing overlapping collect requests. Would this require a
                 // spec change?
                 //
-                // TODO(MVP) Figure out if we can avoid re-serializing by passing raw bytes here instead
+                // TODO Figure out if we can avoid re-serializing by passing raw bytes here instead
                 // of JSON.
                 let collect_req: CollectReq = req.json().await?;
                 let collect_req_digest = digest(&SHA256, &collect_req.get_encoded());
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 203930a..d3ffb83 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -90,12 +90,9 @@ pub async fn main(req: Request, env: Env, _ctx: worker::Context) -> Result<Respo
         .get_async("/hpke_config", |req, ctx| async move {
             let req = worker_request_to_dap(req).await?;
             let config = DaphneWorkerConfig::from_worker_context(ctx)?;
-            // TODO(MVP) Have this method return a DapResponse.
+            // TODO(cjpatton) Have this method return a DapResponse.
             match config.http_get_hpke_config(&req).await {
-                Ok(hpke_config_data) => dap_response_to_worker(DapResponse {
-                    media_type: Some(constants::MEDIA_TYPE_HPKE_CONFIG),
-                    payload: hpke_config_data,
-                }),
+                Ok(req) => dap_response_to_worker(req),
                 Err(e) => abort(e),
             }
         })
@@ -144,7 +141,6 @@ pub async fn main(req: Request, env: Env, _ctx: worker::Context) -> Result<Respo
                         let task_id = parse_id!(ctx.param("task_id"));
                         let collect_id = parse_id!(ctx.param("collect_id"));
                         let config = DaphneWorkerConfig::from_worker_context(ctx)?;
-                        // TODO(MVP) Consider deriving the logic below.
                         match config.poll_collect_job(&task_id, &collect_id).await {
                             Ok(DapCollectJob::Done(collect_resp)) => {
                                 dap_response_to_worker(DapResponse {
diff --git a/daphne_worker/wrangler.toml b/daphne_worker/wrangler.toml
index e8cea4b..a0e5f76 100644
--- a/daphne_worker/wrangler.toml
+++ b/daphne_worker/wrangler.toml
@@ -5,7 +5,7 @@
 name = "daphne-worker-dev"
 type = "javascript"
 workers_dev = true
-# TODO(MVP) Figure out the semantics of this field and update it.
+# TODO(cjpatton) Figure out the semantics of this field and update it.
 compatibility_date = "2021-11-12"
 
 [vars]


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 1682ecb..b1ac239 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -174,10 +174,7 @@ impl DapHelper<BearerToken> for MockAggregator {
 impl DapLeader<BearerToken> for MockAggregator {
     type ReportSelector = ();
 
-    async fn put_reports<I: IntoIterator<Item = Report>>(
-        &self,
-        _reports: I,
-    ) -> Result<(), DapError> {
+    async fn put_report(&self, _report: &Report) -> Result<(), DapError> {
         unreachable!("not implemented");
     }
 
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 5ca3c4d..b8facef 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -18,7 +18,7 @@ use prio::vdaf::{
     PrepareTransition,
 };
 use rand::prelude::*;
-use std::{collections::HashSet, fmt::Debug, time::SystemTime};
+use std::{collections::HashMap, fmt::Debug, time::SystemTime};
 
 impl<M: Debug> DapLeaderTransition<M> {
     pub(crate) fn unwrap_continue(self) -> (DapLeaderState, M) {
@@ -222,7 +222,8 @@ fn agg_resp_fail_report_replayed() {
     let reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
     // Simulate the leader replaying a report.
-    t.helper_reports_aggregated.insert(reports[0].nonce.clone());
+    t.early_rejects
+        .insert(reports[0].nonce.clone(), TransitionFailure::ReportReplayed);
 
     let (_, agg_req) = t.produce_agg_init_req(reports).unwrap_continue();
     let (_, agg_resp) = t.handle_agg_init_req(agg_req).unwrap_continue();
@@ -241,7 +242,8 @@ fn agg_resp_fail_batch_collected() {
 
     // Simulate the leader requesting aggregation of a report for a batch that has already been
     // collected.
-    t.helper_batch_collected = true;
+    t.early_rejects
+        .insert(reports[0].nonce.clone(), TransitionFailure::BatchCollected);
 
     let (_, agg_req) = t.produce_agg_init_req(reports).unwrap_continue();
     let (_, agg_resp) = t.handle_agg_init_req(agg_req).unwrap_continue();
@@ -549,8 +551,7 @@ struct Test<'a> {
     vdaf_verify_key: VdafVerifyKey,
     leader_hpke_secret_key: HpkeSecretKey,
     helper_hpke_secret_key: HpkeSecretKey,
-    helper_reports_aggregated: HashSet<Nonce>,
-    helper_batch_collected: bool,
+    early_rejects: HashMap<Nonce, TransitionFailure>,
     client_hpke_config_list: Vec<HpkeConfig>,
     collector_hpke_config: HpkeConfig,
     collector_hpke_secret_key: HpkeSecretKey,
@@ -578,8 +579,7 @@ impl<'a> Test<'a> {
             vdaf_verify_key,
             leader_hpke_secret_key,
             helper_hpke_secret_key,
-            helper_reports_aggregated: HashSet::new(),
-            helper_batch_collected: false,
+            early_rejects: HashMap::default(),
             client_hpke_config_list: vec![leader_hpke_config, helper_hpke_config],
             collector_hpke_config,
             collector_hpke_secret_key,
@@ -629,13 +629,16 @@ impl<'a> Test<'a> {
                 &self.helper_hpke_secret_key,
                 &self.vdaf_verify_key,
                 &agg_init_req,
-                |nonce| self.helper_early_tran_fail_for(nonce),
+                &self.early_rejects,
             )
             .unwrap();
 
         for report_share in agg_init_req.report_shares {
-            self.helper_reports_aggregated
-                .insert(report_share.nonce.clone());
+            // Make sure the Leader doesn't try to aggregate these reports again.
+            self.early_rejects.insert(
+                report_share.nonce.clone(),
+                TransitionFailure::ReportReplayed,
+            );
         }
 
         agg_resp
@@ -737,13 +740,4 @@ impl<'a> Test<'a> {
             )
             .unwrap()
     }
-
-    fn helper_early_tran_fail_for(&self, nonce: &Nonce) -> Option<TransitionFailure> {
-        if !self.helper_batch_collected && self.helper_reports_aggregated.contains(nonce) {
-            return Some(TransitionFailure::ReportReplayed);
-        } else if self.helper_batch_collected && !self.helper_reports_aggregated.contains(nonce) {
-            return Some(TransitionFailure::BatchCollected);
-        }
-        None
-    }
 }
diff --git a/daphne/src/vdaf/prio3_test.rs b/daphne/src/vdaf/prio3_test.rs
index 9afa012..87970ad 100644
--- a/daphne/src/vdaf/prio3_test.rs
+++ b/daphne/src/vdaf/prio3_test.rs
@@ -4,7 +4,7 @@
 use crate::{
     vdaf::prio3::{
         prio3_encode_prepare_message, prio3_helper_prepare_finish, prio3_leader_prepare_finish,
-        prio3_prepare_start, prio3_shard, prio3_unshard, Prio3Error,
+        prio3_prepare_init, prio3_shard, prio3_unshard, Prio3Error,
     },
     DapAggregateResult, DapMeasurement, Prio3Config,
 };
@@ -58,10 +58,10 @@ fn test_prepare(
 
     // Prepare
     let (leader_state, leader_share) =
-        prio3_prepare_start(&config, &verify_key, 0, nonce, &encoded_input_shares[0])?;
+        prio3_prepare_init(&config, &verify_key, 0, nonce, &encoded_input_shares[0])?;
 
     let (helper_state, helper_share) =
-        prio3_prepare_start(&config, &verify_key, 1, nonce, &encoded_input_shares[1])?;
+        prio3_prepare_init(&config, &verify_key, 1, nonce, &encoded_input_shares[1])?;
 
     let helper_share_data = prio3_encode_prepare_message(&helper_share);
 
diff --git a/daphne_worker/tests/e2e.rs b/daphne_worker/tests/e2e.rs
index e8dee8f..e849761 100644
--- a/daphne_worker/tests/e2e.rs
+++ b/daphne_worker/tests/e2e.rs
@@ -143,18 +143,6 @@ async fn e2e_leader_upload() {
     );
 }
 
-// TODO Add tests for the following constants:
-//
-//  - Leader (resp. Helper) can't decrypt its input share (report gets dropped)
-//      - nonce/extension mismatch
-//      - unknown HPKE config ID
-//      - incorrect public key
-//  - Leader skips on init request (no valid reports)
-//  - Helper skips on init aggregate request
-//  - Input is invalid (report gets dropped)
-//  - Helper doesn't recognize the task ID
-//  - Helper can't parse (initial) aggregate request
-//  - Helper got repeated nonce
 #[tokio::test]
 #[cfg_attr(not(feature = "test_e2e"), ignore)]
 async fn e2e_internal_leader_process() {
diff --git a/daphne_worker/tests/janus.rs b/daphne_worker/tests/janus.rs
index ee7b311..90d8e10 100644
--- a/daphne_worker/tests/janus.rs
+++ b/daphne_worker/tests/janus.rs
@@ -8,7 +8,10 @@ mod test_runner;
 use assert_matches::assert_matches;
 use daphne::constants;
 use daphne_worker::InternalAggregateInfo;
-use prio::{codec::Decode, vdaf::prio3::Prio3};
+use prio::{
+    codec::{Decode, Encode},
+    vdaf::prio3::Prio3,
+};
 use rand::prelude::*;
 use test_runner::{TestRunner, COLLECTOR_HPKE_SECRET_KEY};
 
@@ -66,14 +69,9 @@ async fn janus_client() {
 // Test that daphne can run the aggregation sub-protocol with a Janus helper.
 //
 // For tracing, run `cargo test` with RUST_LOG=aggregator=trace,janus=trace,warning.
-//
-// TODO(MVP) Update Daphne to latest version of the draft to make this test pass.
 #[tokio::test]
 #[cfg_attr(not(feature = "test_janus"), ignore)]
 async fn janus_helper() {
-    // TODO Use the same version of prio as Janus.
-    use prio::codec::{Decode, Encode};
-
     janus_server::trace::test_util::install_test_trace_subscriber();
     let (t, janus_helper) = TestRunner::janus_helper().await;
     let client = t.http_client();