Repository: cloudflare/daphne
Pull request #727 (2024-11-27 14:35:59)

Document capnp requirements


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 569ec39..85a76ff 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -846,7 +846,6 @@ name = "dapf"
 version = "0.3.0"
 dependencies = [
  "anyhow",
- "async-trait",
  "chrono",
  "clap",
  "daphne",
diff --git a/README.md b/README.md
index 36fc569..d89f97f 100644
--- a/README.md
+++ b/README.md
@@ -31,6 +31,19 @@ The remaining crates are not intended for general use:
   deployments. Some of its features are specific to Cloudflare's own
   deployment.
 
+
+## Requirements
+
+[Cap'n Proto](https://capnproto.org/) is required to build DAP.
+
+```sh
+# debian
+apt install capnproto
+
+# macos
+brew install capnp
+```
+
 ## Testing
 
 The `daphne` crate relies on unit tests. To test integration with Cloudflare's
diff --git a/crates/dapf/Cargo.toml b/crates/dapf/Cargo.toml
index 153d71a..afbf7bf 100644
--- a/crates/dapf/Cargo.toml
+++ b/crates/dapf/Cargo.toml
@@ -13,7 +13,6 @@ repository.workspace = true
 
 [dependencies]
 anyhow.workspace = true
-async-trait.workspace = true
 chrono.workspace = true
 clap = { workspace = true, features = ["env"] }
 daphne = { path = "../daphne", features = ["report-generator", "prometheus"] }
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index 49df4db..4edf036 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -20,7 +20,6 @@ pub mod load_testing;
 
 use crate::{deduce_dap_version_from_url, functions, test_durations::TestDurations, HttpClient};
 use anyhow::{anyhow, bail, Context, Result};
-use async_trait::async_trait;
 use daphne::{
     hpke::{HpkeConfig, HpkeKemId, HpkeReceiverConfig},
     messages::{
@@ -28,7 +27,6 @@ use daphne::{
         BatchId, BatchSelector, PartialBatchSelector, TaskId,
     },
     metrics::DaphneMetrics,
-    roles::DapReportInitializer,
     testing::report_generator::ReportGenerator,
     vdaf::VdafConfig,
     DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapMeasurement, DapQueryConfig,
@@ -41,7 +39,7 @@ use rand::{rngs, Rng};
 use std::{
     convert::TryFrom,
     env,
-    ops::Range,
+    num::NonZeroU32,
     path::PathBuf,
     sync::atomic::{AtomicUsize, Ordering},
     time::{Duration, Instant, SystemTime},
@@ -383,7 +381,7 @@ impl Test {
             lifetime: 60,
             min_batch_size: reports_per_batch.try_into().unwrap(),
             query: DapQueryConfig::FixedSize {
-                max_batch_size: Some(reports_per_batch.try_into().unwrap()),
+                max_batch_size: NonZeroU32::new(reports_per_batch.try_into().unwrap()),
             },
             vdaf: self.vdaf_config,
             ..Default::default()
@@ -488,7 +486,7 @@ impl Test {
         let (agg_job_state, agg_job_init_req) = task_config
             .test_produce_agg_job_req(
                 fake_leader_hpke_receiver_config,
-                self,
+                messages::Time::MIN..messages::Time::MAX,
                 task_id,
                 part_batch_sel,
                 &DapAggregationParam::Empty,
@@ -696,14 +694,6 @@ impl Test {
     }
 }
 
-#[async_trait]
-impl DapReportInitializer for Test {
-    fn valid_report_time_range(&self) -> Range<messages::Time> {
-        // Accept reports with any timestmap.
-        0..u64::MAX
-    }
-}
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Now(u64);
 pub fn now() -> Now {
diff --git a/crates/daphne-server/src/roles/aggregator.rs b/crates/daphne-server/src/roles/aggregator.rs
index a4a3192..e2b2564 100644
--- a/crates/daphne-server/src/roles/aggregator.rs
+++ b/crates/daphne-server/src/roles/aggregator.rs
@@ -13,7 +13,7 @@ use daphne::{
     metrics::DaphneMetrics,
     roles::{
         aggregator::{MergeAggShareError, TaskprovConfig},
-        DapAggregator, DapReportInitializer,
+        DapAggregator,
     },
     taskprov, DapAggregateShare, DapAggregateSpan, DapError, DapGlobalConfig, DapTaskConfig,
     DapVersion,
@@ -344,10 +344,7 @@ impl DapAggregator for crate::App {
     fn audit_log(&self) -> &dyn AuditLog {
         &*self.audit_log
     }
-}
 
-#[async_trait]
-impl DapReportInitializer for crate::App {
     fn valid_report_time_range(&self) -> Range<messages::Time> {
         let now = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
diff --git a/crates/daphne/src/lib.rs b/crates/daphne/src/lib.rs
index 2ccd227..e081434 100644
--- a/crates/daphne/src/lib.rs
+++ b/crates/daphne/src/lib.rs
@@ -83,7 +83,7 @@ use std::{
     cmp::{max, min},
     collections::{HashMap, HashSet},
     fmt::Debug,
-    num::NonZeroUsize,
+    num::{NonZeroU32, NonZeroUsize},
     str::FromStr,
 };
 use url::Url;
@@ -237,7 +237,7 @@ pub enum DapQueryConfig {
     /// Aggregators are meant to stop aggregating reports when this limit is reached.
     FixedSize {
         #[serde(default)]
-        max_batch_size: Option<u64>,
+        max_batch_size: Option<NonZeroU32>,
     },
 }
 
@@ -730,7 +730,7 @@ impl DapTaskConfig {
             DapQueryConfig::FixedSize {
                 max_batch_size: Some(max_batch_size),
             } => {
-                if report_count > max_batch_size {
+                if report_count > u64::from(max_batch_size.get()) {
                     return Err(DapAbort::InvalidBatchSize {
                         detail: format!(
                             "Report count ({report_count}) exceeds maximum ({max_batch_size})"
diff --git a/crates/daphne/src/messages/taskprov.rs b/crates/daphne/src/messages/taskprov.rs
index 8d26cd2..d833cbe 100644
--- a/crates/daphne/src/messages/taskprov.rs
+++ b/crates/daphne/src/messages/taskprov.rs
@@ -15,6 +15,7 @@ use prio::codec::{
 };
 use serde::{Deserialize, Serialize};
 use std::io::{Cursor, Read};
+use std::num::NonZeroU32;
 
 use super::{
     decode_base64url_vec, decode_u16_prefixed, encode_base64url, encode_u16_prefixed, TaskId,
@@ -318,7 +319,7 @@ impl Decode for UrlBytes {
 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
 pub enum QueryConfigVar {
     TimeInterval,
-    FixedSize { max_batch_size: u32 },
+    FixedSize { max_batch_size: Option<NonZeroU32> },
     NotImplemented { typ: u8, param: Vec<u8> },
 }
 
@@ -363,7 +364,13 @@ impl ParameterizedEncode<DapVersion> for QueryConfig {
         match &self.var {
             QueryConfigVar::TimeInterval => (),
             QueryConfigVar::FixedSize { max_batch_size } => {
-                max_batch_size.encode(bytes)?;
+                match version {
+                    DapVersion::Draft09 => match max_batch_size {
+                        Some(x) => x.get().encode(bytes)?,
+                        None => 0u32.encode(bytes)?,
+                    },
+                    DapVersion::Latest => (),
+                };
             }
             QueryConfigVar::NotImplemented { typ: _, param } => {
                 bytes.extend_from_slice(param);
@@ -393,7 +400,10 @@ impl ParameterizedDecode<(DapVersion, Option<usize>)> for QueryConfig {
             match (bytes_left, query_type) {
                 (.., QUERY_TYPE_TIME_INTERVAL) => QueryConfigVar::TimeInterval,
                 (.., QUERY_TYPE_FIXED_SIZE) => QueryConfigVar::FixedSize {
-                    max_batch_size: u32::decode(bytes)?,
+                    max_batch_size: match version {
+                        DapVersion::Draft09 => NonZeroU32::new(u32::decode(bytes)?),
+                        DapVersion::Latest => None,
+                    },
                 },
                 (Some(bytes_left), ..) => {
                     let mut param = vec![0; bytes_left - fixed_size];
@@ -549,7 +559,12 @@ mod tests {
                     DapVersion::Latest => 1,
                 },
                 min_batch_size: 55,
-                var: QueryConfigVar::FixedSize { max_batch_size: 57 },
+                var: QueryConfigVar::FixedSize {
+                    max_batch_size: match version {
+                        DapVersion::Draft09 => Some(NonZeroU32::new(57).unwrap()),
+                        DapVersion::Latest => None,
+                    },
+                },
             },
             task_expiration: 23_232_232_232,
             vdaf_config: VdafConfig {
@@ -576,8 +591,8 @@ mod tests {
                 101, 46, 99, 111, 109, 47, 118, 48, 50, 0, 42, 104, 116, 116, 112, 115, 58, 47, 47,
                 115, 111, 109, 101, 115, 101, 114, 118, 105, 99, 101, 46, 99, 108, 111, 117, 100,
                 102, 108, 97, 114, 101, 114, 101, 115, 101, 97, 114, 99, 104, 46, 99, 111, 109, 0,
-                17, 0, 0, 0, 0, 0, 188, 79, 242, 0, 0, 0, 55, 2, 0, 0, 0, 57, 0, 0, 0, 5, 104, 191,
-                187, 40, 0, 11, 0, 1, 1, 255, 255, 0, 0, 0, 1, 134, 159,
+                13, 0, 0, 0, 0, 0, 188, 79, 242, 0, 0, 0, 55, 2, 0, 0, 0, 5, 104, 191, 187, 40, 0,
+                11, 0, 1, 1, 255, 255, 0, 0, 0, 1, 134, 159,
             ]
             .as_slice(),
         };
@@ -613,7 +628,10 @@ mod tests {
             },
             min_batch_size: 12_345_678,
             var: QueryConfigVar::FixedSize {
-                max_batch_size: 12_345_678,
+                max_batch_size: match version {
+                    DapVersion::Draft09 => Some(NonZeroU32::new(12_345_678).unwrap()),
+                    DapVersion::Latest => None,
+                },
             },
         };
         let encoded = query_config.get_encoded_with_param(&version).unwrap();
@@ -783,7 +801,12 @@ mod tests {
                     DapVersion::Latest => 1,
                 },
                 min_batch_size: 55,
-                var: QueryConfigVar::FixedSize { max_batch_size: 57 },
+                var: QueryConfigVar::FixedSize {
+                    max_batch_size: match version {
+                        DapVersion::Draft09 => Some(NonZeroU32::new(57).unwrap()),
+                        DapVersion::Latest => None,
+                    },
+                },
             },
             task_expiration: 23_232_232_232,
             vdaf_config: VdafConfig {
diff --git a/crates/daphne/src/protocol/aggregator.rs b/crates/daphne/src/protocol/aggregator.rs
index b2d4123..f5a9e23 100644
--- a/crates/daphne/src/protocol/aggregator.rs
+++ b/crates/daphne/src/protocol/aggregator.rs
@@ -14,7 +14,6 @@ use crate::{
         Transition, TransitionFailure, TransitionVar,
     },
     metrics::{DaphneMetrics, ReportStatus},
-    roles::DapReportInitializer,
     vdaf::{
         prio2::{prio2_prep_finish, prio2_prep_finish_from_shares, prio2_prep_init},
         prio3::{prio3_prep_finish, prio3_prep_finish_from_shares, prio3_prep_init},
@@ -294,7 +293,7 @@ impl DapTaskConfig {
     pub fn produce_agg_job_req<S>(
         &self,
         decrypter: impl HpkeDecrypter,
-        initializer: &impl DapReportInitializer,
+        valid_report_time_range: Range<messages::Time>,
         task_id: &TaskId,
         part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
@@ -306,7 +305,7 @@ impl DapTaskConfig {
     {
         self.produce_agg_job_req_impl(
             decrypter,
-            initializer,
+            valid_report_time_range,
             task_id,
             part_batch_sel,
             agg_param,
@@ -320,7 +319,7 @@ impl DapTaskConfig {
     fn produce_agg_job_req_impl<S>(
         &self,
         decrypter: impl HpkeDecrypter,
-        initializer: &impl DapReportInitializer,
+        valid_report_time_range: Range<messages::Time>,
         task_id: &TaskId,
         part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
@@ -354,7 +353,7 @@ impl DapTaskConfig {
 
             let initialized_report = InitializedReport::new(
                 &decrypter,
-                initializer.valid_report_time_range(),
+                valid_report_time_range.clone(),
                 true,
                 task_id,
                 self,
@@ -430,7 +429,7 @@ impl DapTaskConfig {
     pub fn test_produce_agg_job_req<S>(
         &self,
         decrypter: impl HpkeDecrypter,
-        initializer: &impl DapReportInitializer,
+        valid_report_time_range: Range<messages::Time>,
         task_id: &TaskId,
         part_batch_sel: &PartialBatchSelector,
         agg_param: &DapAggregationParam,
@@ -443,7 +442,7 @@ impl DapTaskConfig {
     {
         self.produce_agg_job_req_impl(
             decrypter,
-            initializer,
+            valid_report_time_range,
             task_id,
             part_batch_sel,
             agg_param,
@@ -456,17 +455,16 @@ impl DapTaskConfig {
     /// Helper: Consume the `AggregationJobInitReq` sent by the Leader and return the initialized
     /// reports.
     #[tracing::instrument(skip_all)]
-    pub fn consume_agg_job_req<H, I>(
+    pub fn consume_agg_job_req<H>(
         &self,
         decrypter: &H,
-        initializer: &I,
+        valid_report_time_range: Range<messages::Time>,
         task_id: &TaskId,
         agg_job_init_req: AggregationJobInitReq,
         replay_protection: ReplayProtection,
     ) -> Result<Vec<InitializedReport>, DapError>
     where
         H: HpkeDecrypter + Sync,
-        I: DapReportInitializer + Sync,
     {
         let agg_param =
             DapAggregationParam::get_decoded_with_param(&self.vdaf, &agg_job_init_req.agg_param)
@@ -490,7 +488,7 @@ impl DapTaskConfig {
             .map(|prep_init| {
                 InitializedReport::new(
                     decrypter,
-                    initializer.valid_report_time_range(),
+                    valid_report_time_range.clone(),
                     false,
                     task_id,
                     self,
diff --git a/crates/daphne/src/protocol/mod.rs b/crates/daphne/src/protocol/mod.rs
index fecd887..46e08a1 100644
--- a/crates/daphne/src/protocol/mod.rs
+++ b/crates/daphne/src/protocol/mod.rs
@@ -24,7 +24,6 @@ mod test {
             TransitionVar,
         },
         protocol::aggregator::InitializedReport,
-        roles::DapReportInitializer,
         test_versions,
         testing::AggregationJobTest,
         vdaf::{Prio3Config, VdafConfig},
diff --git a/crates/daphne/src/roles/aggregator.rs b/crates/daphne/src/roles/aggregator.rs
index 7c8e3af..2bb3997 100644
--- a/crates/daphne/src/roles/aggregator.rs
+++ b/crates/daphne/src/roles/aggregator.rs
@@ -17,14 +17,6 @@ use crate::{
     DapTaskConfig, DapVersion,
 };
 
-/// Report initializer. Used by a DAP Aggregator [`DapAggregator`] when initializing an aggregation
-/// job.
-#[async_trait]
-pub trait DapReportInitializer {
-    /// Return the time range in which a report must appear in order to be considered valid.
-    fn valid_report_time_range(&self) -> Range<Time>;
-}
-
 #[derive(Debug)]
 pub enum MergeAggShareError {
     AlreadyCollected,
@@ -42,7 +34,7 @@ pub struct TaskprovConfig<'s> {
 
 /// DAP Aggregator functionality.
 #[async_trait]
-pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
+pub trait DapAggregator: HpkeProvider + Sized {
     /// Look up the DAP global configuration.
     async fn get_global_config(&self) -> Result<DapGlobalConfig, DapError>;
 
@@ -137,6 +129,9 @@ pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
 
     /// Access the audit log.
     fn audit_log(&self) -> &dyn AuditLog;
+
+    /// Return the time range in which a report must appear in order to be considered valid.
+    fn valid_report_time_range(&self) -> Range<Time>;
 }
 
 /// Handle request for the Aggregator's HPKE configuration.
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index 23235b3..bfd5a9f 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -48,7 +48,7 @@ pub async fn handle_agg_job_init_req<A: DapHelper + Sync>(
     let version = req.version;
     let initialized_reports = task_config.consume_agg_job_req(
         &aggregator.get_receiver_configs(task_config.version).await?,
-        aggregator,
+        aggregator.valid_report_time_range(),
         &task_id,
         req.payload,
         replay_protection,
diff --git a/crates/daphne/src/roles/leader/mod.rs b/crates/daphne/src/roles/leader/mod.rs
index eff1752..c0076cb 100644
--- a/crates/daphne/src/roles/leader/mod.rs
+++ b/crates/daphne/src/roles/leader/mod.rs
@@ -301,7 +301,7 @@ async fn run_agg_job<A: DapLeader>(
     let agg_job_id = AggregationJobId(thread_rng().gen());
     let (agg_job_state, agg_job_init_req) = task_config.produce_agg_job_req(
         aggregator.get_receiver_configs(task_config.version).await?,
-        aggregator,
+        aggregator.valid_report_time_range(),
         task_id,
         part_batch_sel,
         agg_param,
diff --git a/crates/daphne/src/roles/mod.rs b/crates/daphne/src/roles/mod.rs
index 3afc2ca..15153d1 100644
--- a/crates/daphne/src/roles/mod.rs
+++ b/crates/daphne/src/roles/mod.rs
@@ -13,7 +13,7 @@ use crate::{
     DapAbort, DapError, DapGlobalConfig, DapQueryConfig, DapRequestMeta, DapTaskConfig,
 };
 
-pub use aggregator::{DapAggregator, DapReportInitializer};
+pub use aggregator::DapAggregator;
 pub use helper::DapHelper;
 pub use leader::DapLeader;
 
@@ -153,7 +153,13 @@ mod test {
     #[cfg(feature = "experimental")]
     use prio::{idpf::IdpfInput, vdaf::poplar1::Poplar1AggregationParam};
     use rand::{thread_rng, Rng};
-    use std::{collections::HashMap, num::NonZeroUsize, sync::Arc, time::SystemTime, vec};
+    use std::{
+        collections::HashMap,
+        num::{NonZeroU32, NonZeroUsize},
+        sync::Arc,
+        time::SystemTime,
+        vec,
+    };
     use url::Url;
 
     pub(super) struct TestData {
@@ -235,7 +241,7 @@ mod test {
                     not_after: now + Self::TASK_TIME_PRECISION,
                     min_batch_size: 1,
                     query: DapQueryConfig::FixedSize {
-                        max_batch_size: Some(2),
+                        max_batch_size: Some(NonZeroU32::new(2).unwrap()),
                     },
                     vdaf: vdaf_config,
                     vdaf_verify_key: vdaf_config.gen_verify_key(),
@@ -453,7 +459,7 @@ mod test {
             let (leader_state, agg_job_init_req) = task_config
                 .produce_agg_job_req(
                     &*self.leader.hpke_receiver_config_list,
-                    &*self.leader,
+                    self.leader.valid_report_time_range(),
                     task_id,
                     &part_batch_sel,
                     &agg_param,
@@ -1416,7 +1422,7 @@ mod test {
             version,
             min_batch_size: 1,
             query: DapQueryConfig::FixedSize {
-                max_batch_size: Some(2),
+                max_batch_size: Some(NonZeroU32::new(2).unwrap()),
             },
             vdaf: vdaf_config,
             ..Default::default()
diff --git a/crates/daphne/src/taskprov.rs b/crates/daphne/src/taskprov.rs
index bb8cb4b..5d70cf3 100644
--- a/crates/daphne/src/taskprov.rs
+++ b/crates/daphne/src/taskprov.rs
@@ -122,12 +122,9 @@ fn url_from_bytes(task_id: &TaskId, url_bytes: &[u8]) -> Result<Url, DapAbort> {
 impl DapQueryConfig {
     fn try_from_taskprov(task_id: &TaskId, var: QueryConfigVar) -> Result<Self, DapAbort> {
         match var {
-            QueryConfigVar::FixedSize { max_batch_size: 0 } => Ok(DapQueryConfig::FixedSize {
-                max_batch_size: None,
-            }),
-            QueryConfigVar::FixedSize { max_batch_size } => Ok(DapQueryConfig::FixedSize {
-                max_batch_size: Some(max_batch_size.into()),
-            }),
+            QueryConfigVar::FixedSize { max_batch_size } => {
+                Ok(DapQueryConfig::FixedSize { max_batch_size })
+            }
             QueryConfigVar::TimeInterval => Ok(DapQueryConfig::TimeInterval),
             QueryConfigVar::NotImplemented { typ, .. } => Err(DapAbort::InvalidTask {
                 detail: format!("unimplemented query type ({typ})"),
@@ -359,9 +356,7 @@ impl TryFrom<&DapQueryConfig> for messages::taskprov::QueryConfigVar {
             DapQueryConfig::TimeInterval => messages::taskprov::QueryConfigVar::TimeInterval,
             DapQueryConfig::FixedSize { max_batch_size } => {
                 messages::taskprov::QueryConfigVar::FixedSize {
-                    max_batch_size: max_batch_size.unwrap_or(0).try_into().map_err(|_| {
-                        fatal_error!(err = "task max batch size is too large for taskprov")
-                    })?,
+                    max_batch_size: *max_batch_size,
                 }
             }
         })
@@ -448,7 +443,7 @@ impl TryFrom<&DapTaskConfig> for messages::taskprov::TaskprovAdvertisement {
 
 #[cfg(test)]
 mod test {
-    use std::num::NonZeroUsize;
+    use std::num::{NonZeroU32, NonZeroUsize};
 
     use prio::codec::ParameterizedEncode;
 
@@ -477,7 +472,9 @@ mod test {
                 time_precision: 3600,
                 max_batch_query_count: 1,
                 min_batch_size: 1,
-                var: messages::taskprov::QueryConfigVar::FixedSize { max_batch_size: 2 },
+                var: messages::taskprov::QueryConfigVar::FixedSize {
+                    max_batch_size: Some(NonZeroU32::new(2).unwrap()),
+                },
             },
             task_expiration: 1337,
             vdaf_config: messages::taskprov::VdafConfig {
@@ -557,7 +554,9 @@ mod test {
                     time_precision: 3600,
                     max_batch_query_count: 1,
                     min_batch_size: 1,
-                    var: messages::taskprov::QueryConfigVar::FixedSize { max_batch_size: 2 },
+                    var: messages::taskprov::QueryConfigVar::FixedSize {
+                        max_batch_size: Some(NonZeroU32::new(2).unwrap()),
+                    },
                 },
                 task_expiration: 0,
                 vdaf_config: messages::taskprov::VdafConfig {
@@ -622,7 +621,9 @@ mod test {
                     time_precision: 3600,
                     max_batch_query_count: 1,
                     min_batch_size: 1,
-                    var: messages::taskprov::QueryConfigVar::FixedSize { max_batch_size: 2 },
+                    var: messages::taskprov::QueryConfigVar::FixedSize {
+                        max_batch_size: Some(NonZeroU32::new(2).unwrap()),
+                    },
                 },
                 task_expiration: 0,
                 vdaf_config: messages::taskprov::VdafConfig {


Test patch:
diff --git a/crates/daphne-server/tests/e2e/test_runner.rs b/crates/daphne-server/tests/e2e/test_runner.rs
index 1fabcde..aeb6bad 100644
--- a/crates/daphne-server/tests/e2e/test_runner.rs
+++ b/crates/daphne-server/tests/e2e/test_runner.rs
@@ -23,7 +23,7 @@ use serde_json::json;
 use std::time::SystemTime;
 use std::{
     any::{self, Any},
-    num::NonZeroUsize,
+    num::{NonZeroU32, NonZeroUsize},
     ops::Range,
 };
 use tokio::time::timeout;
@@ -31,7 +31,7 @@ use url::Url;
 
 const VDAF_CONFIG: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Sum { bits: 10 });
 pub(crate) const MIN_BATCH_SIZE: u64 = 10;
-pub(crate) const MAX_BATCH_SIZE: u64 = 12;
+pub(crate) const MAX_BATCH_SIZE: u32 = 12;
 pub(crate) const TIME_PRECISION: Duration = 3600; // seconds
 
 #[derive(Deserialize)]
@@ -66,7 +66,7 @@ impl TestRunner {
         Self::with(
             version,
             &DapQueryConfig::FixedSize {
-                max_batch_size: Some(MAX_BATCH_SIZE),
+                max_batch_size: Some(NonZeroU32::new(MAX_BATCH_SIZE).unwrap()),
             },
         )
         .await
diff --git a/crates/daphne-service-utils/src/test_route_types.rs b/crates/daphne-service-utils/src/test_route_types.rs
index 6591188..904a929 100644
--- a/crates/daphne-service-utils/src/test_route_types.rs
+++ b/crates/daphne-service-utils/src/test_route_types.rs
@@ -10,6 +10,7 @@ use daphne::{
     vdaf::{Prio3Config, VdafConfig},
 };
 use serde::{Deserialize, Serialize};
+use std::num::NonZeroU32;
 use url::Url;
 
 #[derive(Deserialize)]
@@ -86,7 +87,7 @@ pub struct InternalTestAddTask {
     pub query_type: u8,
     pub min_batch_size: u64,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub max_batch_size: Option<u64>,
+    pub max_batch_size: Option<NonZeroU32>,
     pub time_precision: Duration,
     pub collector_hpke_config: String, // base64url
     pub task_expiration: Time,
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index fb666b1..c9ddb61 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -21,7 +21,7 @@ use crate::{
         aggregator::{MergeAggShareError, TaskprovConfig},
         helper,
         leader::{in_memory_leader::InMemoryLeaderState, WorkItem},
-        DapAggregator, DapHelper, DapLeader, DapReportInitializer,
+        DapAggregator, DapHelper, DapLeader,
     },
     taskprov, DapAbort, DapAggregateResult, DapAggregateShare, DapAggregateSpan,
     DapAggregationJobState, DapAggregationParam, DapBatchBucket, DapCollectionJob, DapError,
@@ -68,14 +68,11 @@ pub struct AggregationJobTest {
     pub(crate) leader_metrics: DaphnePromMetrics,
 }
 
-#[async_trait]
-impl DapReportInitializer for AggregationJobTest {
-    fn valid_report_time_range(&self) -> Range<Time> {
+impl AggregationJobTest {
+    pub fn valid_report_time_range(&self) -> Range<Time> {
         self.valid_report_range.clone()
     }
-}
 
-impl AggregationJobTest {
     /// Create an aggregation job test with the given VDAF config, HPKE KEM algorithm, DAP protocol
     /// version. The KEM algorithm is used to generate an HPKE config for each party.
     pub fn new(vdaf: &VdafConfig, kem_id: HpkeKemId, version: DapVersion) -> Self {
@@ -193,7 +190,7 @@ impl AggregationJobTest {
         self.task_config
             .test_produce_agg_job_req(
                 &self.leader_hpke_receiver_config,
-                self,
+                self.valid_report_time_range(),
                 &self.task_id,
                 &PartialBatchSelector::TimeInterval,
                 agg_param,
@@ -219,7 +216,7 @@ impl AggregationJobTest {
                     .task_config
                     .consume_agg_job_req(
                         &self.helper_hpke_receiver_config,
-                        self,
+                        self.valid_report_time_range(),
                         &self.task_id,
                         agg_job_init_req,
                         self.replay_protection,
@@ -654,14 +651,6 @@ impl HpkeProvider for InMemoryAggregator {
     }
 }
 
-#[async_trait]
-impl DapReportInitializer for InMemoryAggregator {
-    fn valid_report_time_range(&self) -> Range<messages::Time> {
-        // Accept reports with any timestmap.
-        0..u64::MAX
-    }
-}
-
 #[async_trait]
 impl DapAggregator for InMemoryAggregator {
     async fn get_global_config(&self) -> Result<DapGlobalConfig, DapError> {
@@ -866,6 +855,11 @@ impl DapAggregator for InMemoryAggregator {
     fn audit_log(&self) -> &dyn AuditLog {
         &self.audit_log
     }
+
+    fn valid_report_time_range(&self) -> Range<messages::Time> {
+        // Accept reports with any timestmap.
+        messages::Time::MIN..messages::Time::MAX
+    }
 }
 
 #[async_trait]