Refactors to prepare for moving the aggregation logic back into the workers

crates/dapf/src/functions/test_routes.rs

@@ -5,7 +5,7 @@
 //!
 //! [interop]: https://divergentdave.github.io/draft-dcook-ppm-dap-interop-test-design/draft-dcook-ppm-dap-interop-test-design.html
 
-use anyhow::Context;
+use anyhow::{bail, Context};
 use daphne::{
     hpke::{HpkeKemId, HpkeReceiverConfig},
     messages::HpkeConfigList,
@@ -53,12 +53,18 @@ impl HttpClient {
     }
 
     pub async fn delete_all_storage(&self, aggregator_url: &Url) -> anyhow::Result<()> {
-        self.post(aggregator_url.join("/internal/delete_all").unwrap())
+        let resp = self
+            .post(aggregator_url.join("/internal/delete_all").unwrap())
             .send()
             .await
-            .context("deleting storage")?
-            .error_for_status()
             .context("deleting storage")?;
-        Ok(())
+        if resp.status().is_success() {
+            return Ok(());
+        }
+        bail!(
+            "delete storage request failed. {} {}",
+            resp.status(),
+            resp.text().await?
+        );
     }
 }

crates/daphne-server/src/roles/aggregator.rs

@@ -216,19 +216,22 @@ impl DapAggregator for crate::App {
     ) -> Result<(), DapError> {
         let expiration_time = task_config.not_after;
 
-        if self.service_config.role.is_leader() {
-            self.kv()
-                .put_with_expiration::<kv::prefix::TaskConfig>(
-                    task_id,
-                    task_config,
-                    expiration_time,
-                )
-                .await
-                .map_err(|e| fatal_error!(err = ?e, "failed to put the a task config in kv"))?;
-        } else {
-            self.kv()
-                .only_cache_put::<kv::prefix::TaskConfig>(task_id, task_config)
-                .await;
+        match self.service_config.role {
+            daphne::constants::DapAggregatorRole::Leader => {
+                self.kv()
+                    .put_with_expiration::<kv::prefix::TaskConfig>(
+                        task_id,
+                        task_config,
+                        expiration_time,
+                    )
+                    .await
+                    .map_err(|e| fatal_error!(err = ?e, "failed to put the a task config in kv"))?;
+            }
+            daphne::constants::DapAggregatorRole::Helper => {
+                self.kv()
+                    .only_cache_put::<kv::prefix::TaskConfig>(task_id, task_config)
+                    .await;
+            }
         }
         Ok(())
     }

crates/daphne-server/src/storage_proxy_connection/mod.rs

@@ -6,9 +6,9 @@ pub(crate) mod kv;
 use std::fmt::Debug;
 
 use axum::http::StatusCode;
-use daphne_service_utils::durable_requests::{
-    bindings::{DurableMethod, DurableRequestPayload, DurableRequestPayloadExt},
-    DurableRequest, ObjectIdFrom, DO_PATH_PREFIX,
+use daphne_service_utils::{
+    capnproto_payload::{CapnprotoPayloadEncode, CapnprotoPayloadEncodeExt as _},
+    durable_requests::{bindings::DurableMethod, DurableRequest, ObjectIdFrom, DO_PATH_PREFIX},
 };
 use serde::de::DeserializeOwned;
 
@@ -95,7 +95,7 @@ impl<'d, B: DurableMethod + Debug, P: AsRef<[u8]>> RequestBuilder<'d, B, P> {
 }
 
 impl<'d, B: DurableMethod> RequestBuilder<'d, B, [u8; 0]> {
-    pub fn encode<T: DurableRequestPayload>(self, payload: &T) -> RequestBuilder<'d, B, Vec<u8>> {
+    pub fn encode<T: CapnprotoPayloadEncode>(self, payload: &T) -> RequestBuilder<'d, B, Vec<u8>> {
         self.with_body(payload.encode_to_bytes().unwrap())
     }
 

crates/daphne-service-utils/src/capnproto_payload.rs

@@ -0,0 +1,54 @@
+// Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+pub trait CapnprotoPayloadEncode {
+    fn encode_to_builder(&self) -> capnp::message::Builder<capnp::message::HeapAllocator>;
+}
+
+pub trait CapnprotoPayloadEncodeExt {
+    fn encode_to_bytes(&self) -> capnp::Result<Vec<u8>>;
+}
+
+pub trait CapnprotoPayloadDecode {
+    fn decode_from_reader(
+        reader: capnp::message::Reader<capnp::serialize::OwnedSegments>,
+    ) -> capnp::Result<Self>
+    where
+        Self: Sized;
+}
+
+pub trait CapnprotoPayloadDecodeExt {
+    fn decode_from_bytes(bytes: &[u8]) -> capnp::Result<Self>
+    where
+        Self: Sized;
+}
+
+impl<T> CapnprotoPayloadEncodeExt for T
+where
+    T: CapnprotoPayloadEncode,
+{
+    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)
+    }
+}
+
+impl<T> CapnprotoPayloadDecodeExt for T
+where
+    T: CapnprotoPayloadDecode,
+{
+    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)
+    }
+}

crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs

@@ -11,11 +11,11 @@ use daphne::{
 use serde::{Deserialize, Serialize};
 
 use crate::{
+    capnproto_payload::{CapnprotoPayloadDecode, CapnprotoPayloadEncode},
     durable_request_capnp::{aggregate_store_merge_req, dap_aggregate_share},
     durable_requests::ObjectIdFrom,
 };
 
-use super::DurableRequestPayload;
 use prio::{
     field::{FieldElement, FieldError},
     vdaf::AggregateShare,
@@ -68,7 +68,7 @@ pub struct AggregateStoreMergeOptions {
     pub skip_replay_protection: bool,
 }
 
-impl DurableRequestPayload for AggregateStoreMergeReq {
+impl CapnprotoPayloadEncode for AggregateStoreMergeReq {
     fn encode_to_builder(&self) -> capnp::message::Builder<capnp::message::HeapAllocator> {
         let Self {
             contained_reports,
@@ -162,7 +162,9 @@ impl DurableRequestPayload for AggregateStoreMergeReq {
         }
         message
     }
+}
 
+impl CapnprotoPayloadDecode for AggregateStoreMergeReq {
     fn decode_from_reader(
         reader: capnp::message::Reader<capnp::serialize::OwnedSegments>,
     ) -> capnp::Result<Self> {
@@ -285,7 +287,9 @@ mod test {
     };
     use rand::{thread_rng, Rng};
 
-    use crate::durable_requests::bindings::DurableRequestPayloadExt;
+    use crate::capnproto_payload::{
+        CapnprotoPayloadDecodeExt as _, CapnprotoPayloadEncodeExt as _,
+    };
 
     use super::*;
 

crates/daphne-service-utils/src/durable_requests/bindings/mod.rs

@@ -37,48 +37,6 @@ 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;

crates/daphne-service-utils/src/lib.rs

@@ -5,6 +5,8 @@
 
 pub mod bearer_token;
 #[cfg(feature = "durable_requests")]
+pub mod capnproto_payload;
+#[cfg(feature = "durable_requests")]
 pub mod durable_requests;
 pub mod http_headers;
 #[cfg(feature = "test-utils")]

crates/daphne-worker/src/durable/mod.rs

@@ -24,8 +24,9 @@ pub(crate) mod aggregate_store;
 pub(crate) mod test_state_cleaner;
 
 use crate::tracing_utils::shorten_paths;
-use daphne_service_utils::durable_requests::bindings::{
-    DurableMethod, DurableRequestPayload, DurableRequestPayloadExt,
+use daphne_service_utils::{
+    capnproto_payload::{CapnprotoPayloadDecode, CapnprotoPayloadDecodeExt},
+    durable_requests::bindings::DurableMethod,
 };
 use serde::{Deserialize, Serialize};
 use tracing::info_span;
@@ -209,7 +210,7 @@ pub(crate) async fn state_set_if_not_exists<T: for<'a> Deserialize<'a> + Seriali
 
 async fn req_parse<T>(req: &mut Request) -> Result<T>
 where
-    T: DurableRequestPayload,
+    T: CapnprotoPayloadDecode,
 {
     T::decode_from_bytes(&req.bytes().await?).map_err(|e| Error::RustError(e.to_string()))
 }

crates/daphne/src/constants.rs

@@ -75,16 +75,6 @@ pub enum DapRole {
     Helper,
 }
 
-impl DapRole {
-    pub fn is_leader(self) -> bool {
-        self == Self::Leader
-    }
-
-    pub fn is_helper(self) -> bool {
-        self == Self::Helper
-    }
-}
-
 impl FromStr for DapRole {
     type Err = String;
 
@@ -119,16 +109,6 @@ pub enum DapAggregatorRole {
     Helper,
 }
 
-impl DapAggregatorRole {
-    pub fn is_leader(self) -> bool {
-        self == Self::Leader
-    }
-
-    pub fn is_helper(self) -> bool {
-        self == Self::Helper
-    }
-}
-
 impl FromStr for DapAggregatorRole {
     type Err = String;
 

crates/daphne/src/lib.rs

@@ -93,7 +93,7 @@ use url::Url;
 use vdaf::mastic::MasticWeight;
 
 pub use messages::request::{DapRequest, DapRequestMeta, DapResponse};
-pub use protocol::report_init::{InitializedReport, WithPeerPrepShare};
+pub use protocol::report_init::{InitializedReport, PartialDapTaskConfig, WithPeerPrepShare};
 
 /// DAP version used for a task.
 #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]

crates/daphne/src/pine/vdaf.rs

@@ -5,7 +5,10 @@
 
 use std::{borrow::Cow, iter};
 
-use crate::pine::{dst, USAGE_QUERY_RAND};
+use crate::{
+    constants::DapAggregatorRole,
+    pine::{dst, USAGE_QUERY_RAND},
+};
 use prio_draft09::{
     codec::{CodecError, Decode, Encode, ParameterizedDecode},
     field::FftFriendlyFieldElement,
@@ -376,28 +379,29 @@ impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> Encode for PinePrepStat
     }
 }
 
-impl<F, X, const SEED_SIZE: usize> ParameterizedDecode<(&Pine<F, X, SEED_SIZE>, bool)>
+impl<F, X, const SEED_SIZE: usize> ParameterizedDecode<(&Pine<F, X, SEED_SIZE>, DapAggregatorRole)>
     for PinePrepState<F, SEED_SIZE>
 where
     F: FftFriendlyFieldElement,
     X: Xof<SEED_SIZE>,
 {
     fn decode_with_param(
-        (pine, is_leader): &(&Pine<F, X, SEED_SIZE>, bool),
+        (pine, role): &(&Pine<F, X, SEED_SIZE>, DapAggregatorRole),
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
-        let (gradient_share, meas_share_seed) = if *is_leader {
-            (
+        let (gradient_share, meas_share_seed) = match role {
+            DapAggregatorRole::Leader => (
                 std::iter::repeat_with(|| F::decode(bytes))
                     .take(pine.flp.dimension)
                     .collect::<Result<Vec<_>, _>>()?,
                 None,
-            )
-        } else {
-            let seed = Seed::decode(bytes)?;
-            let mut gradient_share = pine.helper_meas_share(seed.as_ref());
-            gradient_share.truncate(pine.flp.dimension);
-            (gradient_share, Some(seed))
+            ),
+            DapAggregatorRole::Helper => {
+                let seed = Seed::decode(bytes)?;
+                let mut gradient_share = pine.helper_meas_share(seed.as_ref());
+                gradient_share.truncate(pine.flp.dimension);
+                (gradient_share, Some(seed))
+            }
         };
 
         Ok(Self {
@@ -765,7 +769,10 @@ mod tests {
         },
     };
 
-    use crate::pine::{msg, norm_bound_f64_to_u64, vdaf::PineVec, Pine, PineParam};
+    use crate::{
+        constants::DapAggregatorRole,
+        pine::{msg, norm_bound_f64_to_u64, vdaf::PineVec, Pine, PineParam},
+    };
 
     use assert_matches::assert_matches;
 
@@ -981,14 +988,14 @@ mod tests {
         };
 
         let got = PinePrepState::get_decoded_with_param(
-            &(&pine, true),
+            &(&pine, DapAggregatorRole::Leader),
             &leader_prep_state.get_encoded().unwrap(),
         )
         .unwrap();
         assert_eq!(got, leader_prep_state);
 
         let got = PinePrepState::get_decoded_with_param(
-            &(&pine, false),
+            &(&pine, DapAggregatorRole::Helper),
             &helper_prep_state.get_encoded().unwrap(),
         )
         .unwrap();