Repository: cloudflare/daphne
Pull request #599 (2024-05-20 09:32:22)

Drop bincode in favor of capnproto for DO message serialization


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 2c7b3e1..357101b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -796,7 +796,6 @@ dependencies = [
  "anyhow",
  "assert_matches",
  "axum",
- "bincode",
  "clap",
  "config",
  "daphne",
@@ -839,6 +838,7 @@ dependencies = [
  "hex",
  "itertools 0.12.1",
  "p256",
+ "prio",
  "prometheus",
  "rand",
  "rayon",
diff --git a/Cargo.toml b/Cargo.toml
index 0b79c6b..bbefea6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -34,6 +34,7 @@ assert_matches = "1.5.0"
 async-trait = "0.1.77"
 axum = "0.6"
 base64 = "0.21.7"
+# TODO(mendess): delete. Left here for legacy_req_parse.
 bincode = "1.3.3"
 cap = "0.1.2"
 capnp = "0.18.13"
diff --git a/crates/daphne-server/Cargo.toml b/crates/daphne-server/Cargo.toml
index b74025e..30a964c 100644
--- a/crates/daphne-server/Cargo.toml
+++ b/crates/daphne-server/Cargo.toml
@@ -14,7 +14,6 @@ description = "Workers backend for Daphne"
 
 [dependencies]
 axum.workspace = true
-bincode.workspace = true
 daphne = { path = "../daphne" }
 daphne-service-utils = { path = "../daphne-service-utils", features = ["durable_requests"] }
 futures.workspace = true
diff --git a/crates/daphne-server/src/roles/aggregator.rs b/crates/daphne-server/src/roles/aggregator.rs
index 0823324..cad93fb 100644
--- a/crates/daphne-server/src/roles/aggregator.rs
+++ b/crates/daphne-server/src/roles/aggregator.rs
@@ -45,7 +45,7 @@ impl DapAggregator<DaphneAuth> for crate::App {
                         bindings::AggregateStore::Merge,
                         (task_config.version, &task_id_hex, &bucket),
                     )
-                    .encode_bincode(AggregateStoreMergeReq {
+                    .encode(&AggregateStoreMergeReq {
                         contained_reports: report_metadatas.iter().map(|(id, _)| *id).collect(),
                         agg_share_delta: agg_share,
                     })
diff --git a/crates/daphne-server/src/roles/helper.rs b/crates/daphne-server/src/roles/helper.rs
index c4c816f..5d07a16 100644
--- a/crates/daphne-server/src/roles/helper.rs
+++ b/crates/daphne-server/src/roles/helper.rs
@@ -2,73 +2,8 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use axum::async_trait;
-use daphne::{
-    error::DapAbort,
-    fatal_error,
-    messages::{AggregationJobId, TaskId},
-    roles::{DapAggregator, DapHelper},
-    DapAggregationJobState, DapError,
-};
-use daphne_service_utils::{auth::DaphneAuth, durable_requests::bindings};
-use prio::codec::Encode;
+use daphne::roles::DapHelper;
+use daphne_service_utils::auth::DaphneAuth;
 
 #[async_trait]
-impl DapHelper<DaphneAuth> for crate::App {
-    async fn put_helper_state_if_not_exists(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-        helper_state: &DapAggregationJobState,
-    ) -> Result<bool, DapError> {
-        let task_config = self
-            .get_task_config_for(task_id)
-            .await?
-            .ok_or(DapError::Abort(DapAbort::UnrecognizedTask))?;
-        let helper_state_hex = hex::encode(helper_state.get_encoded().map_err(DapError::encoding)?);
-        Ok(self
-            .durable()
-            .with_retry()
-            .request(
-                bindings::HelperState::PutIfNotExists,
-                (task_config.as_ref().version, task_id, agg_job_id),
-            )
-            .encode_bincode(helper_state_hex)
-            .send()
-            .await
-            .map_err(|e| fatal_error!(err = ?e))?)
-    }
-
-    async fn get_helper_state(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-    ) -> Result<Option<DapAggregationJobState>, DapError> {
-        let task_config = self
-            .get_task_config_for(task_id)
-            .await?
-            .ok_or(DapError::Abort(DapAbort::UnrecognizedTask))?;
-        // TODO(cjpatton) Figure out if retry is safe, since the request is not actually
-        // idempotent. (It removes the helper's state from storage if it exists.)
-        let res: Option<String> = self
-            .durable()
-            .with_retry()
-            .request(
-                bindings::HelperState::Get,
-                (task_config.as_ref().version, task_id, agg_job_id),
-            )
-            .send()
-            .await
-            .map_err(|e| fatal_error!(err = ?e))?;
-
-        match res {
-            Some(helper_state_hex) => {
-                let data = hex::decode(helper_state_hex)
-                    .map_err(|e| DapAbort::from_hex_error(e, *task_id))?;
-                let helper_state =
-                    DapAggregationJobState::get_decoded(&task_config.as_ref().vdaf, &data)?;
-                Ok(Some(helper_state))
-            }
-            None => Ok(None),
-        }
-    }
-}
+impl DapHelper<DaphneAuth> for crate::App {}
diff --git a/crates/daphne-server/src/storage_proxy_connection/mod.rs b/crates/daphne-server/src/storage_proxy_connection/mod.rs
index 130940d..5f4c6b5 100644
--- a/crates/daphne-server/src/storage_proxy_connection/mod.rs
+++ b/crates/daphne-server/src/storage_proxy_connection/mod.rs
@@ -11,9 +11,10 @@ use std::fmt::Debug;
 
 use axum::http::{Method, StatusCode};
 use daphne_service_utils::durable_requests::{
-    bindings::DurableMethod, DurableRequest, ObjectIdFrom, DO_PATH_PREFIX,
+    bindings::{DurableMethod, DurableRequestPayload, DurableRequestPayloadExt},
+    DurableRequest, ObjectIdFrom, DO_PATH_PREFIX,
 };
-use serde::{de::DeserializeOwned, Serialize};
+use serde::de::DeserializeOwned;
 
 pub(crate) use kv::Kv;
 
@@ -96,8 +97,8 @@ impl<'d, B: DurableMethod + Debug, P: AsRef<[u8]>> RequestBuilder<'d, B, P> {
 }
 
 impl<'d, B: DurableMethod> RequestBuilder<'d, B, [u8; 0]> {
-    pub fn encode_bincode<T: Serialize>(self, payload: T) -> RequestBuilder<'d, B, Vec<u8>> {
-        self.with_body(bincode::serialize(&payload).unwrap())
+    pub fn encode<T: DurableRequestPayload>(self, payload: &T) -> RequestBuilder<'d, B, Vec<u8>> {
+        self.with_body(payload.encode_to_bytes().unwrap())
     }
 
     pub fn with_body<T: AsRef<[u8]>>(self, payload: T) -> RequestBuilder<'d, B, T> {
diff --git a/crates/daphne-service-utils/Cargo.toml b/crates/daphne-service-utils/Cargo.toml
index 3cedbba..f0f421e 100644
--- a/crates/daphne-service-utils/Cargo.toml
+++ b/crates/daphne-service-utils/Cargo.toml
@@ -20,6 +20,7 @@ itertools.workspace = true
 hex.workspace = true
 p256.workspace = true
 prometheus = { workspace = true, optional = true }
+prio.workspace = true
 ring.workspace = true
 serde.workspace = true
 serde_json.workspace = true
diff --git a/crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs b/crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs
new file mode 100644
index 0000000..e7eb98e
--- /dev/null
+++ b/crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs
@@ -0,0 +1,265 @@
+// Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use std::collections::HashSet;
+
+use daphne::{
+    messages::ReportId, vdaf::VdafAggregateShare, DapAggregateShare, DapBatchBucket, DapVersion,
+};
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    durable_request_capnp::{aggregate_store_merge_req, dap_aggregate_share},
+    durable_requests::ObjectIdFrom,
+};
+
+use super::DurableRequestPayload;
+
+super::define_do_binding! {
+    const BINDING = "DAP_AGGREGATE_STORE";
+    enum AggregateStore {
+        GetMerged = "/internal/do/aggregate_store/get_merged",
+        Get = "/internal/do/aggregate_store/get",
+        Merge = "/internal/do/aggregate_store/merge",
+        MarkCollected = "/internal/do/aggregate_store/mark_collected",
+        CheckCollected = "/internal/do/aggregate_store/check_collected",
+    }
+
+    fn name((version, task_id_hex, bucket): (DapVersion, &'n str, &'n DapBatchBucket)) -> ObjectIdFrom {
+        fn durable_name_bucket(bucket: &DapBatchBucket) -> String {
+            format!("{bucket}")
+        }
+        ObjectIdFrom::Name(format!(
+            "{}/{}",
+            durable_name_task(version, task_id_hex),
+            durable_name_bucket(bucket),
+        ))
+    }
+}
+
+fn durable_name_task(version: DapVersion, task_id_hex: &str) -> String {
+    format!("{}/task/{}", version.as_ref(), task_id_hex)
+}
+
+#[derive(Debug, PartialEq, Eq)]
+// TODO(mendess): delete. Only here to suport deserialization with bincode for backwards
+// compatibility.
+#[derive(serde::Deserialize)]
+pub struct AggregateStoreMergeReq {
+    pub contained_reports: Vec<ReportId>,
+    pub agg_share_delta: DapAggregateShare,
+}
+
+impl DurableRequestPayload for AggregateStoreMergeReq {
+    fn encode_to_builder(&self) -> capnp::message::Builder<capnp::message::HeapAllocator> {
+        let mut message = capnp::message::Builder::new_default();
+        let mut request = message.init_root::<aggregate_store_merge_req::Builder>();
+        {
+            let mut contained_reports = request.reborrow().init_contained_reports(
+                self.contained_reports
+                    .len()
+                    .try_into()
+                    .expect("can't serialize more than u32::MAX reports"),
+            );
+            for (i, r) in self.contained_reports.iter().enumerate() {
+                let mut report = contained_reports.reborrow().get(i.try_into().unwrap());
+                let bytes: &[u8; 16] = r.as_ref();
+                let low = u64::from_le_bytes(bytes[..8].try_into().unwrap());
+                report.set_low(low);
+                let high = u64::from_le_bytes(bytes[8..].try_into().unwrap());
+                report.set_high(high);
+            }
+        }
+        {
+            let mut agg_share_delta = request.init_agg_share_delta();
+            agg_share_delta.set_report_count(self.agg_share_delta.report_count);
+            agg_share_delta.set_min_time(self.agg_share_delta.min_time);
+            agg_share_delta.set_max_time(self.agg_share_delta.max_time);
+            {
+                let checksum = agg_share_delta
+                    .reborrow()
+                    .init_checksum(self.agg_share_delta.checksum.len().try_into().unwrap());
+                checksum.copy_from_slice(&self.agg_share_delta.checksum);
+            }
+            {
+                fn encode<'b, F, B, const ENCODED_SIZE: usize>(
+                    field: &prio::vdaf::AggregateShare<F>,
+                    get_bytes: B,
+                ) where
+                    F: prio::field::FieldElement + Into<[u8; ENCODED_SIZE]>,
+                    B: FnOnce(u32) -> &'b mut [u8],
+                {
+                    let mut bytes = get_bytes(
+                        (F::ENCODED_SIZE * field.as_ref().len())
+                            .try_into()
+                            .expect("trying to encode a buffer longer than u32::MAX"),
+                    );
+                    for f in field.as_ref() {
+                        let f: [u8; ENCODED_SIZE] = (*f).into();
+                        bytes[..ENCODED_SIZE].copy_from_slice(&f);
+                        bytes = &mut bytes[ENCODED_SIZE..];
+                    }
+                }
+                let mut data = agg_share_delta.init_data();
+                match &self.agg_share_delta.data {
+                    Some(VdafAggregateShare::Field64(field)) => {
+                        encode(field, |len| data.init_field64(len));
+                    }
+                    Some(VdafAggregateShare::Field128(field)) => {
+                        encode(field, |len| data.init_field128(len));
+                    }
+                    Some(VdafAggregateShare::FieldPrio2(field)) => {
+                        encode(field, |len| data.init_field_prio2(len));
+                    }
+                    None => data.set_none(()),
+                };
+            }
+        }
+        message
+    }
+
+    fn decode_from_reader(
+        reader: capnp::message::Reader<capnp::serialize::OwnedSegments>,
+    ) -> capnp::Result<Self> {
+        let request = reader.get_root::<aggregate_store_merge_req::Reader>()?;
+        let agg_share_delta = {
+            let agg_share_delta = request.get_agg_share_delta()?;
+            let data = {
+                fn decode<F>(fields: &[u8]) -> capnp::Result<prio::vdaf::AggregateShare<F>>
+                where
+                    F: prio::field::FieldElement
+                        + for<'s> TryFrom<&'s [u8], Error = prio::field::FieldError>,
+                {
+                    let iter = fields.chunks_exact(F::ENCODED_SIZE);
+                    if let length @ 1.. = iter.remainder().len() {
+                        return Err(capnp::Error {
+                            kind: capnp::ErrorKind::Failed,
+                            extra: format!("leftover bytes still present in buffer: {length}"),
+                        });
+                    }
+                    Ok(prio::vdaf::AggregateShare::from(
+                        iter.map(|f| f.try_into().unwrap()).collect::<Vec<_>>(),
+                    ))
+                }
+                match agg_share_delta.get_data().which()? {
+                    dap_aggregate_share::data::Which::Field64(field) => {
+                        Some(VdafAggregateShare::Field64(decode(field?)?))
+                    }
+                    dap_aggregate_share::data::Which::Field128(field) => {
+                        Some(VdafAggregateShare::Field128(decode(field?)?))
+                    }
+                    dap_aggregate_share::data::Which::FieldPrio2(field) => {
+                        Some(VdafAggregateShare::FieldPrio2(decode(field?)?))
+                    }
+                    dap_aggregate_share::data::Which::None(()) => None,
+                }
+            };
+            DapAggregateShare {
+                report_count: agg_share_delta.get_report_count(),
+                min_time: agg_share_delta.get_min_time(),
+                max_time: agg_share_delta.get_max_time(),
+                checksum: agg_share_delta
+                    .get_checksum()?
+                    .try_into()
+                    .map_err(|_| capnp::Error {
+                        kind: capnp::ErrorKind::Failed,
+                        extra: "checksum had unexpected size".into(),
+                    })?,
+                data,
+            }
+        };
+        Ok(Self {
+            contained_reports: request
+                .get_contained_reports()?
+                .into_iter()
+                .map(|report| {
+                    let low = report.get_low();
+                    let high = report.get_high();
+
+                    let mut buffer = [0; 16];
+                    buffer[..8].copy_from_slice(&low.to_le_bytes());
+                    buffer[8..].copy_from_slice(&high.to_le_bytes());
+                    ReportId(buffer)
+                })
+                .collect(),
+            agg_share_delta,
+        })
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub enum AggregateStoreMergeResp {
+    Ok,
+    ReplaysDetected(HashSet<ReportId>),
+    AlreadyCollected,
+}
+
+#[cfg(test)]
+mod test {
+    use prio::{
+        codec::Decode,
+        field::{Field128, Field64, FieldElement, FieldPrio2},
+        vdaf::AggregateShare,
+    };
+    use rand::{thread_rng, Rng};
+
+    use crate::durable_requests::bindings::DurableRequestPayloadExt;
+
+    use super::*;
+
+    #[test]
+    fn serialization_deserialization_round_trip() {
+        let mut rng = thread_rng();
+        for len in 0..20 {
+            let test_data = [
+                VdafAggregateShare::Field64(AggregateShare::from(
+                    (0..len)
+                        .map(|_| {
+                            Field64::get_decoded(&rng.gen::<[_; Field64::ENCODED_SIZE]>()).unwrap()
+                        })
+                        .collect::<Vec<_>>(),
+                )),
+                VdafAggregateShare::Field128(AggregateShare::from(
+                    (0..len)
+                        .map(|_| {
+                            Field128::get_decoded(&rng.gen::<[_; Field128::ENCODED_SIZE]>())
+                                .unwrap()
+                        })
+                        .collect::<Vec<_>>(),
+                )),
+                VdafAggregateShare::FieldPrio2(AggregateShare::from(
+                    (0..len)
+                        .map(|_| {
+                            // idk how to consistently generate a valid FieldPrio2 value, so I just
+                            // retry until I hit a valid one. Doesn't usualy take too long.
+                            (0..)
+                                .find_map(|_| FieldPrio2::get_decoded(&rng.gen::<[_; 4]>()).ok())
+                                .unwrap()
+                        })
+                        .collect::<Vec<_>>(),
+                )),
+            ]
+            .map(Some)
+            .into_iter()
+            .chain([None]);
+            for data in test_data {
+                let this = AggregateStoreMergeReq {
+                    contained_reports: (0..len)
+                        .map(|_| ReportId::get_decoded(&rng.gen::<[_; 16]>()).unwrap())
+                        .collect(),
+                    agg_share_delta: DapAggregateShare {
+                        report_count: rng.gen(),
+                        min_time: rng.gen(),
+                        max_time: rng.gen(),
+                        checksum: rng.gen(),
+                        data,
+                    },
+                };
+                let other =
+                    AggregateStoreMergeReq::decode_from_bytes(&this.encode_to_bytes().unwrap())
+                        .unwrap();
+                assert_eq!(this, other);
+            }
+        }
+    }
+}
diff --git a/crates/daphne-service-utils/src/durable_requests/bindings.rs b/crates/daphne-service-utils/src/durable_requests/bindings/mod.rs
similarity index 58%
rename from crates/daphne-service-utils/src/durable_requests/bindings.rs
rename to crates/daphne-service-utils/src/durable_requests/bindings/mod.rs
index d5a5db7..7115a18 100644
--- a/crates/daphne-service-utils/src/durable_requests/bindings.rs
+++ b/crates/daphne-service-utils/src/durable_requests/bindings/mod.rs
@@ -6,16 +6,16 @@
 //!
 //! It also defines types that are used as the body of requests sent to these objects.
 
-use std::collections::HashSet;
-
-use daphne::{
-    messages::{AggregationJobId, ReportId, TaskId},
-    DapAggregateShare, DapBatchBucket, DapVersion,
-};
-use serde::{Deserialize, Serialize};
+mod aggregate_store;
+#[cfg(feature = "test-utils")]
+mod test_state_cleaner;
 
 use super::ObjectIdFrom;
 
+pub use aggregate_store::{AggregateStore, AggregateStoreMergeReq, AggregateStoreMergeResp};
+#[cfg(feature = "test-utils")]
+pub use test_state_cleaner::TestStateCleaner;
+
 /// A durable object method.
 pub trait DurableMethod {
     /// The binding of the object this method belongs to as configured in the wrangler.toml
@@ -35,6 +35,48 @@ pub trait DurableMethod {
     fn name(params: Self::NameParameters<'_>) -> ObjectIdFrom;
 }
 
+pub trait DurableRequestPayload {
+    fn decode_from_reader(
+        reader: capnp::message::Reader<capnp::serialize::OwnedSegments>,
+    ) -> capnp::Result<Self>
+    where
+        Self: Sized;
+
+    fn encode_to_builder(&self) -> capnp::message::Builder<capnp::message::HeapAllocator>;
+}
+
+pub trait DurableRequestPayloadExt {
+    fn decode_from_bytes(bytes: &[u8]) -> capnp::Result<Self>
+    where
+        Self: Sized;
+    fn encode_to_bytes(&self) -> capnp::Result<Vec<u8>>;
+}
+
+impl<T> DurableRequestPayloadExt for T
+where
+    T: DurableRequestPayload,
+{
+    fn encode_to_bytes(&self) -> capnp::Result<Vec<u8>> {
+        let mut buf = Vec::new();
+        let message = self.encode_to_builder();
+        capnp::serialize_packed::write_message(&mut buf, &message)?;
+        Ok(buf)
+    }
+
+    fn decode_from_bytes(bytes: &[u8]) -> capnp::Result<Self>
+    where
+        Self: Sized,
+    {
+        let mut cursor = std::io::Cursor::new(bytes);
+        let reader = capnp::serialize_packed::read_message(
+            &mut cursor,
+            capnp::message::ReaderOptions::new(),
+        )?;
+
+        T::decode_from_reader(reader)
+    }
+}
+
 macro_rules! define_do_binding {
     (
         const BINDING = $binding:literal;
@@ -61,7 +103,7 @@ macro_rules! define_do_binding {
             $($op),*
         }
 
-        impl DurableMethod for $name {
+        impl $crate::durable_requests::bindings::DurableMethod for $name {
             const BINDING: &'static str = $binding;
 
             type NameParameters<'n> = $params_ty;
@@ -79,87 +121,14 @@ macro_rules! define_do_binding {
                 }
             }
 
-            fn name($params: Self::NameParameters<'_>) -> ObjectIdFrom {
+            fn name($params: Self::NameParameters<'_>) -> $crate::durable_requests::bindings::ObjectIdFrom {
                 $name_impl
             }
         }
     };
 }
 
-define_do_binding! {
-    const BINDING = "DAP_AGGREGATE_STORE";
-    enum AggregateStore {
-        GetMerged = "/internal/do/aggregate_store/get_merged",
-        Get = "/internal/do/aggregate_store/get",
-        Merge = "/internal/do/aggregate_store/merge",
-        MarkCollected = "/internal/do/aggregate_store/mark_collected",
-        CheckCollected = "/internal/do/aggregate_store/check_collected",
-    }
-
-    fn name((version, task_id_hex, bucket): (DapVersion, &'n str, &'n DapBatchBucket)) -> ObjectIdFrom {
-        fn durable_name_bucket(bucket: &DapBatchBucket) -> String {
-            format!("{bucket}")
-        }
-        ObjectIdFrom::Name(format!(
-            "{}/{}",
-            durable_name_task(version, task_id_hex),
-            durable_name_bucket(bucket),
-        ))
-    }
-}
-
-fn durable_name_task(version: DapVersion, task_id_hex: &str) -> String {
-    format!("{}/task/{}", version.as_ref(), task_id_hex)
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct AggregateStoreMergeReq {
-    pub contained_reports: Vec<ReportId>,
-    pub agg_share_delta: DapAggregateShare,
-}
-
-#[derive(Debug, Serialize, Deserialize)]
-pub enum AggregateStoreMergeResp {
-    Ok,
-    ReplaysDetected(HashSet<ReportId>),
-    AlreadyCollected,
-}
-
-define_do_binding! {
-    const BINDING = "DAP_HELPER_STATE_STORE";
-    enum HelperState {
-        PutIfNotExists = "/internal/do/helper_state/put_if_not_exists",
-        Get = "/internal/do/helper_state/get",
-    }
-
-    fn name((version, task_id, agg_job_id): (DapVersion, &'n TaskId, &'n AggregationJobId)) -> ObjectIdFrom {
-        ObjectIdFrom::Name(format!(
-            "{}/task/{}/agg_job/{}",
-            version.as_ref(),
-            task_id.to_hex(),
-            agg_job_id.to_hex()
-        ))
-    }
-
-}
-
-#[cfg(feature = "test-utils")]
-define_do_binding! {
-    const BINDING = "DAP_TEST_STATE_CLEANER";
-    enum TestStateCleaner {
-        Put = "/internal/do/test_state_cleaner/put",
-        DeleteAll = "/internal/do/delete_all",
-    }
-
-    fn name((): ()) -> ObjectIdFrom {
-        ObjectIdFrom::Name(Self::NAME_STR.into())
-    }
-}
-
-#[cfg(feature = "test-utils")]
-impl TestStateCleaner {
-    pub const NAME_STR: &'static str = "test_do_cleaner";
-}
+pub(crate) use define_do_binding;
 
 #[cfg(test)]
 mod tests {
diff --git a/crates/daphne-service-utils/src/durable_requests/durable_request.capnp b/crates/daphne-service-utils/src/durable_requests/durable_request.capnp
index d1aac34..dc32e94 100644
--- a/crates/daphne-service-utils/src/durable_requests/durable_request.capnp
+++ b/crates/daphne-service-utils/src/durable_requests/durable_request.capnp
@@ -3,7 +3,7 @@
 
 @0xd076f8051f8de41a;
 
-enum Method {
+enum Method @0xdd078556311145a1 {
     get @0;
     post @1;
     put @2;
@@ -23,3 +23,28 @@ struct DurableRequest @0xfbd55b93d47690b9 {
     }
     retry @3 :Bool;
 }
+
+struct UInt128 @0x8f8f75793b90ee0c {
+    low @0 :UInt64;
+    high @1 :UInt64;
+}
+
+struct DapAggregateShare @0xb34ce529a4a66aed {
+    reportCount @0 :UInt64;
+    minTime @1 :UInt64;
+    maxTime @2 :UInt64;
+    checksum @3 :Data;
+    data :union {
+        field64 @4 :Data;
+        field128 @5 :Data;
+        fieldPrio2 @6 :Data;
+        none @7 :Void;
+    }
+}
+
+struct AggregateStoreMergeReq @0xbaad7bdeb4b06161 {
+    using ReportId = UInt128;
+
+    containedReports @0 :List(ReportId);
+    aggShareDelta @1 :DapAggregateShare;
+}
diff --git a/crates/daphne-worker/Cargo.toml b/crates/daphne-worker/Cargo.toml
index f72904a..87bbdce 100644
--- a/crates/daphne-worker/Cargo.toml
+++ b/crates/daphne-worker/Cargo.toml
@@ -16,6 +16,7 @@ description = "Workers backend for Daphne"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
+# TODO(mendess): delete
 bincode.workspace = true
 chrono = { workspace = true, default-features = false, features = ["clock", "wasmbind"] }
 daphne = { path = "../daphne", features = ["prometheus"] }
diff --git a/crates/daphne-worker/src/durable/aggregate_store.rs b/crates/daphne-worker/src/durable/aggregate_store.rs
index 55105d0..f41da52 100644
--- a/crates/daphne-worker/src/durable/aggregate_store.rs
+++ b/crates/daphne-worker/src/durable/aggregate_store.rs
@@ -306,7 +306,7 @@ fn shard_bytes_to_object(
 
 super::mk_durable_object! {
     /// Where the aggregate share is stored. For the binding name see its
-    /// [`BINDING`](bindings::HelperState::BINDING)
+    /// [`BINDING`](bindings::AggregateStore::BINDING)
     struct AggregateStore {
         state: State,
         env: Env,
diff --git a/crates/daphne-worker/src/durable/helper_state_store.rs b/crates/daphne-worker/src/durable/helper_state_store.rs
deleted file mode 100644
index cac4b3b..0000000
--- a/crates/daphne-worker/src/durable/helper_state_store.rs
+++ /dev/null
@@ -1,90 +0,0 @@
-// Copyright (c) 2022 Cloudflare, Inc. All rights reserved.
-// SPDX-License-Identifier: BSD-3-Clause
-
-//! Durable Object (DO) for storing the Helper's state for a given aggregation job.
-//!
-//! This object implements the following API endpoints:
-//!
-//! - `DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS`: Stores Helper's hex-encoded state unless the state
-//!    already exists. Returns a boolean indicating whether the operation succeeded.
-//! - `DURABLE_HELPER_STATE_GET`: Drains the Helper's hex-encoded state.
-//!
-//! The state blob is stored in `helper_state`.
-
-use std::{sync::OnceLock, time::Duration};
-
-use crate::int_err;
-use daphne_service_utils::durable_requests::bindings::{self, DurableMethod};
-use worker::{Env, Request, Response, Result, ScheduledTime, State};
-
-use super::{req_parse, GcDurableObject};
-
-super::mk_durable_object! {
-    /// Where the helper state is stored. For the binding name see its
-    /// [`BINDING`](bindings::HelperState::BINDING)
-    struct HelperStateStore {
-        state: State,
-        env: Env,
-    }
-}
-
-impl GcDurableObject for HelperStateStore {
-    type DurableMethod = bindings::HelperState;
-
-    fn with_state_and_env(state: State, env: Env) -> Self {
-        Self { state, env }
-    }
-
-    async fn handle(&mut self, mut req: Request) -> Result<Response> {
-        match bindings::HelperState::try_from_uri(&req.path()) {
-            // Store the Helper's state.
-            //
-            // Non-idempotent
-            // Input: `helper_state_hex: String` (hex-encoded state)
-            // Output: `bool`
-            Some(bindings::HelperState::PutIfNotExists) => {
-                let helper_state_hex: String = req_parse(&mut req).await?;
-                let success = self
-                    .set_if_not_exists("helper_state", &helper_state_hex)
-                    .await?
-                    .is_none();
-                Response::from_json(&success)
-            }
-
-            // Get the Helper's state.
-            //
-            // Idempotent
-            // Output: `String` (hex-encoded state)
-            Some(bindings::HelperState::Get) => {
-                let helper_state: Option<String> = self.get("helper_state").await?;
-                Response::from_json(&helper_state)
-            }
-
-            _ => Err(int_err(format!(
-                "HelperStateStore: unexpected request: method={:?}; path={:?}",
-                req.method(),
-                req.path()
-            ))),
-        }
-    }
-
-    fn should_cleanup_at(&self) -> Option<ScheduledTime> {
-        const VAR_NAME: &str = "DAP_DURABLE_HELPER_STATE_STORE_GC_AFTER_SECS";
-        static SELF_DELETE_AFTER: OnceLock<Duration> = OnceLock::new();
-
-        let duration = SELF_DELETE_AFTER.get_or_init(|| {
-            Duration::from_secs(
-                self.env
-                    .var(VAR_NAME)
-                    .map(|v| {
-                        v.to_string().parse().unwrap_or_else(|e| {
-                            panic!("{VAR_NAME} could not be parsed as a number of seconds: {e}")
-                        })
-                    })
-                    .unwrap_or(60 * 60 * 24 * 7), // one week
-            )
-        });
-
-        Some(ScheduledTime::from(*duration))
-    }
-}
diff --git a/crates/daphne-worker/src/durable/mod.rs b/crates/daphne-worker/src/durable/mod.rs
index bb9fb8d..f39a1fe 100644
--- a/crates/daphne-worker/src/durable/mod.rs
+++ b/crates/daphne-worker/src/durable/mod.rs
@@ -20,18 +20,18 @@
 //! this module as well as the [`instantiate_durable_object`] macro, respectively.
 
 pub(crate) mod aggregate_store;
-pub(crate) mod helper_state_store;
 #[cfg(feature = "test-utils")]
 pub(crate) mod test_state_cleaner;
 
 use crate::tracing_utils::shorten_paths;
-use daphne_service_utils::durable_requests::bindings::DurableMethod;
-use serde::{de::DeserializeOwned, Deserialize, Serialize};
+use daphne_service_utils::durable_requests::bindings::{
+    DurableMethod, DurableRequestPayload, DurableRequestPayloadExt,
+};
+use serde::{Deserialize, Serialize};
 use tracing::info_span;
 use worker::{Env, Error, Request, Response, Result, ScheduledTime, State};
 
 pub use aggregate_store::AggregateStore;
-pub use helper_state_store::HelperStateStore;
 
 const ERR_NO_VALUE: &str = "No such value in storage.";
 
@@ -204,10 +204,22 @@ pub(crate) async fn state_set_if_not_exists<T: for<'a> Deserialize<'a> + Seriali
     Ok(None)
 }
 
-async fn req_parse<T: DeserializeOwned>(req: &mut Request) -> Result<T> {
+async fn req_parse<T>(req: &mut Request) -> Result<T>
+where
+    T: DurableRequestPayload,
+    // TODO(mendess): delete
+    T: serde::de::DeserializeOwned,
+{
     let bytes = req.bytes().await?;
-    bincode::deserialize(&bytes)
-        .map_err(|e| Error::RustError(format!("failed to deserialize bincode: {e:?}")))
+    T::decode_from_bytes(&bytes)
+        .map_err(|e| Error::RustError(e.to_string()))
+        .or_else(|cap_error| {
+            bincode::deserialize(&bytes).map_err(|e| {
+                Error::RustError(format!(
+                    "failed to deserialize capnproto ({cap_error:?}) and bincode ({e:?})"
+                ))
+            })
+        })
 }
 
 fn create_span_from_request(req: &Request) -> tracing::Span {
diff --git a/crates/daphne/src/lib.rs b/crates/daphne/src/lib.rs
index 4a46b4d..e044756 100644
--- a/crates/daphne/src/lib.rs
+++ b/crates/daphne/src/lib.rs
@@ -981,7 +981,7 @@ impl DapAggregationJobState {
 }
 
 /// An aggregate share computed by combining a set of output shares.
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
 #[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
 pub struct DapAggregateShare {
     /// Number of reports in the batch.
diff --git a/crates/daphne/src/messages/mod.rs b/crates/daphne/src/messages/mod.rs
index 505e8d0..f446c39 100644
--- a/crates/daphne/src/messages/mod.rs
+++ b/crates/daphne/src/messages/mod.rs
@@ -93,6 +93,12 @@ macro_rules! id_struct {
             }
         }
 
+        impl AsRef<[u8; $len]> for $sname {
+            fn as_ref(&self) -> &[u8; $len] {
+                &self.0
+            }
+        }
+
         impl fmt::Display for $sname {
             fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                 write!(f, "{}", self.to_hex())
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index 9776592..cefc902 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -13,41 +13,19 @@ use crate::{
     constants::DapMediaType,
     error::DapAbort,
     messages::{
-        constant_time_eq, AggregateShare, AggregateShareReq, AggregationJobId,
-        AggregationJobInitReq, AggregationJobResp, PartialBatchSelector, TaskId, TransitionFailure,
-        TransitionVar,
+        constant_time_eq, AggregateShare, AggregateShareReq, AggregationJobInitReq,
+        AggregationJobResp, PartialBatchSelector, TaskId, TransitionFailure, TransitionVar,
     },
     metrics::{DaphneMetrics, DaphneRequestType, ReportStatus},
     protocol::aggregator::ReportProcessedStatus,
     roles::aggregator::MergeAggShareError,
-    DapAggregationJobState, DapAggregationParam, DapError, DapRequest, DapResource, DapResponse,
-    DapTaskConfig, EarlyReportStateInitialized,
+    DapAggregationParam, DapError, DapRequest, DapResource, DapResponse, DapTaskConfig,
+    EarlyReportStateInitialized,
 };
 
 /// DAP Helper functionality.
-//
-// TODO draft02 cleanup: Remove methods for mutating aggregation job state. This is not needed
-// until we add support for multi-round VDAFs. Note that we will likely still need this trait for
-// features coming in the next draft.
 #[async_trait]
-pub trait DapHelper<S: Sync>: DapAggregator<S> {
-    /// Store the Helper's aggregation-flow state unless it already exists. Returns a boolean
-    /// indicating if the operation succeeded.
-    async fn put_helper_state_if_not_exists(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-        helper_state: &DapAggregationJobState,
-    ) -> Result<bool, DapError>;
-
-    /// Fetch the Helper's aggregation-flow state. `None` is returned if the Helper has no state
-    /// associated with the given task and aggregation job.
-    async fn get_helper_state(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-    ) -> Result<Option<DapAggregationJobState>, DapError>;
-}
+pub trait DapHelper<S: Sync>: DapAggregator<S> {}
 
 pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     aggregator: &A,
diff --git a/crates/daphne/src/vdaf/mod.rs b/crates/daphne/src/vdaf/mod.rs
index f9d8e0c..bc4dd20 100644
--- a/crates/daphne/src/vdaf/mod.rs
+++ b/crates/daphne/src/vdaf/mod.rs
@@ -324,7 +324,7 @@ impl ParameterizedDecode<VdafPrepState> for VdafPrepMessage {
     }
 }
 
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
 #[serde(rename_all = "snake_case")]
 pub enum VdafAggregateShare {
     Field64(prio::vdaf::AggregateShare<Field64>),


Test patch:
diff --git a/crates/daphne-service-utils/src/durable_requests/bindings/test_state_cleaner.rs b/crates/daphne-service-utils/src/durable_requests/bindings/test_state_cleaner.rs
new file mode 100644
index 0000000..312fd97
--- /dev/null
+++ b/crates/daphne-service-utils/src/durable_requests/bindings/test_state_cleaner.rs
@@ -0,0 +1,20 @@
+// Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use crate::durable_requests::ObjectIdFrom;
+
+super::define_do_binding! {
+    const BINDING = "DAP_TEST_STATE_CLEANER";
+    enum TestStateCleaner {
+        Put = "/internal/do/test_state_cleaner/put",
+        DeleteAll = "/internal/do/delete_all",
+    }
+
+    fn name((): ()) -> ObjectIdFrom {
+        ObjectIdFrom::Name(Self::NAME_STR.into())
+    }
+}
+
+impl TestStateCleaner {
+    pub const NAME_STR: &'static str = "test_do_cleaner";
+}
diff --git a/crates/daphne-worker-test/src/durable.rs b/crates/daphne-worker-test/src/durable.rs
index 56896a6..72b50ba 100644
--- a/crates/daphne-worker-test/src/durable.rs
+++ b/crates/daphne-worker-test/src/durable.rs
@@ -10,11 +10,3 @@ instantiate_durable_object! {
         daphne_worker::tracing_utils::initialize_tracing(env);
     }
 }
-
-instantiate_durable_object! {
-    struct HelperStateStore < durable::HelperStateStore;
-
-    fn pre_init(_state: State, env: Env) {
-        daphne_worker::tracing_utils::initialize_tracing(env);
-    }
-}
diff --git a/crates/daphne-worker-test/wrangler.storage-proxy.toml b/crates/daphne-worker-test/wrangler.storage-proxy.toml
index 77c616c..e46c0f0 100644
--- a/crates/daphne-worker-test/wrangler.storage-proxy.toml
+++ b/crates/daphne-worker-test/wrangler.storage-proxy.toml
@@ -35,7 +35,6 @@ ip = "0.0.0.0"
 bindings = [
     { name = "DAP_AGGREGATE_STORE", class_name = "AggregateStore" },
     { name = "DAP_TEST_STATE_CLEANER", class_name = "TestStateCleaner" },
-    { name = "DAP_HELPER_STATE_STORE", class_name = "HelperStateStore" },
 ]
 
 
diff --git a/crates/daphne-worker/src/durable/test_state_cleaner.rs b/crates/daphne-worker/src/durable/test_state_cleaner.rs
index 0a24398..af1ae06 100644
--- a/crates/daphne-worker/src/durable/test_state_cleaner.rs
+++ b/crates/daphne-worker/src/durable/test_state_cleaner.rs
@@ -3,10 +3,7 @@
 
 use std::{cmp::min, ops::ControlFlow};
 
-use crate::{
-    durable::{create_span_from_request, req_parse},
-    initialize_tracing, int_err,
-};
+use crate::{durable::create_span_from_request, initialize_tracing, int_err};
 use daphne::messages::TaskId;
 use daphne_service_utils::durable_requests::bindings::{self, DurableMethod};
 use rand::{thread_rng, Rng};
@@ -46,9 +43,10 @@ impl TestStateCleaner {
         match bindings::TestStateCleaner::try_from_uri(&req.path()) {
             // Schedule a durable object (DO) instance for deletion.
             Some(bindings::TestStateCleaner::Put) => {
-                let durable_ref: DurableReference = req_parse(&mut req).await?;
+                let durable_ref: DurableReference =
+                    serde_json::from_slice(&req.bytes().await?).unwrap();
                 match durable_ref.binding.as_ref() {
-                    bindings::AggregateStore::BINDING | bindings::HelperState::BINDING => (),
+                    bindings::AggregateStore::BINDING => (),
                     s => {
                         let message = format!("GarbageCollector: unrecognized binding: {s}");
                         error!("{}", message);
@@ -358,7 +356,7 @@ impl<'srv> DurableConnector<'srv> {
     {
         let req = match (&method, &data) {
             (Method::Post, Some(data)) => {
-                let data = bincode::serialize(&data).map_err(|e| {
+                let data = serde_json::to_vec(&data).map_err(|e| {
                     worker::Error::RustError(format!("failed to serialize data: {e:?}"))
                 })?;
                 let buffer =
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index cfbfe1b..a65a2a7 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -1009,54 +1009,7 @@ impl DapAggregator<BearerToken> for InMemoryAggregator {
 }
 
 #[async_trait]
-impl DapHelper<BearerToken> for InMemoryAggregator {
-    async fn put_helper_state_if_not_exists(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-        helper_state: &DapAggregationJobState,
-    ) -> Result<bool, DapError> {
-        let helper_state_info = HelperStateInfo {
-            task_id: *task_id,
-            agg_job_id: *agg_job_id,
-        };
-
-        let mut helper_state_store = self
-            .helper_state_store
-            .lock()
-            .map_err(|e| fatal_error!(err = ?e))?;
-
-        if helper_state_store.contains_key(&helper_state_info) {
-            return Ok(false);
-        }
-
-        // NOTE: This code is only correct for VDAFs with exactly one round of preparation.
-        // For VDAFs with more rounds, the helper state blob will need to be updated here.
-        helper_state_store.insert(helper_state_info, helper_state.clone());
-
-        Ok(true)
-    }
-
-    async fn get_helper_state(
-        &self,
-        task_id: &TaskId,
-        agg_job_id: &AggregationJobId,
-    ) -> Result<Option<DapAggregationJobState>, DapError> {
-        let helper_state_info = HelperStateInfo {
-            task_id: *task_id,
-            agg_job_id: *agg_job_id,
-        };
-
-        let helper_state_store = self
-            .helper_state_store
-            .lock()
-            .map_err(|e| fatal_error!(err = ?e))?;
-
-        // NOTE: This code is only correct for VDAFs with exactly one round of preparation.
-        // For VDAFs with more rounds, the helper state blob will need to be updated here.
-        Ok(helper_state_store.get(&helper_state_info).cloned())
-    }
-}
+impl DapHelper<BearerToken> for InMemoryAggregator {}
 
 #[async_trait]
 impl DapLeader<BearerToken> for InMemoryAggregator {