Repository: cloudflare/daphne
Pull request #454 (2023-12-05 14:05:29)

Remove url from DapRequest
This is necessary for building a native service, which won't have access to the
domain name



Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 7b95f77..7a95bfe 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -613,6 +613,7 @@ dependencies = [
  "tracing",
  "tracing-core",
  "tracing-subscriber",
+ "url",
  "worker",
 ]
 
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index d812c1a..0c4ec7f 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -456,13 +456,14 @@ pub struct DapTaskParameters {
 impl DapTaskParameters {
     /// Construct a new task config using the taskprov extension. Return the task ID, the taskprov
     /// advertisement (if applicable), and the payload of the report extension.
+    #[allow(clippy::type_complexity)]
     pub fn to_config_with_taskprov(
         &self,
         task_info: Vec<u8>,
         now: Time,
         vdaf_verify_key_init: &[u8; 32],
         collector_hpke_config: &HpkeConfig,
-    ) -> Result<(DapTaskConfig, TaskId, Option<String>, Vec<u8>), DapError> {
+    ) -> Result<(DapTaskConfig, TaskId, Option<String>, Option<Vec<u8>>), DapError> {
         let taskprov_config = messages::taskprov::TaskConfig {
             task_info,
             leader_url: messages::taskprov::UrlBytes {
@@ -498,12 +499,10 @@ impl DapTaskParameters {
         .unwrap();
 
         let (taskprov_advertisement, taskprov_report_extension_payload) = match self.version {
-            DapVersion::DraftLatest => {
-                (Some(encode_base64url(&encoded_taskprov_config)), Vec::new())
-            }
+            DapVersion::DraftLatest => (Some(encode_base64url(&encoded_taskprov_config)), None),
             // draft02 compatibility: The taskprov config is advertised in an HTTP header in
             // the latest draft. In draft02, it is carried by a report extension.
-            DapVersion::Draft02 => (None, encoded_taskprov_config),
+            DapVersion::Draft02 => (None, Some(encoded_taskprov_config)),
         };
 
         Ok((
@@ -1136,9 +1135,6 @@ pub struct DapRequest<S> {
     /// Request payload.
     pub payload: Vec<u8>,
 
-    /// Requst path (i.e., URL).
-    pub url: Url,
-
     /// Sender authorization, e.g., a bearer token.
     pub sender_auth: Option<S>,
 
@@ -1155,7 +1151,6 @@ impl<S> Default for DapRequest<S> {
             task_id: Default::default(),
             resource: Default::default(),
             payload: Default::default(),
-            url: Url::parse("http://example.com").unwrap(),
             sender_auth: Default::default(),
             taskprov: Default::default(),
         }
@@ -1197,12 +1192,6 @@ impl<S> DapRequest<S> {
             ))
         }
     }
-
-    /// Return the hostname of the request URL. The value is "unspecified-host" if the URL does not
-    /// indicate a hostname.
-    pub fn host(&self) -> &str {
-        self.url.host_str().unwrap_or("unspecified-host")
-    }
 }
 
 /// DAP response.
diff --git a/daphne/src/messages/mod.rs b/daphne/src/messages/mod.rs
index ee3aaf8..dfba589 100644
--- a/daphne/src/messages/mod.rs
+++ b/daphne/src/messages/mod.rs
@@ -1221,9 +1221,7 @@ pub fn encode_base64url<T: AsRef<[u8]>>(input: T) -> String {
 }
 
 /// Decode the input as a URL-safe, base64 encoding of an `OUT_LEN`-length byte string.
-pub(crate) fn decode_base64url<T: AsRef<[u8]>, const OUT_LEN: usize>(
-    input: T,
-) -> Option<[u8; OUT_LEN]> {
+pub fn decode_base64url<T: AsRef<[u8]>, const OUT_LEN: usize>(input: T) -> Option<[u8; OUT_LEN]> {
     let mut bytes = [0; OUT_LEN];
     // NOTE(cjpatton) It would be better to use `decode_slice` here, but this function uses a
     // conservative estimate of the decoded length (`decoded_len_estimate`). See
diff --git a/daphne/src/messages/taskprov.rs b/daphne/src/messages/taskprov.rs
index 41a1791..783c1b8 100644
--- a/daphne/src/messages/taskprov.rs
+++ b/daphne/src/messages/taskprov.rs
@@ -248,19 +248,19 @@ impl ParameterizedEncode<DapVersion> for QueryConfig {
         self.time_precision.encode(bytes);
         self.max_batch_query_count.encode(bytes);
         self.min_batch_size.encode(bytes);
+        if *version == DapVersion::DraftLatest {
+            self.encode_query_type(bytes);
+        }
         match &self.var {
             QueryConfigVar::TimeInterval => {
-                QUERY_TYPE_TIME_INTERVAL.encode(bytes);
                 taskprov_encode_u16_prefixed(*version, bytes, |_, _| ());
             }
             QueryConfigVar::FixedSize { max_batch_size } => {
-                QUERY_TYPE_FIXED_SIZE.encode(bytes);
                 taskprov_encode_u16_prefixed(*version, bytes, |_version, inner| {
                     max_batch_size.encode(inner);
                 });
             }
-            QueryConfigVar::NotImplemented { typ, param } => {
-                typ.encode(bytes);
+            QueryConfigVar::NotImplemented { typ: _, param } => {
                 taskprov_encode_u16_prefixed(*version, bytes, |_version, inner| {
                     inner.extend_from_slice(param);
                 });
@@ -276,12 +276,12 @@ impl ParameterizedDecode<DapVersion> for QueryConfig {
     ) -> Result<Self, CodecError> {
         let query_type = match version {
             DapVersion::DraftLatest => None,
-            DapVersion::Draft02 => Some(u8::decode(bytes)?),
+            DapVersion::Draft02 => Some(Ok(u8::decode(bytes)?)),
         };
         let time_precision = Duration::decode(bytes)?;
         let max_batch_query_count = u16::decode(bytes)?;
         let min_batch_size = u32::decode(bytes)?;
-        let query_type = query_type.unwrap_or(u8::decode(bytes)?);
+        let query_type = query_type.unwrap_or_else(|| u8::decode(bytes))?;
         let var = match (version, query_type) {
             (.., QUERY_TYPE_TIME_INTERVAL) => {
                 taskprov_decode_u16_prefixed::<()>(*version, bytes, |_version, inner| {
diff --git a/daphne/src/roles/aggregator.rs b/daphne/src/roles/aggregator.rs
index ad891c6..e5de0a5 100644
--- a/daphne/src/roles/aggregator.rs
+++ b/daphne/src/roles/aggregator.rs
@@ -12,8 +12,7 @@ use crate::{
     error::DapAbort,
     hpke::{HpkeConfig, HpkeDecrypter},
     messages::{
-        decode_base64url, BatchId, BatchSelector, HpkeConfigList, PartialBatchSelector, ReportId,
-        TaskId, Time,
+        BatchId, BatchSelector, HpkeConfigList, PartialBatchSelector, ReportId, TaskId, Time,
     },
     metrics::{DaphneMetrics, DaphneRequestType},
     vdaf::{EarlyReportStateConsumed, EarlyReportStateInitialized},
@@ -159,12 +158,17 @@ pub trait DapAggregator<S: Sync>: HpkeDecrypter + DapReportInitializer + Sized {
 
     /// Access the audit log.
     fn audit_log(&self) -> &dyn AuditLog;
+
+    /// Return the hostname of the request URL. The value is "unspecified-host" if the URL does not
+    /// indicate a hostname.
+    fn host(&self) -> &str;
 }
 
 /// Handle request for the Aggregator's HPKE configuration.
 pub async fn handle_hpke_config_req<S, A>(
     aggregator: &A,
     req: &DapRequest<S>,
+    task_id: Option<TaskId>,
 ) -> Result<DapResponse, DapError>
 where
     S: Sync,
@@ -172,25 +176,11 @@ where
 {
     let metrics = aggregator.metrics();
 
-    // Parse the task ID from the query string, ensuring that it is the only query parameter.
-    let mut id = None;
-    for (k, v) in req.url.query_pairs() {
-        if k != "task_id" {
-            return Err(DapAbort::BadRequest("unexpected query parameter".into()).into());
-        }
-
-        let bytes = decode_base64url(v.as_bytes()).ok_or(DapAbort::BadRequest(
-            "failed to parse query parameter as URL-safe Base64".into(),
-        ))?;
-
-        id = Some(TaskId(bytes));
-    }
-
     let hpke_config = aggregator
-        .get_hpke_config_for(req.version, id.as_ref())
+        .get_hpke_config_for(req.version, task_id.as_ref())
         .await?;
 
-    if let Some(task_id) = id {
+    if let Some(task_id) = task_id {
         let task_config = aggregator
             .get_task_config_for(&task_id)
             .await?
diff --git a/daphne/src/roles/helper.rs b/daphne/src/roles/helper.rs
index d9f14d5..d0f3061 100644
--- a/daphne/src/roles/helper.rs
+++ b/daphne/src/roles/helper.rs
@@ -68,7 +68,7 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     if aggregator.get_global_config().allow_taskprov {
         // draft02 compatibility: We also need to ensure that all of the reports include the task
         // config in the report extensions. (See draft-wang-ppm-dap-taskprov-02, Section 6.)
-        let first_metadata = if req.version == DapVersion::default() {
+        let first_metadata = if req.version == DapVersion::Draft02 {
             let using_taskprov = agg_job_init_req
                 .prep_inits
                 .iter()
@@ -208,7 +208,7 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     };
 
     aggregator.audit_log().on_aggregation_job(
-        req.host(),
+        aggregator.host(),
         task_id,
         task_config,
         agg_job_init_req.prep_inits.len() as u64,
@@ -288,7 +288,7 @@ pub async fn handle_agg_job_cont_req<'req, S: Sync, A: DapHelper<S>>(
         .try_into()
         .expect("usize to fit in u64");
     aggregator.audit_log().on_aggregation_job(
-        req.host(),
+        aggregator.host(),
         task_id,
         task_config,
         out_shares_count,
diff --git a/daphne/src/roles/leader.rs b/daphne/src/roles/leader.rs
index b2ca54e..84e5a1e 100644
--- a/daphne/src/roles/leader.rs
+++ b/daphne/src/roles/leader.rs
@@ -66,7 +66,6 @@ async fn leader_send_http_request<S: Sync>(
         media_type: req_media_type.clone(),
         task_id: Some(*task_id),
         resource,
-        url,
         sender_auth: Some(
             role.authorize(task_id, task_config, &req_media_type, &req_data)
                 .await?,
@@ -76,8 +75,8 @@ async fn leader_send_http_request<S: Sync>(
     };
 
     let resp = match method {
-        LeaderHttpRequestMethod::Put => role.send_http_put(req).await?,
-        LeaderHttpRequestMethod::Post => role.send_http_post(req).await?,
+        LeaderHttpRequestMethod::Put => role.send_http_put(req, url).await?,
+        LeaderHttpRequestMethod::Post => role.send_http_post(req, url).await?,
     };
 
     check_response_content_type(&resp, resp_media_type)?;
@@ -147,10 +146,10 @@ pub trait DapLeader<S: Sync>: DapAuthorizedSender<S> + DapAggregator<S> {
     ) -> Result<(), DapError>;
 
     /// Send an HTTP POST request.
-    async fn send_http_post(&self, req: DapRequest<S>) -> Result<DapResponse, DapError>;
+    async fn send_http_post(&self, req: DapRequest<S>, url: Url) -> Result<DapResponse, DapError>;
 
     /// Send an HTTP PUT request.
-    async fn send_http_put(&self, req: DapRequest<S>) -> Result<DapResponse, DapError>;
+    async fn send_http_put(&self, req: DapRequest<S>, url: Url) -> Result<DapResponse, DapError>;
 }
 
 /// Handle a report from a Client.
diff --git a/daphne/src/roles/mod.rs b/daphne/src/roles/mod.rs
index fe58f50..eb43e9d 100644
--- a/daphne/src/roles/mod.rs
+++ b/daphne/src/roles/mod.rs
@@ -490,7 +490,6 @@ mod test {
                 task_id: Some(*task_id),
                 resource: DapResource::Undefined,
                 payload: report.get_encoded_with_param(&version),
-                url: task_config.leader_url.join("upload").unwrap(),
                 ..Default::default()
             }
         }
@@ -539,7 +538,6 @@ mod test {
                     Some(&agg_job_id),
                     DapMediaType::AggregationJobInitReq,
                     agg_job_init_req,
-                    task_config.helper_url.join("aggregate").unwrap(),
                 )
                 .await,
             )
@@ -565,7 +563,6 @@ mod test {
                     round,
                     transitions,
                 },
-                task_config.helper_url.join("aggregate").unwrap(),
             )
             .await
         }
@@ -594,12 +591,6 @@ mod test {
             let task_id = &self.time_interval_task_id;
             let task_config = self.leader.unchecked_get_task_config(task_id).await;
 
-            let url_path = if task_config.version == DapVersion::Draft02 {
-                "aggregate_shares".to_string()
-            } else {
-                format!("tasks/{}/aggregate_shares", task_id.to_base64url())
-            };
-
             self.leader_authorized_req(
                 task_id,
                 &task_config,
@@ -612,7 +603,6 @@ mod test {
                     report_count,
                     checksum,
                 },
-                task_config.helper_url.join(&url_path).unwrap(),
             )
             .await
         }
@@ -682,7 +672,6 @@ mod test {
                     query: query.clone(),
                     agg_param: Vec::default(),
                 },
-                task_config.helper_url.join("collect").unwrap(),
             );
 
             // Leader: Handle request from Collector.
@@ -709,7 +698,6 @@ mod test {
             agg_job_id: Option<&MetaAggregationJobId>,
             media_type: DapMediaType,
             msg: M,
-            url: Url,
         ) -> DapRequest<BearerToken> {
             let payload = msg.get_encoded_with_param(&task_config.version);
             let sender_auth = Some(
@@ -724,7 +712,6 @@ mod test {
                 task_id: Some(*task_id),
                 resource: agg_job_id.map_or(DapResource::Undefined, |id| id.for_request_path()),
                 payload,
-                url,
                 sender_auth,
                 ..Default::default()
             }
@@ -736,7 +723,6 @@ mod test {
             task_config: &DapTaskConfig,
             media_type: DapMediaType,
             msg: M,
-            url: Url,
         ) -> DapRequest<BearerToken> {
             let mut rng = thread_rng();
             let collect_job_id = CollectionJobId(rng.gen());
@@ -756,7 +742,6 @@ mod test {
                     DapResource::CollectionJob(collect_job_id)
                 },
                 payload: msg.get_encoded_with_param(&task_config.version),
-                url,
                 sender_auth,
                 ..Default::default()
             }
@@ -787,7 +772,6 @@ mod test {
                     },
                     prep_inits: Vec::default(),
                 },
-                task_config.helper_url.join("aggregate").unwrap(),
             )
             .await;
         assert_matches!(
@@ -868,16 +852,11 @@ mod test {
             payload: Vec::new(),
             task_id: Some(task_id),
             resource: DapResource::Undefined,
-            url: Url::parse(&format!(
-                "http://aggregator.biz/v02/hpke_config?task_id={}",
-                task_id.to_base64url()
-            ))
-            .unwrap(),
             ..Default::default()
         };
 
         assert_matches!(
-            aggregator::handle_hpke_config_req(&*t.leader, &req).await,
+            aggregator::handle_hpke_config_req(&*t.leader, &req, Some(task_id)).await,
             Err(DapError::Abort(DapAbort::UnrecognizedTask))
         );
     }
@@ -892,7 +871,6 @@ mod test {
             task_id: Some(t.time_interval_task_id),
             resource: DapResource::Undefined,
             payload: Vec::new(),
-            url: Url::parse("http://aggregator.biz/v02/hpke_config").unwrap(),
             ..Default::default()
         };
 
@@ -900,7 +878,7 @@ mod test {
         // that Daphne-Workder does not implement this behavior. Instead it returns the HPKE config
         // used for all tasks.
         assert_matches!(
-            aggregator::handle_hpke_config_req(&*t.leader, &req).await,
+            aggregator::handle_hpke_config_req(&*t.leader, &req, None).await,
             Err(DapError::Abort(DapAbort::MissingTaskId))
         );
     }
@@ -984,7 +962,6 @@ mod test {
                     report_count: 0,
                     checksum: [0; 32],
                 },
-                task_config.helper_url.join("aggregate_share").unwrap(),
             )
             .await;
         assert_matches!(
@@ -1014,7 +991,6 @@ mod test {
                     report_count: 0,
                     checksum: [0; 32],
                 },
-                task_config.helper_url.join("aggregate_share").unwrap(),
             )
             .await;
         assert_matches!(
@@ -1033,15 +1009,6 @@ mod test {
         let task_id = &t.time_interval_task_id;
         let task_config = t.leader.unchecked_get_task_config(task_id).await;
         let collect_job_id = CollectionJobId(rng.gen());
-        let url_path = if task_config.version == DapVersion::Draft02 {
-            "collect".to_string()
-        } else {
-            format!(
-                "tasks/{}/collection_jobs/{}",
-                task_id.to_base64url(),
-                collect_job_id.to_base64url()
-            )
-        };
         let mut req = DapRequest {
             version: task_config.version,
             media_type: DapMediaType::CollectReq,
@@ -1057,7 +1024,6 @@ mod test {
                 agg_param: Vec::default(),
             }
             .get_encoded_with_param(&task_config.version),
-            url: task_config.leader_url.join(&url_path).unwrap(),
             ..Default::default() // Unauthorized request.
         };
 
@@ -1348,7 +1314,6 @@ mod test {
             task_id: Some(TaskId([0; 32])),
             resource: DapResource::Undefined,
             payload: report_invalid_task_id.get_encoded_with_param(&task_config.version),
-            url: task_config.leader_url.join("upload").unwrap(),
             ..Default::default()
         };
 
@@ -1374,7 +1339,6 @@ mod test {
             task_id: Some(*task_id),
             resource: DapResource::Undefined,
             payload: report.get_encoded_with_param(&version),
-            url: task_config.leader_url.join("upload").unwrap(),
             ..Default::default()
         };
 
@@ -1434,7 +1398,6 @@ mod test {
                 query: task_config.query_for_current_batch_window(t.now),
                 agg_param: Vec::default(),
             },
-            task_config.helper_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1527,7 +1490,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.helper_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1555,7 +1517,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.helper_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1583,7 +1544,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.leader_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1618,7 +1578,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.leader_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1676,7 +1635,6 @@ mod test {
             &task_config,
             DapMediaType::CollectReq,
             collector_collect_req.clone(),
-            task_config.leader_url.join("collect").unwrap(),
         );
 
         // Leader: Handle the CollectReq received from Collector.
@@ -1727,7 +1685,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.leader_url.join("collect").unwrap(),
         );
         assert_matches!(
             leader::handle_collect_job_req(&*t.leader, &req)
@@ -1752,7 +1709,6 @@ mod test {
                 },
                 agg_param: Vec::default(),
             },
-            task_config.leader_url.join("collect").unwrap(),
         );
         assert_matches!(
             leader::handle_collect_job_req(&*t.leader, &req)
@@ -1909,7 +1865,7 @@ mod test {
                     vec![Extension::Taskprov {
                         draft02_payload: match version {
                             DapVersion::DraftLatest => None,
-                            DapVersion::Draft02 => Some(taskprov_report_extension_payload.clone()),
+                            DapVersion::Draft02 => taskprov_report_extension_payload.clone(),
                         },
                     }],
                     version,
@@ -1922,7 +1878,6 @@ mod test {
                 task_id: Some(task_id),
                 resource: DapResource::Undefined,
                 payload: report.get_encoded_with_param(&version),
-                url: Url::parse("https://leader.com/upload").unwrap(),
                 taskprov: taskprov_advertisement.clone(),
                 ..Default::default()
             };
diff --git a/daphne/src/taskprov.rs b/daphne/src/taskprov.rs
index 9fc0e65..78d9fa4 100644
--- a/daphne/src/taskprov.rs
+++ b/daphne/src/taskprov.rs
@@ -423,7 +423,6 @@ impl ReportMetadata {
 #[cfg(test)]
 mod test {
     use prio::codec::ParameterizedEncode;
-    use url::Url;
 
     use super::{compute_task_id, compute_vdaf_verify_key, resolve_advertised_task_config};
     use crate::{
@@ -649,7 +648,6 @@ mod test {
                 task_id: Some(task_id),
                 resource: DapResource::Undefined, // ignored by test
                 payload: Vec::default(),          // ignored by test
-                url: Url::parse("https://example.com/").unwrap(), // ignored by test
                 sender_auth: None,                // ignored by test
                 taskprov: Some(taskprov_task_config_base64url),
             };
diff --git a/daphne_worker/Cargo.toml b/daphne_worker/Cargo.toml
index 280aeb7..26e1c11 100644
--- a/daphne_worker/Cargo.toml
+++ b/daphne_worker/Cargo.toml
@@ -20,6 +20,7 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 arc-swap = "1.6.0"
 async-trait.workspace = true
+bincode = "1.3.3"
 chrono = { version = "0.4.31", default-features = false, features = ["clock", "wasmbind"] }
 daphne = { path = "../daphne" }
 futures.workspace = true
@@ -36,8 +37,8 @@ serde_json.workspace = true
 tracing-core = "0.1.32"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"]}
 tracing.workspace = true
+url.workspace = true
 worker.workspace = true
-bincode = "1.3.3"
 
 [dev-dependencies]
 daphne = { path = "../daphne", features = ["test-utils"] }
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index 2be49c2..0938804 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -410,7 +410,7 @@ pub(crate) struct DaphneWorkerRequestState<'srv> {
     /// Metrics.
     pub(crate) metrics: DaphneWorkerMetrics,
 
-    /// Hostname parsed from the HTTP request URL. Set to "unspecified-daphne-worker-hsot" if the
+    /// Hostname parsed from the HTTP request URL. Set to "unspecified-host" if the
     /// hostname is not part of the URL.
     pub(crate) host: String,
 
@@ -431,7 +431,7 @@ impl<'srv> DaphneWorkerRequestState<'srv> {
         let host = req
             .url()?
             .host_str()
-            .unwrap_or("unspecified-daphne-worker-host")
+            .unwrap_or("unspecified-host")
             .to_string();
 
         let prometheus_registry = Registry::new_custom(
@@ -1160,7 +1160,6 @@ impl<'srv> DaphneWorker<'srv> {
             task_id,
             resource,
             payload,
-            url: req.url().map_err(|e| fatal_error!(err = ?e))?,
             media_type,
             sender_auth,
             taskprov: req
@@ -1183,8 +1182,9 @@ impl<'srv> DaphneWorker<'srv> {
         &self,
         req: DapRequest<DaphneWorkerAuth>,
         is_put: bool,
+        url: Url,
     ) -> std::result::Result<DapResponse, DapError> {
-        let (payload, url) = (req.payload, req.url);
+        let payload = req.payload;
 
         let mut headers = reqwest_wasm::header::HeaderMap::new();
 
diff --git a/daphne_worker/src/roles/aggregator.rs b/daphne_worker/src/roles/aggregator.rs
index 5d382bc..da8c0cd 100644
--- a/daphne_worker/src/roles/aggregator.rs
+++ b/daphne_worker/src/roles/aggregator.rs
@@ -412,4 +412,8 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
     fn audit_log(&self) -> &dyn AuditLog {
         self.state.audit_log
     }
+
+    fn host(&self) -> &str {
+        &self.state.host
+    }
 }
diff --git a/daphne_worker/src/roles/leader.rs b/daphne_worker/src/roles/leader.rs
index def60fa..1661be8 100644
--- a/daphne_worker/src/roles/leader.rs
+++ b/daphne_worker/src/roles/leader.rs
@@ -42,6 +42,7 @@ use daphne::{
 use prio::codec::{ParameterizedDecode, ParameterizedEncode};
 use std::collections::HashMap;
 use tracing::debug;
+use url::Url;
 
 #[async_trait(?Send)]
 impl DapAuthorizedSender<DaphneWorkerAuth> for DaphneWorker<'_> {
@@ -338,14 +339,16 @@ impl<'srv> DapLeader<DaphneWorkerAuth> for DaphneWorker<'srv> {
     async fn send_http_post(
         &self,
         req: DapRequest<DaphneWorkerAuth>,
+        url: Url,
     ) -> std::result::Result<DapResponse, DapError> {
-        self.send_http(req, false).await
+        self.send_http(req, false, url).await
     }
 
     async fn send_http_put(
         &self,
         req: DapRequest<DaphneWorkerAuth>,
+        url: Url,
     ) -> std::result::Result<DapResponse, DapError> {
-        self.send_http(req, true).await
+        self.send_http(req, true, url).await
     }
 }
diff --git a/daphne_worker/src/router/aggregator.rs b/daphne_worker/src/router/aggregator.rs
index 14460c9..6a9b2c2 100644
--- a/daphne_worker/src/router/aggregator.rs
+++ b/daphne_worker/src/router/aggregator.rs
@@ -1,4 +1,8 @@
-use daphne::roles::aggregator;
+use daphne::{
+    error::DapAbort,
+    messages::{decode_base64url, TaskId},
+    roles::aggregator,
+};
 use tracing::Instrument;
 
 use crate::info_span_from_dap_request;
@@ -8,6 +12,28 @@ use super::{dap_response_to_worker, DapRouter};
 pub(super) fn add_aggregator_routes(router: DapRouter<'_>) -> DapRouter<'_> {
     router.get_async("/:version/hpke_config", |req, ctx| async move {
         let daph = ctx.data.handler(&ctx.env);
+
+        // Parse the task ID from the query string, ensuring that it is the only query parameter.
+        let task_id = {
+            let url = req.url()?;
+            let mut query = url.query_pairs();
+            let id = match (query.next(), query.next()) {
+                (Some((k, v)), None) if k == "task_id" => decode_base64url(v.as_bytes())
+                    .map(|bytes| Some(TaskId(bytes)))
+                    .ok_or_else(|| {
+                        DapAbort::BadRequest(
+                            "failed to parse query parameter as URL-safe Base64".into(),
+                        )
+                    }),
+                (None, None) => Ok(None),
+                _ => Err(DapAbort::BadRequest("unexpected query parameter".into())),
+            };
+            match id {
+                Ok(maybe_id) => maybe_id,
+                Err(e) => return daph.state.dap_abort_to_worker_response(e),
+            }
+        };
+
         let req = match daph.worker_request_to_dap(req, &ctx).await {
             Ok(req) => req,
             Err(e) => return daph.state.dap_abort_to_worker_response(e),
@@ -15,7 +41,7 @@ pub(super) fn add_aggregator_routes(router: DapRouter<'_>) -> DapRouter<'_> {
 
         let span = info_span_from_dap_request!("hpke_config", req);
 
-        match aggregator::handle_hpke_config_req(&daph, &req)
+        match aggregator::handle_hpke_config_req(&daph, &req, task_id)
             .instrument(span)
             .await
         {


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 6d79c7f..e201e21 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -1162,6 +1162,10 @@ impl DapAggregator<BearerToken> for MockAggregator {
     fn audit_log(&self) -> &dyn AuditLog {
         &self.audit_log
     }
+
+    fn host(&self) -> &str {
+        "unspecified-host"
+    }
 }
 
 #[cfg_attr(not(feature = "send-traits"), async_trait(?Send))]
@@ -1435,7 +1439,11 @@ impl DapLeader<BearerToken> for MockAggregator {
         }
     }
 
-    async fn send_http_post(&self, req: DapRequest<BearerToken>) -> Result<DapResponse, DapError> {
+    async fn send_http_post(
+        &self,
+        req: DapRequest<BearerToken>,
+        _url: Url,
+    ) -> Result<DapResponse, DapError> {
         match req.media_type {
             DapMediaType::AggregationJobInitReq | DapMediaType::AggregationJobContinueReq => {
                 Ok(helper::handle_agg_job_req(
@@ -1455,7 +1463,11 @@ impl DapLeader<BearerToken> for MockAggregator {
         }
     }
 
-    async fn send_http_put(&self, req: DapRequest<BearerToken>) -> Result<DapResponse, DapError> {
+    async fn send_http_put(
+        &self,
+        req: DapRequest<BearerToken>,
+        _url: Url,
+    ) -> Result<DapResponse, DapError> {
         if req.media_type == DapMediaType::AggregationJobInitReq {
             Ok(helper::handle_agg_job_req(
                 &**self.peer.as_ref().expect("peer not configured"),
diff --git a/daphne_worker_test/tests/e2e/e2e.rs b/daphne_worker_test/tests/e2e/e2e.rs
index 85c9bb6..2c8295a 100644
--- a/daphne_worker_test/tests/e2e/e2e.rs
+++ b/daphne_worker_test/tests/e2e/e2e.rs
@@ -1313,7 +1313,7 @@ async fn leader_collect_taskprov_ok(version: DapVersion) {
         let extensions = vec![Extension::Taskprov {
             draft02_payload: match version {
                 DapVersion::DraftLatest => None,
-                DapVersion::Draft02 => Some(taskprov_report_extension_payload.clone()),
+                DapVersion::Draft02 => taskprov_report_extension_payload.clone(),
             },
         }];
         let now = rng.gen_range(t.report_interval(&batch_interval));