Add tests to capnp decoding/encoding

Author: mendess

crates/daphne-service-utils/src/capnproto/mod.rs

@@ -268,3 +268,80 @@ pub fn usize_to_capnp_len(u: usize) -> u32 {
     u.try_into()
         .expect("capnp can't encode more that u32::MAX of something")
 }
+
+#[cfg(test)]
+pub fn roundtrip_test<T>(before: T)
+where
+    T: CapnprotoPayloadDecode + CapnprotoPayloadEncode + PartialEq + std::fmt::Debug,
+{
+    assert_eq!(
+        before,
+        T::decode_from_bytes(&before.encode_to_bytes()).unwrap()
+    );
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::Rng;
+
+    #[test]
+    fn test_u8_array_serialize_deserialize() {
+        roundtrip_test(rand::thread_rng().gen::<[u8; 32]>());
+        roundtrip_test(rand::thread_rng().gen::<[u8; 16]>());
+    }
+
+    #[test]
+    fn test_partial_batch_selector_serialize_deserialize() {
+        roundtrip_test(PartialBatchSelector::TimeInterval);
+
+        roundtrip_test(PartialBatchSelector::LeaderSelectedByBatchId {
+            batch_id: BatchId(rand::thread_rng().gen()),
+        });
+    }
+
+    #[test]
+    fn test_report_error_conversion() {
+        // cause a compilation error if the variants change
+        const _: () = {
+            #[allow(clippy::match_same_arms)]
+            match messages::ReportError::Reserved {
+                messages::ReportError::Reserved => (),
+                messages::ReportError::BatchCollected => (),
+                messages::ReportError::ReportReplayed => (),
+                messages::ReportError::ReportDropped => (),
+                messages::ReportError::HpkeUnknownConfigId => (),
+                messages::ReportError::HpkeDecryptError => (),
+                messages::ReportError::VdafPrepError => (),
+                messages::ReportError::BatchSaturated => (),
+                messages::ReportError::TaskExpired => (),
+                messages::ReportError::InvalidMessage => (),
+                messages::ReportError::ReportTooEarly => (),
+                messages::ReportError::TaskNotStarted => (),
+            }
+        };
+        let all_errors = vec![
+            messages::ReportError::Reserved,
+            messages::ReportError::BatchCollected,
+            messages::ReportError::ReportReplayed,
+            messages::ReportError::ReportDropped,
+            messages::ReportError::HpkeUnknownConfigId,
+            messages::ReportError::HpkeDecryptError,
+            messages::ReportError::VdafPrepError,
+            messages::ReportError::BatchSaturated,
+            messages::ReportError::TaskExpired,
+            messages::ReportError::InvalidMessage,
+            messages::ReportError::ReportTooEarly,
+            messages::ReportError::TaskNotStarted,
+        ];
+
+        for error in all_errors {
+            let converted: base_capnp::ReportError = error.into();
+            let back_converted: messages::ReportError = converted.into();
+            assert_eq!(
+                error, back_converted,
+                "Conversion symmetry failed for {error:?}",
+            );
+        }
+    }
+}

crates/daphne-service-utils/src/compute_offload/mod.rs

@@ -23,6 +23,7 @@ use daphne::{
 use prio::codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode};
 use std::{borrow::Cow, ops::Range};
 
+#[cfg_attr(any(test, feature = "test-utils"), derive(PartialEq, Debug))]
 pub struct InitializeReports<'s> {
     pub hpke_keys: Cow<'s, [HpkeReceiverConfig]>,
     /// Output of [`DapAggregator::valid_report_time_range`](daphne::roles::DapAggregator) at the
@@ -489,11 +490,108 @@ fn to_capnp<E: ToString>(e: E) -> capnp::Error {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::capnproto::{CapnprotoPayloadDecodeExt, CapnprotoPayloadEncodeExt};
+    use crate::capnproto::roundtrip_test;
+    use daphne::{vdaf::VdafVerifyKey, DapVersion};
+    use rand::Rng;
 
     #[test]
     fn report_metadata_roundtrip() {
-        let report_metadata = ReportMetadata {
+        roundtrip_test(generate_random_report_metadata());
+    }
+
+    #[test]
+    fn test_encode_decode_range() {
+        let mut rng = rand::thread_rng();
+
+        let start = rng.gen::<messages::Time>();
+        let end = rng.gen::<messages::Time>();
+
+        roundtrip_test(start..end);
+    }
+
+    #[test]
+    fn test_encode_decode_partial_dap_task_config_for_report_init() {
+        roundtrip_test(generate_random_partial_task_config());
+    }
+
+    #[test]
+    fn test_encode_decode_prepare_init() {
+        roundtrip_test(generate_random_prepare_init());
+    }
+
+    #[test]
+    fn test_encode_decode_hpke_config() {
+        roundtrip_test(generate_random_hpke_config());
+    }
+
+    #[test]
+    fn test_encode_decode_hpke_receiver_config() {
+        roundtrip_test(HpkeReceiverConfig {
+            config: generate_random_hpke_config(),
+            private_key: generate_random_bytes().into(),
+        });
+    }
+
+    #[test]
+    fn test_encode_decode_initialize_reports() {
+        let initialize_reports = InitializeReports {
+            hpke_keys: generate_random_hpke_receiver_configs().into(),
+            valid_report_range: generate_random_valid_report_range(),
+            task_id: TaskId(rand::thread_rng().gen()),
+            task_config: generate_random_partial_task_config(),
+            agg_param: generate_random_bytes().into(),
+            prep_inits: generate_random_prep_inits(),
+        };
+
+        roundtrip_test(initialize_reports);
+    }
+
+    fn generate_random_partial_task_config() -> PartialDapTaskConfigForReportInit<'static> {
+        let mut rng = rand::thread_rng();
+        PartialDapTaskConfigForReportInit {
+            not_before: rng.gen::<u64>(),
+            not_after: rng.gen::<u64>(),
+            method_is_taskprov: rng.gen::<bool>(),
+            version: [DapVersion::Latest, DapVersion::Draft09][rng.gen_range(0..2)],
+            vdaf: Cow::Owned(generate_random_vdaf()),
+            vdaf_verify_key: VdafVerifyKey::from(rng.gen::<[u8; 32]>()),
+        }
+    }
+
+    fn generate_random_hpke_receiver_configs() -> Vec<HpkeReceiverConfig> {
+        (0..3)
+            .map(|_| generate_random_hpke_receiver_config())
+            .collect()
+    }
+
+    fn generate_random_hpke_receiver_config() -> HpkeReceiverConfig {
+        HpkeReceiverConfig {
+            config: generate_random_hpke_config(),
+            private_key: generate_random_bytes().into(),
+        }
+    }
+
+    fn generate_random_valid_report_range() -> std::ops::Range<u64> {
+        let mut rng = rand::thread_rng();
+        rng.gen_range(0..1000)..rng.gen_range(1000..2000)
+    }
+
+    fn generate_random_prep_inits() -> Vec<PrepareInit> {
+        (0..2).map(|_| generate_random_prepare_init()).collect()
+    }
+
+    fn generate_random_hpke_config() -> HpkeConfig {
+        HpkeConfig {
+            id: rand::thread_rng().gen(),
+            kem_id: daphne::hpke::HpkeKemId::P256HkdfSha256,
+            kdf_id: daphne::hpke::HpkeKdfId::HkdfSha256,
+            aead_id: daphne::hpke::HpkeAeadId::Aes128Gcm,
+            public_key: generate_random_bytes().into(),
+        }
+    }
+
+    fn generate_random_report_metadata() -> ReportMetadata {
+        ReportMetadata {
             id: messages::ReportId(rand::random()),
             time: rand::random(),
             public_extensions: Some(vec![
@@ -503,27 +601,35 @@ mod test {
                     payload: b"some extension payload".to_vec(),
                 },
             ]),
-        };
+        }
+    }
 
-        assert_eq!(
-            report_metadata,
-            ReportMetadata::decode_from_bytes(&report_metadata.encode_to_bytes()).unwrap()
-        );
+    fn generate_random_bytes() -> Vec<u8> {
+        let mut rng = rand::thread_rng();
+        (0..rng.gen_range(0..32)).map(|_| rng.gen()).collect()
     }
 
-    #[test]
-    fn report_metadata_roundtrip_draft09() {
-        let report_metadata = ReportMetadata {
-            id: messages::ReportId(rand::random()),
-            time: rand::random(),
-            // draft09 compatibility: Previously there was no extensions field in the report
-            // metadata.
-            public_extensions: None,
-        };
+    fn generate_random_vdaf() -> VdafConfig {
+        // Replace with the actual logic to generate random VdafConfig
+        VdafConfig::Prio2 { dimension: 1234 }
+    }
 
-        assert_eq!(
-            report_metadata,
-            ReportMetadata::decode_from_bytes(&report_metadata.encode_to_bytes()).unwrap()
-        );
+    fn generate_random_prepare_init() -> PrepareInit {
+        PrepareInit {
+            report_share: generate_random_report_share(),
+            payload: generate_random_bytes(),
+        }
+    }
+
+    fn generate_random_report_share() -> ReportShare {
+        ReportShare {
+            report_metadata: generate_random_report_metadata(),
+            public_share: generate_random_bytes(),
+            encrypted_input_share: HpkeCiphertext {
+                config_id: rand::thread_rng().gen::<u8>(),
+                enc: generate_random_bytes(),
+                payload: generate_random_bytes(),
+            },
+        }
     }
 }

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

@@ -31,6 +31,11 @@ super::mk_durable_object! {
     ///
     /// - The report id has been seen but belongs to a different agg job id:
     ///     => We return it in duplicates set
+    ///
+    ///
+    /// This is only correct for as long as two invariants are upheld somewhere else in the code:
+    /// - We stricly require aggregation jobs to not change once submited (this is required per the DAP spec)
+    /// - All storage operations after this point that rely on replay protection are idempotent.
     struct ReplayChecker {
         state: State,
         env: Env,
@@ -59,6 +64,8 @@ impl GcDurableObject for ReplayChecker {
 
                 let mut duplicates = HashSet::new();
 
+                // Check the cache for duplicates and compute the set of report IDs we need to read
+                // from the disk.
                 let report_ids_as_string = report_ids
                     .iter()
                     .filter(|r| match self.seen.get(r) {
@@ -73,7 +80,7 @@ impl GcDurableObject for ReplayChecker {
                     .map(ToString::to_string)
                     .collect::<Vec<_>>();
 
-                let aggregation_job_id_as_str = agg_job_id.to_string();
+                let agg_job_id_as_str = agg_job_id.to_string();
 
                 let result = self
                     .state
@@ -87,14 +94,14 @@ impl GcDurableObject for ReplayChecker {
 
                     let v = result.get(&JsValue::from_str(as_str));
                     if let Some(stored_agg_job_id) = v.as_string() {
-                        if stored_agg_job_id != aggregation_job_id_as_str {
+                        if stored_agg_job_id != agg_job_id_as_str {
                             duplicates.insert(*id);
                         }
                     } else {
                         js_sys::Reflect::set(
                             &obj_to_update,
                             &JsValue::from_str(as_str),
-                            &JsValue::from_str(aggregation_job_id_as_str.as_ref()),
+                            &JsValue::from_str(agg_job_id_as_str.as_ref()),
                         )?;
                     }
                 }

crates/daphne/src/protocol/report_init.rs

@@ -123,6 +123,8 @@ impl<'s> From<&'s PartialDapTaskConfigForReportInit<'_>> for PartialDapTaskConfi
     }
 }
 
+#[derive(Clone)]
+#[cfg_attr(any(test, feature = "test-utils"), derive(PartialEq, Debug))]
 pub struct PartialDapTaskConfigForReportInit<'s> {
     pub not_before: messages::Time,
     pub not_after: messages::Time,