Repository: cloudflare/daphne
Pull request #750 (2024-12-20 02:11:52)

Unify VDAF dispatching logic
Partially addresses #698.
Stacked on #749.

Add dispatchers for sharding, preparation, and unsharding to `VdafConfig`. Pass the DAP version to the dispatcher so that we can properly resolve supported VDAFs.

Remove support for the following, unspecified combinations:

* DAP-13/Pine32HmacSha256Aes128
* DAP-13/Pine64HmacSha256Aes128
* DAP-09/Mastic

Technically, since Mastic is still a shim at the moment and not a functional VDAF, we could support it in DAP-09 as well. When the implementation is complete (https://github.com/divviup/libprio-rs/issues/947), it will only be VDAF-13-compatible and thus only supported in DAP-13.

cjpatton: Rebased.

Patch:
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index d781da6..dc82028 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -532,7 +532,6 @@ impl Test {
             agg_job_state,
             agg_job_resp,
             self.metrics(),
-            task_config.version,
         )?;
 
         let aggregated_report_count = agg_share_span
diff --git a/crates/daphne/src/protocol/aggregator.rs b/crates/daphne/src/protocol/aggregator.rs
index 124863d..479d423 100644
--- a/crates/daphne/src/protocol/aggregator.rs
+++ b/crates/daphne/src/protocol/aggregator.rs
@@ -5,8 +5,6 @@ use super::{
     no_duplicates,
     report_init::{InitializedReport, WithPeerPrepShare},
 };
-#[cfg(feature = "experimental")]
-use crate::vdaf::mastic::{mastic_prep_finish, mastic_prep_finish_from_shares};
 use crate::{
     constants::DapAggregatorRole,
     error::DapAbort,
@@ -19,12 +17,9 @@ use crate::{
     },
     metrics::{DaphneMetrics, ReportStatus},
     protocol::{decode_ping_pong_framed, PingPongMessageType},
-    vdaf::{
-        prio2::{prio2_prep_finish, prio2_prep_finish_from_shares},
-        VdafError,
-    },
+    vdaf::VdafError,
     AggregationJobReportState, DapAggregateShare, DapAggregateSpan, DapAggregationJobState,
-    DapAggregationParam, DapError, DapTaskConfig, DapVersion, VdafConfig,
+    DapAggregationParam, DapError, DapTaskConfig, DapVersion,
 };
 
 use prio::codec::{encode_u32_items, Encode, ParameterizedDecode, ParameterizedEncode};
@@ -284,7 +279,6 @@ impl DapTaskConfig {
         report_status: &HashMap<ReportId, ReportProcessedStatus>,
         part_batch_sel: &PartialBatchSelector,
         initialized_reports: &[InitializedReport<WithPeerPrepShare>],
-        version: DapVersion,
     ) -> Result<(DapAggregateSpan<DapAggregateShare>, AggregationJobResp), DapError> {
         let num_reports = initialized_reports.len();
         let mut agg_span = DapAggregateSpan::default();
@@ -302,39 +296,14 @@ impl DapTaskConfig {
                         prep_share: helper_prep_share,
                         prep_state: helper_prep_state,
                     } => {
-                        let res = match &self.vdaf {
-                            VdafConfig::Prio3(prio3_config) => prio3_config
-                                .prep_finish_from_shares(
-                                    version,
-                                    1,
-                                    task_id,
-                                    helper_prep_state.clone(),
-                                    helper_prep_share.clone(),
-                                    leader_prep_share,
-                                ),
-                            VdafConfig::Prio2 { dimension } => prio2_prep_finish_from_shares(
-                                *dimension,
-                                helper_prep_state.clone(),
-                                helper_prep_share.clone(),
-                                leader_prep_share,
-                            ),
-                            #[cfg(feature = "experimental")]
-                            VdafConfig::Mastic {
-                                input_size: _,
-                                weight_config,
-                            } => mastic_prep_finish_from_shares(
-                                *weight_config,
-                                helper_prep_state.clone(),
-                                helper_prep_share.clone(),
-                                leader_prep_share,
-                            ),
-                            VdafConfig::Pine(pine) => pine.prep_finish_from_shares(
-                                1,
-                                helper_prep_state.clone(),
-                                helper_prep_share.clone(),
-                                leader_prep_share,
-                            ),
-                        };
+                        let res = self.vdaf.prep_finish_from_shares(
+                            self.version,
+                            1,
+                            task_id,
+                            helper_prep_state.clone(),
+                            helper_prep_share.clone(),
+                            leader_prep_share,
+                        );
 
                         match res {
                             Ok((data, prep_msg)) => {
@@ -397,7 +366,6 @@ impl DapTaskConfig {
         state: DapAggregationJobState,
         agg_job_resp: AggregationJobResp,
         metrics: &dyn DaphneMetrics,
-        version: DapVersion,
     ) -> Result<DapAggregateSpan<DapAggregateShare>, DapError> {
         if agg_job_resp.transitions.len() != state.seq.len() {
             return Err(DapAbort::InvalidMessage {
@@ -450,17 +418,9 @@ impl DapTaskConfig {
                 }
             };
 
-            let res = match &self.vdaf {
-                VdafConfig::Prio3(prio3_config) => {
-                    prio3_config.prep_finish(leader.prep_state, prep_msg, *task_id, version)
-                }
-                VdafConfig::Prio2 { dimension } => {
-                    prio2_prep_finish(*dimension, leader.prep_state, prep_msg)
-                }
-                #[cfg(feature = "experimental")]
-                VdafConfig::Mastic { .. } => mastic_prep_finish(leader.prep_state, prep_msg),
-                VdafConfig::Pine(pine) => pine.prep_finish(leader.prep_state, prep_msg),
-            };
+            let res = self
+                .vdaf
+                .prep_finish(leader.prep_state, prep_msg, *task_id, self.version);
 
             match res {
                 Ok(data) => {
diff --git a/crates/daphne/src/protocol/client.rs b/crates/daphne/src/protocol/client.rs
index 349c380..253bd8b 100644
--- a/crates/daphne/src/protocol/client.rs
+++ b/crates/daphne/src/protocol/client.rs
@@ -1,13 +1,10 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-#[cfg(feature = "experimental")]
-use crate::vdaf::mastic::mastic_shard;
 use crate::{
     constants::DapAggregatorRole,
     hpke::{info_and_aad, HpkeConfig},
     messages::{Extension, PlaintextInputShare, Report, ReportId, ReportMetadata, TaskId, Time},
-    vdaf::{prio2::prio2_shard, VdafError},
     DapError, DapMeasurement, DapVersion, VdafConfig,
 };
 use prio::codec::ParameterizedEncode;
@@ -45,7 +42,7 @@ impl VdafConfig {
         let mut rng = thread_rng();
         let report_id = ReportId(rng.gen());
         let (public_share, input_shares) = self
-            .produce_input_shares(measurement, &report_id.0, task_id, version)
+            .shard(measurement, &report_id.0, *task_id, version)
             .map_err(DapError::from_vdaf)?;
         Self::produce_report_with_extensions_for_shares(
             public_share,
@@ -116,28 +113,6 @@ impl VdafConfig {
         })
     }
 
-    /// Generate shares for a measurement.
-    pub(crate) fn produce_input_shares(
-        &self,
-        measurement: DapMeasurement,
-        nonce: &[u8; 16],
-        task_id: &TaskId,
-        version: DapVersion,
-    ) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
-        match self {
-            Self::Prio3(prio3_config) => {
-                Ok(prio3_config.shard(version, measurement, nonce, *task_id)?)
-            }
-            Self::Prio2 { dimension } => Ok(prio2_shard(*dimension, measurement, nonce)?),
-            #[cfg(feature = "experimental")]
-            VdafConfig::Mastic {
-                input_size,
-                weight_config,
-            } => Ok(mastic_shard(*input_size, *weight_config, measurement)?),
-            VdafConfig::Pine(pine) => Ok(pine.shard(measurement, nonce)?),
-        }
-    }
-
     /// Generate a report for a measurement. This method is run by the Client.
     ///
     /// # Inputs
diff --git a/crates/daphne/src/protocol/collector.rs b/crates/daphne/src/protocol/collector.rs
index f6fc6a8..8bc10d0 100644
--- a/crates/daphne/src/protocol/collector.rs
+++ b/crates/daphne/src/protocol/collector.rs
@@ -1,14 +1,11 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-#[cfg(feature = "experimental")]
-use crate::vdaf::mastic::mastic_unshard;
 use crate::{
     constants::DapAggregatorRole,
     fatal_error,
     hpke::{info_and_aad, HpkeDecrypter},
     messages::{BatchSelector, HpkeCiphertext, TaskId},
-    vdaf::prio2::prio2_unshard,
     DapAggregateResult, DapAggregationParam, DapError, DapVersion, VdafConfig,
 };
 
@@ -72,18 +69,7 @@ impl VdafConfig {
         }
 
         let num_measurements = usize::try_from(report_count).unwrap();
-        match self {
-            Self::Prio3(prio3_config) => {
-                prio3_config.unshard(version, num_measurements, agg_shares)
-            }
-            Self::Prio2 { dimension } => prio2_unshard(*dimension, num_measurements, agg_shares),
-            #[cfg(feature = "experimental")]
-            Self::Mastic {
-                input_size: _,
-                weight_config,
-            } => mastic_unshard(*weight_config, agg_param, agg_shares),
-            Self::Pine(pine) => pine.unshard(num_measurements, agg_shares),
-        }
-        .map_err(|e| fatal_error!(err = ?e, "failed to unshard agg_shares"))
+        self.unshard(version, agg_param, num_measurements, agg_shares)
+            .map_err(|e| fatal_error!(err = ?e, "failed to unshard agg_shares"))
     }
 }
diff --git a/crates/daphne/src/protocol/mod.rs b/crates/daphne/src/protocol/mod.rs
index 47627e7..d768614 100644
--- a/crates/daphne/src/protocol/mod.rs
+++ b/crates/daphne/src/protocol/mod.rs
@@ -841,10 +841,10 @@ mod test {
             let (invalid_public_share, mut invalid_input_shares) = self
                 .task_config
                 .vdaf
-                .produce_input_shares(
+                .shard(
                     measurement,
                     &report_id.0,
-                    &self.task_id,
+                    self.task_id,
                     self.task_config.version,
                 )
                 .unwrap();
@@ -872,10 +872,10 @@ mod test {
             let (mut invalid_public_share, invalid_input_shares) = self
                 .task_config
                 .vdaf
-                .produce_input_shares(
+                .shard(
                     measurement,
                     &report_id.0,
-                    &self.task_id,
+                    self.task_id,
                     self.task_config.version,
                 )
                 .unwrap();
@@ -903,10 +903,10 @@ mod test {
             let (invalid_public_share, mut invalid_input_shares) = self
                 .task_config
                 .vdaf
-                .produce_input_shares(
+                .shard(
                     measurement,
                     &report_id.0,
-                    &self.task_id,
+                    self.task_id,
                     self.task_config.version,
                 )
                 .unwrap();
diff --git a/crates/daphne/src/protocol/report_init.rs b/crates/daphne/src/protocol/report_init.rs
index 0d1b880..fd4a9a8 100644
--- a/crates/daphne/src/protocol/report_init.rs
+++ b/crates/daphne/src/protocol/report_init.rs
@@ -1,8 +1,6 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-#[cfg(feature = "experimental")]
-use crate::vdaf::mastic::mastic_prep_init;
 use crate::{
     constants::DapAggregatorRole,
     hpke::{info_and_aad, HpkeDecrypter},
@@ -10,7 +8,7 @@ use crate::{
         self, Extension, PlaintextInputShare, ReportError, ReportMetadata, ReportShare, TaskId,
     },
     protocol::{decode_ping_pong_framed, no_duplicates, PingPongMessageType},
-    vdaf::{prio2::prio2_prep_init, VdafConfig, VdafPrepShare, VdafPrepState},
+    vdaf::{VdafPrepShare, VdafPrepState},
     DapAggregationParam, DapError, DapTaskConfig,
 };
 use prio::codec::{CodecError, ParameterizedDecode as _};
@@ -98,9 +96,6 @@ impl<P> InitializedReport<P> {
         task_config: &DapTaskConfig,
         report_share: ReportShare,
         prep_init_payload: S,
-        // We need to use this variable for Mastic, which is currently fenced by the
-        // "experimental" feature.
-        #[cfg_attr(not(feature = "experimental"), expect(unused_variables))]
         agg_param: &DapAggregationParam,
     ) -> Result<Self, DapError>
     where
@@ -193,44 +188,16 @@ impl<P> InitializedReport<P> {
             DapAggregatorRole::Leader => 0,
             DapAggregatorRole::Helper => 1,
         };
-        let res = match &task_config.vdaf {
-            VdafConfig::Prio3(prio3_config) => prio3_config.prep_init(
-                task_config.version,
-                &task_config.vdaf_verify_key,
-                *task_id,
-                agg_id,
-                &report_share.report_metadata.id.0,
-                &report_share.public_share,
-                &input_share,
-            ),
-            VdafConfig::Prio2 { dimension } => prio2_prep_init(
-                *dimension,
-                &task_config.vdaf_verify_key,
-                agg_id,
-                &report_share.report_metadata.id.0,
-                &report_share.public_share,
-                &input_share,
-            ),
-            #[cfg(feature = "experimental")]
-            VdafConfig::Mastic {
-                input_size,
-                weight_config,
-            } => mastic_prep_init(
-                *input_size,
-                *weight_config,
-                &task_config.vdaf_verify_key,
-                agg_param,
-                &report_share.public_share,
-                input_share.as_ref(),
-            ),
-            VdafConfig::Pine(pine) => pine.prep_init(
-                &task_config.vdaf_verify_key,
-                agg_id,
-                &report_share.report_metadata.id.0,
-                &report_share.public_share,
-                &input_share,
-            ),
-        };
+        let res = task_config.vdaf.prep_init(
+            task_config.version,
+            &task_config.vdaf_verify_key,
+            *task_id,
+            agg_id,
+            agg_param,
+            report_share.report_metadata.id.as_ref(),
+            &report_share.public_share,
+            &input_share,
+        );
 
         match res {
             Ok((prep_state, prep_share)) => Ok(InitializedReport::Ready {
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index 2d4e3b1..39f9998 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -232,7 +232,6 @@ async fn finish_agg_job_and_aggregate(
             &report_status,
             part_batch_sel,
             initialized_reports,
-            task_config.version,
         )?;
 
         let put_shares_result = helper
diff --git a/crates/daphne/src/roles/leader/mod.rs b/crates/daphne/src/roles/leader/mod.rs
index a1fb093..5034618 100644
--- a/crates/daphne/src/roles/leader/mod.rs
+++ b/crates/daphne/src/roles/leader/mod.rs
@@ -341,13 +341,8 @@ async fn run_agg_job<A: DapLeader>(
             .map_err(|e| DapAbort::from_codec_error(e, *task_id))?;
 
     // Handle AggregationJobResp.
-    let agg_span = task_config.consume_agg_job_resp(
-        task_id,
-        agg_job_state,
-        agg_job_resp,
-        metrics,
-        task_config.version,
-    )?;
+    let agg_span =
+        task_config.consume_agg_job_resp(task_id, agg_job_state, agg_job_resp, metrics)?;
 
     let out_shares_count = agg_span.report_count() as u64;
     if out_shares_count == 0 {
diff --git a/crates/daphne/src/vdaf/mastic.rs b/crates/daphne/src/vdaf/mastic.rs
index 0918998..3178e6d 100644
--- a/crates/daphne/src/vdaf/mastic.rs
+++ b/crates/daphne/src/vdaf/mastic.rs
@@ -232,18 +232,19 @@ mod test {
 
     use super::*;
     use crate::{
-        hpke::HpkeKemId, test_versions, testing::AggregationJobTest, vdaf::VdafConfig,
-        DapAggregateResult, DapMeasurement, DapVersion,
+        hpke::HpkeKemId, testing::AggregationJobTest, vdaf::VdafConfig, DapAggregateResult,
+        DapMeasurement, DapVersion,
     };
 
-    fn roundtrip_count(version: DapVersion) {
+    #[test]
+    fn roundtrip_count() {
         let mut t = AggregationJobTest::new(
             &VdafConfig::Mastic {
                 input_size: 4,
                 weight_config: MasticWeightConfig::Count,
             },
             HpkeKemId::X25519HkdfSha256,
-            version,
+            DapVersion::Latest,
         );
         let got = t.roundtrip(
             DapAggregationParam::Mastic(
@@ -279,6 +280,4 @@ mod test {
 
         assert_eq!(got, DapAggregateResult::U64Vec(vec![1, 2]));
     }
-
-    test_versions! { roundtrip_count }
 }
diff --git a/crates/daphne/src/vdaf/mod.rs b/crates/daphne/src/vdaf/mod.rs
index 9c95f53..2762c7d 100644
--- a/crates/daphne/src/vdaf/mod.rs
+++ b/crates/daphne/src/vdaf/mod.rs
@@ -13,6 +13,7 @@ pub(crate) mod prio3_draft09;
 
 use crate::pine::vdaf::PinePrepState;
 use crate::{fatal_error, messages::TaskId, DapError};
+use crate::{DapAggregateResult, DapAggregationParam, DapMeasurement, DapVersion};
 use pine::PineConfig;
 use prio::{
     codec::{CodecError, Encode, ParameterizedDecode},
@@ -123,6 +124,205 @@ impl std::fmt::Display for VdafConfig {
     }
 }
 
+impl VdafConfig {
+    pub(crate) fn shard(
+        &self,
+        measurement: DapMeasurement,
+        nonce: &[u8; 16],
+        task_id: TaskId,
+        version: DapVersion,
+    ) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
+        match (version, self) {
+            (_, Self::Prio2 { dimension }) => prio2::prio2_shard(*dimension, measurement, nonce),
+            (_, Self::Prio3(prio3_config)) => {
+                prio3_config.shard(version, measurement, nonce, task_id)
+            }
+            (DapVersion::Draft09, Self::Pine(pine_config)) => pine_config.shard(measurement, nonce),
+            #[cfg(feature = "experimental")]
+            (
+                DapVersion::Latest,
+                VdafConfig::Mastic {
+                    input_size,
+                    weight_config,
+                },
+            ) => Ok(mastic::mastic_shard(
+                *input_size,
+                *weight_config,
+                measurement,
+            )?),
+            _ => Err(VdafError::Dap(fatal_error!(
+                err = format!("{self:?} is not supported in DAP {version}")
+            ))),
+        }
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn prep_init(
+        &self,
+        version: DapVersion,
+        verify_key: &VdafVerifyKey,
+        task_id: TaskId,
+        agg_id: usize,
+        // This is used by Mastic, which for the moment is behind the "experimental" flag.
+        #[cfg_attr(not(feature = "experimental"), expect(unused_variables))]
+        agg_param: &DapAggregationParam,
+        nonce: &[u8; 16],
+        public_share_data: &[u8],
+        input_share_data: &[u8],
+    ) -> Result<(VdafPrepState, VdafPrepShare), VdafError> {
+        match (version, self) {
+            (_, VdafConfig::Prio2 { dimension }) => prio2::prio2_prep_init(
+                *dimension,
+                verify_key,
+                agg_id,
+                nonce,
+                public_share_data,
+                input_share_data,
+            ),
+            (_, Self::Prio3(prio3_config)) => prio3_config.prep_init(
+                version,
+                verify_key,
+                task_id,
+                agg_id,
+                nonce,
+                public_share_data,
+                input_share_data,
+            ),
+            (DapVersion::Draft09, Self::Pine(pine)) => pine.prep_init(
+                verify_key,
+                agg_id,
+                nonce,
+                public_share_data,
+                input_share_data,
+            ),
+            #[cfg(feature = "experimental")]
+            (
+                DapVersion::Latest,
+                Self::Mastic {
+                    input_size,
+                    weight_config,
+                },
+            ) => mastic::mastic_prep_init(
+                *input_size,
+                *weight_config,
+                verify_key,
+                agg_param,
+                public_share_data,
+                input_share_data,
+            ),
+            _ => Err(VdafError::Dap(fatal_error!(
+                err = format!("{self:?} is not supported in DAP {version}")
+            ))),
+        }
+    }
+
+    pub(crate) fn prep_finish_from_shares(
+        &self,
+        version: DapVersion,
+        agg_id: usize,
+        task_id: TaskId,
+        host_state: VdafPrepState,
+        host_share: VdafPrepShare,
+        peer_share_data: &[u8],
+    ) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
+        match (version, self) {
+            (_, Self::Prio2 { dimension }) => prio2::prio2_prep_finish_from_shares(
+                *dimension,
+                host_state,
+                host_share,
+                peer_share_data,
+            ),
+            (_, Self::Prio3(prio3_config)) => prio3_config.prep_finish_from_shares(
+                version,
+                agg_id,
+                task_id,
+                host_state,
+                host_share,
+                peer_share_data,
+            ),
+            (DapVersion::Draft09, Self::Pine(pine_config)) => {
+                pine_config.prep_finish_from_shares(agg_id, host_state, host_share, peer_share_data)
+            }
+            #[cfg(feature = "experimental")]
+            (
+                DapVersion::Latest,
+                Self::Mastic {
+                    input_size: _,
+                    weight_config,
+                },
+            ) => mastic::mastic_prep_finish_from_shares(
+                *weight_config,
+                host_state,
+                host_share,
+                peer_share_data,
+            ),
+            _ => Err(VdafError::Dap(fatal_error!(
+                err = format!("{self:?} is not supported in DAP {version}")
+            ))),
+        }
+    }
+
+    pub(crate) fn prep_finish(
+        &self,
+        host_state: VdafPrepState,
+        peer_message_data: &[u8],
+        task_id: TaskId,
+        version: DapVersion,
+    ) -> Result<VdafAggregateShare, VdafError> {
+        match (version, self) {
+            (_, Self::Prio2 { dimension }) => {
+                prio2::prio2_prep_finish(*dimension, host_state, peer_message_data)
+            }
+            (_, Self::Prio3(prio3_config)) => {
+                prio3_config.prep_finish(host_state, peer_message_data, task_id, version)
+            }
+            (DapVersion::Draft09, Self::Pine(pine_config)) => {
+                pine_config.prep_finish(host_state, peer_message_data)
+            }
+            #[cfg(feature = "experimental")]
+            (DapVersion::Latest, Self::Mastic { .. }) => {
+                mastic::mastic_prep_finish(host_state, peer_message_data)
+            }
+            _ => Err(VdafError::Dap(fatal_error!(
+                err = format!("{self:?} is not supported in DAP {version}")
+            ))),
+        }
+    }
+
+    pub(crate) fn unshard<M: IntoIterator<Item = Vec<u8>>>(
+        &self,
+        version: DapVersion,
+        // This is used by Mastic, which for the moment is behind the "experimental" flag.
+        #[cfg_attr(not(feature = "experimental"), expect(unused_variables))]
+        agg_param: &DapAggregationParam,
+        num_measurements: usize,
+        agg_shares: M,
+    ) -> Result<DapAggregateResult, VdafError> {
+        match (version, self) {
+            (_, Self::Prio2 { dimension }) => {
+                prio2::prio2_unshard(*dimension, num_measurements, agg_shares)
+            }
+            (_, Self::Prio3(prio3_config)) => {
+                prio3_config.unshard(version, num_measurements, agg_shares)
+            }
+            (DapVersion::Draft09, Self::Pine(pine_config)) => {
+                pine_config.unshard(num_measurements, agg_shares)
+            }
+            #[cfg(feature = "experimental")]
+            (
+                DapVersion::Latest,
+                Self::Mastic {
+                    input_size: _,
+                    weight_config,
+                },
+            ) => mastic::mastic_unshard(*weight_config, agg_param, agg_shares),
+            _ => Err(VdafError::Dap(fatal_error!(
+                err = format!("{self:?} is not supported in DAP {version}")
+            ))),
+        }
+    }
+}
+
 /// Supported data types for prio3.
 #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize)]
 #[serde(rename_all = "snake_case")]
diff --git a/crates/daphne/src/vdaf/pine.rs b/crates/daphne/src/vdaf/pine.rs
index 9230f3a..2f83914 100644
--- a/crates/daphne/src/vdaf/pine.rs
+++ b/crates/daphne/src/vdaf/pine.rs
@@ -240,13 +240,14 @@ fn prep_init<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usi
 #[cfg(test)]
 mod test {
     use crate::{
-        hpke::HpkeKemId, pine::PineParam, test_versions, testing::AggregationJobTest,
-        vdaf::VdafConfig, DapAggregateResult, DapAggregationParam, DapMeasurement, DapVersion,
+        hpke::HpkeKemId, pine::PineParam, testing::AggregationJobTest, vdaf::VdafConfig,
+        DapAggregateResult, DapAggregationParam, DapMeasurement, DapVersion,
     };
 
     use super::PineConfig;
 
-    fn roundtrip_pine32_hmac_sha256_aes128(version: DapVersion) {
+    #[test]
+    fn roundtrip_pine32_hmac_sha256_aes128_draft09() {
         let mut t = AggregationJobTest::new(
             &VdafConfig::Pine(PineConfig::Field32HmacSha256Aes128 {
                 param: PineParam {
@@ -262,7 +263,7 @@ mod test {
                 },
             }),
             HpkeKemId::X25519HkdfSha256,
-            version,
+            DapVersion::Draft09,
         );
         let DapAggregateResult::F64Vec(got) = t.roundtrip(
             DapAggregationParam::Empty,
@@ -282,9 +283,8 @@ mod test {
         }
     }
 
-    test_versions! { roundtrip_pine32_hmac_sha256_aes128 }
-
-    fn roundtrip_pine64_hmac_sha256_aes128(version: DapVersion) {
+    #[test]
+    fn roundtrip_pine64_hmac_sha256_aes128_draft09() {
         let mut t = AggregationJobTest::new(
             &VdafConfig::Pine(PineConfig::Field64HmacSha256Aes128 {
                 param: PineParam {
@@ -300,7 +300,7 @@ mod test {
                 },
             }),
             HpkeKemId::X25519HkdfSha256,
-            version,
+            DapVersion::Draft09,
         );
         let DapAggregateResult::F64Vec(got) = t.roundtrip(
             DapAggregationParam::Empty,
@@ -319,6 +319,4 @@ mod test {
             );
         }
     }
-
-    test_versions! { roundtrip_pine64_hmac_sha256_aes128 }
 }


Test patch:
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index 5b32c30..ce074d9 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -223,7 +223,6 @@ impl AggregationJobTest {
                         self.replay_protection,
                     )
                     .unwrap(),
-                self.task_config.version,
             )
             .unwrap()
     }
@@ -242,7 +241,6 @@ impl AggregationJobTest {
                 leader_state,
                 agg_job_resp,
                 &self.leader_metrics,
-                self.task_config.version,
             )
             .unwrap()
     }
@@ -255,13 +253,7 @@ impl AggregationJobTest {
     ) -> DapError {
         let metrics = &self.leader_metrics;
         self.task_config
-            .consume_agg_job_resp(
-                &self.task_id,
-                leader_state,
-                agg_job_resp,
-                metrics,
-                self.task_config.version,
-            )
+            .consume_agg_job_resp(&self.task_id, leader_state, agg_job_resp, metrics)
             .expect_err("consume_agg_job_resp() succeeded; expected failure")
     }