Repository: cloudflare/daphne
Pull request #769 (2025-01-27 18:26:12)

Add support for the Mastic VDAF
Mastic (https://datatracker.ietf.org/doc/draft-mouris-cfrg-mastic/) is a
VDAF that enables a richer set of functionalities than the VDAFs we
support so so far. The `daphne::vdaf::mastic` module contains a "dummy"
version of Mastic intended to exercise the DAP protocol logic we would
need in order to fully support this VDAF. The `prio` crate now
implements Mastic, so upgrade to a version of the crate that supports it
and replace the dummy VDAF with the real one.

In addition, to complete aggregation of a report, it is necessary to
know the aggregation parameter, which currently is only plumbed to
report initialization. In particular,
`DapTaskConfig::produce_agg_job_resp()` needs the aggregation parameter
from the aggregation job request message. (Likewise,
`ToInitializedReportsTransition::with_initialized_reports()` needs the
aggregation parameter.)

Finally, clean up some API cruft in `daphne::vdaf`:

1. Encapsulate variants of Mastic behind a `MasticConfig` as we've done
   for other VDAFs.

2. Modify `prep_finish_from_shares()` to not take the aggregator ID.
   This is a relic of when we supported DAP-02, when this method may
   have been called by either the Leader or the Helper. Now it's always
   called by the Helper.

3. Implement state encoding for Mastic, as required by the new async
   Helper implementation.

4. Generalize the `prep_init()` function in `daphne::prio3` to be used
   for Mastic as well.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 10a881b..55c7043 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2528,7 +2528,7 @@ dependencies = [
 [[package]]
 name = "prio"
 version = "0.17.0-alpha.0"
-source = "git+https://github.com/divviup/libprio-rs.git?rev=c50bb9a47b396ad6a08a3fec36b98bcc2d9217a1#c50bb9a47b396ad6a08a3fec36b98bcc2d9217a1"
+source = "git+https://github.com/divviup/libprio-rs.git?rev=e5e8a47ee4567f7588d0b5c8d20f75dde4061b2f#e5e8a47ee4567f7588d0b5c8d20f75dde4061b2f"
 dependencies = [
  "aes",
  "bitvec",
@@ -2536,7 +2536,6 @@ dependencies = [
  "ctr",
  "fiat-crypto",
  "fixed",
- "getrandom",
  "hex",
  "hmac",
  "num-bigint",
diff --git a/Cargo.toml b/Cargo.toml
index 59f098e..51cefce 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -63,7 +63,8 @@ matchit = "0.7.3"
 p256 = { version = "0.13.2", features = ["ecdsa-core", "ecdsa", "pem"] }
 paste = "1.0.15"
 prio_draft09 = { package = "prio", version = "0.16.7" }
-prio = { git = "https://github.com/divviup/libprio-rs.git", rev = "c50bb9a47b396ad6a08a3fec36b98bcc2d9217a1" }
+# TODO Point to version `0.17.0` once release. This revision is one commit ahead of `0.17.0-alpha.0`.
+prio = { git = "https://github.com/divviup/libprio-rs.git", rev = "e5e8a47ee4567f7588d0b5c8d20f75dde4061b2f" }
 prometheus = "0.13.4"
 rand = "0.8.5"
 rayon = "1.10.0"
diff --git a/crates/daphne-worker/src/aggregator/router/helper.rs b/crates/daphne-worker/src/aggregator/router/helper.rs
index 65ae20a..848739a 100644
--- a/crates/daphne-worker/src/aggregator/router/helper.rs
+++ b/crates/daphne-worker/src/aggregator/router/helper.rs
@@ -11,6 +11,7 @@ use axum::{
     routing::{post, put},
 };
 use daphne::{
+    error::DapAbort,
     fatal_error,
     hpke::HpkeProvider,
     messages::AggregateShareReq,
@@ -18,11 +19,11 @@ use daphne::{
         helper::{self, HashedAggregationJobReq},
         DapAggregator, DapHelper,
     },
-    DapError, DapResponse,
+    DapAggregationParam, DapError, DapResponse,
 };
 use daphne_service_utils::compute_offload;
 use http::StatusCode;
-use prio::codec::ParameterizedEncode;
+use prio::codec::{ParameterizedDecode, ParameterizedEncode};
 use std::{borrow::Cow, sync::Arc};
 
 pub(super) fn add_helper_routes(router: super::Router<App>) -> super::Router<App> {
@@ -60,6 +61,12 @@ async fn agg_job(
 
         let hpke_receiver_configs = app.get_hpke_receiver_configs(req.version).await?;
 
+        let agg_param = DapAggregationParam::get_decoded_with_param(
+            &transition.task_config.vdaf,
+            &req.payload.agg_param,
+        )
+        .map_err(|e| DapAbort::from_codec_error(e, req.task_id))?;
+
         let initialized_reports: compute_offload::InitializedReports = app
             .compute_offload
             .compute(
@@ -77,7 +84,7 @@ async fn agg_job(
             .map_err(|e| fatal_error!(err = ?e, "failed to offload report initialization"))?;
 
         transition
-            .with_initialized_reports(initialized_reports.reports)
+            .with_initialized_reports(agg_param, initialized_reports.reports)
             .finish_and_aggregate(&*app)
             .await
     }
diff --git a/crates/daphne/src/hpke.rs b/crates/daphne/src/hpke.rs
index aac757d..3291f66 100644
--- a/crates/daphne/src/hpke.rs
+++ b/crates/daphne/src/hpke.rs
@@ -505,7 +505,7 @@ pub mod info_and_aad {
         }
     }
 
-    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    #[derive(Debug, Clone, Copy, PartialEq)]
     pub struct AggregateShare<'s> {
         // info
         pub version: DapVersion,
diff --git a/crates/daphne/src/lib.rs b/crates/daphne/src/lib.rs
index 2dce27a..a3ad3fc 100644
--- a/crates/daphne/src/lib.rs
+++ b/crates/daphne/src/lib.rs
@@ -70,7 +70,7 @@ use error::FatalDapError;
 use hpke::{HpkeConfig, HpkeKemId};
 use messages::taskprov::TaskprovAdvertisement;
 #[cfg(feature = "experimental")]
-use prio::{codec::Decode, vdaf::poplar1::Poplar1AggregationParam};
+use prio::{codec::Decode, vdaf::mastic::MasticAggregationParam};
 use prio::{
     codec::{CodecError, Encode, ParameterizedDecode},
     vdaf::Aggregatable as AggregatableTrait,
@@ -813,11 +813,11 @@ pub enum DapMeasurement {
 }
 
 /// An aggregation parameter.
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum DapAggregationParam {
     Empty,
     #[cfg(feature = "experimental")]
-    Mastic(Poplar1AggregationParam),
+    Mastic(MasticAggregationParam),
 }
 
 #[cfg(any(test, feature = "test-utils"))]
@@ -877,7 +877,7 @@ impl ParameterizedDecode<VdafConfig> for DapAggregationParam {
         let _ = bytes;
         match vdaf_config {
             #[cfg(feature = "experimental")]
-            VdafConfig::Mastic { .. } => Ok(Self::Mastic(Poplar1AggregationParam::decode(bytes)?)),
+            VdafConfig::Mastic(_) => Ok(Self::Mastic(MasticAggregationParam::decode(bytes)?)),
             _ => Ok(Self::Empty),
         }
     }
diff --git a/crates/daphne/src/protocol/aggregator.rs b/crates/daphne/src/protocol/aggregator.rs
index a2ce008..0925c3b 100644
--- a/crates/daphne/src/protocol/aggregator.rs
+++ b/crates/daphne/src/protocol/aggregator.rs
@@ -236,7 +236,13 @@ impl DapTaskConfig {
         task_id: &TaskId,
         agg_job_init_req: AggregationJobInitReq,
         replay_protection: ReplayProtection,
-    ) -> Result<Vec<InitializedReport<WithPeerPrepShare>>, DapError>
+    ) -> Result<
+        (
+            DapAggregationParam,
+            Vec<InitializedReport<WithPeerPrepShare>>,
+        ),
+        DapError,
+    >
     where
         H: HpkeDecrypter + Sync,
     {
@@ -260,7 +266,7 @@ impl DapTaskConfig {
             agg_job_init_req.prep_inits.len()
         );
 
-        agg_job_init_req
+        let initialized_reports = agg_job_init_req
             .prep_inits
             .into_par_iter()
             .map(|prep_init| {
@@ -274,7 +280,8 @@ impl DapTaskConfig {
                     &agg_param,
                 )
             })
-            .collect()
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok((agg_param, initialized_reports))
     }
 
     /// Helper -> Leader: Produce the `AggregationJobResp` message to send to the Leader and
@@ -283,6 +290,7 @@ impl DapTaskConfig {
     pub(crate) fn produce_agg_job_resp(
         &self,
         task_id: TaskId,
+        agg_param: &DapAggregationParam,
         report_status: &HashMap<ReportId, ReportProcessedStatus>,
         part_batch_sel: &PartialBatchSelector,
         initialized_reports: &[InitializedReport<WithPeerPrepShare>],
@@ -305,8 +313,8 @@ impl DapTaskConfig {
                     } => {
                         let res = self.vdaf.prep_finish_from_shares(
                             self.version,
-                            1,
                             task_id,
+                            agg_param,
                             helper_prep_state.clone(),
                             helper_prep_share.clone(),
                             leader_prep_share,
diff --git a/crates/daphne/src/roles/helper/handle_agg_job.rs b/crates/daphne/src/roles/helper/handle_agg_job.rs
index b8f7e1d..dad6f04 100644
--- a/crates/daphne/src/roles/helper/handle_agg_job.rs
+++ b/crates/daphne/src/roles/helper/handle_agg_job.rs
@@ -11,7 +11,7 @@ use crate::{
     metrics::ReportStatus,
     protocol::aggregator::ReportProcessedStatus,
     roles::{aggregator::MergeAggShareError, resolve_task_config},
-    DapError, DapRequest, DapTaskConfig, InitializedReport, WithPeerPrepShare,
+    DapAggregationParam, DapError, DapRequest, DapTaskConfig, InitializedReport, WithPeerPrepShare,
 };
 use std::{collections::HashMap, sync::Once};
 
@@ -48,6 +48,7 @@ pub struct ToInitializedReportsTransition {
 /// The reports have been initialized and are ready for aggregation.
 pub struct InitializedReports {
     task_id: TaskId,
+    agg_param: DapAggregationParam,
     part_batch_sel: PartialBatchSelector,
     task_config: DapTaskConfig,
     reports: Vec<InitializedReport<WithPeerPrepShare>>,
@@ -142,7 +143,7 @@ impl HandleAggJob<WithTaskConfig> {
         } = self.state;
         let task_id = request.task_id;
         let part_batch_sel = request.payload.part_batch_sel.clone();
-        let initialized_reports = task_config.consume_agg_job_req(
+        let (agg_param, initialized_reports) = task_config.consume_agg_job_req(
             &aggregator
                 .get_hpke_receiver_configs(task_config.version)
                 .await?,
@@ -155,6 +156,7 @@ impl HandleAggJob<WithTaskConfig> {
         Ok(HandleAggJob {
             state: InitializedReports {
                 task_id,
+                agg_param,
                 task_config,
                 part_batch_sel,
                 reports: initialized_reports,
@@ -207,6 +209,7 @@ impl ToInitializedReportsTransition {
     /// Provide the initialized reports that should be aggregated.
     pub fn with_initialized_reports(
         self,
+        agg_param: DapAggregationParam,
         reports: Vec<InitializedReport<WithPeerPrepShare>>,
     ) -> HandleAggJob<InitializedReports> {
         let Self {
@@ -217,6 +220,7 @@ impl ToInitializedReportsTransition {
         HandleAggJob {
             state: InitializedReports {
                 task_id,
+                agg_param,
                 part_batch_sel,
                 task_config,
                 reports,
@@ -236,6 +240,7 @@ impl HandleAggJob<InitializedReports> {
             state:
                 InitializedReports {
                     task_id,
+                    agg_param,
                     part_batch_sel,
                     task_config,
                     reports,
@@ -257,6 +262,7 @@ impl HandleAggJob<InitializedReports> {
         for _ in 0..RETRY_COUNT {
             let (agg_span, agg_job_resp) = task_config.produce_agg_job_resp(
                 task_id,
+                &agg_param,
                 &report_status,
                 &part_batch_sel,
                 &reports,
diff --git a/crates/daphne/src/roles/mod.rs b/crates/daphne/src/roles/mod.rs
index 091f46e..2ab83e8 100644
--- a/crates/daphne/src/roles/mod.rs
+++ b/crates/daphne/src/roles/mod.rs
@@ -130,7 +130,7 @@ async fn resolve_task_config(
 mod test {
     use super::{aggregator, helper, leader, DapLeader};
     #[cfg(feature = "experimental")]
-    use crate::vdaf::{mastic::MasticWeight, MasticWeightConfig};
+    use crate::vdaf::mastic::{MasticConfig, MasticWeight, MasticWeightConfig};
     use crate::{
         assert_metrics_include, async_test_versions,
         constants::DapMediaType,
@@ -151,7 +151,7 @@ mod test {
     use assert_matches::assert_matches;
     use prio::codec::{Encode, ParameterizedDecode};
     #[cfg(feature = "experimental")]
-    use prio::{idpf::IdpfInput, vdaf::poplar1::Poplar1AggregationParam};
+    use prio::idpf::IdpfInput;
     use rand::{thread_rng, Rng};
     use std::{
         collections::HashMap,
@@ -275,10 +275,10 @@ mod test {
 
             #[cfg(feature = "experimental")]
             {
-                let mastic = VdafConfig::Mastic {
-                    input_size: 1,
+                let mastic = VdafConfig::Mastic(MasticConfig {
+                    bits: 8,
                     weight_config: MasticWeightConfig::Count,
-                };
+                });
                 tasks.insert(
                     mastic_task_id,
                     DapTaskConfig {
@@ -2018,6 +2018,8 @@ mod test {
     #[cfg(feature = "experimental")]
     #[tokio::test]
     async fn mastic() {
+        use prio::vdaf::mastic::MasticAggregationParam;
+
         let t = Test::new(DapVersion::Latest);
         let task_id = &t.mastic_task_id;
         let task_config = t
@@ -2043,11 +2045,14 @@ mod test {
         // Collector: Request result from the Leader.
         let query = task_config.query_for_current_batch_window(t.now);
         let agg_param = DapAggregationParam::Mastic(
-            Poplar1AggregationParam::try_from_prefixes(vec![
-                IdpfInput::from_bytes(&[0]),
-                IdpfInput::from_bytes(&[1]),
-                IdpfInput::from_bytes(&[7]),
-            ])
+            MasticAggregationParam::new(
+                vec![
+                    IdpfInput::from_bytes(&[0]),
+                    IdpfInput::from_bytes(&[1]),
+                    IdpfInput::from_bytes(&[7]),
+                ],
+                true,
+            )
             .unwrap(),
         );
         leader::handle_coll_job_req(
diff --git a/crates/daphne/src/vdaf/draft09.rs b/crates/daphne/src/vdaf/draft09.rs
index 8e5bfff..4be5f3c 100644
--- a/crates/daphne/src/vdaf/draft09.rs
+++ b/crates/daphne/src/vdaf/draft09.rs
@@ -97,7 +97,6 @@ where
 
 pub(crate) fn prep_finish_from_shares<V, const VERIFY_KEY_SIZE: usize, const NONCE_SIZE: usize>(
     vdaf: &V,
-    agg_id: usize,
     host_state: V::PrepareState,
     host_share: V::PrepareShare,
     peer_share_data: &[u8],
@@ -105,18 +104,11 @@ pub(crate) fn prep_finish_from_shares<V, const VERIFY_KEY_SIZE: usize, const NON
 where
     V: Vdaf<AggregationParam = ()> + Aggregator<VERIFY_KEY_SIZE, NONCE_SIZE>,
 {
-    // Decode the Helper's inbound message.
+    // Decode the peer's inbound message.
     let peer_share = V::PrepareShare::get_decoded_with_param(&host_state, peer_share_data)?;
 
     // Preprocess the inbound messages.
-    let message = vdaf.prepare_shares_to_prepare_message(
-        &(),
-        if agg_id == 0 {
-            [host_share, peer_share]
-        } else {
-            [peer_share, host_share]
-        },
-    )?;
+    let message = vdaf.prepare_shares_to_prepare_message(&(), [peer_share, host_share])?;
     let message_data = message.get_encoded()?;
 
     // Compute the host's output share.
diff --git a/crates/daphne/src/vdaf/mastic.rs b/crates/daphne/src/vdaf/mastic.rs
index a81c081..e3b854e 100644
--- a/crates/daphne/src/vdaf/mastic.rs
+++ b/crates/daphne/src/vdaf/mastic.rs
@@ -7,21 +7,31 @@
 //!
 //! [draft-mouris-cfrg-mastic]: https://datatracker.ietf.org/doc/draft-mouris-cfrg-mastic/
 
-use std::array;
-
-use crate::{fatal_error, DapAggregateResult, DapAggregationParam, DapMeasurement};
-
-use super::{
-    decode_field_vec, VdafAggregateShare, VdafError, VdafPrepShare, VdafPrepState, VdafVerifyKey,
+use crate::{
+    fatal_error,
+    messages::TaskId,
+    vdaf::{prep_finish, prep_finish_from_shares, prep_init, shard_then_encode},
+    DapAggregateResult, DapAggregationParam, DapMeasurement,
 };
 
-use prio::{
-    codec::{CodecError, Decode},
-    field::{Field64, FieldElement},
-    vdaf::AggregateShare,
-};
+use super::{unshard, VdafAggregateShare, VdafError, VdafPrepShare, VdafPrepState, VdafVerifyKey};
+
+use prio::{vdaf::mastic::Mastic, vidpf::VidpfInput};
 use serde::{Deserialize, Serialize};
 
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize)]
+#[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
+pub struct MasticConfig {
+    pub bits: usize,
+    pub weight_config: MasticWeightConfig,
+}
+
+impl std::fmt::Display for MasticConfig {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}, Weight({})", self.bits, self.weight_config)
+    }
+}
+
 #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize)]
 #[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
 /// The type of each input's weight.
@@ -44,185 +54,144 @@ pub enum MasticWeight {
     Bool(bool),
 }
 
-pub(crate) fn mastic_shard(
-    input_size: usize,
-    weight_config: MasticWeightConfig,
-    measurement: DapMeasurement,
-) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
-    match (weight_config, measurement) {
-        (
-            MasticWeightConfig::Count,
-            DapMeasurement::Mastic {
-                input,
-                weight: MasticWeight::Bool(counter),
-            },
-        ) if input.len() == input_size => {
-            // Simulate Mastic, insecurely. Set the public share to the input and each input share
-            // to the weight.
-            Ok((input, array::from_fn(|_| vec![u8::from(counter)])))
-        }
-        _ => Err(VdafError::Dap(fatal_error!(
-            err = "mastic: unexpected measurement type"
-        ))),
-    }
-}
-
-pub(crate) fn mastic_prep_init(
-    input_size: usize,
-    weight_config: MasticWeightConfig,
-    _verify_key: &VdafVerifyKey,
-    agg_param: &DapAggregationParam,
-    public_share_bytes: &[u8],
-    input_share_bytes: &[u8],
-) -> Result<(VdafPrepState, VdafPrepShare), VdafError> {
-    match (weight_config, agg_param) {
-        (MasticWeightConfig::Count, DapAggregationParam::Mastic(agg_param)) => {
-            // Simulate Mastic, insecurely. The public share encodes the plaintext input; the input
-            // share encodes the plaintext weight.
-            if public_share_bytes.len() != input_size {
-                return Err(VdafError::Codec(prio::codec::CodecError::Other(
-                    "mastic: malformed public share".into(),
-                )));
-            }
+impl MasticConfig {
+    pub(crate) fn shard(
+        self,
+        measurement: DapMeasurement,
+        nonce: &[u8; 16],
+        task_id: TaskId,
+    ) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
+        match (self.weight_config, measurement) {
+            (
+                MasticWeightConfig::Count,
+                DapMeasurement::Mastic {
+                    input,
+                    weight: MasticWeight::Bool(counter),
+                },
+            ) => {
+                let vdaf = Mastic::new_count(self.bits).map_err(|e| {
+                    VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
+                })?;
 
-            if input_share_bytes.len() != 1 {
-                return Err(VdafError::Codec(CodecError::Other(
-                    "mastic: malformed input share".into(),
-                )));
+                let alpha = VidpfInput::from_bytes(&input);
+                shard_then_encode(&vdaf, task_id, &(alpha, counter), nonce)
             }
-
-            let weight = Field64::from(u64::from(input_share_bytes[0]));
-            let out_share = agg_param
-                .prefixes()
-                .iter()
-                .map(|prefix| {
-                    let prefix_bytes = prefix.to_bytes();
-                    if prefix_bytes.len() > input_size {
-                        return Err(VdafError::Codec(CodecError::Other(
-                            "mastic: malformed agg param: path with invalid length".into(),
-                        )));
-                    }
-
-                    // If the path is a prefix of the input, then the value is the
-                    // weight; otherwise the value is 0.
-                    let value = if prefix_bytes == public_share_bytes[..prefix_bytes.len()] {
-                        weight
-                    } else {
-                        Field64::zero()
-                    };
-
-                    // Each Aggregator computes a share of the value, so divide by 2.
-                    Ok(value / Field64::from(2))
-                })
-                .collect::<Result<Vec<Field64>, _>>()?;
-
-            Ok((
-                VdafPrepState::Mastic { out_share },
-                VdafPrepShare::Mastic(weight),
-            ))
+            _ => todo!(),
         }
-        _ => Err(VdafError::Dap(fatal_error!(
-            err = "mastic: unexpected agg param type"
-        ))),
     }
-}
 
-pub(crate) fn mastic_prep_finish_from_shares(
-    weight_config: MasticWeightConfig,
-    host_state: VdafPrepState,
-    host_share: VdafPrepShare,
-    peer_share_bytes: &[u8],
-) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
-    match (weight_config, host_state, host_share) {
-        (
-            MasticWeightConfig::Count,
-            VdafPrepState::Mastic { out_share },
-            VdafPrepShare::Mastic(host_weight),
-        ) => {
-            // Simulate Mastic. Check that both Aggregators got the same weight, and the weight is
-            // valid. This is not secure because the weight is revealed to the caller.
-            let peer_weight = Field64::get_decoded(peer_share_bytes)?;
-            if peer_weight != host_weight {
-                return Err(VdafError::Vdaf(prio::vdaf::VdafError::Uncategorized(
-                    "mastic: weights do not match".into(),
-                )));
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn prep_init(
+        self,
+        VdafVerifyKey(verify_key): &VdafVerifyKey,
+        task_id: TaskId,
+        agg_id: usize,
+        agg_param: &DapAggregationParam,
+        nonce: &[u8; 16],
+        public_share_data: &[u8],
+        input_share_data: &[u8],
+    ) -> Result<(VdafPrepState, VdafPrepShare), VdafError> {
+        match (self.weight_config, agg_param) {
+            (MasticWeightConfig::Count, DapAggregationParam::Mastic(agg_param)) => {
+                let vdaf = Mastic::new_count(self.bits).map_err(|e| {
+                    VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
+                })?;
+
+                let (prep_state, prep_share) = prep_init(
+                    &vdaf,
+                    task_id,
+                    verify_key,
+                    agg_id,
+                    agg_param,
+                    nonce,
+                    public_share_data,
+                    input_share_data,
+                )?;
+
+                Ok((
+                    VdafPrepState::MasticField64(prep_state),
+                    VdafPrepShare::MasticField64(prep_share),
+                ))
             }
+            (_, DapAggregationParam::Empty) => todo!(),
+        }
+    }
 
-            if peer_weight != Field64::one() && peer_weight != Field64::zero() {
-                return Err(VdafError::Vdaf(prio::vdaf::VdafError::Uncategorized(
-                    "mastic: weight is out of range".into(),
-                )));
+    pub(crate) fn prep_finish_from_shares(
+        self,
+        task_id: TaskId,
+        agg_param: &DapAggregationParam,
+        host_state: VdafPrepState,
+        host_share: VdafPrepShare,
+        peer_share_data: &[u8],
+    ) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
+        match (self.weight_config, agg_param, host_state, host_share) {
+            (
+                MasticWeightConfig::Count,
+                DapAggregationParam::Mastic(agg_param),
+                VdafPrepState::MasticField64(host_state),
+                VdafPrepShare::MasticField64(host_share),
+            ) => {
+                let vdaf = Mastic::new_count(self.bits).map_err(|e| {
+                    VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
+                })?;
+
+                let (out_share, prep_msg_data) = prep_finish_from_shares(
+                    &vdaf,
+                    task_id,
+                    agg_param,
+                    host_state,
+                    host_share,
+                    peer_share_data,
+                )?;
+                Ok((VdafAggregateShare::Field64(out_share.into()), prep_msg_data))
             }
-
-            Ok((
-                VdafAggregateShare::Field64(AggregateShare::from(out_share)),
-                // Empty prep message for now.
-                Vec::new(),
-            ))
+            _ => todo!(),
         }
-        _ => Err(VdafError::Dap(fatal_error!(
-            err = "mastic: unexpected prep state"
-        ))),
     }
-}
 
-pub(crate) fn mastic_prep_finish(
-    host_state: VdafPrepState,
-    peer_message_bytes: &[u8],
-) -> Result<VdafAggregateShare, VdafError> {
-    match host_state {
-        VdafPrepState::Mastic { out_share } => {
-            // Simulate Mastic: If the prep message is empty, then accept the output share.
-            if !peer_message_bytes.is_empty() {
-                return Err(VdafError::Vdaf(prio::vdaf::VdafError::Uncategorized(
-                    "mastic: invalid prep message".into(),
-                )));
+    pub(crate) fn prep_finish(
+        self,
+        host_state: VdafPrepState,
+        peer_message_data: &[u8],
+        task_id: TaskId,
+    ) -> Result<VdafAggregateShare, VdafError> {
+        match (self.weight_config, host_state) {
+            (MasticWeightConfig::Count, VdafPrepState::MasticField64(host_state)) => {
+                let vdaf = Mastic::new_count(self.bits).map_err(|e| {
+                    VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
+                })?;
+
+                let out_share = prep_finish(&vdaf, task_id, host_state, peer_message_data)?;
+                Ok(VdafAggregateShare::Field64(out_share.into()))
             }
-
-            Ok(VdafAggregateShare::Field64(AggregateShare::from(out_share)))
+            _ => todo!(),
         }
-        _ => Err(VdafError::Dap(fatal_error!(
-            err = "mastic: unexpected prep state"
-        ))),
     }
-}
 
-pub(crate) fn mastic_unshard<M: IntoIterator<Item = Vec<u8>>>(
-    weight_config: MasticWeightConfig,
-    agg_param: &DapAggregationParam,
-    agg_share_bytes: M,
-) -> Result<DapAggregateResult, VdafError> {
-    match (weight_config, agg_param) {
-        (MasticWeightConfig::Count, DapAggregationParam::Mastic(agg_param)) => {
-            let agg: Vec<Field64> = agg_share_bytes
-                .into_iter()
-                .map(|bytes| decode_field_vec(&bytes, agg_param.prefixes().len()))
-                .reduce(|r, agg_share| {
-                    let mut agg = r?;
-                    for (x, y) in agg.iter_mut().zip(agg_share?.into_iter()) {
-                        *x += y;
-                    }
-                    Ok(agg)
-                })
-                .ok_or_else(|| {
-                    VdafError::Dap(fatal_error!(
-                        err = "mastic: unexpected number of agg shares"
-                    ))
-                })??;
-
-            Ok(DapAggregateResult::U64Vec(
-                agg.into_iter().map(u64::from).collect(),
-            ))
+    pub(crate) fn unshard<M: IntoIterator<Item = Vec<u8>>>(
+        self,
+        agg_param: &DapAggregationParam,
+        agg_shares: M,
+        num_measurements: usize,
+    ) -> Result<DapAggregateResult, VdafError> {
+        match (self.weight_config, agg_param) {
+            (MasticWeightConfig::Count, DapAggregationParam::Mastic(agg_param)) => {
+                let vdaf = Mastic::new_count(self.bits).map_err(|e| {
+                    VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
+                })?;
+
+                let agg_result = unshard(&vdaf, agg_param, num_measurements, agg_shares)?;
+                Ok(DapAggregateResult::U64Vec(agg_result))
+            }
+            _ => todo!(),
         }
-        _ => Err(VdafError::Dap(fatal_error!(
-            err = "mastic: unexpected agg param type"
-        ))),
     }
 }
 
 #[cfg(test)]
 mod test {
-    use prio::{idpf::IdpfInput, vdaf::poplar1::Poplar1AggregationParam};
+    use prio::{idpf::IdpfInput, vdaf::mastic::MasticAggregationParam};
 
     use super::*;
     use crate::{
@@ -233,19 +202,22 @@ mod test {
     #[test]
     fn roundtrip_count() {
         let mut t = AggregationJobTest::new(
-            &VdafConfig::Mastic {
-                input_size: 4,
+            &VdafConfig::Mastic(MasticConfig {
+                bits: 32,
                 weight_config: MasticWeightConfig::Count,
-            },
+            }),
             HpkeKemId::X25519HkdfSha256,
             DapVersion::Latest,
         );
         let got = t.roundtrip(
             DapAggregationParam::Mastic(
-                Poplar1AggregationParam::try_from_prefixes(vec![
-                    IdpfInput::from_bytes(b"cool"),
-                    IdpfInput::from_bytes(b"trip"),
-                ])
+                MasticAggregationParam::new(
+                    vec![
+                        IdpfInput::from_bytes(b"cool"),
+                        IdpfInput::from_bytes(b"trip"),
+                    ],
+                    true,
+                )
                 .unwrap(),
             ),
             vec![
diff --git a/crates/daphne/src/vdaf/mod.rs b/crates/daphne/src/vdaf/mod.rs
index adb190a..1af845a 100644
--- a/crates/daphne/src/vdaf/mod.rs
+++ b/crates/daphne/src/vdaf/mod.rs
@@ -14,7 +14,11 @@ pub(crate) mod prio3;
 use crate::pine::vdaf::PinePrepState;
 use crate::{fatal_error, messages::TaskId, DapError};
 use crate::{DapAggregateResult, DapAggregationParam, DapMeasurement, DapVersion};
+#[cfg(feature = "experimental")]
+use mastic::MasticConfig;
 use pine::PineConfig;
+#[cfg(feature = "experimental")]
+use prio::vdaf::mastic::{MasticPrepareShare, MasticPrepareState};
 use prio::{
     codec::{CodecError, Encode, ParameterizedDecode},
     field::{Field128, Field64, FieldPrio2},
@@ -25,7 +29,7 @@ use prio::{
     },
 };
 
-#[cfg(any(test, feature = "test-utils", feature = "experimental"))]
+#[cfg(any(test, feature = "test-utils"))]
 use prio::field::FieldElement;
 use prio_draft09::{
     codec::{CodecError as CodecErrorDraft09, Encode as EncodeDraft09},
@@ -43,8 +47,6 @@ use prio_draft09::{
 use rand::prelude::*;
 use ring::hkdf::KeyType;
 use serde::{Deserialize, Serialize};
-#[cfg(feature = "experimental")]
-use std::io::Read;
 
 #[cfg(feature = "experimental")]
 pub use self::mastic::MasticWeightConfig;
@@ -87,13 +89,7 @@ pub enum VdafConfig {
         dimension: usize,
     },
     #[cfg(feature = "experimental")]
-    Mastic {
-        /// Length of each input, in number of bytes.
-        input_size: usize,
-
-        /// The type of each weight.
-        weight_config: MasticWeightConfig,
-    },
+    Mastic(MasticConfig),
     Pine(PineConfig),
 }
 
@@ -111,10 +107,7 @@ impl std::fmt::Display for VdafConfig {
             VdafConfig::Prio3(prio3_config) => write!(f, "Prio3({prio3_config})"),
             VdafConfig::Prio2 { dimension } => write!(f, "Prio2({dimension})"),
             #[cfg(feature = "experimental")]
-            VdafConfig::Mastic {
-                input_size,
-                weight_config,
-            } => write!(f, "Mastic({input_size}, {weight_config})"),
+            VdafConfig::Mastic(mastic_config) => write!(f, "Mastic({mastic_config})"),
             VdafConfig::Pine(pine_config) => write!(f, "{pine_config}"),
         }
     }
@@ -135,17 +128,9 @@ impl VdafConfig {
             }
             (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,
-            )?),
+            (DapVersion::Latest, VdafConfig::Mastic(mastic_config)) => {
+                mastic_config.shard(measurement, nonce, task_id)
+            }
             _ => Err(VdafError::Dap(fatal_error!(
                 err = format!("{self:?} is not supported in DAP {version}")
             ))),
@@ -192,17 +177,12 @@ impl VdafConfig {
                 input_share_data,
             ),
             #[cfg(feature = "experimental")]
-            (
-                DapVersion::Latest,
-                Self::Mastic {
-                    input_size,
-                    weight_config,
-                },
-            ) => mastic::mastic_prep_init(
-                *input_size,
-                *weight_config,
+            (DapVersion::Latest, Self::Mastic(mastic_config)) => mastic_config.prep_init(
                 verify_key,
+                task_id,
+                agg_id,
                 agg_param,
+                nonce,
                 public_share_data,
                 input_share_data,
             ),
@@ -215,12 +195,15 @@ impl VdafConfig {
     pub(crate) fn prep_finish_from_shares(
         &self,
         version: DapVersion,
-        agg_id: usize,
         task_id: TaskId,
+        agg_param: &DapAggregationParam,
         host_state: VdafPrepState,
         host_share: VdafPrepShare,
         peer_share_data: &[u8],
     ) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
+        #[cfg(not(feature = "experimental"))]
+        let _ = agg_param;
+
         match (version, self) {
             (_, Self::Prio2 { dimension }) => prio2::prio2_prep_finish_from_shares(
                 *dimension,
@@ -230,28 +213,23 @@ impl VdafConfig {
             ),
             (_, 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)
+                pine_config.prep_finish_from_shares(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,
-            ),
+            (DapVersion::Latest, Self::Mastic(mastic_config)) => mastic_config
+                .prep_finish_from_shares(
+                    task_id,
+                    agg_param,
+                    host_state,
+                    host_share,
+                    peer_share_data,
+                ),
             _ => Err(VdafError::Dap(fatal_error!(
                 err = format!("{self:?} is not supported in DAP {version}")
             ))),
@@ -276,8 +254,8 @@ impl VdafConfig {
                 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)
+            (DapVersion::Latest, Self::Mastic(mastic_config)) => {
+                mastic_config.prep_finish(host_state, peer_message_data, task_id)
             }
             _ => Err(VdafError::Dap(fatal_error!(
                 err = format!("{self:?} is not supported in DAP {version}")
@@ -305,13 +283,9 @@ impl VdafConfig {
                 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),
+            (DapVersion::Latest, Self::Mastic(mastic_config)) => {
+                mastic_config.unshard(agg_param, agg_shares, num_measurements)
+            }
             _ => Err(VdafError::Dap(fatal_error!(
                 err = format!("{self:?} is not supported in DAP {version}")
             ))),
@@ -433,16 +407,14 @@ fn upgrade_codec_error(error: CodecErrorDraft09) -> CodecError {
 
 /// VDAF preparation state.
 #[derive(Clone)]
-#[cfg_attr(any(test, feature = "test-utils"), derive(Debug, Eq, PartialEq))]
+#[cfg_attr(any(test, feature = "test-utils"), derive(Debug, PartialEq))]
 pub enum VdafPrepState {
     Prio2(Prio2PrepareState),
     Prio3Draft09Field64HmacSha256Aes128(Prio3Draft09PrepareState<Field64Draft09, 32>),
     Prio3Field64(Prio3PrepareState<Field64, 32>),
     Prio3Field128(Prio3PrepareState<Field128, 32>),
     #[cfg(feature = "experimental")]
-    Mastic {
-        out_share: Vec<Field64>,
-    },
+    MasticField64(MasticPrepareState<Field64>),
     Pine64HmacSha256Aes128(PinePrepState<Field64Draft09, 32>),
     Pine32HmacSha256Aes128(PinePrepState<FieldPrio2Draft09, 32>),
 }
@@ -461,7 +433,7 @@ impl Encode for VdafPrepState {
             Self::Pine64HmacSha256Aes128(state) => state.encode(bytes).map_err(upgrade_codec_error),
             Self::Pine32HmacSha256Aes128(state) => state.encode(bytes).map_err(upgrade_codec_error),
             #[cfg(feature = "experimental")]
-            Self::Mastic { .. } => todo!("encoding of prep state is not implemented"),
+            Self::MasticField64(state) => state.encode(bytes),
         }
     }
 }
@@ -499,7 +471,7 @@ impl deepsize::DeepSizeOf for VdafPrepState {
             | Self::Pine64HmacSha256Aes128(_)
             | Self::Pine32HmacSha256Aes128(_) => 0,
             #[cfg(feature = "experimental")]
-            Self::Mastic { .. } => 0,
+            Self::MasticField64(_) => 0,
         }
     }
 }
@@ -513,7 +485,7 @@ pub enum VdafPrepShare {
     Prio3Field64(Prio3PrepareShare<Field64, 32>),
     Prio3Field128(Prio3PrepareShare<Field128, 32>),
     #[cfg(feature = "experimental")]
-    Mastic(Field64),
+    MasticField64(MasticPrepareShare<Field64>),
     Pine64HmacSha256Aes128(crate::pine::msg::PrepShare<Field64Draft09, 32>),
     Pine32HmacSha256Aes128(crate::pine::msg::PrepShare<FieldPrio2Draft09, 32>),
 }
@@ -535,7 +507,7 @@ impl deepsize::DeepSizeOf for VdafPrepShare {
             | Self::Pine64HmacSha256Aes128(_)
             | Self::Pine32HmacSha256Aes128(_) => 0,
             #[cfg(feature = "experimental")]
-            Self::Mastic(..) => 0,
+            Self::MasticField64(..) => 0,
         }
     }
 }
@@ -550,7 +522,7 @@ impl Encode for VdafPrepShare {
             Self::Prio3Field128(share) => share.encode(bytes),
             Self::Prio2(share) => share.encode(bytes),
             #[cfg(feature = "experimental")]
-            Self::Mastic(share) => share.encode(bytes),
+            Self::MasticField64(share) => share.encode(bytes),
             Self::Pine64HmacSha256Aes128(share) => share.encode(bytes).map_err(upgrade_codec_error),
             Self::Pine32HmacSha256Aes128(share) => share.encode(bytes).map_err(upgrade_codec_error),
         }
@@ -579,9 +551,9 @@ impl ParameterizedDecode<VdafPrepState> for VdafPrepShare {
                 Prio2PrepareShare::decode_with_param(state, bytes)?,
             )),
             #[cfg(feature = "experimental")]
-            VdafPrepState::Mastic { .. } => {
-                todo!("mastic: decoding of prep messages is not implemented")
-            }
+            VdafPrepState::MasticField64(state) => Ok(VdafPrepShare::MasticField64(
+                MasticPrepareShare::decode_with_param(state, bytes)?,
+            )),
             VdafPrepState::Pine64HmacSha256Aes128(state) => {
                 Ok(VdafPrepShare::Pine64HmacSha256Aes128(
                     crate::pine::msg::PrepShare::decode_with_param(state, bytes)
@@ -670,22 +642,6 @@ impl VdafConfig {
     }
 }
 
-#[cfg(feature = "experimental")]
-pub(crate) fn decode_field_vec<F: FieldElement>(
-    bytes: &[u8],
-    len: usize,
-) -> Result<Vec<F>, CodecError> {
-    debug_assert!(F::ENCODED_SIZE < 64);
-    let mut buf = [0; 64];
-    let mut r = std::io::Cursor::new(bytes);
-    let mut v = Vec::with_capacity(len);
-    for _ in 0..len {
-        r.read_exact(&mut buf[..F::ENCODED_SIZE])?;
-        v.push(F::get_decoded(&buf[..F::ENCODED_SIZE])?);
-    }
-    Ok(v)
-}
-
 fn shard_then_encode<V: Vdaf + Client<16>>(
     vdaf: &V,
     task_id: TaskId,
@@ -712,18 +668,54 @@ fn shard_then_encode<V: Vdaf + Client<16>>(
     ))
 }
 
-fn prep_finish_from_shares<V, const VERIFY_KEY_SIZE: usize, const NONCE_SIZE: usize>(
+#[allow(clippy::too_many_arguments)]
+fn prep_init<V, const VERIFY_KEY_SIZE: usize, const NONCE_SIZE: usize>(
     vdaf: &V,
     task_id: TaskId,
+    verify_key: &[u8; VERIFY_KEY_SIZE],
     agg_id: usize,
+    agg_param: &V::AggregationParam,
+    nonce: &[u8; NONCE_SIZE],
+    public_share_data: &[u8],
+    input_share_data: &[u8],
+) -> Result<(V::PrepareState, V::PrepareShare), VdafError>
+where
+    V: Vdaf + Aggregator<VERIFY_KEY_SIZE, NONCE_SIZE>,
+{
+    // Parse the public share.
+    let public_share = V::PublicShare::get_decoded_with_param(vdaf, public_share_data)?;
+
+    // Parse the input share.
+    let input_share = V::InputShare::get_decoded_with_param(&(vdaf, agg_id), input_share_data)?;
+
+    let mut ctx = [0; CTX_STRING_PREFIX.len() + 32];
+    ctx[..CTX_STRING_PREFIX.len()].copy_from_slice(CTX_STRING_PREFIX);
+    ctx[CTX_STRING_PREFIX.len()..].copy_from_slice(&task_id.0);
+    let (prep_state, prep_share) = vdaf.prepare_init(
+        verify_key,
+        &ctx,
+        agg_id,
+        agg_param,
+        nonce,
+        &public_share,
+        &input_share,
+    )?;
+
+    Ok((prep_state, prep_share))
+}
+
+fn prep_finish_from_shares<V, const VERIFY_KEY_SIZE: usize, const NONCE_SIZE: usize>(
+    vdaf: &V,
+    task_id: TaskId,
+    agg_param: &V::AggregationParam,
     host_state: V::PrepareState,
     host_share: V::PrepareShare,
     peer_share_data: &[u8],
 ) -> Result<(V::OutputShare, Vec<u8>), VdafError>
 where
-    V: Vdaf<AggregationParam = ()> + Aggregator<VERIFY_KEY_SIZE, NONCE_SIZE>,
+    V: Vdaf + Aggregator<VERIFY_KEY_SIZE, NONCE_SIZE>,
 {
-    // Decode the Helper's inbound message.
+    // Decode the peer's inbound message.
     let peer_share = V::PrepareShare::get_decoded_with_param(&host_state, peer_share_data)?;
 
     let mut ctx = [0; CTX_STRING_PREFIX.len() + 32];
@@ -731,15 +723,8 @@ where
     ctx[CTX_STRING_PREFIX.len()..].copy_from_slice(&task_id.0);
 
     // Preprocess the inbound messages.
-    let message = vdaf.prepare_shares_to_prepare_message(
-        &ctx,
-        &(),
-        if agg_id == 0 {
-            [host_share, peer_share]
-        } else {
-            [peer_share, host_share]
-        },
-    )?;
+    let message =
+        vdaf.prepare_shares_to_prepare_message(&ctx, agg_param, [peer_share, host_share])?;
     let message_data = message.get_encoded()?;
 
     // Compute the host's output share.
@@ -778,17 +763,19 @@ where
 
 fn unshard<V, M>(
     vdaf: &V,
+    agg_param: &V::AggregationParam,
     num_measurements: usize,
     agg_shares: M,
 ) -> Result<V::AggregateResult, VdafError>
 where
-    V: Vdaf<AggregationParam = ()> + Collector,
+    V: Vdaf + Collector,
     M: IntoIterator<Item = Vec<u8>>,
 {
     let mut agg_shares_vec = Vec::with_capacity(vdaf.num_aggregators());
     for data in agg_shares {
-        let agg_share = V::AggregateShare::get_decoded_with_param(&(vdaf, &()), data.as_ref())?;
+        let agg_share =
+            V::AggregateShare::get_decoded_with_param(&(vdaf, agg_param), data.as_ref())?;
         agg_shares_vec.push(agg_share);
     }
-    Ok(vdaf.unshard(&(), agg_shares_vec, num_measurements)?)
+    Ok(vdaf.unshard(agg_param, agg_shares_vec, num_measurements)?)
 }
diff --git a/crates/daphne/src/vdaf/pine.rs b/crates/daphne/src/vdaf/pine.rs
index ae3f691..5bf2935 100644
--- a/crates/daphne/src/vdaf/pine.rs
+++ b/crates/daphne/src/vdaf/pine.rs
@@ -123,7 +123,6 @@ impl PineConfig {
 
     pub(crate) fn prep_finish_from_shares(
         &self,
-        agg_id: usize,
         host_state: VdafPrepState,
         host_share: VdafPrepShare,
         peer_share_data: &[u8],
@@ -136,7 +135,7 @@ impl PineConfig {
             ) => {
                 let vdaf = pine32_hmac_sha256_aes128(param)?;
                 let (out_share, outbound) =
-                    draft09::prep_finish_from_shares(&vdaf, agg_id, state, share, peer_share_data)?;
+                    draft09::prep_finish_from_shares(&vdaf, state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field32Draft09(AggregateShare::from(
                     OutputShare::from(out_share.0),
                 ));
@@ -149,7 +148,7 @@ impl PineConfig {
             ) => {
                 let vdaf = pine64_hmac_sha256_aes128(param)?;
                 let (out_share, outbound) =
-                    draft09::prep_finish_from_shares(&vdaf, agg_id, state, share, peer_share_data)?;
+                    draft09::prep_finish_from_shares(&vdaf, state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field64Draft09(AggregateShare::from(
                     OutputShare::from(out_share.0),
                 ));
diff --git a/crates/daphne/src/vdaf/prio3.rs b/crates/daphne/src/vdaf/prio3.rs
index 8cf57ad..461f8af 100644
--- a/crates/daphne/src/vdaf/prio3.rs
+++ b/crates/daphne/src/vdaf/prio3.rs
@@ -12,21 +12,17 @@ use crate::{
     VdafPrepState,
 };
 
-use super::{prep_finish, prep_finish_from_shares, shard_then_encode, unshard};
+use super::{prep_finish, prep_finish_from_shares, prep_init, shard_then_encode, unshard};
 
 use prio::{
     codec::ParameterizedDecode,
-    flp::Type,
     vdaf::{
-        prio3::{Prio3, Prio3InputShare, Prio3PrepareShare, Prio3PrepareState, Prio3PublicShare},
-        xof::Xof,
+        prio3::{Prio3, Prio3PrepareState},
         Aggregator,
     },
 };
 use std::io::Cursor;
 
-const CTX_STRING_PREFIX: &[u8] = b"dap-13";
-
 impl Prio3Config {
     pub(crate) fn shard(
         &self,
@@ -118,16 +114,17 @@ impl Prio3Config {
         public_share_data: &[u8],
         input_share_data: &[u8],
     ) -> Result<(VdafPrepState, VdafPrepShare), VdafError> {
-        return match (version, self) {
+        match (version, self) {
             (DapVersion::Latest, Prio3Config::Count) => {
                 let vdaf = Prio3::new_count(2).map_err(|e| {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (state, share) = prep_init(
-                    vdaf,
+                    &vdaf,
                     task_id,
                     verify_key,
                     agg_id,
+                    &(),
                     nonce,
                     public_share_data,
                     input_share_data,
@@ -142,10 +139,11 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (state, share) = prep_init(
-                    vdaf,
+                    &vdaf,
                     task_id,
                     verify_key,
                     agg_id,
+                    &(),
                     nonce,
                     public_share_data,
                     input_share_data,
@@ -166,10 +164,11 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (state, share) = prep_init(
-                    vdaf,
+                    &vdaf,
                     task_id,
                     verify_key,
                     agg_id,
+                    &(),
                     nonce,
                     public_share_data,
                     input_share_data,
@@ -191,10 +190,11 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (state, share) = prep_init(
-                    vdaf,
+                    &vdaf,
                     task_id,
                     verify_key,
                     agg_id,
+                    &(),
                     nonce,
                     public_share_data,
                     input_share_data,
@@ -236,53 +236,12 @@ impl Prio3Config {
                 err =
                     format!("unexpected verify key or {self:?} is not supported in DAP {version}")
             ))),
-        };
-
-        type Prio3Prepared<T, const SEED_SIZE: usize> = (
-            Prio3PrepareState<<T as Type>::Field, SEED_SIZE>,
-            Prio3PrepareShare<<T as Type>::Field, SEED_SIZE>,
-        );
-
-        fn prep_init<T, P, const SEED_SIZE: usize>(
-            vdaf: Prio3<T, P, SEED_SIZE>,
-            task_id: TaskId,
-            verify_key: &[u8; SEED_SIZE],
-            agg_id: usize,
-            nonce: &[u8; 16],
-            public_share_data: &[u8],
-            input_share_data: &[u8],
-        ) -> Result<Prio3Prepared<T, SEED_SIZE>, VdafError>
-        where
-            T: Type,
-            P: Xof<SEED_SIZE>,
-        {
-            // Parse the public share.
-            let public_share = Prio3PublicShare::get_decoded_with_param(&vdaf, public_share_data)?;
-
-            // Parse the input share.
-            let input_share =
-                Prio3InputShare::get_decoded_with_param(&(&vdaf, agg_id), input_share_data)?;
-
-            let mut ctx = [0; CTX_STRING_PREFIX.len() + 32];
-            ctx[..CTX_STRING_PREFIX.len()].copy_from_slice(CTX_STRING_PREFIX);
-            ctx[CTX_STRING_PREFIX.len()..].copy_from_slice(&task_id.0);
-            // Run the prepare-init algorithm, returning the initial state.
-            Ok(vdaf.prepare_init(
-                verify_key,
-                &ctx,
-                agg_id,
-                &(),
-                nonce,
-                &public_share,
-                &input_share,
-            )?)
         }
     }
 
     pub(crate) fn prep_finish_from_shares(
         &self,
         version: DapVersion,
-        agg_id: usize,
         task_id: TaskId,
         host_state: VdafPrepState,
         host_share: VdafPrepShare,
@@ -299,7 +258,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (out_share, outbound) =
-                    prep_finish_from_shares(&vdaf, task_id, agg_id, state, share, peer_share_data)?;
+                    prep_finish_from_shares(&vdaf, task_id, &(), state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field64(vdaf.aggregate(&(), [out_share])?);
                 (agg_share, outbound)
             }
@@ -313,7 +272,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (out_share, outbound) =
-                    prep_finish_from_shares(&vdaf, task_id, agg_id, state, share, peer_share_data)?;
+                    prep_finish_from_shares(&vdaf, task_id, &(), state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field64(vdaf.aggregate(&(), [out_share])?);
                 (agg_share, outbound)
             }
@@ -330,7 +289,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (out_share, outbound) =
-                    prep_finish_from_shares(&vdaf, task_id, agg_id, state, share, peer_share_data)?;
+                    prep_finish_from_shares(&vdaf, task_id, &(), state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?);
                 (agg_share, outbound)
             }
@@ -348,7 +307,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
                 let (out_share, outbound) =
-                    prep_finish_from_shares(&vdaf, task_id, agg_id, state, share, peer_share_data)?;
+                    prep_finish_from_shares(&vdaf, task_id, &(), state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?);
                 (agg_share, outbound)
             }
@@ -370,7 +329,7 @@ impl Prio3Config {
                     *num_proofs,
                 )?;
                 let (out_share, outbound) =
-                    draft09::prep_finish_from_shares(&vdaf, agg_id, state, share, peer_share_data)?;
+                    draft09::prep_finish_from_shares(&vdaf, state, share, peer_share_data)?;
                 let agg_share = VdafAggregateShare::Field64Draft09(
                     prio_draft09::vdaf::Aggregator::aggregate(&vdaf, &(), [out_share])?,
                 );
@@ -490,7 +449,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
 
-                let agg_res = unshard(&vdaf, num_measurements, agg_shares)?;
+                let agg_res = unshard(&vdaf, &(), num_measurements, agg_shares)?;
                 Ok(DapAggregateResult::U64(agg_res))
             }
             (DapVersion::Latest, Prio3Config::Sum { max_measurement }) => {
@@ -498,7 +457,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
 
-                let agg_res = unshard(&vdaf, num_measurements, agg_shares)?;
+                let agg_res = unshard(&vdaf, &(), num_measurements, agg_shares)?;
                 Ok(DapAggregateResult::U64(agg_res))
             }
             (
@@ -512,7 +471,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
 
-                let agg_res = unshard(&vdaf, num_measurements, agg_shares)?;
+                let agg_res = unshard(&vdaf, &(), num_measurements, agg_shares)?;
                 Ok(DapAggregateResult::U128Vec(agg_res))
             }
             (
@@ -527,7 +486,7 @@ impl Prio3Config {
                     VdafError::Dap(fatal_error!(err = ?e, "initializing {self:?} failed"))
                 })?;
 
-                let agg_res = unshard(&vdaf, num_measurements, agg_shares)?;
+                let agg_res = unshard(&vdaf, &(), num_measurements, agg_shares)?;
                 Ok(DapAggregateResult::U128Vec(agg_res))
             }
             (


Test patch:
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index 4eaab89..fd3f1fb 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -209,22 +209,26 @@ impl AggregationJobTest {
         &self,
         agg_job_init_req: AggregationJobInitReq,
     ) -> (DapAggregateSpan<DapAggregateShare>, AggregationJobResp) {
+        let part_batch_sel = agg_job_init_req.part_batch_sel.clone();
+        let (agg_param, initialized_reports) = self
+            .task_config
+            .consume_agg_job_req(
+                &self.helper_hpke_receiver_config,
+                self.valid_report_time_range(),
+                &self.task_id,
+                agg_job_init_req,
+                self.replay_protection,
+            )
+            .unwrap();
+
         let (span, resp) = self
             .task_config
             .produce_agg_job_resp(
                 self.task_id,
+                &agg_param,
                 &HashMap::default(),
-                &agg_job_init_req.part_batch_sel.clone(),
-                &self
-                    .task_config
-                    .consume_agg_job_req(
-                        &self.helper_hpke_receiver_config,
-                        self.valid_report_time_range(),
-                        &self.task_id,
-                        agg_job_init_req,
-                        self.replay_protection,
-                    )
-                    .unwrap(),
+                &part_batch_sel,
+                &initialized_reports,
             )
             .unwrap();
         (span, resp.into())