Repository: cloudflare/daphne
Pull request #389 (2023-09-06 00:25:05)

Clean up bearer token resolution logic
Decouple taskprov from the `BearerTokenProvider` by getting rid of the
`is_taskprov_*token()` calls and forcing the `get_*token()` calls to
resolve the taskprov tokens.

Accordingly, pass the `DapTaskConfig` so that the implementer knows if
the task was configured via taskprov. This requires resolving the
`DapTaskConfig` before authorizing the request, so re-order the logic in
`DapLeader` and `DapHelper` accordingly.

While at it, pick up an efficiency improvement for DaphneWorker: if
taskprov is enabled for the task, then we use the token from the global
taskprov config and avoid a KV lookup.


Patch:
diff --git a/daphne/src/auth.rs b/daphne/src/auth.rs
index a2ad7da..1344de1 100644
--- a/daphne/src/auth.rs
+++ b/daphne/src/auth.rs
@@ -7,7 +7,7 @@ use crate::{
     constants::DapMediaType,
     fatal_error,
     messages::{constant_time_eq, TaskId},
-    DapError, DapRequest, DapSender,
+    DapError, DapRequest, DapSender, DapTaskConfig,
 };
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
@@ -62,30 +62,27 @@ pub trait BearerTokenProvider {
     async fn get_leader_bearer_token_for<'s>(
         &'s self,
         task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
     ) -> Result<Option<Self::WrappedBearerToken<'s>>, DapError>;
 
     /// Fetch the Collector's bearer token for the given task, if the task is recognized.
     async fn get_collector_bearer_token_for<'s>(
         &'s self,
         task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
     ) -> Result<Option<Self::WrappedBearerToken<'s>>, DapError>;
 
-    /// Returns true if the given bearer token matches the leader token configured for the "taskprov" extension.
-    fn is_taskprov_leader_bearer_token(&self, token: &BearerToken) -> bool;
-
-    /// Returns true if the given bearer token matches the collector token configured for the "taskprov" extension.
-    fn is_taskprov_collector_bearer_token(&self, token: &BearerToken) -> bool;
-
     /// Return a bearer token that can be used to authorize a request with the given task ID and
     /// media type.
     async fn authorize_with_bearer_token<'s>(
         &'s self,
         task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
         media_type: &DapMediaType,
     ) -> Result<Self::WrappedBearerToken<'s>, DapError> {
         if matches!(media_type.sender(), Some(DapSender::Leader)) {
             let token = self
-                .get_leader_bearer_token_for(task_id)
+                .get_leader_bearer_token_for(task_id, task_config)
                 .await?
                 .ok_or_else(|| {
                     fatal_error!(err = "attempted to authorize request with unknown task ID")
@@ -105,6 +102,7 @@ pub trait BearerTokenProvider {
     /// is the reason for the failure.
     async fn bearer_token_authorized<T: AsRef<BearerToken>>(
         &self,
+        task_config: &DapTaskConfig,
         req: &DapRequest<T>,
     ) -> Result<Option<String>, DapError> {
         if req.task_id.is_none() {
@@ -121,38 +119,31 @@ pub trait BearerTokenProvider {
         // token is not formatted properly.
         if matches!(req.media_type.sender(), Some(DapSender::Leader)) {
             if let Some(ref got) = req.sender_auth {
-                if let Some(expected) = self.get_leader_bearer_token_for(task_id).await? {
+                if let Some(expected) = self
+                    .get_leader_bearer_token_for(task_id, task_config)
+                    .await?
+                {
                     return Ok(if got.as_ref() == expected.as_ref() {
                         None
                     } else {
-                        Some("The indicated beareer token is incorrect for the Leader.".into())
+                        Some("The indicated bearer token is incorrect for the Leader.".into())
                     });
                 }
-                return Ok(if self.is_taskprov_leader_bearer_token(got.as_ref()) {
-                    None
-                } else {
-                    Some("The indicated beaer token is incorrect for Taskprov Leader.".into())
-                });
             }
         }
 
         if matches!(req.media_type.sender(), Some(DapSender::Collector)) {
             if let Some(ref got) = req.sender_auth {
-                if let Some(expected) = self.get_collector_bearer_token_for(task_id).await? {
+                if let Some(expected) = self
+                    .get_collector_bearer_token_for(task_id, task_config)
+                    .await?
+                {
                     return Ok(if got.as_ref() == expected.as_ref() {
                         None
                     } else {
                         Some("The indicated bearer token is incorrect for the Collector.".into())
                     });
                 }
-                return Ok(if self.is_taskprov_collector_bearer_token(got.as_ref()) {
-                    None
-                } else {
-                    Some(
-                        "The indicated bearer token is incorrect for the Taskprov Collector."
-                            .into(),
-                    )
-                });
             }
         }
 
diff --git a/daphne/src/roles/aggregator.rs b/daphne/src/roles/aggregator.rs
index e367f8c..7960ac5 100644
--- a/daphne/src/roles/aggregator.rs
+++ b/daphne/src/roles/aggregator.rs
@@ -49,7 +49,11 @@ pub trait DapAggregator<S>: HpkeDecrypter + DapReportInitializer + Sized {
     /// If the return value is `None`, then the request is authorized. If the return value is
     /// `Some(reason)`, then the request is denied and `reason` conveys details about how the
     /// decision was reached.
-    async fn unauthorized_reason(&self, req: &DapRequest<S>) -> Result<Option<String>, DapError>;
+    async fn unauthorized_reason(
+        &self,
+        task_config: &DapTaskConfig,
+        req: &DapRequest<S>,
+    ) -> Result<Option<String>, DapError>;
 
     /// Look up the DAP global configuration.
     fn get_global_config(&self) -> &DapGlobalConfig;
diff --git a/daphne/src/roles/helper.rs b/daphne/src/roles/helper.rs
index 5e51a85..9fc6a52 100644
--- a/daphne/src/roles/helper.rs
+++ b/daphne/src/roles/helper.rs
@@ -54,14 +54,6 @@ pub trait DapHelper<S>: DapAggregator<S> {
             return Err(DapAbort::version_unknown());
         }
 
-        if let Some(reason) = self.unauthorized_reason(req).await? {
-            error!("aborted unauthorized collect request: {reason}");
-            return Err(DapAbort::UnauthorizedRequest {
-                detail: reason,
-                task_id: task_id.clone(),
-            });
-        }
-
         match req.media_type {
             DapMediaType::AggregationJobInitReq => {
                 let agg_job_init_req =
@@ -110,6 +102,14 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     .ok_or(DapAbort::UnrecognizedTask)?;
                 let task_config = wrapped_task_config.as_ref();
 
+                if let Some(reason) = self.unauthorized_reason(task_config, req).await? {
+                    error!("aborted unauthorized collect request: {reason}");
+                    return Err(DapAbort::UnauthorizedRequest {
+                        detail: reason,
+                        task_id: task_id.clone(),
+                    });
+                }
+
                 // draft02 compatibility: In draft02, the aggregation job ID is parsed from the
                 // HTTP request payload; in the latest draft, the aggregation job ID is parsed from
                 // the request path.
@@ -198,6 +198,14 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     .ok_or(DapAbort::UnrecognizedTask)?;
                 let task_config = wrapped_task_config.as_ref();
 
+                if let Some(reason) = self.unauthorized_reason(task_config, req).await? {
+                    error!("aborted unauthorized collect request: {reason}");
+                    return Err(DapAbort::UnauthorizedRequest {
+                        detail: reason,
+                        task_id: task_id.clone(),
+                    });
+                }
+
                 // 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));
@@ -306,7 +314,13 @@ pub trait DapHelper<S>: DapAggregator<S> {
 
         resolve_taskprov(self, task_id, req, None).await?;
 
-        if let Some(reason) = self.unauthorized_reason(req).await? {
+        let wrapped_task_config = self
+            .get_task_config_for(Cow::Borrowed(req.task_id()?))
+            .await?
+            .ok_or(DapAbort::UnrecognizedTask)?;
+        let task_config = wrapped_task_config.as_ref();
+
+        if let Some(reason) = self.unauthorized_reason(task_config, req).await? {
             error!("aborted unauthorized collect request: {reason}");
             return Err(DapAbort::UnauthorizedRequest {
                 detail: reason,
@@ -314,12 +328,6 @@ pub trait DapHelper<S>: DapAggregator<S> {
             });
         }
 
-        let wrapped_task_config = self
-            .get_task_config_for(Cow::Borrowed(req.task_id()?))
-            .await?
-            .ok_or(DapAbort::UnrecognizedTask)?;
-        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));
diff --git a/daphne/src/roles/leader.rs b/daphne/src/roles/leader.rs
index 996c298..1811b49 100644
--- a/daphne/src/roles/leader.rs
+++ b/daphne/src/roles/leader.rs
@@ -62,7 +62,10 @@ async fn leader_send_http_request<S>(
         task_id: Some(task_id.clone()),
         resource,
         url,
-        sender_auth: Some(role.authorize(task_id, &req_media_type, &req_data).await?),
+        sender_auth: Some(
+            role.authorize(task_id, task_config, &req_media_type, &req_data)
+                .await?,
+        ),
         payload: req_data,
         taskprov: None,
     };
@@ -83,6 +86,7 @@ pub trait DapAuthorizedSender<S> {
     async fn authorize(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         media_type: &DapMediaType,
         payload: &[u8],
     ) -> Result<S, DapError>;
@@ -226,7 +230,13 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
 
         resolve_taskprov(self, task_id, req, None).await?;
 
-        if let Some(reason) = self.unauthorized_reason(req).await? {
+        let wrapped_task_config = self
+            .get_task_config_for(Cow::Borrowed(req.task_id()?))
+            .await?
+            .ok_or(DapAbort::UnrecognizedTask)?;
+        let task_config = wrapped_task_config.as_ref();
+
+        if let Some(reason) = self.unauthorized_reason(task_config, req).await? {
             error!("aborted unauthorized collect request: {reason}");
             return Err(DapAbort::UnauthorizedRequest {
                 detail: reason,
@@ -237,11 +247,6 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         let mut collect_req =
             CollectionReq::get_decoded_with_param(&req.version, req.payload.as_ref())
                 .map_err(|e| DapAbort::from_codec_error(e, task_id.clone()))?;
-        let wrapped_task_config = self
-            .get_task_config_for(Cow::Borrowed(req.task_id()?))
-            .await?
-            .ok_or(DapAbort::UnrecognizedTask)?;
-        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 {
diff --git a/daphne/src/roles/mod.rs b/daphne/src/roles/mod.rs
index 5ec3721..a3a2033 100644
--- a/daphne/src/roles/mod.rs
+++ b/daphne/src/roles/mod.rs
@@ -255,6 +255,7 @@ mod test {
         leader: Arc<MockAggregator>,
         helper: Arc<MockAggregator>,
         collector_token: BearerToken,
+        taskprov_collector_token: BearerToken,
         time_interval_task_id: TaskId,
         fixed_size_task_id: TaskId,
         expired_task_id: TaskId,
@@ -345,12 +346,14 @@ mod test {
                 },
             );
 
-            // Authorization tokens, used for all tasks.
-            let leader_token = BearerToken::from("this is a bearer token!");
-            let collector_token = BearerToken::from("This is a DIFFERENT token.");
+            // Authorization tokens. These are normally chosen at random.
+            let leader_token = BearerToken::from("leader_token");
+            let collector_token = BearerToken::from("collector_token");
 
-            // taskprov: VDAF verification key.
+            // taskprov
             let taskprov_vdaf_verify_key_init = rng.gen::<[u8; 32]>();
+            let taskprov_leader_token = BearerToken::from("taskprov_leader_token");
+            let taskprov_collector_token = BearerToken::from("taskprov_collector_token");
 
             let prometheus_registry = prometheus::Registry::new();
 
@@ -370,6 +373,8 @@ mod test {
                 agg_store: Arc::new(Mutex::new(HashMap::new())),
                 collector_hpke_config: collector_hpke_receiver_config.config.clone(),
                 taskprov_vdaf_verify_key_init,
+                taskprov_leader_token: taskprov_leader_token.clone(),
+                taskprov_collector_token: None,
                 metrics: DaphneMetrics::register(&prometheus_registry, Some("test_helper"))
                     .unwrap(),
                 audit_log: MockAuditLog::default(),
@@ -392,6 +397,8 @@ mod test {
                 agg_store: Arc::new(Mutex::new(HashMap::new())),
                 collector_hpke_config: collector_hpke_receiver_config.config,
                 taskprov_vdaf_verify_key_init,
+                taskprov_leader_token,
+                taskprov_collector_token: Some(taskprov_collector_token.clone()),
                 metrics: DaphneMetrics::register(&prometheus_registry, Some("test_leader"))
                     .unwrap(),
                 audit_log: MockAuditLog::default(),
@@ -403,6 +410,7 @@ mod test {
                 leader,
                 helper,
                 collector_token,
+                taskprov_collector_token,
                 time_interval_task_id,
                 fixed_size_task_id,
                 expired_task_id,
@@ -446,10 +454,10 @@ mod test {
             };
 
             let agg_job_id = MetaAggregationJobId::gen_for_version(&version);
-            self.leader_authorized_req_with_version(
+            self.leader_authorized_req(
                 task_id,
+                &task_config,
                 Some(&agg_job_id),
-                task_config.version,
                 DapMediaType::AggregationJobInitReq,
                 AggregationJobInitReq {
                     draft02_task_id: task_id.for_request_payload(&version),
@@ -474,8 +482,8 @@ mod test {
 
             self.leader_authorized_req(
                 task_id,
+                &task_config,
                 Some(agg_job_id),
-                task_config.version,
                 DapMediaType::AggregationJobContinueReq,
                 AggregationJobContinueReq {
                     draft02_task_id: task_id.for_request_payload(&task_config.version),
@@ -517,10 +525,10 @@ mod test {
                 format!("tasks/{}/aggregate_shares", task_id.to_base64url())
             };
 
-            self.leader_authorized_req_with_version(
+            self.leader_authorized_req(
                 task_id,
+                &task_config,
                 None,
-                task_config.version,
                 DapMediaType::AggregateShareReq,
                 AggregateShareReq {
                     draft02_task_id: task_id.for_request_payload(&task_config.version),
@@ -602,9 +610,9 @@ mod test {
             // Collector->Leader: Initialize collection job.
             let req = self
                 .collector_authorized_req(
-                    task_config.version,
-                    DapMediaType::CollectReq,
                     task_id,
+                    task_config,
+                    DapMediaType::CollectReq,
                     CollectionReq {
                         draft02_task_id: task_id.for_request_payload(&task_config.version),
                         query: query.clone(),
@@ -636,49 +644,21 @@ mod test {
         async fn leader_authorized_req<M: ParameterizedEncode<DapVersion>>(
             &self,
             task_id: &TaskId,
+            task_config: &DapTaskConfig,
             agg_job_id: Option<&MetaAggregationJobId<'_>>,
-            version: DapVersion,
-            media_type: DapMediaType,
-            msg: M,
-            url: Url,
-        ) -> DapRequest<BearerToken> {
-            let payload = msg.get_encoded_with_param(&version);
-            let sender_auth = Some(
-                self.leader
-                    .authorize(task_id, &media_type, &payload)
-                    .await
-                    .unwrap(),
-            );
-            DapRequest {
-                version,
-                media_type,
-                task_id: Some(task_id.clone()),
-                resource: agg_job_id.map_or(DapResource::Undefined, |id| id.for_request_path()),
-                payload,
-                url,
-                sender_auth,
-                ..Default::default()
-            }
-        }
-
-        async fn leader_authorized_req_with_version<M: ParameterizedEncode<DapVersion>>(
-            &self,
-            task_id: &TaskId,
-            agg_job_id: Option<&MetaAggregationJobId<'_>>,
-            version: DapVersion,
             media_type: DapMediaType,
             msg: M,
             url: Url,
         ) -> DapRequest<BearerToken> {
-            let payload = msg.get_encoded_with_param(&version);
+            let payload = msg.get_encoded_with_param(&task_config.version);
             let sender_auth = Some(
                 self.leader
-                    .authorize(task_id, &media_type, &payload)
+                    .authorize(task_id, task_config, &media_type, &payload)
                     .await
                     .unwrap(),
             );
             DapRequest {
-                version,
+                version: task_config.version,
                 media_type,
                 task_id: Some(task_id.clone()),
                 resource: agg_job_id.map_or(DapResource::Undefined, |id| id.for_request_path()),
@@ -691,26 +671,32 @@ mod test {
 
         async fn collector_authorized_req<M: ParameterizedEncode<DapVersion>>(
             &self,
-            version: DapVersion,
-            media_type: DapMediaType,
             task_id: &TaskId,
+            task_config: &DapTaskConfig,
+            media_type: DapMediaType,
             msg: M,
             url: Url,
         ) -> DapRequest<BearerToken> {
             let mut rng = thread_rng();
             let collect_job_id = CollectionJobId(rng.gen());
+            let sender_auth = if task_config.taskprov {
+                Some(self.taskprov_collector_token.clone())
+            } else {
+                Some(self.collector_token.clone())
+            };
+
             DapRequest {
-                version,
+                version: task_config.version,
                 media_type,
                 task_id: Some(task_id.clone()),
-                resource: if version == DapVersion::Draft02 {
+                resource: if task_config.version == DapVersion::Draft02 {
                     DapResource::Undefined
                 } else {
                     DapResource::CollectionJob(collect_job_id)
                 },
-                payload: msg.get_encoded_with_param(&version),
+                payload: msg.get_encoded_with_param(&task_config.version),
                 url,
-                sender_auth: Some(self.collector_token.clone()),
+                sender_auth,
                 ..Default::default()
             }
         }
@@ -726,10 +712,10 @@ mod test {
 
         // Helper expects "time_interval" query, but Leader indicates "fixed_size".
         let req = t
-            .leader_authorized_req_with_version(
+            .leader_authorized_req(
                 task_id,
+                &task_config,
                 Some(&agg_job_id),
-                task_config.version,
                 DapMediaType::AggregationJobInitReq,
                 AggregationJobInitReq {
                     draft02_task_id: task_id.for_request_payload(&version),
@@ -1008,10 +994,10 @@ mod test {
             .unchecked_get_task_config(&t.time_interval_task_id)
             .await;
         let req = t
-            .leader_authorized_req_with_version(
+            .leader_authorized_req(
                 &t.time_interval_task_id,
+                &task_config,
                 None,
-                task_config.version,
                 DapMediaType::AggregateShareReq,
                 AggregateShareReq {
                     draft02_task_id: t.time_interval_task_id.for_request_payload(&version),
@@ -1036,10 +1022,10 @@ mod test {
             .unchecked_get_task_config(&t.fixed_size_task_id)
             .await;
         let req = t
-            .leader_authorized_req_with_version(
+            .leader_authorized_req(
                 &t.fixed_size_task_id,
+                &task_config,
                 None,
-                task_config.version,
                 DapMediaType::AggregateShareReq,
                 AggregateShareReq {
                     draft02_task_id: t.fixed_size_task_id.for_request_payload(&version),
@@ -1441,9 +1427,9 @@ mod test {
         let version = task_config.version;
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: task_id.for_request_payload(&version),
                     query: task_config.query_for_current_batch_window(t.now),
@@ -1517,9 +1503,9 @@ mod test {
         // Collector: Create a CollectReq with a very large batch interval.
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: task_id.for_request_payload(&version),
                     query: Query::TimeInterval {
@@ -1544,9 +1530,9 @@ mod test {
         // Collector: Create a CollectReq with a batch interval in the past.
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: task_id.for_request_payload(&version),
                     query: Query::TimeInterval {
@@ -1572,9 +1558,9 @@ mod test {
         // Collector: Create a CollectReq with a batch interval in the future.
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: task_id.for_request_payload(&version),
                     query: Query::TimeInterval {
@@ -1608,9 +1594,9 @@ mod test {
         // Collector: Create a CollectReq with a very large batch interval.
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: task_id.for_request_payload(&version),
                     query: Query::TimeInterval {
@@ -1676,9 +1662,9 @@ mod test {
         };
         let req = t
             .collector_authorized_req(
-                version,
-                DapMediaType::CollectReq,
                 task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 collector_collect_req.clone(),
                 task_config.leader_url.join("collect").unwrap(),
             )
@@ -1721,9 +1707,9 @@ mod test {
             .await;
         let req = t
             .collector_authorized_req(
-                task_config.version,
-                DapMediaType::CollectReq,
                 &t.time_interval_task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: t.time_interval_task_id.for_request_payload(&version),
                     query: Query::FixedSizeByBatchId {
@@ -1746,9 +1732,9 @@ mod test {
             .await;
         let req = t
             .collector_authorized_req(
-                task_config.version,
-                DapMediaType::CollectReq,
                 &t.fixed_size_task_id,
+                &task_config,
+                DapMediaType::CollectReq,
                 CollectionReq {
                     draft02_task_id: t.fixed_size_task_id.for_request_payload(&version),
                     query: Query::FixedSizeByBatchId {
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index d773fdc..060127b 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -1283,6 +1283,15 @@ impl<K: Clone + Eq + std::hash::Hash, V> KvPair<'_, K, V> {
 
 pub(crate) type BearerTokenKvPair<'a> = KvPair<'a, TaskId, BearerToken>;
 
+impl<'a> BearerTokenKvPair<'a> {
+    pub(crate) fn new(task_id: &'a TaskId, bearer_token: &BearerToken) -> Self {
+        Self {
+            key: Cow::Borrowed(task_id),
+            value: bearer_token.clone(),
+        }
+    }
+}
+
 impl AsRef<BearerToken> for BearerTokenKvPair<'_> {
     fn as_ref(&self) -> &BearerToken {
         self.value()
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 09eb778..d0ffe1a 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -46,7 +46,7 @@ use crate::{
 use async_trait::async_trait;
 use daphne::{
     audit_log::AuditLog,
-    auth::{BearerToken, BearerTokenProvider},
+    auth::BearerTokenProvider,
     constants::DapMediaType,
     error::DapAbort,
     fatal_error,
@@ -163,7 +163,17 @@ impl<'srv> BearerTokenProvider for DaphneWorker<'srv> {
     async fn get_leader_bearer_token_for<'s>(
         &'s self,
         task_id: &'s TaskId,
-    ) -> std::result::Result<Option<Self::WrappedBearerToken<'s>>, DapError> {
+        task_config: &DapTaskConfig,
+    ) -> std::result::Result<Option<BearerTokenKvPair<'s>>, DapError> {
+        if let Some(ref taskprov_config) = self.config().taskprov {
+            if self.get_global_config().allow_taskprov && task_config.taskprov {
+                return Ok(Some(BearerTokenKvPair::new(
+                    task_id,
+                    taskprov_config.leader_auth.as_ref(),
+                )));
+            }
+        }
+
         self.get_leader_bearer_token(task_id)
             .await
             .map_err(|e| fatal_error!(err = e))
@@ -172,33 +182,24 @@ impl<'srv> BearerTokenProvider for DaphneWorker<'srv> {
     async fn get_collector_bearer_token_for<'s>(
         &'s self,
         task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
     ) -> std::result::Result<Option<Self::WrappedBearerToken<'s>>, DapError> {
-        self.get_collector_bearer_token(task_id)
-            .await
-            .map_err(|e| fatal_error!(err = e))
-    }
-
-    fn is_taskprov_leader_bearer_token(&self, token: &BearerToken) -> bool {
-        self.get_global_config().allow_taskprov
-            && match &self.config().taskprov {
-                Some(config) => config.leader_auth.as_ref() == token,
-                None => false,
-            }
-    }
-
-    fn is_taskprov_collector_bearer_token(&self, token: &BearerToken) -> bool {
-        self.get_global_config().allow_taskprov
-            && match &self.config().taskprov {
-                Some(config) => {
-                    config
+        if let Some(ref taskprov_config) = self.config().taskprov {
+            if self.get_global_config().allow_taskprov && task_config.taskprov {
+                return Ok(Some(BearerTokenKvPair::new(
+                    task_id,
+                    taskprov_config
                         .collector_auth
                         .as_ref()
                         .expect("collector authorization method not set")
-                        .as_ref()
-                        == token
-                }
-                None => false,
+                        .as_ref(),
+                )));
             }
+        }
+
+        self.get_collector_bearer_token(task_id)
+            .await
+            .map_err(|e| fatal_error!(err = e))
     }
 }
 
@@ -207,12 +208,13 @@ impl DapAuthorizedSender<DaphneWorkerAuth> for DaphneWorker<'_> {
     async fn authorize(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         media_type: &DapMediaType,
         _payload: &[u8],
     ) -> std::result::Result<DaphneWorkerAuth, DapError> {
         Ok(DaphneWorkerAuth {
             bearer_token: Some(
-                self.authorize_with_bearer_token(task_id, media_type)
+                self.authorize_with_bearer_token(task_id, task_config, media_type)
                     .await?
                     .value()
                     .clone(),
@@ -370,6 +372,7 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
 
     async fn unauthorized_reason(
         &self,
+        task_config: &DapTaskConfig,
         req: &DapRequest<DaphneWorkerAuth>,
     ) -> std::result::Result<Option<String>, DapError> {
         let mut authorized = false;
@@ -380,7 +383,9 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
 
         // If a bearer token is present, verify that it can be used to authorize the request.
         if sender_auth.bearer_token.is_some() {
-            if let Some(unauthorized_reason) = self.bearer_token_authorized(req).await? {
+            if let Some(unauthorized_reason) =
+                self.bearer_token_authorized(task_config, req).await?
+            {
                 return Ok(Some(unauthorized_reason));
             }
             authorized = true;


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index e25c3a6..6890c29 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -127,10 +127,14 @@ pub(crate) struct MockAggregator {
     pub(crate) helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapHelperState>>>,
     pub(crate) agg_store: Arc<Mutex<HashMap<TaskId, HashMap<DapBatchBucketOwned, AggStore>>>>,
     pub(crate) collector_hpke_config: HpkeConfig,
-    pub(crate) taskprov_vdaf_verify_key_init: [u8; 32],
     pub(crate) metrics: DaphneMetrics,
     pub(crate) audit_log: MockAuditLog,
 
+    // taskprov
+    pub(crate) taskprov_vdaf_verify_key_init: [u8; 32],
+    pub(crate) taskprov_leader_token: BearerToken,
+    pub(crate) taskprov_collector_token: Option<BearerToken>, // Not set by Helper
+
     // Leader: Reference to peer. Used to simulate HTTP requests from Leader to Helper, i.e.,
     // implement `DapLeader::send_http_post()` for `MockAggregator`. Not set by the Helper.
     pub(crate) peer: Option<Arc<MockAggregator>>,
@@ -263,36 +267,30 @@ impl BearerTokenProvider for MockAggregator {
     async fn get_leader_bearer_token_for<'s>(
         &'s self,
         _task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
     ) -> Result<Option<Self::WrappedBearerToken<'s>>, DapError> {
-        Ok(Some(&self.leader_token))
+        if task_config.taskprov {
+            Ok(Some(&self.taskprov_leader_token))
+        } else {
+            Ok(Some(&self.leader_token))
+        }
     }
 
     async fn get_collector_bearer_token_for<'s>(
         &'s self,
         _task_id: &'s TaskId,
+        task_config: &DapTaskConfig,
     ) -> Result<Option<Self::WrappedBearerToken<'s>>, DapError> {
-        if let Some(ref collector_token) = self.collector_token {
-            Ok(Some(collector_token))
+        if task_config.taskprov {
+            Ok(Some(self.taskprov_collector_token.as_ref().expect(
+                "MockAggregator not configured with taskprov collector token",
+            )))
         } else {
-            Err(fatal_error!(
-                err = "MockAggregator not configured with Collector bearer token",
-            ))
+            Ok(Some(self.collector_token.as_ref().expect(
+                "MockAggregator not configured with collector token",
+            )))
         }
     }
-
-    fn is_taskprov_leader_bearer_token(&self, _token: &BearerToken) -> bool {
-        // MockAggregator currently uses the same token for all tasks, regardless of how the task
-        // is configured. As a result, we don't expect BearerTokenProver::bearer_token_authorized()
-        // to ever reach this point.
-        unreachable!("did not expect to check bearer token");
-    }
-
-    fn is_taskprov_collector_bearer_token(&self, _token: &BearerToken) -> bool {
-        // MockAggregator currently uses the same token for all tasks, regardless of how the task
-        // is configured. As a result, we don't expect BearerTokenProver::bearer_token_authorized()
-        // to ever reach this point.
-        unreachable!("did not expect to check bearer token");
-    }
 }
 
 #[async_trait(?Send)]
@@ -346,11 +344,12 @@ impl DapAuthorizedSender<BearerToken> for MockAggregator {
     async fn authorize(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         media_type: &DapMediaType,
         _payload: &[u8],
     ) -> Result<BearerToken, DapError> {
         Ok(self
-            .authorize_with_bearer_token(task_id, media_type)
+            .authorize_with_bearer_token(task_id, task_config, media_type)
             .await?
             .clone())
     }
@@ -417,9 +416,10 @@ impl DapAggregator<BearerToken> for MockAggregator {
 
     async fn unauthorized_reason(
         &self,
+        task_config: &DapTaskConfig,
         req: &DapRequest<BearerToken>,
     ) -> Result<Option<String>, DapError> {
-        self.bearer_token_authorized(req).await
+        self.bearer_token_authorized(task_config, req).await
     }
 
     fn get_global_config(&self) -> &DapGlobalConfig {