Split CapnprotoPayload into Encode/Decode

Author: mendess

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

@@ -7,7 +7,7 @@ use std::fmt::Debug;
 
 use axum::http::StatusCode;
 use daphne_service_utils::{
-    capnproto_payload::{CapnprotoPayload, CapnprotoPayloadExt as _},
+    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: CapnprotoPayload>(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

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

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

@@ -11,7 +11,7 @@ use daphne::{
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    capnproto_payload::CapnprotoPayload,
+    capnproto_payload::{CapnprotoPayloadDecode, CapnprotoPayloadEncode},
     durable_request_capnp::{aggregate_store_merge_req, dap_aggregate_share},
     durable_requests::ObjectIdFrom,
 };
@@ -68,7 +68,7 @@ pub struct AggregateStoreMergeOptions {
     pub skip_replay_protection: bool,
 }
 
-impl CapnprotoPayload 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 CapnprotoPayload 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::capnproto_payload::CapnprotoPayloadExt as _;
+    use crate::capnproto_payload::{
+        CapnprotoPayloadDecodeExt as _, CapnprotoPayloadEncodeExt as _,
+    };
 
     use super::*;
 

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

@@ -25,7 +25,7 @@ pub(crate) mod test_state_cleaner;
 
 use crate::tracing_utils::shorten_paths;
 use daphne_service_utils::{
-    capnproto_payload::{CapnprotoPayload, CapnprotoPayloadExt},
+    capnproto_payload::{CapnprotoPayloadDecode, CapnprotoPayloadDecodeExt},
     durable_requests::bindings::DurableMethod,
 };
 use serde::{Deserialize, Serialize};
@@ -210,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: CapnprotoPayload,
+    T: CapnprotoPayloadDecode,
 {
     T::decode_from_bytes(&req.bytes().await?).map_err(|e| Error::RustError(e.to_string()))
 }