Repository: cloudflare/daphne
Pull request #621 (2024-06-18 11:07:04)

Allow disabling replay protection in daphne-server


Patch:
diff --git a/crates/dapf/src/main.rs b/crates/dapf/src/main.rs
index e08e78f..25c85ab 100644
--- a/crates/dapf/src/main.rs
+++ b/crates/dapf/src/main.rs
@@ -370,6 +370,9 @@ async fn main() -> Result<()> {
         .build()
         .with_context(|| "failed to create HTTP client")?;
 
+    if std::env::var("REPLAY_REPORTS").unwrap_or_default() == "1" {
+        daphne::testing::report_generator::replay_reports(true);
+    }
     match cli.action {
         Action::GetHpkeConfig {
             aggregator_url,
diff --git a/crates/daphne-server/src/roles/aggregator.rs b/crates/daphne-server/src/roles/aggregator.rs
index 92ee4e6..f734859 100644
--- a/crates/daphne-server/src/roles/aggregator.rs
+++ b/crates/daphne-server/src/roles/aggregator.rs
@@ -27,7 +27,10 @@ use futures::{future::try_join_all, StreamExt, TryStreamExt};
 use mappable_rc::Marc;
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
-use crate::storage_proxy_connection::kv::{self, KvGetOptions};
+use crate::{
+    roles::fetch_replay_protection_override,
+    storage_proxy_connection::kv::{self, KvGetOptions},
+};
 
 #[async_trait]
 impl DapAggregator<DaphneAuth> for crate::App {
@@ -41,21 +44,7 @@ impl DapAggregator<DaphneAuth> for crate::App {
         let task_id_hex = task_id.to_hex();
         let durable = self.durable();
 
-        let skip_replay_protection = self
-            .kv()
-            .get_cloned::<kv::prefix::GlobalConfigOverride<bool>>(
-                &kv::prefix::GlobalOverrides::SkipReplayProtection,
-                &KvGetOptions {
-                    cache_not_found: true,
-                },
-            )
-            .await
-            .inspect_err(
-                |e| tracing::error!(error = ?e, "failed to fetch skip_replay_protection from kv"),
-            )
-            .ok() // treat error as false
-            .flatten()
-            .unwrap_or_default(); // treat missing as false
+        let replay_protection = fetch_replay_protection_override(self.kv()).await;
 
         futures::stream::iter(agg_share_span)
             .map(|(bucket, (agg_share, report_metadatas))| async {
@@ -68,7 +57,7 @@ impl DapAggregator<DaphneAuth> for crate::App {
                         contained_reports: report_metadatas.iter().map(|(id, _)| *id).collect(),
                         agg_share_delta: agg_share,
                         options: AggregateStoreMergeOptions {
-                            skip_replay_protection,
+                            skip_replay_protection: replay_protection.disabled(),
                         },
                     })
                     .send::<AggregateStoreMergeResp>()
diff --git a/crates/daphne-server/src/roles/mod.rs b/crates/daphne-server/src/roles/mod.rs
index c275fad..a55d594 100644
--- a/crates/daphne-server/src/roles/mod.rs
+++ b/crates/daphne-server/src/roles/mod.rs
@@ -1,10 +1,37 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
+use daphne::ReplayProtection;
+
+use crate::storage_proxy_connection::kv::{self, Kv, KvGetOptions};
+
 mod aggregator;
 mod helper;
 mod leader;
 
+pub async fn fetch_replay_protection_override(kv: Kv<'_>) -> ReplayProtection {
+    let skip_replay_protection = kv
+        .get_cloned::<kv::prefix::GlobalConfigOverride<bool>>(
+            &kv::prefix::GlobalOverrides::SkipReplayProtection,
+            &KvGetOptions {
+                cache_not_found: true,
+            },
+        )
+        .await
+        .inspect_err(
+            |e| tracing::error!(error = ?e, "failed to fetch skip_replay_protection from kv"),
+        )
+        .ok() // treat error as false
+        .flatten()
+        .unwrap_or_default(); // treat missing as false
+    if skip_replay_protection {
+        tracing::debug!("replay protection is disabled");
+        ReplayProtection::Disabled
+    } else {
+        ReplayProtection::Enabled
+    }
+}
+
 #[cfg(feature = "test-utils")]
 mod test_utils {
     use daphne::{
diff --git a/crates/daphne-server/src/router/helper.rs b/crates/daphne-server/src/router/helper.rs
index a53945b..087c537 100644
--- a/crates/daphne-server/src/router/helper.rs
+++ b/crates/daphne-server/src/router/helper.rs
@@ -16,11 +16,12 @@ use daphne::{
 use daphne_service_utils::auth::DaphneAuth;
 use http::StatusCode;
 
+use crate::{roles::fetch_replay_protection_override, App};
+
 use super::{AxumDapResponse, DapRequestExtractor, DaphneService};
 
-pub(super) fn add_helper_routes<A, B>(router: super::Router<A, B>) -> super::Router<A, B>
+pub(super) fn add_helper_routes<B>(router: super::Router<App, B>) -> super::Router<App, B>
 where
-    A: DapHelper<DaphneAuth> + DaphneService + Send + Sync + 'static,
     B: Send + HttpBody + 'static,
     B::Data: Send,
     B::Error: Send + Sync,
@@ -41,16 +42,18 @@ where
         version = ?req.version
     )
 )]
-async fn agg_job<A>(
-    State(app): State<Arc<A>>,
+async fn agg_job(
+    State(app): State<Arc<App>>,
     DapRequestExtractor(req): DapRequestExtractor,
-) -> AxumDapResponse
-where
-    A: DapHelper<DaphneAuth> + DaphneService + Send + Sync,
-{
+) -> AxumDapResponse {
     match req.media_type {
         Some(DapMediaType::AggregationJobInitReq) => {
-            let resp = helper::handle_agg_job_init_req(&*app, &req).await;
+            let resp = helper::handle_agg_job_init_req(
+                &*app,
+                &req,
+                fetch_replay_protection_override(app.kv()).await,
+            )
+            .await;
             AxumDapResponse::from_result_with_success_code(
                 resp,
                 app.server_metrics(),
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
index 964fb49..e03d727 100644
--- a/crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs
+++ b/crates/daphne-service-utils/src/durable_requests/bindings/aggregate_store.rs
@@ -65,11 +65,16 @@ pub struct AggregateStoreMergeOptions {
 
 impl DurableRequestPayload for AggregateStoreMergeReq {
     fn encode_to_builder(&self) -> capnp::message::Builder<capnp::message::HeapAllocator> {
+        let Self {
+            contained_reports,
+            agg_share_delta,
+            options,
+        } = self;
         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
+                contained_reports
                     .len()
                     .try_into()
                     .expect("can't serialize more than u32::MAX reports"),
@@ -84,15 +89,15 @@ impl DurableRequestPayload for AggregateStoreMergeReq {
             }
         }
         {
-            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 mut agg_share_delta_packet = request.reborrow().init_agg_share_delta();
+            agg_share_delta_packet.set_report_count(agg_share_delta.report_count);
+            agg_share_delta_packet.set_min_time(agg_share_delta.min_time);
+            agg_share_delta_packet.set_max_time(agg_share_delta.max_time);
             {
-                let checksum = agg_share_delta
+                let checksum = agg_share_delta_packet
                     .reborrow()
-                    .init_checksum(self.agg_share_delta.checksum.len().try_into().unwrap());
-                checksum.copy_from_slice(&self.agg_share_delta.checksum);
+                    .init_checksum(agg_share_delta.checksum.len().try_into().unwrap());
+                checksum.copy_from_slice(&agg_share_delta.checksum);
             }
             {
                 fn encode<'b, F, B, const ENCODED_SIZE: usize>(
@@ -113,7 +118,7 @@ impl DurableRequestPayload for AggregateStoreMergeReq {
                         bytes = &mut bytes[ENCODED_SIZE..];
                     }
                 }
-                let mut data = agg_share_delta.init_data();
+                let mut data = agg_share_delta_packet.init_data();
                 match &self.agg_share_delta.data {
                     Some(VdafAggregateShare::Field64(field)) => {
                         encode(field, |len| data.init_field64(len));
@@ -128,6 +133,13 @@ impl DurableRequestPayload for AggregateStoreMergeReq {
                 };
             }
         }
+        {
+            let AggregateStoreMergeOptions {
+                skip_replay_protection,
+            } = options;
+            let mut options_packet = request.init_options();
+            options_packet.set_skip_replay_protection(*skip_replay_protection);
+        }
         message
     }
 
@@ -262,7 +274,7 @@ mod test {
             .map(Some)
             .into_iter()
             .chain([None]);
-            for data in test_data {
+            for (i, data) in test_data.enumerate() {
                 let this = AggregateStoreMergeReq {
                     contained_reports: (0..len)
                         .map(|_| ReportId::get_decoded(&rng.gen::<[_; 16]>()).unwrap())
@@ -275,7 +287,7 @@ mod test {
                         data,
                     },
                     options: AggregateStoreMergeOptions {
-                        skip_replay_protection: false,
+                        skip_replay_protection: i % 2 == 0,
                     },
                 };
                 let other =
diff --git a/crates/daphne/benches/aggregation.rs b/crates/daphne/benches/aggregation.rs
index 6e65e3d..3098589 100644
--- a/crates/daphne/benches/aggregation.rs
+++ b/crates/daphne/benches/aggregation.rs
@@ -46,7 +46,9 @@ fn aggregate(c: &mut Criterion) {
                     let mut total = Duration::ZERO;
                     for _ in 0..iters {
                         let now = Instant::now();
-                        let ret = black_box(helper::handle_agg_job_req(aggregator, req).await);
+                        let ret = black_box(
+                            helper::handle_agg_job_req(aggregator, req, Default::default()).await,
+                        );
                         total += now.elapsed();
                         aggregator.clear_storage();
                         drop(ret.unwrap());
diff --git a/crates/daphne/src/lib.rs b/crates/daphne/src/lib.rs
index e873ce5..879dcc9 100644
--- a/crates/daphne/src/lib.rs
+++ b/crates/daphne/src/lib.rs
@@ -77,6 +77,7 @@ use prio::{
     codec::{CodecError, Decode, Encode, ParameterizedDecode, ParameterizedEncode},
     vdaf::Aggregatable as AggregatableTrait,
 };
+pub use protocol::aggregator::ReplayProtection;
 use serde::{Deserialize, Serialize};
 use std::{
     cmp::{max, min},
diff --git a/crates/daphne/src/protocol/aggregator.rs b/crates/daphne/src/protocol/aggregator.rs
index 0f1b54f..311f82d 100644
--- a/crates/daphne/src/protocol/aggregator.rs
+++ b/crates/daphne/src/protocol/aggregator.rs
@@ -374,6 +374,23 @@ pub(crate) enum ReportProcessedStatus {
     Rejected(TransitionFailure),
 }
 
+#[derive(Default, Debug)]
+pub enum ReplayProtection {
+    #[default]
+    Enabled,
+    Disabled,
+}
+
+impl ReplayProtection {
+    pub const fn enabled(&self) -> bool {
+        matches!(self, ReplayProtection::Enabled)
+    }
+
+    pub const fn disabled(&self) -> bool {
+        matches!(self, ReplayProtection::Disabled)
+    }
+}
+
 impl DapTaskConfig {
     /// Leader -> Helper: Initialize the aggregation flow for a sequence of reports. The outputs are the Leader's
     /// state for the aggregation flow and the outbound `AggregationJobInitReq` message.
@@ -398,7 +415,15 @@ impl DapTaskConfig {
             let mut processed = HashSet::with_capacity(report_count_hint);
             let mut reports = pin!(reports);
             while let Some(report) = reports.next().await {
-                if processed.contains(&report.report_metadata.id) {
+                let enforce_replay_check = {
+                    #[cfg(feature = "report-generator")]
+                    {
+                        !crate::testing::report_generator::generator_replaying_reports()
+                    }
+                    #[cfg(not(feature = "report-generator"))]
+                    true
+                };
+                if enforce_replay_check && processed.contains(&report.report_metadata.id) {
                     return Err(fatal_error!(
                         err = "tried to process report sequence with non-unique report IDs",
                         non_unique_id = %report.report_metadata.id,
@@ -503,13 +528,16 @@ impl DapTaskConfig {
         initializer: &impl DapReportInitializer,
         task_id: &TaskId,
         agg_job_init_req: AggregationJobInitReq,
+        replay_protection: ReplayProtection,
     ) -> Result<Vec<EarlyReportStateInitialized>, DapError> {
         let num_reports = agg_job_init_req.prep_inits.len();
         let mut consumed_reports = Vec::with_capacity(num_reports);
         async {
             let mut processed = HashSet::with_capacity(num_reports);
             for prep_init in agg_job_init_req.prep_inits {
-                if processed.contains(&prep_init.report_share.report_metadata.id) {
+                if replay_protection.enabled()
+                    && processed.contains(&prep_init.report_share.report_metadata.id)
+                {
                     return Err(DapAbort::InvalidMessage {
                         detail: format!(
                             "report ID {} appears twice in the same aggregation job",
@@ -552,6 +580,7 @@ impl DapTaskConfig {
 
     /// Helper -> Leader: Produce the `AggregationJobResp` message to send to the Leader and
     /// compute Helper's aggregate share span.
+    #[tracing::instrument(skip_all, fields(report_count = report_status.len()))]
     pub(crate) fn produce_agg_job_resp(
         &self,
         report_status: &HashMap<ReportId, ReportProcessedStatus>,
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index cefc902..64d5bc9 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -17,7 +17,7 @@ use crate::{
         AggregationJobResp, PartialBatchSelector, TaskId, TransitionFailure, TransitionVar,
     },
     metrics::{DaphneMetrics, DaphneRequestType, ReportStatus},
-    protocol::aggregator::ReportProcessedStatus,
+    protocol::aggregator::{ReplayProtection, ReportProcessedStatus},
     roles::aggregator::MergeAggShareError,
     DapAggregationParam, DapError, DapRequest, DapResource, DapResponse, DapTaskConfig,
     EarlyReportStateInitialized,
@@ -30,6 +30,7 @@ pub trait DapHelper<S: Sync>: DapAggregator<S> {}
 pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     aggregator: &A,
     req: &'req DapRequest<S>,
+    replay_protection: ReplayProtection,
 ) -> Result<DapResponse, DapError> {
     let global_config = aggregator.get_global_config().await?;
     let task_id = req.task_id()?;
@@ -80,7 +81,13 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     let prep_init_count = agg_job_init_req.prep_inits.len();
     let part_batch_sel = agg_job_init_req.part_batch_sel.clone();
     let initialized_reports = task_config
-        .consume_agg_job_req(aggregator, aggregator, task_id, agg_job_init_req)
+        .consume_agg_job_req(
+            aggregator,
+            aggregator,
+            task_id,
+            agg_job_init_req,
+            replay_protection,
+        )
         .await?;
 
     let agg_job_resp = {
@@ -119,9 +126,12 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
 pub async fn handle_agg_job_req<'req, S: Sync, A: DapHelper<S>>(
     aggregator: &A,
     req: &DapRequest<S>,
+    replay_protection: ReplayProtection,
 ) -> Result<DapResponse, DapError> {
     match req.media_type {
-        Some(DapMediaType::AggregationJobInitReq) => handle_agg_job_init_req(aggregator, req).await,
+        Some(DapMediaType::AggregationJobInitReq) => {
+            handle_agg_job_init_req(aggregator, req, replay_protection).await
+        }
         _ => Err(DapAbort::BadRequest("unexpected media type".into()).into()),
     }
 }
diff --git a/crates/daphne/src/roles/mod.rs b/crates/daphne/src/roles/mod.rs
index 3511655..9c2925e 100644
--- a/crates/daphne/src/roles/mod.rs
+++ b/crates/daphne/src/roles/mod.rs
@@ -657,7 +657,7 @@ mod test {
             )
             .await;
         assert_matches!(
-            helper::handle_agg_job_req(&*t.helper, &req)
+            helper::handle_agg_job_req(&*t.helper, &req, Default::default())
                 .await
                 .unwrap_err(),
             DapError::Abort(DapAbort::QueryMismatch { .. })
@@ -712,14 +712,14 @@ mod test {
 
         // Expect failure due to missing bearer token.
         assert_matches!(
-            helper::handle_agg_job_req(&*t.helper, &req).await,
+            helper::handle_agg_job_req(&*t.helper, &req, Default::default()).await,
             Err(DapError::Abort(DapAbort::UnauthorizedRequest { .. }))
         );
 
         // Expect failure due to incorrect bearer token.
         req.sender_auth = Some(BearerToken::from("incorrect auth token!".to_string()));
         assert_matches!(
-            helper::handle_agg_job_req(&*t.helper, &req).await,
+            helper::handle_agg_job_req(&*t.helper, &req, Default::default()).await,
             Err(DapError::Abort(DapAbort::UnauthorizedRequest { .. }))
         );
 
@@ -904,7 +904,7 @@ mod test {
 
         // Get AggregationJobResp and then extract the transition data from inside.
         let agg_job_resp = AggregationJobResp::get_decoded(
-            &helper::handle_agg_job_req(&*t.helper, &req)
+            &helper::handle_agg_job_req(&*t.helper, &req, Default::default())
                 .await
                 .unwrap()
                 .payload,
@@ -932,7 +932,7 @@ mod test {
 
         // Get AggregationJobResp and then extract the transition data from inside.
         let agg_job_resp = AggregationJobResp::get_decoded(
-            &helper::handle_agg_job_req(&*t.helper, &req)
+            &helper::handle_agg_job_req(&*t.helper, &req, Default::default())
                 .await
                 .unwrap()
                 .payload,


Test patch:
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index 91046d7..c0c51e2 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -262,6 +262,7 @@ impl AggregationJobTest {
                         self,
                         &self.task_id,
                         agg_job_init_req,
+                        Default::default(),
                     )
                     .await
                     .unwrap(),
@@ -1107,6 +1108,7 @@ impl DapLeader<BearerToken> for InMemoryAggregator {
             Some(DapMediaType::AggregationJobInitReq) => Ok(helper::handle_agg_job_req(
                 &**self.peer.as_ref().expect("peer not configured"),
                 &req,
+                Default::default(),
             )
             .await
             .expect("peer aborted unexpectedly")),
@@ -1129,6 +1131,7 @@ impl DapLeader<BearerToken> for InMemoryAggregator {
             Ok(helper::handle_agg_job_req(
                 &**self.peer.as_ref().expect("peer not configured"),
                 &req,
+                Default::default(),
             )
             .await
             .expect("peer aborted unexpectedly"))
diff --git a/crates/daphne/src/testing/report_generator.rs b/crates/daphne/src/testing/report_generator.rs
index 2dfa57d..0bd254e 100644
--- a/crates/daphne/src/testing/report_generator.rs
+++ b/crates/daphne/src/testing/report_generator.rs
@@ -4,8 +4,8 @@
 use std::{
     pin::Pin,
     sync::{
-        atomic::{AtomicUsize, Ordering},
-        Mutex,
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Mutex, OnceLock,
     },
     task::Context,
     time::Instant,
@@ -14,7 +14,7 @@ use std::{
 use crate::{
     fatal_error,
     hpke::HpkeConfig,
-    messages::{self, TaskId, Time},
+    messages::{self, Report, TaskId, Time},
     vdaf::VdafConfig,
     DapError, DapMeasurement, DapVersion,
 };
@@ -56,6 +56,14 @@ impl Stream for ReportGenerator {
     }
 }
 
+static REPLAY_REPORTS: AtomicBool = AtomicBool::new(false);
+pub fn replay_reports(b: bool) {
+    REPLAY_REPORTS.store(b, Ordering::SeqCst);
+}
+pub fn generator_replaying_reports() -> bool {
+    REPLAY_REPORTS.load(Ordering::SeqCst)
+}
+
 impl ReportGenerator {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
@@ -89,14 +97,31 @@ impl ReportGenerator {
                     let now = Instant::now();
                     // ----
 
-                    let report = vdaf.produce_report_with_extensions(
-                        &hpke_config_list,
-                        report_time_dist.sample(&mut thread_rng()),
-                        &task_id,
-                        measurement.clone(),
-                        extensions.clone(),
-                        version,
-                    )?;
+                    static LAST_REPORT: OnceLock<Report> = OnceLock::new();
+                    let report = if REPLAY_REPORTS.load(Ordering::SeqCst) {
+                        LAST_REPORT
+                            .get_or_init(|| {
+                                vdaf.produce_report_with_extensions(
+                                    &hpke_config_list,
+                                    report_time_dist.sample(&mut thread_rng()),
+                                    &task_id,
+                                    measurement.clone(),
+                                    extensions.clone(),
+                                    version,
+                                )
+                                .expect("we have to panic here since we can't return the error")
+                            })
+                            .clone()
+                    } else {
+                        vdaf.produce_report_with_extensions(
+                            &hpke_config_list,
+                            report_time_dist.sample(&mut thread_rng()),
+                            &task_id,
+                            measurement.clone(),
+                            extensions.clone(),
+                            version,
+                        )?
+                    };
 
                     // perf measurements
                     let count = GENERATED_REPORT_COUNTER.fetch_add(1, Ordering::SeqCst) + 1;