Repository: cloudflare/daphne
Pull request #122 (2022-09-20 16:52:13)

daphne: Upgrade prio to 0.10.0
Partially addresses #100.

Contains wire-breaking changes.

DAP-02 uses VDAF-03, which is implemented by version 0.10.0 of the `prio` crate. Upgrade the dependency and plumb syntax changes into the protocol accordingly.

While at, it use prio's serialization cdoe isntead of rolling our own.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 1d92666..7f6d029 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -112,6 +112,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "az"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b7e4c2464d97fe331d41de9d5db0def0a96f4d823b8b32a2efd503578988973"
+
 [[package]]
 name = "base16ct"
 version = "0.1.1"
@@ -175,6 +181,12 @@ version = "3.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c1ad822118d20d2c234f427000d5acc36eabe1e29a348c89b63dd60b13f28e5d"
 
+[[package]]
+name = "bytemuck"
+version = "1.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f5715e491b5a1598fc2bef5a606847b5dc1d48ea625bd3c02c00de8285591da"
+
 [[package]]
 name = "byteorder"
 version = "1.4.3"
@@ -413,6 +425,12 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "crunchy"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
+
 [[package]]
 name = "crypto-bigint"
 version = "0.2.11"
@@ -497,7 +515,7 @@ dependencies = [
  "base64",
  "clap 3.2.19",
  "daphne 0.1.2",
- "prio",
+ "prio 0.8.2",
  "reqwest",
  "serde",
  "serde_json",
@@ -518,7 +536,7 @@ dependencies = [
  "hex",
  "hpke 0.8.0",
  "matchit 0.6.0",
- "prio",
+ "prio 0.8.2",
  "rand",
  "ring",
  "serde",
@@ -538,7 +556,7 @@ dependencies = [
  "hex",
  "hpke 0.8.0",
  "matchit 0.6.0",
- "prio",
+ "prio 0.10.0",
  "rand",
  "ring",
  "serde",
@@ -565,7 +583,7 @@ dependencies = [
  "janus_core",
  "janus_server",
  "lazy_static",
- "prio",
+ "prio 0.8.2",
  "rand",
  "reqwest",
  "ring",
@@ -590,7 +608,7 @@ dependencies = [
  "hex",
  "hpke 0.8.0",
  "matchit 0.6.0",
- "prio",
+ "prio 0.8.2",
  "rand",
  "reqwest-wasm",
  "ring",
@@ -812,6 +830,18 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "fixed"
+version = "1.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e0371cd413fb63f8ec1b9eb4dff47fa2c88b21abc681771234c84808b9920991"
+dependencies = [
+ "az",
+ "bytemuck",
+ "half",
+ "typenum",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -1007,6 +1037,15 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "half"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad6a9459c9c30b177b925162351f97e7d967c7ea8bab3b8352805327daf45554"
+dependencies = [
+ "crunchy",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
@@ -1341,7 +1380,7 @@ source = "git+https://github.com/divviup/janus?rev=e5ed423cfc96051894de2997aa013
 dependencies = [
  "http",
  "janus_core",
- "prio",
+ "prio 0.8.2",
  "reqwest",
  "thiserror",
  "tokio",
@@ -1365,7 +1404,7 @@ dependencies = [
  "num_enum",
  "postgres-protocol",
  "postgres-types",
- "prio",
+ "prio 0.8.2",
  "rand",
  "ring",
  "serde",
@@ -1400,7 +1439,7 @@ dependencies = [
  "num_enum",
  "opentelemetry",
  "postgres-types",
- "prio",
+ "prio 0.8.2",
  "rand",
  "reqwest",
  "ring",
@@ -1951,6 +1990,26 @@ dependencies = [
  "thiserror",
 ]
 
+[[package]]
+name = "prio"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0453e1ec5f2f48af9a67aab1d812f9aec48619dda0d9a59e2f79ebd235448ec"
+dependencies = [
+ "aes 0.8.1",
+ "aes-gcm",
+ "base64",
+ "byteorder",
+ "cmac",
+ "ctr 0.9.1",
+ "fixed",
+ "getrandom",
+ "ring",
+ "serde",
+ "static_assertions",
+ "thiserror",
+]
+
 [[package]]
 name = "proc-macro-crate"
 version = "1.2.1"
diff --git a/daphne/Cargo.toml b/daphne/Cargo.toml
index 65f03a9..76d8c51 100644
--- a/daphne/Cargo.toml
+++ b/daphne/Cargo.toml
@@ -23,7 +23,7 @@ async-trait = "0.1.56"
 base64 = "0.13.0"
 getrandom = { version = "0.2", features = ["js"] } # Required for prio
 serde_json = "1.0.82"
-prio = { version = "=0.8.2", features = ["prio2"] }
+prio = { version = "0.10.0", features = ["prio2"] }
 hpke = { version = "0.8.0", features = ["std", "serde_impls"] }
 ring = "0.16.20"
 hex = { version = "0.4.3", features = ["serde"] }
diff --git a/daphne/dapf/Cargo.toml b/daphne/dapf/Cargo.toml
index 46ed62c..cc0cb94 100644
--- a/daphne/dapf/Cargo.toml
+++ b/daphne/dapf/Cargo.toml
@@ -14,7 +14,7 @@ license = "BSD-3-Clause"
 daphne = "0.1.2" # TODO(issue #100) Use path = "../daphne" instead.
 assert_matches = "1.5.0"
 base64 = "0.13.0"
-prio = { version = "=0.8.2" }
+prio = "=0.8.2"
 serde = { version = "^1.0", features = ["derive"] }
 serde_json = "1.0.82"
 url = { version = "2.2.2", features = ["serde"] }
diff --git a/daphne/src/hpke.rs b/daphne/src/hpke.rs
index 025b055..6d3e7df 100644
--- a/daphne/src/hpke.rs
+++ b/daphne/src/hpke.rs
@@ -5,8 +5,8 @@
 
 use crate::{
     messages::{
-        decode_prefixed_bytes, encode_prefixed_bytes, HpkeAeadId, HpkeCiphertext, HpkeConfig,
-        HpkeKdfId, HpkeKemId, Id,
+        decode_u16_bytes, encode_u16_bytes, HpkeAeadId, HpkeCiphertext, HpkeConfig, HpkeKdfId,
+        HpkeKemId, Id,
     },
     DapError,
 };
@@ -286,7 +286,7 @@ impl<'a> HpkeDecrypter<'a> for HpkeReceiverConfig {
 impl Encode for HpkeReceiverConfig {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.config.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.secret_key);
+        encode_u16_bytes(bytes, &self.secret_key);
     }
 }
 
@@ -294,7 +294,7 @@ impl Decode for HpkeReceiverConfig {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
             config: HpkeConfig::decode(bytes)?,
-            secret_key: decode_prefixed_bytes::<2>(bytes)?,
+            secret_key: decode_u16_bytes(bytes)?,
         })
     }
 }
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index 5b1d08d..39f3e85 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -4,7 +4,10 @@
 //! Messages in the DAP protocol.
 
 use crate::DapTaskConfig;
-use prio::codec::{CodecError, Decode, Encode};
+use prio::codec::{
+    decode_u16_items, decode_u32_items, encode_u16_items, encode_u32_items, CodecError, Decode,
+    Encode,
+};
 use serde::{Deserialize, Serialize};
 use std::{
     convert::{TryFrom, TryInto},
@@ -97,7 +100,7 @@ impl Encode for Extension {
         match self {
             Self::Unhandled { typ, payload } => {
                 typ.encode(bytes);
-                encode_prefixed_bytes::<2>(bytes, payload);
+                encode_u16_bytes(bytes, payload);
             }
         }
     }
@@ -107,7 +110,7 @@ impl Decode for Extension {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self::Unhandled {
             typ: u16::decode(bytes)?,
-            payload: decode_prefixed_bytes::<2>(bytes)?,
+            payload: decode_u16_bytes(bytes)?,
         })
     }
 }
@@ -125,7 +128,7 @@ impl Encode for ReportMetadata {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.time.encode(bytes);
         self.nonce.encode(bytes);
-        encode_prefixed_items::<2, _>(bytes, &self.extensions);
+        encode_u16_items(bytes, &(), &self.extensions);
     }
 }
 
@@ -134,7 +137,7 @@ impl Decode for ReportMetadata {
         Ok(Self {
             time: Time::decode(bytes)?,
             nonce: Nonce::decode(bytes)?,
-            extensions: decode_prefixed_items::<2, _>(bytes)?,
+            extensions: decode_u16_items(&(), bytes)?,
         })
     }
 }
@@ -153,8 +156,8 @@ impl Encode for Report {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
         self.metadata.encode(bytes);
-        encode_prefixed_bytes::<4>(bytes, &self.public_share);
-        encode_prefixed_items::<4, _>(bytes, &self.encrypted_input_shares);
+        encode_u32_bytes(bytes, &self.public_share);
+        encode_u32_items(bytes, &(), &self.encrypted_input_shares);
     }
 }
 
@@ -163,8 +166,8 @@ impl Decode for Report {
         Ok(Self {
             task_id: Id::decode(bytes)?,
             metadata: ReportMetadata::decode(bytes)?,
-            public_share: decode_prefixed_bytes::<4>(bytes)?,
-            encrypted_input_shares: decode_prefixed_items::<4, _>(bytes)?,
+            public_share: decode_u32_bytes(bytes)?,
+            encrypted_input_shares: decode_u32_items(&(), bytes)?,
         })
     }
 }
@@ -182,7 +185,7 @@ pub struct ReportShare {
 impl Encode for ReportShare {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.metadata.encode(bytes);
-        encode_prefixed_bytes::<4>(bytes, &self.public_share);
+        encode_u32_bytes(bytes, &self.public_share);
         self.encrypted_input_share.encode(bytes);
     }
 }
@@ -191,7 +194,7 @@ impl Decode for ReportShare {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
             metadata: ReportMetadata::decode(bytes)?,
-            public_share: decode_prefixed_bytes::<4>(bytes)?,
+            public_share: decode_u32_bytes(bytes)?,
             encrypted_input_share: HpkeCiphertext::decode(bytes)?,
         })
     }
@@ -241,9 +244,9 @@ impl Encode for AggregateInitializeReq {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
         self.agg_job_id.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.agg_param);
+        encode_u16_bytes(bytes, &self.agg_param);
         self.batch_param.encode(bytes);
-        encode_prefixed_items::<4, _>(bytes, &self.report_shares);
+        encode_u32_items(bytes, &(), &self.report_shares);
     }
 }
 
@@ -252,9 +255,9 @@ impl Decode for AggregateInitializeReq {
         Ok(Self {
             task_id: Id::decode(bytes)?,
             agg_job_id: Id::decode(bytes)?,
-            agg_param: decode_prefixed_bytes::<2>(bytes)?,
+            agg_param: decode_u16_bytes(bytes)?,
             batch_param: BatchParameter::decode(bytes)?,
-            report_shares: decode_prefixed_items::<4, _>(bytes)?,
+            report_shares: decode_u32_items(&(), bytes)?,
         })
     }
 }
@@ -271,7 +274,7 @@ impl Encode for AggregateContinueReq {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
         self.agg_job_id.encode(bytes);
-        encode_prefixed_items::<4, _>(bytes, &self.transitions);
+        encode_u32_items(bytes, &(), &self.transitions);
     }
 }
 
@@ -280,7 +283,7 @@ impl Decode for AggregateContinueReq {
         Ok(Self {
             task_id: Id::decode(bytes)?,
             agg_job_id: Id::decode(bytes)?,
-            transitions: decode_prefixed_items::<4, _>(bytes)?,
+            transitions: decode_u32_items(&(), bytes)?,
         })
     }
 }
@@ -325,7 +328,7 @@ impl Encode for TransitionVar {
         match self {
             TransitionVar::Continued(vdaf_message) => {
                 0_u8.encode(bytes);
-                encode_prefixed_bytes::<4>(bytes, vdaf_message);
+                encode_u32_bytes(bytes, vdaf_message);
             }
             TransitionVar::Finished => {
                 1_u8.encode(bytes);
@@ -341,7 +344,7 @@ impl Encode for TransitionVar {
 impl Decode for TransitionVar {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         match u8::decode(bytes)? {
-            0 => Ok(Self::Continued(decode_prefixed_bytes::<4>(bytes)?)),
+            0 => Ok(Self::Continued(decode_u32_bytes(bytes)?)),
             1 => Ok(Self::Finished),
             2 => Ok(Self::Failed(TransitionFailure::decode(bytes)?)),
             _ => Err(CodecError::UnexpectedValue),
@@ -413,14 +416,14 @@ pub struct AggregateResp {
 
 impl Encode for AggregateResp {
     fn encode(&self, bytes: &mut Vec<u8>) {
-        encode_prefixed_items::<4, _>(bytes, &self.transitions);
+        encode_u32_items(bytes, &(), &self.transitions);
     }
 }
 
 impl Decode for AggregateResp {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
-            transitions: decode_prefixed_items::<4, _>(bytes)?,
+            transitions: decode_u32_items(&(), bytes)?,
         })
     }
 }
@@ -536,7 +539,7 @@ impl Encode for CollectReq {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
         self.query.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.agg_param);
+        encode_u16_bytes(bytes, &self.agg_param);
     }
 }
 
@@ -545,7 +548,7 @@ impl Decode for CollectReq {
         Ok(Self {
             task_id: Id::decode(bytes)?,
             query: Query::decode(bytes)?,
-            agg_param: decode_prefixed_bytes::<2>(bytes)?,
+            agg_param: decode_u16_bytes(bytes)?,
         })
     }
 }
@@ -562,7 +565,7 @@ pub struct CollectResp {
 impl Encode for CollectResp {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.report_count.encode(bytes);
-        encode_prefixed_items::<4, _>(bytes, &self.encrypted_agg_shares);
+        encode_u32_items(bytes, &(), &self.encrypted_agg_shares);
     }
 }
 
@@ -570,7 +573,7 @@ impl Decode for CollectResp {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
             report_count: u64::decode(bytes)?,
-            encrypted_agg_shares: decode_prefixed_items::<4, _>(bytes)?,
+            encrypted_agg_shares: decode_u32_items(&(), bytes)?,
         })
     }
 }
@@ -598,7 +601,7 @@ impl Encode for AggregateShareReq {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
         self.batch_selector.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.agg_param);
+        encode_u16_bytes(bytes, &self.agg_param);
         self.report_count.encode(bytes);
         bytes.extend_from_slice(&self.checksum);
     }
@@ -609,7 +612,7 @@ impl Decode for AggregateShareReq {
         Ok(Self {
             task_id: Id::decode(bytes)?,
             batch_selector: BatchSelector::decode(bytes)?,
-            agg_param: decode_prefixed_bytes::<2>(bytes)?,
+            agg_param: decode_u16_bytes(bytes)?,
             report_count: u64::decode(bytes)?,
             checksum: {
                 let mut checksum = [0u8; 32];
@@ -765,7 +768,7 @@ impl Encode for HpkeConfig {
         self.kem_id.encode(bytes);
         self.kdf_id.encode(bytes);
         self.aead_id.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.public_key);
+        encode_u16_bytes(bytes, &self.public_key);
     }
 }
 
@@ -776,7 +779,7 @@ impl Decode for HpkeConfig {
             kem_id: HpkeKemId::decode(bytes)?,
             kdf_id: HpkeKdfId::decode(bytes)?,
             aead_id: HpkeAeadId::decode(bytes)?,
-            public_key: decode_prefixed_bytes::<2>(bytes)?,
+            public_key: decode_u16_bytes(bytes)?,
         })
     }
 }
@@ -796,8 +799,8 @@ pub struct HpkeCiphertext {
 impl Encode for HpkeCiphertext {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.config_id.encode(bytes);
-        encode_prefixed_bytes::<2>(bytes, &self.enc);
-        encode_prefixed_bytes::<4>(bytes, &self.payload);
+        encode_u16_bytes(bytes, &self.enc);
+        encode_u32_bytes(bytes, &self.payload);
     }
 }
 
@@ -805,8 +808,8 @@ impl Decode for HpkeCiphertext {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
             config_id: u8::decode(bytes)?,
-            enc: decode_prefixed_bytes::<2>(bytes)?,
-            payload: decode_prefixed_bytes::<4>(bytes)?,
+            enc: decode_u16_bytes(bytes)?,
+            payload: decode_u32_bytes(bytes)?,
         })
     }
 }
@@ -825,83 +828,30 @@ pub(crate) fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
     r == 0
 }
 
-// TODO(issue #100) Replace this stuff with libprio once 0.10.0 is ready.
-fn encode_prefix<const PREFIX_BYTES: usize>(bytes: &mut Vec<u8>, len: usize) {
-    let prefix_start = bytes.len();
-    bytes.extend_from_slice(&[0; PREFIX_BYTES]);
-    let prefix_end = bytes.len();
-    write_prefix::<PREFIX_BYTES>(&mut bytes[prefix_start..prefix_end], len);
-}
-
-fn write_prefix<const PREFIX_BYTES: usize>(buf: &mut [u8], len: usize) {
-    for i in (0..PREFIX_BYTES).rev() {
-        buf[PREFIX_BYTES - i - 1] = ((len >> (i << 3)) & 255) as u8;
-    }
-}
-
-pub(crate) fn decode_prefix<const PREFIX_BYTES: usize>(
-    bytes: &mut Cursor<&[u8]>,
-) -> Result<usize, CodecError> {
-    let mut len = 0;
-    for i in (0..PREFIX_BYTES).rev() {
-        len |= (u8::decode(bytes)? as usize) << (i << 3);
-    }
-    Ok(len)
-}
-
-pub(crate) fn encode_prefixed_bytes<const PREFIX_BYTES: usize>(bytes: &mut Vec<u8>, input: &[u8]) {
-    encode_prefix::<PREFIX_BYTES>(bytes, input.len());
+pub(crate) fn encode_u16_bytes(bytes: &mut Vec<u8>, input: &[u8]) {
+    u16::try_from(input.len())
+        .expect("length too large for u16")
+        .encode(bytes);
     bytes.extend_from_slice(input);
 }
 
-pub(crate) fn decode_prefixed_bytes<const PREFIX_BYTES: usize>(
-    bytes: &mut Cursor<&[u8]>,
-) -> Result<Vec<u8>, CodecError> {
-    let len = decode_prefix::<PREFIX_BYTES>(bytes)?;
+pub(crate) fn decode_u16_bytes(bytes: &mut Cursor<&[u8]>) -> Result<Vec<u8>, CodecError> {
+    let len = u16::decode(bytes)? as usize;
     let mut out = vec![0; len];
     bytes.read_exact(&mut out)?;
     Ok(out)
 }
 
-fn encode_prefixed_items<const PREFIX_BYTES: usize, I: Encode>(bytes: &mut Vec<u8>, items: &[I]) {
-    let prefix_start = bytes.len();
-    bytes.extend_from_slice(&[0; PREFIX_BYTES]);
-    let prefix_end = bytes.len();
-    for item in items {
-        item.encode(bytes);
-    }
-    let len = bytes.len() - prefix_end;
-    write_prefix::<PREFIX_BYTES>(&mut bytes[prefix_start..prefix_end], len);
-}
-
-fn decode_prefixed_items<const PREFIX_BYTES: usize, I: Decode>(
-    bytes: &mut Cursor<&[u8]>,
-) -> Result<Vec<I>, CodecError> {
-    let len = decode_prefix::<PREFIX_BYTES>(bytes)?;
-    decode_items(len, bytes)
+pub(crate) fn encode_u32_bytes(bytes: &mut Vec<u8>, input: &[u8]) {
+    u32::try_from(input.len())
+        .expect("length too large for u32")
+        .encode(bytes);
+    bytes.extend_from_slice(input);
 }
 
-fn decode_items<I: Decode>(len: usize, bytes: &mut Cursor<&[u8]>) -> Result<Vec<I>, CodecError> {
-    let mut decoded = Vec::new();
-    let initial_position = bytes.position() as usize;
-
-    // Create cursor over specified portion of provided cursor to ensure we can't read past length.
-    let inner = bytes.get_ref();
-
-    // Make sure encoded length doesn't overflow usize or go past the end of provided byte buffer.
-    let (items_end, overflowed) = initial_position.overflowing_add(len);
-    if overflowed || items_end > inner.len() {
-        return Err(CodecError::LengthPrefixTooBig(len));
-    }
-
-    let mut sub = Cursor::new(&bytes.get_ref()[initial_position..items_end]);
-
-    while sub.position() < len as u64 {
-        decoded.push(I::decode(&mut sub)?);
-    }
-
-    // Advance outer cursor by the amount read in the inner cursor
-    bytes.set_position(initial_position as u64 + sub.position());
-
-    Ok(decoded)
+pub(crate) fn decode_u32_bytes(bytes: &mut Cursor<&[u8]>) -> Result<Vec<u8>, CodecError> {
+    let len = u32::decode(bytes)? as usize;
+    let mut out = vec![0; len];
+    bytes.read_exact(&mut out)?;
+    Ok(out)
 }
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 0aebe44..7b05971 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -7,7 +7,7 @@
 use crate::{
     hpke::HpkeDecrypter,
     messages::{
-        encode_prefixed_bytes, AggregateContinueReq, AggregateInitializeReq, AggregateResp,
+        encode_u32_bytes, AggregateContinueReq, AggregateInitializeReq, AggregateResp,
         BatchParameter, BatchSelector, HpkeCiphertext, HpkeConfig, Id, Nonce, Report,
         ReportMetadata, ReportShare, Time, Transition, TransitionFailure, TransitionVar,
     },
@@ -170,7 +170,7 @@ impl VdafConfig {
         task_id.encode(&mut aad);
         metadata.encode(&mut aad);
         // TODO spec: Consider folding the public share into a field called "header".
-        encode_prefixed_bytes::<4>(&mut aad, &public_share);
+        encode_u32_bytes(&mut aad, &public_share);
 
         let mut encrypted_input_shares = Vec::with_capacity(encoded_input_shares.len());
         for (i, (hpke_config, input_share_data)) in hpke_config_list
@@ -239,7 +239,7 @@ impl VdafConfig {
         task_id.encode(&mut aad);
         metadata.encode(&mut aad);
         // TODO spec: Consider folding the public share into a field called "header".
-        encode_prefixed_bytes::<4>(&mut aad, public_share);
+        encode_u32_bytes(&mut aad, public_share);
 
         let input_share_data = decrypter
             .hpke_decrypt(task_id, &info, &aad, encrypted_input_share)
@@ -749,6 +749,7 @@ impl VdafConfig {
         decrypter: &impl HpkeDecrypter<'_>,
         task_id: &Id,
         batch_selector: &BatchSelector,
+        report_count: u64,
         encrypted_agg_shares: Vec<HpkeCiphertext>,
     ) -> Result<DapAggregateResult, DapError> {
         const N: usize = CTX_AGG_SHARE.len();
@@ -780,9 +781,14 @@ impl VdafConfig {
             ));
         }
 
+        let num_measurements = usize::try_from(report_count).unwrap();
         match self {
-            Self::Prio3(prio3_config) => Ok(prio3_unshard(prio3_config, agg_shares)?),
-            Self::Prio2 { dimension } => Ok(prio2_unshard(*dimension, agg_shares)?),
+            Self::Prio3(prio3_config) => {
+                Ok(prio3_unshard(prio3_config, num_measurements, agg_shares)?)
+            }
+            Self::Prio2 { dimension } => {
+                Ok(prio2_unshard(*dimension, num_measurements, agg_shares)?)
+            }
         }
     }
 }
diff --git a/daphne/src/vdaf/prio2.rs b/daphne/src/vdaf/prio2.rs
index ec3f445..6aa8cbe 100644
--- a/daphne/src/vdaf/prio2.rs
+++ b/daphne/src/vdaf/prio2.rs
@@ -23,7 +23,7 @@ pub(crate) fn prio2_shard(
     measurement: DapMeasurement,
 ) -> Result<Vec<Vec<u8>>, VdafError> {
     let vdaf = Prio2::new(dimension as usize)?;
-    let input_shares = match measurement {
+    let (_public_share, input_shares) = match measurement {
         DapMeasurement::U32Vec(ref data) => vdaf.shard(data)?,
         _ => panic!("prio2_shard: unexpected measurement type"),
     };
@@ -45,7 +45,8 @@ pub(crate) fn prio2_prepare_init(
     let vdaf = Prio2::new(dimension as usize)?;
     let input_share: Share<FieldPrio2, 32> =
         Share::get_decoded_with_param(&(&vdaf, agg_id), input_share_data)?;
-    let (state, share) = vdaf.prepare_init(verify_key, agg_id, &(), nonce_data, &input_share)?;
+    let (state, share) =
+        vdaf.prepare_init(verify_key, agg_id, &(), nonce_data, &(), &input_share)?;
     Ok((VdafState::Prio2(state), VdafMessage::Prio2Share(share)))
 }
 
@@ -120,6 +121,7 @@ pub(crate) fn prio2_encode_prepare_message(message: &VdafMessage) -> Vec<u8> {
 /// Interpret `encoded_agg_shares` as a sequence of encoded aggregate shares and unshard them.
 pub(crate) fn prio2_unshard<M: IntoIterator<Item = Vec<u8>>>(
     dimension: u32,
+    num_measurements: usize,
     encoded_agg_shares: M,
 ) -> Result<DapAggregateResult, VdafError> {
     let vdaf = Prio2::new(dimension as usize)?;
@@ -129,6 +131,6 @@ pub(crate) fn prio2_unshard<M: IntoIterator<Item = Vec<u8>>>(
             .map_err(|e| CodecError::Other(Box::new(e)))?;
         agg_shares.push(agg_share)
     }
-    let agg_res = vdaf.unshard(&(), agg_shares)?;
+    let agg_res = vdaf.unshard(&(), agg_shares, num_measurements)?;
     Ok(DapAggregateResult::U32Vec(agg_res))
 }
diff --git a/daphne/src/vdaf/prio3.rs b/daphne/src/vdaf/prio3.rs
index 0238cf8..aef1aac 100644
--- a/daphne/src/vdaf/prio3.rs
+++ b/daphne/src/vdaf/prio3.rs
@@ -27,7 +27,7 @@ macro_rules! shard {
         $measurement:expr
     ) => {{
         // Split measurement into input shares.
-        let input_shares = $vdaf.shard($measurement)?;
+        let (_public_share, input_shares) = $vdaf.shard($measurement)?;
 
         // Encode input shares.
         input_shares
@@ -72,7 +72,7 @@ macro_rules! prep_init {
             Prio3InputShare::get_decoded_with_param(&(&$vdaf, $agg_id), $input_share_data)?;
 
         // Run the prepare-init algorithm, returning the initial state.
-        $vdaf.prepare_init($verify_key, $agg_id, &(), $nonce_data, &input_share)?
+        $vdaf.prepare_init($verify_key, $agg_id, &(), $nonce_data, &(), &input_share)?
     }};
 }
 
@@ -290,6 +290,7 @@ pub(crate) fn prio3_encode_prepare_message(message: &VdafMessage) -> Vec<u8> {
 macro_rules! unshard {
     (
         $vdaf:ident,
+        $num_measurements:ident,
         $agg_shares:expr
     ) => {{
         let mut agg_shares = Vec::with_capacity($vdaf.num_aggregators());
@@ -298,29 +299,30 @@ macro_rules! unshard {
                 .map_err(|e| CodecError::Other(Box::new(e)))?;
             agg_shares.push(agg_share)
         }
-        $vdaf.unshard(&(), agg_shares)
+        $vdaf.unshard(&(), agg_shares, $num_measurements)
     }};
 }
 
 /// Interpret `agg_shares` as a sequence of encoded aggregate shares and unshard them.
 pub(crate) fn prio3_unshard<M: IntoIterator<Item = Vec<u8>>>(
     config: &Prio3Config,
+    num_measurements: usize,
     agg_shares: M,
 ) -> Result<DapAggregateResult, VdafError> {
     match &config {
         Prio3Config::Count => {
             let vdaf = Prio3::new_aes128_count(2)?;
-            let agg_res = unshard!(vdaf, agg_shares)?;
+            let agg_res = unshard!(vdaf, num_measurements, agg_shares)?;
             Ok(DapAggregateResult::U64(agg_res))
         }
         Prio3Config::Histogram { buckets } => {
             let vdaf = Prio3::new_aes128_histogram(2, buckets)?;
-            let agg_res = unshard!(vdaf, agg_shares)?;
+            let agg_res = unshard!(vdaf, num_measurements, agg_shares)?;
             Ok(DapAggregateResult::U128Vec(agg_res))
         }
         Prio3Config::Sum { bits } => {
             let vdaf = Prio3::new_aes128_sum(2, *bits)?;
-            let agg_res = unshard!(vdaf, agg_shares)?;
+            let agg_res = unshard!(vdaf, num_measurements, agg_shares)?;
             Ok(DapAggregateResult::U128(agg_res))
         }
     }


Test patch:
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index d077e6f..a738479 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -127,7 +127,7 @@ async fn roundtrip_report() {
             let helper_agg_share = vdaf.aggregate(&(), [helper_out_share]).unwrap();
 
             assert_eq!(
-                vdaf.unshard(&(), vec![leader_agg_share, helper_agg_share])
+                vdaf.unshard(&(), vec![leader_agg_share, helper_agg_share], 1)
                     .unwrap(),
                 1,
             );
@@ -323,6 +323,7 @@ async fn agg_cont_req() {
 
     let leader_out_shares = t.handle_final_agg_resp(leader_uncommitted, agg_resp);
     assert_eq!(leader_out_shares.len(), 5);
+    let num_measurements = leader_out_shares.len();
 
     let leader_agg_share = leader_out_shares
         .into_iter()
@@ -350,7 +351,7 @@ async fn agg_cont_req() {
 
     let vdaf = Prio3::new_aes128_count(2).unwrap();
     assert_eq!(
-        vdaf.unshard(&(), [leader_agg_share, helper_agg_share])
+        vdaf.unshard(&(), [leader_agg_share, helper_agg_share], num_measurements,)
             .unwrap(),
         3,
     );
@@ -487,6 +488,7 @@ async fn encrypted_agg_share() {
     let agg_res = t
         .consume_encrypted_agg_shares(
             &batch_selector,
+            50,
             vec![leader_encrypted_agg_share, helper_encrypted_agg_share],
         )
         .await;
@@ -707,6 +709,7 @@ impl<'a> Test<'a> {
     async fn consume_encrypted_agg_shares(
         &self,
         batch_selector: &BatchSelector,
+        report_count: u64,
         enc_agg_shares: Vec<HpkeCiphertext>,
     ) -> DapAggregateResult {
         self.vdaf
@@ -714,6 +717,7 @@ impl<'a> Test<'a> {
                 &self.collector_hpke_receiver_config,
                 &self.task_id,
                 batch_selector,
+                report_count,
                 enc_agg_shares,
             )
             .await
@@ -753,6 +757,7 @@ impl<'a> Test<'a> {
             .handle_agg_cont_req(helper_state, &agg_cont)
             .unwrap_finish();
         let leader_out_shares = self.handle_final_agg_resp(uncommitted, agg_resp);
+        let report_count = u64::try_from(leader_out_shares.len()).unwrap();
 
         // Leader: Aggregation
         let leader_agg_shares =
@@ -773,6 +778,7 @@ impl<'a> Test<'a> {
         // Collector: Unshard
         self.consume_encrypted_agg_shares(
             &batch_selector,
+            report_count,
             vec![leader_encrypted_agg_share, helper_encrypted_agg_share],
         )
         .await
diff --git a/daphne/src/vdaf/prio3_test.rs b/daphne/src/vdaf/prio3_test.rs
index c638ace..8d2332e 100644
--- a/daphne/src/vdaf/prio3_test.rs
+++ b/daphne/src/vdaf/prio3_test.rs
@@ -76,6 +76,7 @@ fn test_prepare(
     // Unshard
     let agg_res = prio3_unshard(
         config,
+        1,
         [
             leader_out_share.get_encoded(),
             helper_out_share.get_encoded(),