Repository: cloudflare/daphne
Pull request #709 (2024-10-30 16:45:09)

Simplify how protocol functions resolve the current task
- Replace allow_taskprov with method to optionally fetch the taskprov config
- simplify taskprov resolve flow



Patch:
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index a12f973..2f61124 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -393,8 +393,10 @@ impl Test {
         .to_config_with_taskprov(
             b"cool task".to_vec(),
             now.0,
-            &self.vdaf_verify_init,
-            &fake_collector_hpke_receiver_config.config,
+            daphne::roles::aggregator::TaskprovConfig {
+                hpke_collector_config: &fake_collector_hpke_receiver_config.config,
+                vdaf_verify_key_init: &self.vdaf_verify_init,
+            },
         )?;
 
         Ok((
diff --git a/crates/daphne-server/examples/configuration-helper.toml b/crates/daphne-server/examples/configuration-helper.toml
index 7385737..5380bca 100644
--- a/crates/daphne-server/examples/configuration-helper.toml
+++ b/crates/daphne-server/examples/configuration-helper.toml
@@ -18,7 +18,6 @@ supported_hpke_kems = ["x25519_hkdf_sha256"]
 default_version = "v09"
 report_storage_epoch_duration = 300000
 base_url = "http://127.0.0.1:8788"
-allow_taskprov = true
 default_num_agg_span_shards = 4
 
 [service.taskprov]
diff --git a/crates/daphne-server/examples/configuration-leader.toml b/crates/daphne-server/examples/configuration-leader.toml
index 469bd35..8c1fad5 100644
--- a/crates/daphne-server/examples/configuration-leader.toml
+++ b/crates/daphne-server/examples/configuration-leader.toml
@@ -18,7 +18,6 @@ supported_hpke_kems = ["x25519_hkdf_sha256"]
 default_version = "v09"
 report_storage_epoch_duration = 300000
 base_url = "http://127.0.0.1:8787"
-allow_taskprov = true
 default_num_agg_span_shards = 4
 
 [service.taskprov]
diff --git a/crates/daphne-server/src/lib.rs b/crates/daphne-server/src/lib.rs
index 1ab19a0..d0bb6dc 100644
--- a/crates/daphne-server/src/lib.rs
+++ b/crates/daphne-server/src/lib.rs
@@ -64,7 +64,6 @@ mod storage_proxy_connection;
 ///     min_batch_interval_start: 259_200,
 ///     max_batch_interval_end: 259_200,
 ///     supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
-///     allow_taskprov: false,
 ///     default_num_agg_span_shards: NonZeroUsize::new(2).unwrap(),
 /// };
 /// let service_config = DaphneServiceConfig {
@@ -132,7 +131,7 @@ impl router::DaphneService for App {
             .taskprov
             .as_ref()
             // we only use taskprov auth if it's allowed by config and if the request is using taskprov
-            .filter(|_| self.service_config.global.allow_taskprov && is_taskprov)
+            .filter(|_| self.service_config.taskprov.is_some() && is_taskprov)
         {
             match (&taskprov.peer_auth, sender) {
                 (PeerBearerToken::Leader { expected_token }, DapSender::Leader)
diff --git a/crates/daphne-server/src/roles/aggregator.rs b/crates/daphne-server/src/roles/aggregator.rs
index a66e33b..0870795 100644
--- a/crates/daphne-server/src/roles/aggregator.rs
+++ b/crates/daphne-server/src/roles/aggregator.rs
@@ -11,10 +11,12 @@ use daphne::{
     hpke::{self, HpkeConfig, HpkeProvider, HpkeReceiverConfig},
     messages::{self, BatchId, BatchSelector, HpkeCiphertext, TaskId, Time},
     metrics::DaphneMetrics,
-    roles::{aggregator::MergeAggShareError, DapAggregator, DapReportInitializer},
+    roles::{
+        aggregator::{MergeAggShareError, TaskprovConfig},
+        DapAggregator, DapReportInitializer,
+    },
     taskprov, DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapError, DapGlobalConfig,
-    DapRequestMeta, DapTaskConfig, DapVersion, EarlyReportStateConsumed,
-    EarlyReportStateInitialized,
+    DapTaskConfig, DapVersion, EarlyReportStateConsumed, EarlyReportStateInitialized,
 };
 use daphne_service_utils::durable_requests::bindings::{
     self, AggregateStoreMergeOptions, AggregateStoreMergeReq, AggregateStoreMergeResp,
@@ -143,10 +145,6 @@ impl DapAggregator for crate::App {
         Ok(())
     }
 
-    type WrappedDapTaskConfig<'a> = DapTaskConfig
-    where
-        Self: 'a;
-
     async fn get_global_config(&self) -> Result<DapGlobalConfig, DapError> {
         let mut global_config = self.service_config.global.clone();
 
@@ -172,25 +170,20 @@ impl DapAggregator for crate::App {
         Ok(global_config)
     }
 
-    fn taskprov_vdaf_verify_key_init(&self) -> Option<&[u8; 32]> {
-        self.service_config
-            .taskprov
-            .as_ref()
-            .map(|c| &c.vdaf_verify_key_init)
-    }
-
-    fn taskprov_collector_hpke_config(&self) -> Option<&HpkeConfig> {
+    fn get_taskprov_config(&self) -> Option<TaskprovConfig<'_>> {
         self.service_config
             .taskprov
             .as_ref()
-            .map(|c| &c.hpke_collector_config)
+            .map(|t| TaskprovConfig {
+                hpke_collector_config: &t.hpke_collector_config,
+                vdaf_verify_key_init: &t.vdaf_verify_key_init,
+            })
     }
 
     async fn taskprov_opt_in(
         &self,
         task_id: &TaskId,
         task_config: taskprov::DapTaskConfigNeedsOptIn,
-        global_config: &DapGlobalConfig,
     ) -> Result<DapTaskConfig, DapError> {
         if let Some(param) = self
             .kv()
@@ -200,6 +193,7 @@ impl DapAggregator for crate::App {
         {
             Ok(task_config.into_opted_in(&param))
         } else {
+            let global_config = self.get_global_config().await?;
             let param = taskprov::OptInParam {
                 not_before: self.get_current_time(),
                 num_agg_span_shards: global_config.default_num_agg_span_shards,
@@ -226,13 +220,12 @@ impl DapAggregator for crate::App {
 
     async fn taskprov_put(
         &self,
-        req: &DapRequestMeta,
+        task_id: &TaskId,
         task_config: DapTaskConfig,
     ) -> Result<(), DapError> {
-        let task_id = &req.task_id;
         let expiration_time = task_config.not_after;
 
-        if self.service_config.role.is_leader() || req.taskprov.is_none() {
+        if self.service_config.role.is_leader() {
             self.kv()
                 .put_with_expiration::<kv::prefix::TaskConfig>(
                     task_id,
@@ -252,7 +245,7 @@ impl DapAggregator for crate::App {
     async fn get_task_config_for<'req>(
         &'req self,
         task_id: &'req TaskId,
-    ) -> Result<Option<Self::WrappedDapTaskConfig<'req>>, DapError> {
+    ) -> Result<Option<DapTaskConfig>, DapError> {
         self.kv()
             .get_cloned::<kv::prefix::TaskConfig>(task_id, &KvGetOptions::default())
             .await
diff --git a/crates/daphne/src/lib.rs b/crates/daphne/src/lib.rs
index e72bfdc..4e630a6 100644
--- a/crates/daphne/src/lib.rs
+++ b/crates/daphne/src/lib.rs
@@ -156,10 +156,6 @@ pub struct DapGlobalConfig {
     /// receiver config.
     pub supported_hpke_kems: Vec<HpkeKemId>,
 
-    /// draft-wang-ppm-dap-taskprov: Indicates if the taskprov extension is enabled.
-    #[serde(default)]
-    pub allow_taskprov: bool,
-
     /// Default number of aggregate span shards for a task.
     ///
     /// At the end of an aggregation job, each Aggregator produces a [`DapAggregateSpan`] that maps
@@ -195,7 +191,6 @@ impl Default for DapGlobalConfig {
             min_batch_interval_start: 60,
             max_batch_interval_end: 60,
             supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
-            allow_taskprov: false,
             default_num_agg_span_shards: NonZeroUsize::new(1).unwrap(),
         }
     }
@@ -550,10 +545,9 @@ impl DapTaskParameters {
         &self,
         task_info: Vec<u8>,
         now: Time,
-        vdaf_verify_key_init: &[u8; 32],
-        collector_hpke_config: &HpkeConfig,
+        taskprov_config: roles::aggregator::TaskprovConfig<'_>,
     ) -> Result<(DapTaskConfig, TaskId, String), DapError> {
-        let taskprov_config = messages::taskprov::TaskConfig {
+        let task_config = messages::taskprov::TaskConfig {
             task_info,
             leader_url: messages::taskprov::UrlBytes {
                 bytes: self.leader_url.to_string().into_bytes(),
@@ -574,7 +568,7 @@ impl DapTaskParameters {
             },
         };
 
-        let encoded_taskprov_config = taskprov_config
+        let encoded_taskprov_config = task_config
             .get_encoded_with_param(&self.version)
             .map_err(DapError::encoding)?;
         let task_id = taskprov::compute_task_id(&encoded_taskprov_config);
@@ -583,9 +577,8 @@ impl DapTaskParameters {
         let task_config = taskprov::DapTaskConfigNeedsOptIn::try_from_taskprov(
             self.version,
             &task_id,
+            task_config,
             taskprov_config,
-            vdaf_verify_key_init,
-            collector_hpke_config,
         )
         .unwrap()
         .into_opted_in(&taskprov::OptInParam {
diff --git a/crates/daphne/src/roles/aggregator.rs b/crates/daphne/src/roles/aggregator.rs
index 0b27f52..ff66889 100644
--- a/crates/daphne/src/roles/aggregator.rs
+++ b/crates/daphne/src/roles/aggregator.rs
@@ -15,7 +15,7 @@ use crate::{
     metrics::{DaphneMetrics, DaphneRequestType},
     protocol::aggregator::{EarlyReportStateConsumed, EarlyReportStateInitialized},
     taskprov, DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapError, DapGlobalConfig,
-    DapRequestMeta, DapResponse, DapTaskConfig, DapVersion,
+    DapResponse, DapTaskConfig, DapVersion,
 };
 
 /// Report initializer. Used by a DAP Aggregator [`DapAggregator`] when initializing an aggregation
@@ -43,24 +43,23 @@ pub enum MergeAggShareError {
     Other(DapError),
 }
 
+pub struct TaskprovConfig<'s> {
+    /// HPKE collector configuration for all taskprov tasks.
+    pub hpke_collector_config: &'s HpkeConfig,
+
+    /// VDAF verify key init secret, used to generate the VDAF verification key for a taskprov task.
+    pub vdaf_verify_key_init: &'s [u8; 32],
+}
+
 /// DAP Aggregator functionality.
 #[async_trait]
 pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
-    /// A refernce to a task configuration stored by the Aggregator.
-    type WrappedDapTaskConfig<'a>: AsRef<DapTaskConfig> + Send
-    where
-        Self: 'a;
-
     /// Look up the DAP global configuration.
     async fn get_global_config(&self) -> Result<DapGlobalConfig, DapError>;
 
-    /// taskprov: The VDAF verification key initializer. Used to derive the VDAF verify key for all
-    /// tasks configured by this extension.
-    fn taskprov_vdaf_verify_key_init(&self) -> Option<&[u8; 32]>;
-
-    /// taskprov: The Collector's HPKE configuration used for all tasks configured by this
-    /// extension.
-    fn taskprov_collector_hpke_config(&self) -> Option<&HpkeConfig>;
+    /// Get the taskprov configuration. This method also controls whether taskprov is enabled.
+    /// Returning [`None`] disables taskprov.
+    fn get_taskprov_config(&self) -> Option<TaskprovConfig<'_>>;
 
     /// taskprov: Complete the opt-in process.
     ///
@@ -76,7 +75,6 @@ pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
         &self,
         task_id: &TaskId,
         task_config: taskprov::DapTaskConfigNeedsOptIn,
-        global_config: &DapGlobalConfig,
     ) -> Result<DapTaskConfig, DapError>;
 
     /// taskprov: Configure a task. This is called after opting in. If successful, the next call to
@@ -84,7 +82,7 @@ pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
     /// nothing.
     async fn taskprov_put(
         &self,
-        req: &DapRequestMeta,
+        task_id: &TaskId,
         task_config: DapTaskConfig,
     ) -> Result<(), DapError>;
 
@@ -92,7 +90,7 @@ pub trait DapAggregator: HpkeProvider + DapReportInitializer + Sized {
     async fn get_task_config_for<'req>(
         &'req self,
         task_id: &'req TaskId,
-    ) -> Result<Option<Self::WrappedDapTaskConfig<'req>>, DapError>;
+    ) -> Result<Option<DapTaskConfig>, DapError>;
 
     /// Get the current time (number of seconds since the beginning of UNIX time).
     fn get_current_time(&self) -> Time;
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index b5f88a6..98f7d8a 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -6,7 +6,7 @@ use std::{collections::HashMap, sync::Once};
 use async_trait::async_trait;
 use prio::codec::{Encode, ParameterizedDecode};
 
-use super::{check_batch, resolve_taskprov, DapAggregator};
+use super::{check_batch, resolve_task_config, DapAggregator};
 use crate::{
     constants::DapMediaType,
     error::DapAbort,
@@ -31,32 +31,17 @@ pub async fn handle_agg_job_init_req<A: DapHelper>(
     req: DapRequest<AggregationJobInitReq, resource::AggregationJobId>,
     replay_protection: ReplayProtection,
 ) -> Result<DapResponse, DapError> {
-    let global_config = aggregator.get_global_config().await?;
     let task_id = req.task_id;
     let metrics = aggregator.metrics();
 
     metrics.agg_job_observe_batch_size(req.payload.prep_inits.len());
 
-    // taskprov: Resolve the task config to use for the request.
-    if global_config.allow_taskprov {
-        resolve_taskprov(aggregator, &task_id, &req, &global_config).await?;
-    }
-
-    let wrapped_task_config = aggregator
-        .get_task_config_for(&task_id)
-        .await?
-        .ok_or(DapAbort::UnrecognizedTask { task_id })?;
-    let task_config = wrapped_task_config.as_ref();
-
-    // Check whether the DAP version in the request matches the task config.
-    if task_config.version != req.version {
-        return Err(DapAbort::version_mismatch(req.version, task_config.version).into());
-    }
+    let task_config = resolve_task_config(aggregator, &req.meta).await?;
 
     // Ensure we know which batch the request pertains to.
     check_part_batch(
         &task_id,
-        task_config,
+        &task_config,
         &req.payload.part_batch_sel,
         &req.payload.agg_param,
     )?;
@@ -77,7 +62,7 @@ pub async fn handle_agg_job_init_req<A: DapHelper>(
         let agg_job_resp = finish_agg_job_and_aggregate(
             aggregator,
             &task_id,
-            task_config,
+            &task_config,
             &part_batch_sel,
             &initialized_reports,
             metrics,
@@ -91,7 +76,7 @@ pub async fn handle_agg_job_init_req<A: DapHelper>(
 
     aggregator.audit_log().on_aggregation_job(
         &task_id,
-        task_config,
+        &task_config,
         agg_job_resp.transitions.len() as u64,
         0, /* vdaf step */
     );
@@ -115,20 +100,7 @@ pub async fn handle_agg_share_req<'req, A: DapHelper>(
     let metrics = aggregator.metrics();
     let task_id = req.task_id;
 
-    if global_config.allow_taskprov {
-        resolve_taskprov(aggregator, &task_id, &req, &global_config).await?;
-    }
-
-    let wrapped_task_config = aggregator
-        .get_task_config_for(&task_id)
-        .await?
-        .ok_or(DapAbort::UnrecognizedTask { task_id })?;
-    let task_config = wrapped_task_config.as_ref();
-
-    // Check whether the DAP version in the request matches the task config.
-    if task_config.version != req.version {
-        return Err(DapAbort::version_mismatch(req.version, task_config.version).into());
-    }
+    let task_config = resolve_task_config(aggregator, &req.meta).await?;
 
     let agg_param =
         DapAggregationParam::get_decoded_with_param(&task_config.vdaf, &req.payload.agg_param)
@@ -138,7 +110,7 @@ pub async fn handle_agg_share_req<'req, A: DapHelper>(
     // collected batches.
     check_batch(
         aggregator,
-        task_config,
+        &task_config,
         &task_id,
         &req.payload.batch_sel.clone().into(),
         &req.payload.agg_param,
diff --git a/crates/daphne/src/roles/leader/mod.rs b/crates/daphne/src/roles/leader/mod.rs
index a3e71f8..6941139 100644
--- a/crates/daphne/src/roles/leader/mod.rs
+++ b/crates/daphne/src/roles/leader/mod.rs
@@ -12,7 +12,7 @@ use rand::{thread_rng, Rng};
 use tracing::debug;
 use url::Url;
 
-use super::{aggregator::MergeAggShareError, check_batch, resolve_taskprov, DapAggregator};
+use super::{aggregator::MergeAggShareError, check_batch, DapAggregator};
 use crate::{
     constants::DapMediaType,
     error::DapAbort,
@@ -23,6 +23,7 @@ use crate::{
         PartialBatchSelector, Query, Report, TaskId,
     },
     metrics::{DaphneRequestType, ReportStatus},
+    roles::resolve_task_config,
     DapAggregationParam, DapCollectionJob, DapError, DapLeaderProcessTelemetry, DapRequest,
     DapRequestMeta, DapResponse, DapTaskConfig, DapVersion,
 };
@@ -175,7 +176,6 @@ pub async fn handle_upload_req<A: DapLeader>(
     aggregator: &A,
     req: DapRequest<Report>,
 ) -> Result<(), DapError> {
-    let global_config = aggregator.get_global_config().await?;
     let metrics = aggregator.metrics();
     let task_id = req.task_id;
     debug!("upload for task {task_id}");
@@ -183,18 +183,7 @@ pub async fn handle_upload_req<A: DapLeader>(
     let report = &req.payload;
     debug!("report id is {}", report.report_metadata.id);
 
-    if global_config.allow_taskprov {
-        resolve_taskprov(aggregator, &task_id, &req, &global_config).await?;
-    }
-    let task_config = aggregator
-        .get_task_config_for(&task_id)
-        .await?
-        .ok_or(DapAbort::UnrecognizedTask { task_id })?;
-
-    // Check whether the DAP version in the request matches the task config.
-    if task_config.as_ref().version != req.version {
-        return Err(DapAbort::version_mismatch(req.version, task_config.as_ref().version).into());
-    }
+    let task_config = resolve_task_config(aggregator, &req.meta).await?;
 
     if report.encrypted_input_shares.len() != 2 {
         return Err(DapAbort::InvalidMessage {
@@ -255,15 +244,7 @@ pub async fn handle_coll_job_req<A: DapLeader>(
     let task_id = req.task_id;
     debug!("collect for task {task_id}");
 
-    if global_config.allow_taskprov {
-        resolve_taskprov(aggregator, &task_id, req, &global_config).await?;
-    }
-
-    let wrapped_task_config = aggregator
-        .get_task_config_for(&task_id)
-        .await?
-        .ok_or(DapAbort::UnrecognizedTask { task_id })?;
-    let task_config = wrapped_task_config.as_ref();
+    let task_config = resolve_task_config(aggregator, &req.meta).await?;
 
     let coll_job_req = &req.payload;
 
@@ -271,16 +252,11 @@ pub async fn handle_coll_job_req<A: DapLeader>(
         DapAggregationParam::get_decoded_with_param(&task_config.vdaf, &coll_job_req.agg_param)
             .map_err(|e| DapAbort::from_codec_error(e, task_id))?;
 
-    // Check whether the DAP version in the request matches the task config.
-    if task_config.version != req.version {
-        return Err(DapAbort::version_mismatch(req.version, task_config.version).into());
-    }
-
     // Ensure the batch boundaries are valid and that the batch doesn't overlap with previosuly
     // collected batches.
     check_batch(
         aggregator,
-        task_config,
+        &task_config,
         &task_id,
         &coll_job_req.query,
         &coll_job_req.agg_param,
diff --git a/crates/daphne/src/roles/mod.rs b/crates/daphne/src/roles/mod.rs
index 6b6e213..558bc37 100644
--- a/crates/daphne/src/roles/mod.rs
+++ b/crates/daphne/src/roles/mod.rs
@@ -11,7 +11,6 @@ use crate::{
     messages::{Base64Encode, Query, TaskId, Time},
     taskprov, DapAbort, DapError, DapGlobalConfig, DapQueryConfig, DapRequestMeta, DapTaskConfig,
 };
-use tracing::warn;
 
 pub use aggregator::{DapAggregator, DapReportInitializer};
 pub use helper::DapHelper;
@@ -91,45 +90,37 @@ async fn check_batch(
     Ok(())
 }
 
-async fn resolve_taskprov(
+async fn resolve_task_config(
     agg: &impl DapAggregator,
-    task_id: &TaskId,
     req: &DapRequestMeta,
-    global_config: &DapGlobalConfig,
-) -> Result<(), DapError> {
-    if agg.get_task_config_for(task_id).await?.is_some() {
-        // Task already configured, so nothing to do.
-        return Ok(());
-    }
-
-    let Some(vdaf_verify_key_init) = agg.taskprov_vdaf_verify_key_init() else {
-        warn!("Taskprov disabled due to missing VDAF verification key initializer.");
-        return Ok(());
-    };
-
-    let Some(collector_hpke_config) = agg.taskprov_collector_hpke_config() else {
-        warn!("Taskprov disabled due to missing Collector HPKE configuration.");
-        return Ok(());
-    };
-
-    let Some(task_config) = taskprov::resolve_advertised_task_config(
-        req,
-        vdaf_verify_key_init,
-        collector_hpke_config,
-        task_id,
-    )?
-    else {
-        // No task configuration advertised, so nothing to do.
-        return Ok(());
+) -> Result<DapTaskConfig, DapError> {
+    let task_config = if let Some(config) = agg.get_task_config_for(&req.task_id).await? {
+        config
+    } else {
+        let taskprov_config = agg.get_taskprov_config();
+        let (Some(taskprov_msg), Some(taskprov_config)) = (&req.taskprov, taskprov_config) else {
+            return Err(DapAbort::UnrecognizedTask {
+                task_id: req.task_id,
+            }
+            .into());
+        };
+        let task_config = taskprov::resolve_advertised_task_config(
+            taskprov_msg,
+            req.version,
+            taskprov_config,
+            &req.task_id,
+        )?;
+        let task_config = agg.taskprov_opt_in(&req.task_id, task_config).await?;
+        agg.taskprov_put(&req.task_id, task_config.clone()).await?;
+        task_config
     };
 
-    // This is the opt-in / opt-out decision point.
-    let task_config = agg
-        .taskprov_opt_in(task_id, task_config, global_config)
-        .await?;
+    // Check whether the DAP version in the request matches the task config.
+    if task_config.version != req.version {
+        return Err(DapAbort::version_mismatch(req.version, task_config.version).into());
+    }
 
-    agg.taskprov_put(req, task_config).await?;
-    Ok(())
+    Ok(task_config)
 }
 
 #[cfg(test)]
@@ -194,7 +185,6 @@ mod test {
                 min_batch_interval_start: 259_200,
                 max_batch_interval_end: 259_200,
                 supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
-                allow_taskprov: true,
                 default_num_agg_span_shards: NonZeroUsize::new(4).unwrap(),
             };
 
@@ -1433,8 +1423,7 @@ mod test {
         .to_config_with_taskprov(
             b"cool task".to_vec(),
             t.now,
-            t.leader.taskprov_vdaf_verify_key_init().unwrap(),
-            t.leader.taskprov_collector_hpke_config().unwrap(),
+            t.leader.get_taskprov_config().unwrap(),
         )
         .unwrap();
 
diff --git a/crates/daphne/src/taskprov.rs b/crates/daphne/src/taskprov.rs
index ef8f027..88fde8e 100644
--- a/crates/daphne/src/taskprov.rs
+++ b/crates/daphne/src/taskprov.rs
@@ -16,9 +16,9 @@ use crate::{
         taskprov::{QueryConfigVar, TaskConfig, VdafTypeVar},
         Duration, TaskId, Time,
     },
+    roles::aggregator::TaskprovConfig,
     vdaf::VdafVerifyKey,
-    DapError, DapQueryConfig, DapRequestMeta, DapTaskConfig, DapTaskConfigMethod, DapVersion,
-    VdafConfig,
+    DapError, DapQueryConfig, DapTaskConfig, DapTaskConfigMethod, DapVersion, VdafConfig,
 };
 use crate::{
     pine::PineParam,
@@ -110,39 +110,32 @@ fn malformed_task_config(task_id: &TaskId, detail: String) -> DapAbort {
 /// The `task_id` is the task ID indicated by the request; if this does not match the derived task
 /// ID, then we return `Err(DapError::Abort(DapAbort::UnrecognizedTask))`.
 pub(crate) fn resolve_advertised_task_config(
-    req: &DapRequestMeta,
-    verify_key_init: &[u8; 32],
-    collector_hpke_config: &HpkeConfig,
+    taskprov_msg: &str,
+    version: DapVersion,
+    taskprov_config: TaskprovConfig<'_>,
     task_id: &TaskId,
-) -> Result<Option<DapTaskConfigNeedsOptIn>, DapAbort> {
-    get_taskprov_task_config(req, task_id)?
-        .map(|task_config_msg| {
-            DapTaskConfigNeedsOptIn::try_from_taskprov(
-                req.version,
-                task_id, // get_taskprov_task_config() checks that this matches the derived ID
-                task_config_msg,
-                verify_key_init,
-                collector_hpke_config,
-            )
-        })
-        .transpose()
+) -> Result<DapTaskConfigNeedsOptIn, DapAbort> {
+    let task_config = get_taskprov_task_config(taskprov_msg, task_id, version)?;
+    DapTaskConfigNeedsOptIn::try_from_taskprov(
+        version,
+        task_id, // get_taskprov_task_config() checks that this matches the derived ID
+        task_config,
+        taskprov_config,
+    )
 }
 
 /// Check for a taskprov extension in the report, and return it if found.
 fn get_taskprov_task_config(
-    req: &DapRequestMeta,
+    taskprov_msg: &str,
     task_id: &TaskId,
-) -> Result<Option<TaskConfig>, DapAbort> {
-    let taskprov_data = if let Some(ref taskprov_base64url) = req.taskprov {
-        decode_base64url_vec(taskprov_base64url).ok_or_else(|| {
-            DapAbort::BadRequest(
-                r#"Invalid advertisement in "dap-taskprov" header: base64url parsing failed"#
-                    .to_string(),
-            )
-        })?
-    } else {
-        return Ok(None);
-    };
+    version: DapVersion,
+) -> Result<TaskConfig, DapAbort> {
+    let taskprov_data = decode_base64url_vec(taskprov_msg).ok_or_else(|| {
+        DapAbort::BadRequest(
+            r#"Invalid advertisement in "dap-taskprov" header: base64url parsing failed"#
+                .to_string(),
+        )
+    })?;
 
     if compute_task_id(taskprov_data.as_ref()) != *task_id {
         // Return unrecognizedTask following section 5.1 of the taskprov draft.
@@ -150,10 +143,10 @@ fn get_taskprov_task_config(
     }
 
     // Return unrecognizedMessage if parsing fails following section 5.1 of the taskprov draft.
-    let task_config = TaskConfig::get_decoded_with_param(&req.version, taskprov_data.as_ref())
+    let task_config = TaskConfig::get_decoded_with_param(&version, taskprov_data.as_ref())
         .map_err(|e| DapAbort::from_codec_error(e, *task_id))?;
 
-    Ok(Some(task_config))
+    Ok(task_config)
 }
 
 fn url_from_bytes(task_id: &TaskId, url_bytes: &[u8]) -> Result<Url, DapAbort> {
@@ -346,8 +339,7 @@ impl DapTaskConfigNeedsOptIn {
         version: DapVersion,
         task_id: &TaskId,
         task_config: TaskConfig,
-        vdaf_verify_key_init: &[u8; 32],
-        collector_hpke_config: &HpkeConfig,
+        taskprov_config: TaskprovConfig<'_>,
     ) -> Result<Self, DapAbort> {
         // Only one query per batch is currently supported.
         if task_config.query_config.max_batch_query_count != 1 {
@@ -361,8 +353,12 @@ impl DapTaskConfigNeedsOptIn {
         }
 
         let vdaf = VdafConfig::try_from_taskprov(task_id, version, task_config.vdaf_config.var)?;
-        let vdaf_verify_key =
-            compute_vdaf_verify_key(version, vdaf_verify_key_init, task_id, &vdaf);
+        let vdaf_verify_key = compute_vdaf_verify_key(
+            version,
+            taskprov_config.vdaf_verify_key_init,
+            task_id,
+            &vdaf,
+        );
         Ok(Self {
             version,
             leader_url: url_from_bytes(task_id, &task_config.leader_url.bytes)?,
@@ -373,7 +369,7 @@ impl DapTaskConfigNeedsOptIn {
             query: DapQueryConfig::try_from_taskprov(task_id, task_config.query_config.var)?,
             vdaf,
             vdaf_verify_key,
-            collector_hpke_config: collector_hpke_config.clone(),
+            collector_hpke_config: taskprov_config.hpke_collector_config.clone(),
             method: DapTaskConfigMethod::Taskprov {
                 info: task_config.task_info,
             },
@@ -541,10 +537,12 @@ mod test {
             version,
             &task_id,
             taskprov_config.clone(),
-            &[0; 32],
-            &HpkeReceiverConfig::gen(23, HpkeKemId::P256HkdfSha256)
-                .unwrap()
-                .config,
+            crate::roles::aggregator::TaskprovConfig {
+                hpke_collector_config: &HpkeReceiverConfig::gen(23, HpkeKemId::P256HkdfSha256)
+                    .unwrap()
+                    .config,
+                vdaf_verify_key_init: &[0; 32],
+            },
         )
         .unwrap()
         .into_opted_in(&OptInParam {
@@ -636,7 +634,15 @@ mod test {
             .config;
 
         assert_matches::assert_matches!(
-            resolve_advertised_task_config(&req, &[0; 32], &collector_hpke_config, &task_id).unwrap_err(),
+            resolve_advertised_task_config(
+                req.taskprov.as_deref().unwrap(),
+                req.version,
+                crate::roles::aggregator::TaskprovConfig {
+                    vdaf_verify_key_init: &[0; 32],
+                    hpke_collector_config: &collector_hpke_config,
+                },
+                &task_id
+            ).unwrap_err(),
             DapAbort::InvalidTask{ detail, .. } if detail == "unimplemented VDAF type (1337)"
         );
     }
@@ -688,8 +694,16 @@ mod test {
             .unwrap()
             .config;
 
-        let _ = resolve_advertised_task_config(&req, &[0; 32], &collector_hpke_config, &task_id)
-            .unwrap();
+        let _ = resolve_advertised_task_config(
+            req.taskprov.as_deref().unwrap(),
+            req.version,
+            crate::roles::aggregator::TaskprovConfig {
+                hpke_collector_config: &collector_hpke_config,
+                vdaf_verify_key_init: &[0; 32],
+            },
+            &task_id,
+        )
+        .unwrap();
     }
 
     test_versions! { resolve_advertised_task_config_ignore_unimplemented_dp_ocnfig }


Test patch:
diff --git a/crates/daphne-server/tests/e2e/e2e.rs b/crates/daphne-server/tests/e2e/e2e.rs
index e94e5e3..0765d48 100644
--- a/crates/daphne-server/tests/e2e/e2e.rs
+++ b/crates/daphne-server/tests/e2e/e2e.rs
@@ -339,8 +339,10 @@ async fn leader_upload_taskprov() {
     .to_config_with_taskprov(
         b"cool task".to_vec(),
         t.now,
-        &t.taskprov_vdaf_verify_key_init,
-        &t.taskprov_collector_hpke_receiver.config,
+        daphne::roles::aggregator::TaskprovConfig {
+            hpke_collector_config: &t.taskprov_collector_hpke_receiver.config,
+            vdaf_verify_key_init: &t.taskprov_vdaf_verify_key_init,
+        },
     )
     .unwrap();
 
@@ -1103,8 +1105,10 @@ async fn leader_collect_taskprov_ok(version: DapVersion) {
     .to_config_with_taskprov(
         b"cool task".to_vec(),
         t.now,
-        &t.taskprov_vdaf_verify_key_init,
-        &t.taskprov_collector_hpke_receiver.config,
+        daphne::roles::aggregator::TaskprovConfig {
+            hpke_collector_config: &t.taskprov_collector_hpke_receiver.config,
+            vdaf_verify_key_init: &t.taskprov_vdaf_verify_key_init,
+        },
     )
     .unwrap();
 
diff --git a/crates/daphne-server/tests/e2e/test_runner.rs b/crates/daphne-server/tests/e2e/test_runner.rs
index bd64ce2..abae06d 100644
--- a/crates/daphne-server/tests/e2e/test_runner.rs
+++ b/crates/daphne-server/tests/e2e/test_runner.rs
@@ -103,7 +103,6 @@ impl TestRunner {
             min_batch_interval_start: 259_200,
             max_batch_interval_end: 259_200,
             supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
-            allow_taskprov: true,
             default_num_agg_span_shards: NonZeroUsize::new(1).unwrap(),
         };
 
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index f50b9c3..f7fac5d 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -19,7 +19,7 @@ use crate::{
     metrics::{prometheus::DaphnePromMetrics, DaphneMetrics},
     protocol::aggregator::{EarlyReportStateConsumed, EarlyReportStateInitialized},
     roles::{
-        aggregator::MergeAggShareError,
+        aggregator::{MergeAggShareError, TaskprovConfig},
         helper,
         leader::{in_memory_leader::InMemoryLeaderState, WorkItem},
         DapAggregator, DapHelper, DapLeader, DapReportInitializer,
@@ -754,28 +754,21 @@ impl DapReportInitializer for InMemoryAggregator {
 
 #[async_trait]
 impl DapAggregator for InMemoryAggregator {
-    // The lifetimes on the traits ensure that we can return a reference to a task config stored by
-    // the DapAggregator. (See DaphneWorkerConfig for an example.) For simplicity, InMemoryAggregator
-    // clones the task config as needed.
-    type WrappedDapTaskConfig<'a> = DapTaskConfig;
-
     async fn get_global_config(&self) -> Result<DapGlobalConfig, DapError> {
         Ok(self.global_config.clone())
     }
 
-    fn taskprov_vdaf_verify_key_init(&self) -> Option<&[u8; 32]> {
-        Some(&self.taskprov_vdaf_verify_key_init)
-    }
-
-    fn taskprov_collector_hpke_config(&self) -> Option<&HpkeConfig> {
-        Some(&self.collector_hpke_config)
+    fn get_taskprov_config(&self) -> Option<TaskprovConfig<'_>> {
+        Some(TaskprovConfig {
+            hpke_collector_config: &self.collector_hpke_config,
+            vdaf_verify_key_init: &self.taskprov_vdaf_verify_key_init,
+        })
     }
 
     async fn taskprov_opt_in(
         &self,
         _task_id: &TaskId,
         task_config: taskprov::DapTaskConfigNeedsOptIn,
-        _global_config: &DapGlobalConfig,
     ) -> Result<DapTaskConfig, DapError> {
         // Always opt-in with four shards.
         Ok(task_config.into_opted_in(&taskprov::OptInParam {
@@ -786,18 +779,18 @@ impl DapAggregator for InMemoryAggregator {
 
     async fn taskprov_put(
         &self,
-        req: &DapRequestMeta,
+        task_id: &TaskId,
         task_config: DapTaskConfig,
     ) -> Result<(), DapError> {
         let mut tasks = self.tasks.lock().expect("tasks: lock failed");
-        tasks.deref_mut().insert(req.task_id, task_config);
+        tasks.deref_mut().insert(*task_id, task_config);
         Ok(())
     }
 
     async fn get_task_config_for<'req>(
         &'req self,
         task_id: &'req TaskId,
-    ) -> Result<Option<Self::WrappedDapTaskConfig<'req>>, DapError> {
+    ) -> Result<Option<DapTaskConfig>, DapError> {
         let tasks = self.tasks.lock().expect("tasks: lock failed");
         Ok(tasks.get(task_id).cloned())
     }