Repository: cloudflare/daphne
Pull request #164 (2022-11-09 23:17:05)

Implement draft-wang-ppm-dap-taskprov-01.
This implements the changes in draft-wang-ppm-dap-taskprov-01, [#162] while continuing to support draft-wang-ppm-dap-taskprov-00.  The taskprov extension is not versioned, and we are currently using the same extension id for both versions, so which version we implement is controlled by global server configuration.  We may wish to considering using a different number, e.g. 0xff01, as that would allow the server to support both versions concurrently.

This PR also addresses [#161] as the work fit in will with this work.  We now pass the DAP version obtained from the UI when appropriate.  For example, a DAP draft02 API invocation using taskprov will create a DAP draft02 task.

I'm leaving this as a draft for the moment as I'd like to add end-to-end testing for taskprov-01.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 9c188b9..129ca0e 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -50,6 +50,7 @@ use std::{
     collections::{HashMap, HashSet},
     fmt::Debug,
 };
+use taskprov::TaskprovVersion;
 use url::Url;
 
 /// DAP errors.
@@ -145,6 +146,10 @@ pub enum DapAbort {
     #[error("invalidBatchSize")]
     InvalidBatchSize,
 
+    /// Invalid DAP task.  Sent when a server opts out of a taskprov task configuration.
+    #[error("invalidTask")]
+    InvalidTask,
+
     /// Request with missing task ID.
     #[error("missingTaskID")]
     MissingTaskId,
@@ -202,8 +207,9 @@ impl DapAbort {
             Self::BatchInvalid
             | Self::BatchMismatch
             | Self::BatchOverlap
-            | Self::InvalidProtocolVersion
             | Self::InvalidBatchSize
+            | Self::InvalidProtocolVersion
+            | Self::InvalidTask
             | Self::QueryMismatch
             | Self::MissingTaskId
             | Self::ReplayedReport
@@ -350,6 +356,9 @@ pub struct DapGlobalConfig {
 
     /// Is the taskprov extension allowed?
     pub allow_taskprov: bool,
+
+    /// Which taskprov draft should be used?
+    pub taskprov_version: TaskprovVersion,
 }
 
 impl DapGlobalConfig {
diff --git a/daphne/src/messages/taskprov.rs b/daphne/src/messages/taskprov.rs
index 4dcb559..f1ac3f5 100644
--- a/daphne/src/messages/taskprov.rs
+++ b/daphne/src/messages/taskprov.rs
@@ -8,9 +8,10 @@ use crate::messages::{
     decode_u16_bytes, encode_u16_bytes, Duration, Time, QUERY_TYPE_FIXED_SIZE,
     QUERY_TYPE_TIME_INTERVAL,
 };
+use crate::taskprov::TaskprovVersion;
 use prio::codec::{
     decode_u16_items, decode_u24_items, decode_u8_items, encode_u16_items, encode_u24_items,
-    encode_u8_items, CodecError, Decode, Encode,
+    encode_u8_items, CodecError, Decode, Encode, ParameterizedDecode, ParameterizedEncode,
 };
 use ring::hkdf::KeyType;
 use serde::{Deserialize, Serialize};
@@ -233,8 +234,8 @@ pub struct QueryConfig {
     pub var: QueryConfigVar,
 }
 
-impl Encode for QueryConfig {
-    fn encode(&self, bytes: &mut Vec<u8>) {
+impl QueryConfig {
+    fn encode_query_type(&self, bytes: &mut Vec<u8>) {
         match &self.var {
             QueryConfigVar::TimeInterval => {
                 QUERY_TYPE_TIME_INTERVAL.encode(bytes);
@@ -243,9 +244,15 @@ impl Encode for QueryConfig {
                 QUERY_TYPE_FIXED_SIZE.encode(bytes);
             }
         }
+    }
+}
+
+impl ParameterizedEncode<TaskprovVersion> for QueryConfig {
+    fn encode_with_param(&self, _encoding_parameter: &TaskprovVersion, bytes: &mut Vec<u8>) {
         self.time_precision.encode(bytes);
         self.max_batch_query_count.encode(bytes);
         self.min_batch_size.encode(bytes);
+        self.encode_query_type(bytes);
         match &self.var {
             QueryConfigVar::TimeInterval => (),
             QueryConfigVar::FixedSize { max_batch_size } => {
@@ -255,12 +262,15 @@ impl Encode for QueryConfig {
     }
 }
 
-impl Decode for QueryConfig {
-    fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
-        let query_type = u8::decode(bytes)?;
+impl ParameterizedDecode<TaskprovVersion> for QueryConfig {
+    fn decode_with_param(
+        _decoding_parameter: &TaskprovVersion,
+        bytes: &mut Cursor<&[u8]>,
+    ) -> Result<Self, CodecError> {
         let time_precision = Duration::decode(bytes)?;
         let max_batch_query_count = u16::decode(bytes)?;
         let min_batch_size = u32::decode(bytes)?;
+        let query_type = u8::decode(bytes)?;
         let var = match query_type {
             QUERY_TYPE_TIME_INTERVAL => Ok(QueryConfigVar::TimeInterval),
             QUERY_TYPE_FIXED_SIZE => Ok(QueryConfigVar::FixedSize {
@@ -287,22 +297,26 @@ pub struct TaskConfig {
     pub vdaf_config: VdafConfig,
 }
 
-impl Encode for TaskConfig {
-    fn encode(&self, bytes: &mut Vec<u8>) {
+impl ParameterizedEncode<TaskprovVersion> for TaskConfig {
+    fn encode_with_param(&self, encoding_parameter: &TaskprovVersion, bytes: &mut Vec<u8>) {
         encode_u8_items(bytes, &(), &self.task_info);
         encode_u16_items(bytes, &(), &self.aggregator_endpoints);
-        self.query_config.encode(bytes);
+        self.query_config
+            .encode_with_param(encoding_parameter, bytes);
         self.task_expiration.encode(bytes);
         self.vdaf_config.encode(bytes);
     }
 }
 
-impl Decode for TaskConfig {
-    fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
+impl ParameterizedDecode<TaskprovVersion> for TaskConfig {
+    fn decode_with_param(
+        decoding_parameter: &TaskprovVersion,
+        bytes: &mut Cursor<&[u8]>,
+    ) -> Result<Self, CodecError> {
         Ok(TaskConfig {
             task_info: decode_u8_items(&(), bytes)?,
             aggregator_endpoints: decode_u16_items(&(), bytes)?,
-            query_config: QueryConfig::decode(bytes)?,
+            query_config: QueryConfig::decode_with_param(decoding_parameter, bytes)?,
             task_expiration: Time::decode(bytes)?,
             vdaf_config: VdafConfig::decode(bytes)?,
         })
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 7133936..5965a81 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -54,9 +54,24 @@ 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.
+    ///
+    /// Returning `Ok(true)` opts in, returning `Ok(false)` opts out, and any error is
+    /// also an opt out, but that error code is used instead of InvalidTask.
+    fn taskprov_opt_in_decision(&self, task_config: &DapTaskConfig) -> Result<bool, DapError>;
+
     /// Look up the DAP task configuration for the given task ID.
+    ///
+    /// If a `report` has been provided, then look for the draft-wang-ppm-dap-taskprov-<nn> extension
+    /// in the report.  If a taskprov task configuration is successfully read from the report,
+    /// [`DapAggregator::taskprov_opt_in_decision`] will be called, and if it returns Ok(true) the server will opt-in to the task.
+    /// if it returns Ok(false) or an error then the server will opt-out or return an appropriate error.
+    ///
+    /// 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 get_task_config_considering_taskprov(
         &'srv self,
+        version: DapVersion,
         task_id: Cow<'req, Id>,
         report: Option<&ReportMetadata>,
     ) -> Result<Option<Self::WrappedDapTaskConfig>, DapError>;
@@ -66,7 +81,9 @@ where
         &'srv self,
         task_id: Cow<'req, Id>,
     ) -> Result<Option<Self::WrappedDapTaskConfig>, DapError> {
-        self.get_task_config_considering_taskprov(task_id, None)
+        // We use DapVersion::Unknown here as we don't know it and we don't need to
+        // know it as we will not be doing any taskprov task creation.
+        self.get_task_config_considering_taskprov(DapVersion::Unknown, task_id, None)
             .await
     }
 
@@ -244,6 +261,7 @@ where
         let report = Report::get_decoded(req.payload.as_ref())?;
         let task_config = self
             .get_task_config_considering_taskprov(
+                req.version,
                 Cow::Borrowed(req.task_id()?),
                 Some(&report.metadata),
             )
@@ -593,12 +611,17 @@ where
 
                 // If taskprov is allowed, ensure that either all of the shares have it or none of them
                 // do (section 6 of draft-wang-ppm-dap-taskprov-00).
-                if self.get_global_config().allow_taskprov {
+                let global_config = self.get_global_config();
+                if global_config.allow_taskprov {
                     let task_id = req.task_id()?;
                     let using_taskprov = agg_init_req
                         .report_shares
                         .iter()
-                        .filter(|share| share.metadata.is_taskprov(task_id))
+                        .filter(|share| {
+                            share
+                                .metadata
+                                .is_taskprov(global_config.taskprov_version, task_id)
+                        })
                         .count();
 
                     if using_taskprov == agg_init_req.report_shares.len() {
@@ -616,6 +639,7 @@ where
 
                 let wrapped_task_config = self
                     .get_task_config_considering_taskprov(
+                        req.version,
                         Cow::Borrowed(req.task_id()?),
                         first_metadata,
                     )
diff --git a/daphne/src/taskprov.rs b/daphne/src/taskprov.rs
index 07a390a..2fb0b23 100644
--- a/daphne/src/taskprov.rs
+++ b/daphne/src/taskprov.rs
@@ -9,43 +9,64 @@ use crate::{
     vdaf::VdafVerifyKey,
     DapAbort, DapError, DapQueryConfig, DapTaskConfig, DapVersion, Prio3Config, VdafConfig,
 };
-use prio::codec::Decode;
+use prio::codec::ParameterizedDecode;
 use ring::{
+    digest,
     hkdf::{Prk, Salt, HKDF_SHA256},
-    hmac::HMAC_SHA256,
-    hmac::{sign, Key},
 };
+use serde::{Deserialize, Serialize};
 use std::str;
 use url::Url;
 
-/// SHA-256 of "dap-taskprov-00"
+/// DAP taskprov version.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub enum TaskprovVersion {
+    #[serde(rename = "v01")]
+    Draft01,
+
+    #[serde(other)]
+    Unknown,
+}
+
+/// SHA-256 of "dap-taskprov"
 #[allow(dead_code)]
-pub(crate) const TASK_PROV_SALT: [u8; 32] = [
-    0x4d, 0x63, 0x1a, 0xeb, 0xa8, 0xdf, 0xe0, 0x1b, 0x34, 0x4c, 0x29, 0x2d, 0x17, 0xba, 0x34, 0x9a,
-    0x78, 0x97, 0xbf, 0x64, 0x88, 0x00, 0x55, 0x1c, 0x0d, 0x75, 0x32, 0xab, 0x61, 0x4b, 0xe2, 0x21,
+pub(crate) const TASK_PROV_SALT_DRAFT01: [u8; 32] = [
+    0x28, 0xb9, 0xbb, 0x4f, 0x62, 0x4f, 0x67, 0x9a, 0xc1, 0x98, 0xd9, 0x68, 0xf4, 0xb0, 0x9e, 0xec,
+    0x74, 0x01, 0x7a, 0x52, 0xcb, 0x4c, 0xf6, 0x39, 0xfb, 0x83, 0xe0, 0x47, 0x72, 0x3a, 0x0f, 0xfe,
 ];
 
-/// Compute the task id of a serialized task config.
-pub fn compute_task_id(serialized: &[u8]) -> Id {
-    // This is an implementation of HKDF-Expand as the ring::hkdf API does not expose the result
-    // of the expand step.
-    let key = Key::new(HMAC_SHA256, &TASK_PROV_SALT);
-    let digest = sign(&key, serialized);
+fn compute_task_id_draft01(serialized: &[u8]) -> Id {
+    let d = digest::digest(&digest::SHA256, serialized);
+    let dref = d.as_ref();
     let mut b: [u8; 32] = [0; 32];
-    let d = digest.as_ref();
-    b[..32].copy_from_slice(&d[..32]);
+    b[..32].copy_from_slice(&dref[..32]);
     Id(b)
 }
 
+/// Compute the task id of a serialized task config.
+pub fn compute_task_id(version: TaskprovVersion, serialized: &[u8]) -> Id {
+    match version {
+        TaskprovVersion::Draft01 => compute_task_id_draft01(serialized),
+        _ => panic!("unimplemented taskprov version"),
+    }
+}
+
 // The documentation for ring::hkdf says computing the Salt is expensive, and we use the same PRK all the
 // time, so we have structured the code so it would be possible to generate the PRK once and then
 // compute a VDAF verify key many times from the PRK, but we do not actually exploit this currently.
 
 /// Extract pseudorandom key from the pre-shared secret used for the "taskprov" extension.
-pub(crate) fn extract_prk_from_verify_key_init(verify_key_init: &[u8]) -> Prk {
+pub(crate) fn extract_prk_from_verify_key_init(
+    version: TaskprovVersion,
+    verify_key_init: &[u8],
+) -> Prk {
     // The documentation says computing the Salt is expensive, and we use the same PRK all the
     // time, so we compute it once.
-    Salt::new(HKDF_SHA256, &TASK_PROV_SALT).extract(verify_key_init)
+    let value = match version {
+        TaskprovVersion::Draft01 => &TASK_PROV_SALT_DRAFT01,
+        _ => panic!("unimplemented taskprov version"),
+    };
+    Salt::new(HKDF_SHA256, value).extract(verify_key_init)
 }
 
 /// Expand a pseudorandom key into the VDAF verification key for a given task.
@@ -77,12 +98,13 @@ pub(crate) fn expand_prk_into_verify_key(
 /// directly.
 #[allow(dead_code)]
 pub(crate) fn compute_vdaf_verify_key(
+    version: TaskprovVersion,
     verify_key_init: &[u8],
     task_id: &Id,
     vdaf_type: VdafType,
 ) -> VdafVerifyKey {
     expand_prk_into_verify_key(
-        &extract_prk_from_verify_key_init(verify_key_init),
+        &extract_prk_from_verify_key_init(version, verify_key_init),
         task_id,
         vdaf_type,
     )
@@ -94,6 +116,7 @@ pub fn bad_request(detail: &str) -> DapError {
 
 /// Check for a taskprov extension in the report, and return it if found.
 pub fn get_taskprov_task_config(
+    version: TaskprovVersion,
     task_id: &Id,
     metadata: &ReportMetadata,
 ) -> Result<Option<TaskConfig>, DapError> {
@@ -106,12 +129,12 @@ pub fn get_taskprov_task_config(
         0 => Ok(None),
         1 => match &taskprovs[0] {
             Extension::Taskprov { payload } => {
-                if compute_task_id(&payload[..]) != *task_id {
+                if compute_task_id(version, &payload[..]) != *task_id {
                     // Return unrecognizedTask following section 5.1 of the taskprov draft.
                     return Err(DapError::Abort(DapAbort::UnrecognizedTask));
                 }
                 // Return unrecognizedMessage if parsing fails following section 5.1 of the taskprov draft.
-                let task_config = TaskConfig::get_decoded(payload)
+                let task_config = TaskConfig::get_decoded_with_param(&version, payload)
                     .map_err(|_| DapError::Abort(DapAbort::UnrecognizedMessage))?;
                 Ok(Some(task_config))
             }
@@ -157,7 +180,8 @@ impl From<VdafTypeVar> for VdafConfig {
 
 impl DapTaskConfig {
     pub fn try_from_taskprov(
-        version: DapVersion,
+        dap_version: DapVersion,
+        taskprov_version: TaskprovVersion,
         task_id: &Id,
         task_config: TaskConfig,
         vdaf_verify_key_init: &[u8],
@@ -168,7 +192,7 @@ impl DapTaskConfig {
         }
         let vdaf_type = VdafType::from(task_config.vdaf_config.var.clone());
         Ok(DapTaskConfig {
-            version,
+            version: dap_version,
             leader_url: url_from_bytes(&task_config.aggregator_endpoints[0].bytes)?,
             helper_url: url_from_bytes(&task_config.aggregator_endpoints[1].bytes)?,
             time_precision: task_config.query_config.time_precision,
@@ -176,7 +200,12 @@ impl DapTaskConfig {
             min_batch_size: task_config.query_config.min_batch_size.into(),
             query: DapQueryConfig::from(task_config.query_config.var),
             vdaf: VdafConfig::from(task_config.vdaf_config.var),
-            vdaf_verify_key: compute_vdaf_verify_key(vdaf_verify_key_init, task_id, vdaf_type),
+            vdaf_verify_key: compute_vdaf_verify_key(
+                taskprov_version,
+                vdaf_verify_key_init,
+                task_id,
+                vdaf_type,
+            ),
             collector_hpke_config: collector_hpke_config.clone(),
         })
     }
@@ -184,9 +213,9 @@ impl DapTaskConfig {
 
 impl ReportMetadata {
     /// Does this metatdata have a taskprov extension and does it match the specified id?
-    pub fn is_taskprov(&self, task_id: &Id) -> bool {
+    pub fn is_taskprov(&self, version: TaskprovVersion, task_id: &Id) -> bool {
         return self.extensions.iter().any(|x| match x {
-            Extension::Taskprov { payload } => *task_id == compute_task_id(payload),
+            Extension::Taskprov { payload } => *task_id == compute_task_id(version, payload),
             _ => false,
         });
     }
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index dc41674..24404eb 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -268,11 +268,20 @@ where
         &self.global_config
     }
 
+    fn taskprov_opt_in_decision(
+        &self,
+        _task_config: &DapTaskConfig,
+    ) -> std::result::Result<bool, DapError> {
+        // For now we always opt-in.
+        Ok(true)
+    }
+
     /// 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 get_task_config_considering_taskprov(
         &'srv self,
+        version: DapVersion,
         task_id: Cow<'req, Id>,
         metadata: Option<&ReportMetadata>,
     ) -> std::result::Result<Option<GuardedDapTaskConfig<'req>>, DapError> {
@@ -289,7 +298,11 @@ where
             return Ok(None);
         }
         let metadata_ref = metadata.unwrap();
-        let taskprov_task_config = get_taskprov_task_config(task_id.as_ref(), metadata_ref)?;
+        let taskprov_task_config = get_taskprov_task_config(
+            self.global_config.taskprov_version,
+            task_id.as_ref(),
+            metadata_ref,
+        )?;
         if taskprov_task_config.is_some() {
             let global_config = self.get_global_config();
             if !global_config.allow_taskprov {
@@ -301,18 +314,21 @@ where
                 .as_ref()
                 .ok_or_else(|| DapError::fatal("taskprov configuration not found"))?;
 
-            // We currently always opt-in; any additional checks that would lead to
-            // us opting out should be done here.
-
             let taskprov_task_id = task_id.as_ref().clone();
             let task_config = DapTaskConfig::try_from_taskprov(
-                DapVersion::Draft02,
+                version,
+                self.global_config.taskprov_version,
                 &taskprov_task_id,
                 taskprov_task_config.unwrap(),
                 taskprov_config.vdaf_verify_key_init.as_ref(),
                 taskprov_config.hpke_collector_config.as_ref(),
             )?;
 
+            // This is the opt-in / opt-out decision point.
+            if !self.taskprov_opt_in_decision(&task_config)? {
+                return Err(DapError::Abort(daphne::DapAbort::InvalidTask));
+            }
+
             // Write the leader bearer token to the KV.  We do this so authorize_with_bearer_token()
             // finds something.
             //


Test patch:
diff --git a/daphne/src/messages/mod_test.rs b/daphne/src/messages/mod_test.rs
index a76bbea..a975fc3 100644
--- a/daphne/src/messages/mod_test.rs
+++ b/daphne/src/messages/mod_test.rs
@@ -10,8 +10,8 @@ use crate::messages::{
     HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, PartialBatchSelector, Report, ReportId,
     ReportMetadata, ReportShare, Transition, TransitionVar,
 };
-use crate::taskprov::compute_task_id;
-use prio::codec::{Decode, Encode};
+use crate::taskprov::{compute_task_id, TaskprovVersion};
+use prio::codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode};
 
 #[test]
 fn read_report() {
@@ -195,18 +195,18 @@ fn read_vdaf_config() {
 }
 
 #[test]
-fn read_task_config() {
+fn read_task_config_taskprov_draft01() {
     let data = [
         0x02, 0x48, 0x69, 0x00, 0x0e, 0x00, 0x0c, 0x68, 0x74, 0x74, 0x70, 0x73, 0x3a, 0x2f, 0x2f,
-        0x74, 0x65, 0x73, 0x74, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x80,
-        0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x63, 0x52, 0xf9,
+        0x74, 0x65, 0x73, 0x74, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x80, 0x00,
+        0x00, 0x04, 0x00, 0x02, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x63, 0x52, 0xf9,
         0xa5, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x18, 0x01, 0x02, 0x03, 0x04, 0x04, 0x03,
         0x02, 0x01, 0x02, 0x02, 0x03, 0x04, 0x04, 0x03, 0x02, 0x02, 0x03, 0x02, 0x03, 0x04, 0x04,
         0x03, 0x02, 0x03,
     ];
 
     let buckets = vec![0x0102030404030201, 0x0202030404030202, 0x0302030404030203];
-    let task_config = TaskConfig::get_decoded(&data).unwrap();
+    let task_config = TaskConfig::get_decoded_with_param(&TaskprovVersion::Draft01, &data).unwrap();
     assert_eq!(
         task_config,
         TaskConfig {
@@ -233,7 +233,16 @@ fn read_task_config() {
     );
 
     assert_eq!(
-        compute_task_id(&task_config.get_encoded()).to_hex(),
-        "b4769bb063a8b3312af74297f30fdbf8e0b71c2eb2481f591d1d7de66a4ce34f"
+        compute_task_id(
+            TaskprovVersion::Draft01,
+            &task_config.get_encoded_with_param(&TaskprovVersion::Draft01)
+        )
+        .to_hex(),
+        "2e5eb17cb12faaf122fca963b8e4cca3f11737f31fc83a02cf9d72019d2b2081"
+    );
+
+    assert_eq!(
+        task_config.get_encoded_with_param(&TaskprovVersion::Draft01),
+        &data
     );
 }
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 5a6d09e..3e2ae68 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -15,6 +15,7 @@ use crate::{
         TransitionVar,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
+    taskprov::TaskprovVersion,
     testing::{AggStore, DapBatchBucketOwned, MockAggregator, MockAggregatorReportSelector},
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateShare, DapCollectJob, DapGlobalConfig, DapLeaderTransition,
@@ -71,6 +72,7 @@ impl Test {
             max_batch_interval_end: 259200,
             supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
             allow_taskprov: true,
+            taskprov_version: TaskprovVersion::Draft01,
         };
 
         // Task Parameters that the Leader and Helper must agree on.
diff --git a/daphne/src/taskprov_test.rs b/daphne/src/taskprov_test.rs
index 77e4547..eda7891 100644
--- a/daphne/src/taskprov_test.rs
+++ b/daphne/src/taskprov_test.rs
@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use crate::{
-    messages::taskprov::VdafType, messages::Id, taskprov::compute_vdaf_verify_key,
+    messages::taskprov::VdafType,
+    messages::Id,
+    taskprov::{compute_vdaf_verify_key, TaskprovVersion},
     vdaf::VdafVerifyKey,
 };
 
@@ -18,10 +20,15 @@ fn check_vdaf_key_computation() {
         0x32, 0xd7, 0xe1, 0xbc, 0x6c, 0x75, 0x10, 0x05, 0x60, 0x7b, 0x81, 0xda, 0xc3, 0xa7, 0xda,
         0x76, 0x1d,
     ];
-    let vk = compute_vdaf_verify_key(&verify_key_init, &task_id, VdafType::Prio3Aes128Count);
+    let vk = compute_vdaf_verify_key(
+        TaskprovVersion::Draft01,
+        &verify_key_init,
+        &task_id,
+        VdafType::Prio3Aes128Count,
+    );
     let expected: [u8; 16] = [
-        0x57, 0x7a, 0x80, 0x3c, 0xb7, 0xac, 0xa7, 0x5f, 0xdf, 0xb5, 0x86, 0xa9, 0xc1, 0xb6, 0x01,
-        0x75,
+        0xfb, 0xd1, 0x7d, 0xb5, 0x39, 0x0f, 0x94, 0x9e, 0xe3, 0x2d, 0x26, 0x34, 0xdc, 0x49, 0x9f,
+        0x5b,
     ];
     match &vk {
         VdafVerifyKey::Prio3(bytes) => assert_eq!(*bytes, expected),
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 582645a..1a7608e 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -13,6 +13,7 @@ use crate::{
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     DapAbort, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig,
     DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
+    DapVersion,
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
@@ -289,8 +290,13 @@ where
         &self.global_config
     }
 
+    fn taskprov_opt_in_decision(&self, _task_config: &DapTaskConfig) -> Result<bool, DapError> {
+        panic!("MockAggregator does not implement the taskprov extension")
+    }
+
     async fn get_task_config_considering_taskprov(
         &'srv self,
+        _version: DapVersion,
         task_id: Cow<'req, Id>,
         _metadata: Option<&ReportMetadata>,
     ) -> Result<Option<&'req DapTaskConfig>, DapError> {
diff --git a/daphne_worker_test/tests/e2e.rs b/daphne_worker_test/tests/e2e.rs
index a49810d..c085bbd 100644
--- a/daphne_worker_test/tests/e2e.rs
+++ b/daphne_worker_test/tests/e2e.rs
@@ -15,11 +15,11 @@ use daphne::{
         BatchSelector, CollectReq, CollectResp, Extension, HpkeCiphertext, Id, Interval, Query,
         Report, ReportId, ReportMetadata,
     },
-    taskprov::compute_task_id,
+    taskprov::{compute_task_id, TaskprovVersion},
     DapAggregateResult, DapMeasurement, DapTaskConfig, DapVersion,
 };
 use daphne_worker::DaphneWorkerReportSelector;
-use prio::codec::{Decode, Encode};
+use prio::codec::{Decode, Encode, ParameterizedEncode};
 use rand::prelude::*;
 use serde::Deserialize;
 use serde_json::json;
@@ -312,8 +312,8 @@ async fn e2e_leader_upload() {
             var: VdafTypeVar::Prio3Aes128Count,
         },
     };
-    let payload = taskprov_task_config.get_encoded();
-    let task_id = compute_task_id(&payload);
+    let payload = taskprov_task_config.get_encoded_with_param(&TaskprovVersion::Draft01);
+    let task_id = compute_task_id(TaskprovVersion::Draft01, &payload);
     let extensions = vec![Extension::Taskprov { payload }];
     let report = t
         .task_config
@@ -335,10 +335,10 @@ async fn e2e_leader_upload() {
     .await;
 
     // Generate and upload a report with taskprov but with the wrong id
-    let payload = taskprov_task_config.get_encoded();
+    let payload = taskprov_task_config.get_encoded_with_param(&TaskprovVersion::Draft01);
     let mut bad_payload = payload.clone();
     bad_payload[0] = u8::wrapping_add(bad_payload[0], 1);
-    let task_id = compute_task_id(&bad_payload);
+    let task_id = compute_task_id(TaskprovVersion::Draft01, &bad_payload);
     let extensions = vec![Extension::Taskprov { payload }];
     let report = t
         .task_config
@@ -363,8 +363,8 @@ async fn e2e_leader_upload() {
     .await;
 
     // Generate and upload a report with two copies of the taskprov extension
-    let payload = taskprov_task_config.get_encoded();
-    let task_id = compute_task_id(&payload);
+    let payload = taskprov_task_config.get_encoded_with_param(&TaskprovVersion::Draft01);
+    let task_id = compute_task_id(TaskprovVersion::Draft01, &payload);
     let extensions = vec![
         Extension::Taskprov {
             payload: payload.clone(),
@@ -393,24 +393,34 @@ async fn e2e_leader_upload() {
     )
     .await;
 
-    // Generate and upload a report with taskprov but with only one aggregator endpoint (an error).
-    // (This report is also expired, but we notice the endpoint problem first).
-    let payload = [
-        0x02u8, 0x48, 0x69, 0x00, 0x0e, 0x00, 0x0c, 0x68, 0x74, 0x74, 0x70, 0x73, 0x3a, 0x2f, 0x2f,
-        0x74, 0x65, 0x73, 0x74, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x80,
-        0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x63, 0x52, 0xf9,
-        0xa5, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x18, 0x01, 0x02, 0x03, 0x04, 0x04, 0x03,
-        0x02, 0x01, 0x02, 0x02, 0x03, 0x04, 0x04, 0x03, 0x02, 0x02, 0x03, 0x02, 0x03, 0x04, 0x04,
-        0x03, 0x02, 0x03,
-    ];
-    let extensions = vec![Extension::Taskprov {
-        payload: payload.to_vec(),
-    }];
-    let task_id = Id([
-        0xb4, 0x76, 0x9b, 0xb0, 0x63, 0xa8, 0xb3, 0x31, 0x2a, 0xf7, 0x42, 0x97, 0xf3, 0x0f, 0xdb,
-        0xf8, 0xe0, 0xb7, 0x1c, 0x2e, 0xb2, 0x48, 0x1f, 0x59, 0x1d, 0x1d, 0x7d, 0xe6, 0x6a, 0x4c,
-        0xe3, 0x4f,
-    ]);
+    // Generate and upload a report with taskprov but only one endpoint, which is an error.
+    //
+    // We have to make this by hand as if we cut and paste a pre-serialized one it
+    // will have an expiring task.
+    let taskprov_task_config = TaskConfig {
+        task_info: "Hi".as_bytes().to_vec(),
+        aggregator_endpoints: vec![UrlBytes {
+            bytes: "https://test1".as_bytes().to_vec(),
+        }],
+        query_config: QueryConfig {
+            time_precision: 0x01,
+            max_batch_query_count: 128,
+            min_batch_size: 1024,
+            var: QueryConfigVar::FixedSize {
+                max_batch_size: 2048,
+            },
+        },
+        task_expiration: t.now + 86400,
+        vdaf_config: VdafConfig {
+            dp_config: DpConfig {
+                mechanism: DpMechanism::None,
+            },
+            var: VdafTypeVar::Prio3Aes128Count,
+        },
+    };
+    let payload = taskprov_task_config.get_encoded_with_param(&TaskprovVersion::Draft01);
+    let task_id = compute_task_id(TaskprovVersion::Draft01, &payload);
+    let extensions = vec![Extension::Taskprov { payload }];
     let report = t
         .task_config
         .vdaf
@@ -1131,10 +1141,11 @@ async fn e2e_leader_collect_taskprov_ok() {
             var: VdafTypeVar::Prio3Aes128Sum { bit_length: 10 },
         },
     };
-    let payload = taskprov_task_config.get_encoded();
-    let task_id = compute_task_id(&payload);
+    let payload = taskprov_task_config.get_encoded_with_param(&TaskprovVersion::Draft01);
+    let task_id = compute_task_id(TaskprovVersion::Draft01, &payload);
     let task_config = DapTaskConfig::try_from_taskprov(
         DapVersion::Draft02,
+        TaskprovVersion::Draft01,
         &task_id.clone(),
         taskprov_task_config.clone(),
         &t.taskprov_vdaf_verify_key_init,
diff --git a/daphne_worker_test/tests/test_runner.rs b/daphne_worker_test/tests/test_runner.rs
index 0e077ed..86962f3 100644
--- a/daphne_worker_test/tests/test_runner.rs
+++ b/daphne_worker_test/tests/test_runner.rs
@@ -8,6 +8,7 @@ use daphne::{
     constants::MEDIA_TYPE_COLLECT_REQ,
     hpke::HpkeReceiverConfig,
     messages::{Duration, HpkeAeadId, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Interval},
+    taskprov::TaskprovVersion,
     DapGlobalConfig, DapLeaderProcessTelemetry, DapQueryConfig, DapTaskConfig, DapVersion,
     Prio3Config, VdafConfig,
 };
@@ -111,6 +112,7 @@ impl TestRunner {
             max_batch_interval_end: 259200,
             supported_hpke_kems: vec![HpkeKemId::X25519HkdfSha256],
             allow_taskprov: true,
+            taskprov_version: TaskprovVersion::Draft01,
         };
         let taskprov_vdaf_verify_key_init =
             hex::decode("0074a5dd6e9dac501f73f7a961193b2b").unwrap();
diff --git a/daphne_worker_test/wrangler.toml b/daphne_worker_test/wrangler.toml
index 097fd5c..c1bbb25 100644
--- a/daphne_worker_test/wrangler.toml
+++ b/daphne_worker_test/wrangler.toml
@@ -32,7 +32,8 @@ DAP_GLOBAL_CONFIG = """{
      "min_batch_interval_start": 259200,
      "max_batch_interval_end": 259200,
      "supported_hpke_kems": ["x25519_hkdf_sha256"],
-     "allow_taskprov": true
+     "allow_taskprov": true,
+     "taskprov_version": "v01"
 }"""
 DAP_TASKPROV_HPKE_COLLECTOR_CONFIG = """{
   "id": 23,
@@ -87,7 +88,8 @@ DAP_GLOBAL_CONFIG = """{
   "min_batch_interval_start": 259200,
   "max_batch_interval_end": 259200,
   "supported_hpke_kems": ["x25519_hkdf_sha256"],
-  "allow_taskprov": true
+  "allow_taskprov": true,
+  "taskprov_version": "v01"
 }"""
 DAP_TASKPROV_HPKE_COLLECTOR_CONFIG = """{
   "id": 23,