Repository: cloudflare/daphne
Pull request #318 (2023-05-05 06:03:27)

taskprov: Miscellaneous improvements
Based on #317 (merge that first).

Just some minor things I noticed after some recent refactoring:

*    tasprov: Dedup some shared logic

    Both `MockAggregator` and `DaphneWorker` need to parse need to do the
    following:

    1. Resolve advertisement
    2. Decide to opt-in
    3. Store the task

    Step (1.) is backend-agnostic; steps (2.) and (3.) can be handled as
    callbacks. Refactor the code so that they are.

*    daphne: Add field to `DapTaskConfig` indicating if taskprov was used

    For backwards compatibility, JSON deserialization defaults to `false`.
    This means that any tasks configured via taskprov will be treated as
    non-taskprov. However this is currently a noop, since no protocol logic
    uses this flag.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 0fdd2a2..f1d3be2 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -320,6 +320,10 @@ pub struct DapTaskConfig {
 
     /// The Collector's HPKE configuration for this task.
     pub collector_hpke_config: HpkeConfig,
+
+    /// If true, then the taskprov extension was used to configure this task.
+    #[serde(default)]
+    pub taskprov: bool,
 }
 
 impl DapTaskConfig {
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 83898f4..ef7e1f0 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -9,21 +9,21 @@ use crate::{
     messages::{
         constant_time_eq, decode_base64url, AggregateShare, AggregateShareReq,
         AggregationJobContinueReq, AggregationJobInitReq, AggregationJobResp, BatchId,
-        BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeConfigList, Interval,
-        PartialBatchSelector, Query, Report, ReportId, ReportMetadata, TaskId, Time,
+        BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeConfig, HpkeConfigList,
+        Interval, PartialBatchSelector, Query, Report, ReportId, ReportMetadata, TaskId, Time,
         TransitionFailure, TransitionVar,
     },
     metrics::{DaphneMetrics, DaphneRequestType},
-    DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
-    DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition, DapOutputShare,
-    DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig, DapVersion,
-    MetaAggregationJobId,
+    taskprov, DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig,
+    DapHelperState, DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition,
+    DapOutputShare, DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig,
+    DapVersion, MetaAggregationJobId,
 };
 use async_trait::async_trait;
 use prio::codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode};
 use std::borrow::Cow;
 use std::collections::HashMap;
-use tracing::{debug, error};
+use tracing::{debug, error, warn};
 use url::Url;
 
 /// A party in the DAP protocol who is authorized to send requests to another party.
@@ -57,7 +57,15 @@ where
     /// Look up the DAP global configuration.
     fn get_global_config(&self) -> &DapGlobalConfig;
 
-    /// Decide whether to opt-in or out-out of a task provisioned via taskprov.
+    /// 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>;
+
+    /// taskprov: Decide whether to opt-in or out-out of a task provisioned via taskprov.
     ///
     /// If the return value is `None`, then the decision is to opt-in. If the return value is
     /// `Some(reason)`, then the decision is to opt-out; `reason` conveys details about how the
@@ -67,17 +75,13 @@ where
         task_config: &DapTaskConfig,
     ) -> Result<Option<String>, DapError>;
 
-    /// taskprov: Attempt to configure the task advertised in the request. If taskprov is disabled
-    /// then this is a noop. Otherwise, if successful, the next call to `get_task_config_for()`
-    /// will return the configure task. Otherwise this call will return nothing.
-    ///
-    /// The DAP version must be specified because we may create a DapTaskConfig via taskprov, and
-    /// we want it to have the same version as the API entry point the client is using.
-    async fn resolve_taskprov(
+    /// taskprov: Configure a task. This is called after opting in. If successful, the next call to
+    /// `get_task_config_for()` will return the configure task. Otherwise this call will return
+    /// nothing.
+    async fn taskprov_put(
         &'srv self,
-        task_id: &TaskId,
         req: &'req DapRequest<S>,
-        report_metadata_advertisement: Option<&ReportMetadata>,
+        task_config: DapTaskConfig,
     ) -> Result<(), DapError>;
 
     /// Look up the DAP task configuration for the given task ID.
@@ -319,8 +323,7 @@ where
         let report = Report::get_decoded_with_param(&req.version, req.payload.as_ref())?;
         debug!("report id is {}", report.report_metadata.id);
 
-        self.resolve_taskprov(task_id, req, Some(&report.report_metadata))
-            .await?;
+        resolve_taskprov(self, task_id, req, Some(&report.report_metadata)).await?;
         let task_config = self
             .get_task_config_for(Cow::Borrowed(task_id))
             .await?
@@ -381,7 +384,7 @@ where
 
         check_request_content_type(req, DapMediaType::CollectReq)?;
 
-        self.resolve_taskprov(task_id, req, None).await?;
+        resolve_taskprov(self, task_id, req, None).await?;
 
         if let Some(reason) = self.unauthorized_reason(req).await? {
             error!("aborted unauthorized collect request: {reason}");
@@ -834,7 +837,7 @@ where
                         return Err(DapAbort::UnrecognizedMessage);
                     }
                 }
-                self.resolve_taskprov(task_id, req, first_metadata).await?;
+                resolve_taskprov(self, task_id, req, first_metadata).await?;
 
                 let wrapped_task_config = self
                     .get_task_config_for(Cow::Borrowed(task_id))
@@ -963,7 +966,7 @@ where
                 })
             }
             DapMediaType::AggregationJobContinueReq => {
-                self.resolve_taskprov(task_id, req, None).await?;
+                resolve_taskprov(self, task_id, req, None).await?;
                 let wrapped_task_config = self
                     .get_task_config_for(Cow::Borrowed(task_id))
                     .await?
@@ -1058,7 +1061,7 @@ where
 
         check_request_content_type(req, DapMediaType::AggregateShareReq)?;
 
-        self.resolve_taskprov(task_id, req, None).await?;
+        resolve_taskprov(self, task_id, req, None).await?;
 
         if let Some(reason) = self.unauthorized_reason(req).await? {
             error!("aborted unauthorized collect request: {reason}");
@@ -1315,3 +1318,61 @@ fn check_request_content_type<S>(
         Ok(())
     }
 }
+
+async fn resolve_taskprov<'srv, 'req, S>(
+    agg: &'srv impl DapAggregator<'srv, 'req, S>,
+    task_id: &'req TaskId,
+    req: &'req DapRequest<S>,
+    report_metadata_advertisement: Option<&ReportMetadata>,
+) -> Result<(), DapError>
+where
+    'srv: 'req,
+{
+    if agg
+        .get_task_config_for(Cow::Borrowed(task_id))
+        .await?
+        .is_some()
+    {
+        // Task already configured, so nothing to do.
+        return Ok(());
+    }
+
+    let global_config = agg.get_global_config();
+    if !global_config.allow_taskprov {
+        // Taskprov is disabled, 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,
+            global_config.taskprov_version,
+            vdaf_verify_key_init,
+            collector_hpke_config,
+            task_id,
+            report_metadata_advertisement,
+        )? else {
+            // No task configuration advertised, so nothing to do.
+            return Ok(());
+        };
+
+    // This is the opt-in / opt-out decision point.
+    if let Some(reason) = agg.taskprov_opt_out_reason(&task_config)? {
+        return Err(DapError::Abort(DapAbort::InvalidTask {
+            detail: reason,
+            task_id: task_id.clone(),
+        }));
+    }
+
+    agg.taskprov_put(req, task_config).await?;
+    Ok(())
+}
diff --git a/daphne/src/taskprov.rs b/daphne/src/taskprov.rs
index 3cd390c..36f649a 100644
--- a/daphne/src/taskprov.rs
+++ b/daphne/src/taskprov.rs
@@ -288,6 +288,7 @@ impl DapTaskConfig {
                 vdaf_type,
             ),
             collector_hpke_config: collector_hpke_config.clone(),
+            taskprov: true,
         })
     }
 }
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index abba0ac..19fbe24 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -42,7 +42,7 @@ use std::{
     borrow::Cow,
     collections::HashMap,
     io::Cursor,
-    sync::{Arc, RwLock, RwLockReadGuard},
+    sync::{Arc, RwLock},
     time::Duration,
 };
 use tracing::{error, info, trace};
@@ -558,16 +558,39 @@ impl<'srv> DaphneWorker<'srv> {
         Ok(None)
     }
 
+    /// In memory cache on-top of KV to avoid hitting KV API limits.
+    // NOTE: We shouldn't return guards from this function as they could be held across await points.
+    //
+    // Locks in wasm are assumed to be single threaded without concurrency:
+    // https://github.com/rust-lang/rust/blob/6f8c0557e0b73c73a8a7163a15f4a5a3feca7d5c/library/std/src/sys/unsupported/locks/rwlock.rs#L4
+    //
+    // However, with Cloudflare Workers the situation is not so straightforward.
+    //
+    // A single isolate can process multiple requests concurrently (not in parallel), so holding
+    // ReadGuards across await points is not OK.
+    //
+    // There are some official docs that refer to the behaviour:
+    // https://developers.cloudflare.com/workers/learning/how-workers-works/#distributed-execution
+    //
+    // The errors seen are:
+    //
+    //   Error: The script will never generate a response.
+    //
+    //   and
+    //
+    //   Memory access out of bounds
+    //
+    // They only seem to happen under load, likely because the Workers runtime will readily spawn
+    // additional isolates.
     async fn kv_get_cached<'req, K, V>(
         &self,
         map: &'srv Arc<RwLock<HashMap<K, V>>>,
         kv_key_prefix: &str,
         kv_key_suffix: Cow<'req, K>,
-    ) -> Result<Option<Guarded<'req, K, V>>>
+    ) -> Result<Option<KvPair<'req, K, V>>>
     where
         K: Clone + Eq + std::hash::Hash + ToString,
-        V: for<'de> Deserialize<'de>,
-        'srv: 'req,
+        V: Clone + for<'de> Deserialize<'de>,
     {
         // If the value is cached, then return immediately.
         {
@@ -575,9 +598,9 @@ impl<'srv> DaphneWorker<'srv> {
                 .read()
                 .map_err(|e| Error::RustError(format!("Failed to lock map for reading: {e}")))?;
 
-            if guarded_map.get(&kv_key_suffix).is_some() {
-                return Ok(Some(Guarded {
-                    guarded_map,
+            if let Some(value) = guarded_map.get(&kv_key_suffix) {
+                return Ok(Some(KvPair {
+                    value: value.clone(),
                     key: kv_key_suffix,
                 }));
             }
@@ -601,9 +624,9 @@ impl<'srv> DaphneWorker<'srv> {
             .read()
             .map_err(|e| Error::RustError(format!("Failed to lock map for reading: {e}")))?;
 
-        if guarded_map.get(kv_key_suffix.as_ref()).is_some() {
-            Ok(Some(Guarded {
-                guarded_map,
+        if let Some(value) = guarded_map.get(kv_key_suffix.as_ref()) {
+            Ok(Some(KvPair {
+                value: value.clone(),
                 key: kv_key_suffix,
             }))
         } else {
@@ -616,7 +639,7 @@ impl<'srv> DaphneWorker<'srv> {
     pub(crate) async fn get_hpke_receiver_config(
         &self,
         hpke_receiver_kv_key: HpkeReceiverKvKey,
-    ) -> Result<Option<GuardedHpkeReceiverConfig>> {
+    ) -> Result<Option<HpkeReceiverConfigKvPair>> {
         self.kv_get_cached(
             &self.isolate_state().hpke_receiver_configs,
             KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG,
@@ -629,7 +652,7 @@ impl<'srv> DaphneWorker<'srv> {
     pub(crate) async fn get_leader_bearer_token<'a>(
         &'a self,
         task_id: &'a TaskId,
-    ) -> Result<Option<GuardedBearerToken>> {
+    ) -> Result<Option<BearerTokenKvPair>> {
         self.kv_get_cached(
             &self.isolate_state().leader_bearer_tokens,
             KV_KEY_PREFIX_BEARER_TOKEN_LEADER,
@@ -652,7 +675,7 @@ impl<'srv> DaphneWorker<'srv> {
     pub(crate) async fn get_collector_bearer_token<'a>(
         &'a self,
         task_id: &'a TaskId,
-    ) -> Result<Option<GuardedBearerToken>> {
+    ) -> Result<Option<BearerTokenKvPair>> {
         self.kv_get_cached(
             &self.isolate_state().collector_bearer_tokens,
             KV_KEY_PREFIX_BEARER_TOKEN_COLLECTOR,
@@ -665,7 +688,7 @@ impl<'srv> DaphneWorker<'srv> {
     pub(crate) async fn get_task_config<'req>(
         &'srv self,
         task_id: Cow<'req, TaskId>,
-    ) -> Result<Option<GuardedDapTaskConfig<'req>>>
+    ) -> Result<Option<DapTaskConfigKvPair<'req>>>
     where
         'srv: 'req,
     {
@@ -692,7 +715,7 @@ impl<'srv> DaphneWorker<'srv> {
     pub(crate) async fn try_get_task_config<'req>(
         &'srv self,
         task_id: &'req TaskId,
-    ) -> std::result::Result<GuardedDapTaskConfig<'req>, DapError>
+    ) -> std::result::Result<DapTaskConfigKvPair<'req>, DapError>
     where
         'srv: 'req,
     {
@@ -895,6 +918,7 @@ impl<'srv> DaphneWorker<'srv> {
                     vdaf,
                     vdaf_verify_key,
                     collector_hpke_config,
+                    taskprov: false,
                 },
             )
             .await?
@@ -1113,41 +1137,41 @@ impl<'srv> DaphneWorker<'srv> {
     }
 }
 
-/// RwLockReadGuard'ed object, used to catch items fetched from KV.
-pub(crate) struct Guarded<'a, K: Clone, V> {
-    guarded_map: RwLockReadGuard<'a, HashMap<K, V>>,
+/// KV pair
+pub(crate) struct KvPair<'a, K: Clone, V> {
     key: Cow<'a, K>,
+    value: V,
 }
 
-impl<K: Clone + Eq + std::hash::Hash, V> Guarded<'_, K, V> {
+impl<K: Clone + Eq + std::hash::Hash, V> KvPair<'_, K, V> {
     pub(crate) fn key(&self) -> &K {
         self.key.as_ref()
     }
 
     pub(crate) fn value(&self) -> &V {
-        self.guarded_map.get(self.key.as_ref()).unwrap()
+        &self.value
     }
 }
 
-pub(crate) type GuardedHpkeReceiverConfig<'a> = Guarded<'a, HpkeReceiverKvKey, HpkeReceiverConfig>;
+pub(crate) type HpkeReceiverConfigKvPair<'a> = KvPair<'a, HpkeReceiverKvKey, HpkeReceiverConfig>;
 
-impl AsRef<HpkeConfig> for GuardedHpkeReceiverConfig<'_> {
+impl AsRef<HpkeConfig> for HpkeReceiverConfigKvPair<'_> {
     fn as_ref(&self) -> &HpkeConfig {
         &self.value().config
     }
 }
 
-pub(crate) type GuardedBearerToken<'a> = Guarded<'a, TaskId, BearerToken>;
+pub(crate) type BearerTokenKvPair<'a> = KvPair<'a, TaskId, BearerToken>;
 
-impl AsRef<BearerToken> for GuardedBearerToken<'_> {
+impl AsRef<BearerToken> for BearerTokenKvPair<'_> {
     fn as_ref(&self) -> &BearerToken {
         self.value()
     }
 }
 
-pub(crate) type GuardedDapTaskConfig<'a> = Guarded<'a, TaskId, DapTaskConfig>;
+pub(crate) type DapTaskConfigKvPair<'a> = KvPair<'a, TaskId, DapTaskConfig>;
 
-impl AsRef<DapTaskConfig> for GuardedDapTaskConfig<'_> {
+impl AsRef<DapTaskConfig> for DapTaskConfigKvPair<'_> {
     fn as_ref(&self) -> &DapTaskConfig {
         self.value()
     }
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 7b9d2a6..226b6a2 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -9,7 +9,7 @@
 use crate::{
     auth::DaphneWorkerAuth,
     config::{
-        DaphneWorker, GuardedBearerToken, GuardedDapTaskConfig, GuardedHpkeReceiverConfig,
+        BearerTokenKvPair, DapTaskConfigKvPair, DaphneWorker, HpkeReceiverConfigKvPair,
         HpkeReceiverKvKey, KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG,
     },
     dap_err,
@@ -45,19 +45,19 @@ use crate::{
 };
 use async_trait::async_trait;
 use daphne::{
-    aborts::DapAbort,
     auth::{BearerToken, BearerTokenProvider},
     constants::DapMediaType,
     hpke::HpkeDecrypter,
     messages::{
         BatchId, BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeCiphertext,
-        PartialBatchSelector, Report, ReportId, ReportMetadata, TaskId, TransitionFailure,
+        HpkeConfig, PartialBatchSelector, Report, ReportId, ReportMetadata, TaskId,
+        TransitionFailure,
     },
     metrics::DaphneMetrics,
     roles::{early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    taskprov, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig,
-    DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapSender,
-    DapTaskConfig, DapVersion, MetaAggregationJobId,
+    DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
+    DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapSender, DapTaskConfig, DapVersion,
+    MetaAggregationJobId,
 };
 use futures::future::try_join_all;
 use prio::codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode};
@@ -87,13 +87,13 @@ pub(crate) fn dap_response_to_worker(resp: DapResponse) -> Result<Response> {
 
 #[async_trait(?Send)]
 impl<'srv> HpkeDecrypter<'srv> for DaphneWorker<'srv> {
-    type WrappedHpkeConfig = GuardedHpkeReceiverConfig<'srv>;
+    type WrappedHpkeConfig = HpkeReceiverConfigKvPair<'srv>;
 
     async fn get_hpke_config_for(
         &'srv self,
         version: DapVersion,
         _task_id: Option<&TaskId>,
-    ) -> std::result::Result<GuardedHpkeReceiverConfig<'srv>, DapError> {
+    ) -> std::result::Result<HpkeReceiverConfigKvPair<'srv>, DapError> {
         let kv_store = self.kv().map_err(dap_err)?;
         let keys = kv_store
             .list()
@@ -207,19 +207,19 @@ impl<'srv> HpkeDecrypter<'srv> for DaphneWorker<'srv> {
 
 #[async_trait(?Send)]
 impl<'srv> BearerTokenProvider<'srv> for DaphneWorker<'srv> {
-    type WrappedBearerToken = GuardedBearerToken<'srv>;
+    type WrappedBearerToken = BearerTokenKvPair<'srv>;
 
     async fn get_leader_bearer_token_for(
         &'srv self,
         task_id: &'srv TaskId,
-    ) -> std::result::Result<Option<GuardedBearerToken>, DapError> {
+    ) -> std::result::Result<Option<BearerTokenKvPair>, DapError> {
         self.get_leader_bearer_token(task_id).await.map_err(dap_err)
     }
 
     async fn get_collector_bearer_token_for(
         &'srv self,
         task_id: &'srv TaskId,
-    ) -> std::result::Result<Option<GuardedBearerToken>, DapError> {
+    ) -> std::result::Result<Option<BearerTokenKvPair>, DapError> {
         self.get_collector_bearer_token(task_id)
             .await
             .map_err(dap_err)
@@ -276,7 +276,7 @@ impl<'srv, 'req> DapAggregator<'srv, 'req, DaphneWorkerAuth> for DaphneWorker<'s
 where
     'srv: 'req,
 {
-    type WrappedDapTaskConfig = GuardedDapTaskConfig<'req>;
+    type WrappedDapTaskConfig = DapTaskConfigKvPair<'req>;
 
     async fn unauthorized_reason(
         &self,
@@ -355,6 +355,20 @@ where
         &self.config().global
     }
 
+    fn taskprov_vdaf_verify_key_init(&self) -> Option<&[u8; 32]> {
+        self.config()
+            .taskprov
+            .as_ref()
+            .map(|config| &config.vdaf_verify_key_init)
+    }
+
+    fn taskprov_collector_hpke_config(&self) -> Option<&HpkeConfig> {
+        self.config()
+            .taskprov
+            .as_ref()
+            .map(|config| &config.hpke_collector_config)
+    }
+
     fn taskprov_opt_out_reason(
         &self,
         _task_config: &DapTaskConfig,
@@ -363,57 +377,18 @@ where
         Ok(None)
     }
 
-    /// Get an existing task (whether an ordinary task or a previously created
-    /// taskprov task).  If we can't find it, see if there is a taskprov extension
-    /// in the report, and if so create the task.
-    async fn resolve_taskprov(
+    async fn taskprov_put(
         &'srv self,
-        task_id: &TaskId,
         req: &'req DapRequest<DaphneWorkerAuth>,
-        report_metadata_advertisement: Option<&ReportMetadata>,
+        task_config: DapTaskConfig,
     ) -> std::result::Result<(), DapError> {
-        let found = self
-            .get_task_config(Cow::Borrowed(task_id))
-            .await
-            .map_err(dap_err)?;
-        if found.is_some() {
-            return Ok(());
-        }
-
-        let global_config = self.get_global_config();
-        if !global_config.allow_taskprov {
-            // TODO(bhalleycf) if DAP gets a generic denied error, we should use it here.
-            return Err(DapError::Abort(DapAbort::InvalidTask {
-                detail: "Taskprov extension is disabled.".to_string(),
-                task_id: task_id.clone(),
-            }));
-        }
-
+        let task_id = req.task_id().map_err(DapError::Abort)?;
         let taskprov = self
             .config()
             .taskprov
             .as_ref()
             .ok_or_else(|| DapError::fatal("taskprov configuration not found"))?;
 
-        let Some(task_config) = taskprov::resolve_advertised_task_config(
-            req,
-            global_config.taskprov_version,
-            &taskprov.vdaf_verify_key_init,
-            &taskprov.hpke_collector_config,
-            task_id,
-            report_metadata_advertisement,
-        )? else {
-            return Ok(());
-        };
-
-        // This is the opt-in / opt-out decision point.
-        if let Some(reason) = self.taskprov_opt_out_reason(&task_config)? {
-            return Err(DapError::Abort(DapAbort::InvalidTask {
-                detail: reason,
-                task_id: task_id.clone(),
-            }));
-        }
-
         // If `resolve_advertised_task_config()` returned a `TaskConfig` and `req.taskprov` is set,
         // then the task was advertised in the HTTP "dap-taskprov" header. In this case we expect
         // the peer to send the header in every request for this task.
@@ -459,7 +434,7 @@ where
     async fn get_task_config_for(
         &'srv self,
         task_id: Cow<'req, TaskId>,
-    ) -> std::result::Result<Option<GuardedDapTaskConfig<'req>>, DapError> {
+    ) -> std::result::Result<Option<DapTaskConfigKvPair<'req>>, DapError> {
         self.get_task_config(task_id).await.map_err(dap_err)
     }
 


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 045b09d..7890142 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -107,6 +107,7 @@ impl Test {
                 query: DapQueryConfig::TimeInterval,
                 vdaf: vdaf_config.clone(),
                 vdaf_verify_key: VdafVerifyKey::Prio3(rng.gen()),
+                taskprov: false,
             },
         );
         tasks.insert(
@@ -122,6 +123,7 @@ impl Test {
                 query: DapQueryConfig::FixedSize { max_batch_size: 2 },
                 vdaf: vdaf_config.clone(),
                 vdaf_verify_key: VdafVerifyKey::Prio3(rng.gen()),
+                taskprov: false,
             },
         );
         tasks.insert(
@@ -137,6 +139,7 @@ impl Test {
                 query: DapQueryConfig::TimeInterval,
                 vdaf: vdaf_config,
                 vdaf_verify_key: VdafVerifyKey::Prio3(rng.gen()),
+                taskprov: false,
             },
         );
 
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index cd58c28..56b9157 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -14,9 +14,9 @@ use crate::{
     },
     metrics::DaphneMetrics,
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    taskprov, DapAbort, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError,
-    DapGlobalConfig, DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse,
-    DapTaskConfig, DapVersion, MetaAggregationJobId,
+    DapAbort, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig,
+    DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
+    DapVersion, MetaAggregationJobId,
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
@@ -344,6 +344,14 @@ where
         &self.global_config
     }
 
+    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 taskprov_opt_out_reason(
         &self,
         _task_config: &DapTaskConfig,
@@ -352,44 +360,14 @@ where
         Ok(None)
     }
 
-    async fn resolve_taskprov(
+    async fn taskprov_put(
         &'srv self,
-        task_id: &TaskId,
         req: &'req DapRequest<BearerToken>,
-        report_metadata_advertisement: Option<&ReportMetadata>,
+        task_config: DapTaskConfig,
     ) -> Result<(), DapError> {
-        let taskprov_version = self.global_config.taskprov_version;
-
-        // Before looking up the task configuration, first check if it needs to be configured from
-        // the current request.
-        if self.get_global_config().allow_taskprov
-            && report_metadata_advertisement.is_some()
-            && report_metadata_advertisement
-                .unwrap()
-                .is_taskprov(taskprov_version, task_id)
-        {
-            if let Some(task_config) = taskprov::resolve_advertised_task_config(
-                req,
-                taskprov_version,
-                &self.taskprov_vdaf_verify_key_init,
-                &self.collector_hpke_config,
-                task_id,
-                report_metadata_advertisement,
-            )? {
-                let mut tasks = self.tasks.lock().expect("tasks: lock failed");
-                if tasks.get(task_id).is_none() {
-                    // Decide whether to opt-in to the task.
-                    if let Some(reason) = self.taskprov_opt_out_reason(&task_config)? {
-                        return Err(DapError::Abort(DapAbort::InvalidTask {
-                            detail: reason,
-                            task_id: task_id.clone(),
-                        }));
-                    }
-
-                    tasks.deref_mut().insert(task_id.clone(), task_config);
-                }
-            }
-        }
+        let task_id = req.task_id().map_err(DapError::Abort)?;
+        let mut tasks = self.tasks.lock().expect("tasks: lock failed");
+        tasks.deref_mut().insert(task_id.clone(), task_config);
         Ok(())
     }
 
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 3aaf43a..a7085f5 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -773,6 +773,7 @@ impl Test {
                 vdaf: vdaf.clone(),
                 vdaf_verify_key,
                 collector_hpke_config,
+                taskprov: false,
             },
             prometheus_registry,
             leader_metrics,
diff --git a/daphne_worker_test/tests/test_runner.rs b/daphne_worker_test/tests/test_runner.rs
index 2311979..6e345f9 100644
--- a/daphne_worker_test/tests/test_runner.rs
+++ b/daphne_worker_test/tests/test_runner.rs
@@ -115,6 +115,7 @@ impl TestRunner {
             vdaf: VDAF_CONFIG.clone(),
             vdaf_verify_key: VDAF_CONFIG.gen_verify_key(),
             collector_hpke_config: collector_hpke_receiver.config.clone(),
+            taskprov: false,
         };
 
         // This block needs to be kept in-sync with daphne_worker_test/wrangler.toml.