Repository: cloudflare/daphne
Pull request #131 (2022-10-03 16:31:05)

daphne: Check for query mismatch
Partially addresses #100.
Based on #130 (merge that first).

Ensure all requests that encode a query type assert that the correct query type is indicated. Along the way, refactor the batch validation logic so that the Leader and Helper use more of the same code.

In addition, rename "invalidBatchInterval" to "batchInvalid".


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index dc9e766..7814eb2 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -113,6 +113,10 @@ pub enum DapAbort {
     #[error("badRequest")]
     BadRequest(String),
 
+    /// Invalid batch. Sent in response to a CollectReq or AggregateShareReq.
+    #[error("batchInvalid")]
+    BatchInvalid,
+
     /// Batch mismatch. Sent in response to an AggregateShareReq.
     #[error("batchMismatch")]
     BatchMismatch,
@@ -126,10 +130,6 @@ pub enum DapAbort {
     #[error("internalError: {0}")]
     Internal(#[source] Box<dyn std::error::Error + 'static + Send + Sync>),
 
-    /// Invalid batch interval. Sent in response to a CollectReq or AggregateShareReq.
-    #[error("invalidBatchInterval")]
-    InvalidBatchInterval,
-
     /// Invalid DAP version. Sent in response to requests for an unsupported (or unknown) DAP version.
     #[error("invalidProtocolVersion")]
     InvalidProtocolVersion,
@@ -143,6 +143,10 @@ pub enum DapAbort {
     #[error("missingTaskID")]
     MissingTaskId,
 
+    /// Query mismatch. Sent in response to a CollectReq or AggregateShareReq.
+    #[error("queryMismatch")]
+    QueryMismatch,
+
     /// Replayed report. Sent in response to an upload request containing a Report that has been replayed.
     //
     // TODO spec: Define this error type.
@@ -184,11 +188,12 @@ impl DapAbort {
     /// request was targeted and `task_id` is the associated TaskID.
     pub fn to_problem_details(&self) -> ProblemDetails {
         let (typ, detail) = match self {
-            Self::BatchMismatch
+            Self::BatchInvalid
+            | Self::BatchMismatch
             | Self::BatchOverlap
-            | Self::InvalidBatchInterval
             | Self::InvalidProtocolVersion
             | Self::InvalidBatchSize
+            | Self::QueryMismatch
             | Self::MissingTaskId
             | Self::ReplayedReport
             | Self::StaleReport
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index 945eab9..b0b59d9 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -3,7 +3,6 @@
 
 //! Messages in the DAP protocol.
 
-use crate::DapTaskConfig;
 use prio::codec::{
     decode_u16_items, decode_u32_items, encode_u16_items, encode_u32_items, CodecError, Decode,
     Encode,
@@ -443,17 +442,6 @@ impl Interval {
     pub fn end(&self) -> Time {
         self.start + self.duration
     }
-
-    /// Check that the batch interval is valid for the given task configuration.
-    pub fn is_valid_for(&self, task_config: &DapTaskConfig) -> bool {
-        if self.start % task_config.time_precision != 0
-            || self.duration % task_config.time_precision != 0
-            || self.duration < task_config.time_precision
-        {
-            return false;
-        }
-        true
-    }
 }
 
 impl Encode for Interval {
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 600bd9f..11cef7b 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -12,12 +12,12 @@ use crate::{
     hpke::HpkeDecrypter,
     messages::{
         constant_time_eq, AggregateContinueReq, AggregateInitializeReq, AggregateResp,
-        AggregateShareReq, AggregateShareResp, BatchSelector, CollectReq, CollectResp, Id, Nonce,
-        Report, ReportShare, TransitionFailure, TransitionVar,
+        AggregateShareReq, AggregateShareResp, BatchParameter, BatchSelector, CollectReq,
+        CollectResp, Id, Nonce, Report, ReportShare, Time, TransitionFailure, TransitionVar,
     },
     DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
     DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition, DapOutputShare,
-    DapRequest, DapResponse, DapTaskConfig, DapVersion,
+    DapQueryConfig, DapRequest, DapResponse, DapTaskConfig, DapVersion,
 };
 use async_trait::async_trait;
 use prio::codec::{Decode, Encode};
@@ -25,48 +25,6 @@ use rand::prelude::*;
 use std::collections::HashMap;
 use url::Url;
 
-macro_rules! check_batch_param {
-    (
-        $task_config:expr,
-        $batch_selector:expr,
-        $agg_param:expr
-    ) => {{
-        let batch_interval = $batch_selector.unwrap_interval();
-        if !batch_interval.is_valid_for($task_config) {
-            return Err(DapAbort::InvalidBatchInterval);
-        }
-
-        // Check that the aggreation parameter is suitable for the given VDAF.
-        if !$task_config.vdaf.is_valid_agg_param(&$agg_param) {
-            // TODO spec: Define this behavior.
-            return Err(DapAbort::UnrecognizedMessage);
-        }
-    }};
-}
-
-macro_rules! check_batch_boundary {
-    (
-        $global_config:expr,
-        $now:expr,
-        $batch_selector:expr
-    ) => {{
-        let batch_interval = $batch_selector.unwrap_interval();
-        if batch_interval.duration > $global_config.max_batch_duration {
-            return Err(DapAbort::BadRequest("batch interval too large".to_string()));
-        }
-        if $now.abs_diff(batch_interval.start) > $global_config.min_batch_interval_start {
-            return Err(DapAbort::BadRequest(
-                "batch interval too far into past".to_string(),
-            ));
-        }
-        if $now.abs_diff(batch_interval.end()) > $global_config.max_batch_interval_end {
-            return Err(DapAbort::BadRequest(
-                "batch interval too far into future".to_string(),
-            ));
-        }
-    }};
-}
-
 /// A party in the DAP protocol who is authorized to send requests to another party.
 #[async_trait(?Send)]
 pub trait DapAuthorizedSender<S> {
@@ -282,6 +240,8 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
     /// The return value is a URI that the Collector can poll later on to get the corresponding
     /// [`CollectResp`](crate::messages::CollectResp).
     async fn http_post_collect(&'a self, req: &DapRequest<S>) -> Result<Url, DapAbort> {
+        let now = self.get_current_time();
+
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -297,36 +257,25 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
             .await?
             .ok_or(DapAbort::UnrecognizedTask)?;
         let task_config = wrapped_task_config.as_ref();
-        check_batch_param!(task_config, collect_req.query, collect_req.agg_param);
 
         // Check whether the DAP version in the request matches the task config.
         if task_config.version != req.version {
             return Err(DapAbort::InvalidProtocolVersion);
         }
 
-        // Check that the batch boundaries.
-        let global_config = self.get_global_config();
-        let now = self.get_current_time();
-        check_batch_boundary!(global_config, now, collect_req.query);
-
-        // Check that the batch does not overlap with previous batches.
-        if self
-            .is_batch_overlapping(&collect_req.task_id, &collect_req.query)
-            .await?
-        {
-            return Err(DapAbort::BatchOverlap);
-        }
-
-        let batch_interval = collect_req.query.unwrap_interval();
-        if !batch_interval.is_valid_for(task_config) {
-            return Err(DapAbort::InvalidBatchInterval);
-        }
-
-        // Check that the aggreation parameter is suitable for the given VDAF.
-        if !task_config.vdaf.is_valid_agg_param(&collect_req.agg_param) {
-            // TODO spec: Define this behavior.
-            return Err(DapAbort::UnrecognizedMessage);
-        }
+        // Ensure the batch boundaries are valid and that the batch doesn't overlap with previosuly
+        // collected batches.
+        //
+        // TODO(issue #100) For "fixed_size" we need to check if we recognize the batch ID.
+        check_batch(
+            self,
+            task_config,
+            &collect_req.task_id,
+            &collect_req.query,
+            &collect_req.agg_param,
+            now,
+        )
+        .await?;
 
         Ok(self.init_collect_job(&collect_req).await?)
     }
@@ -587,6 +536,13 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
                     return Err(DapAbort::InvalidProtocolVersion);
                 }
 
+                // Ensure we know which batch the request pertains to.
+                check_part_batch(
+                    task_config,
+                    &agg_init_req.batch_param,
+                    &agg_init_req.agg_param,
+                )?;
+
                 // Run mark_aggregated and handle_agg_init_req concurrently.
                 let early_rejects_future =
                     self.mark_aggregated(&agg_init_req.task_id, &agg_init_req.report_shares);
@@ -694,6 +650,8 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
         &'a self,
         req: &DapRequest<S>,
     ) -> Result<DapResponse, DapAbort> {
+        let now = self.get_current_time();
+
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -709,29 +667,26 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
             .await?
             .ok_or(DapAbort::UnrecognizedTask)?;
         let task_config = wrapped_task_config.as_ref();
-        check_batch_param!(
-            task_config,
-            agg_share_req.batch_selector,
-            agg_share_req.agg_param
-        );
 
         // Check whether the DAP version in the request matches the task config.
         if task_config.version != req.version {
             return Err(DapAbort::InvalidProtocolVersion);
         }
 
-        // Check the batch boundaries.
-        let global_config = self.get_global_config();
-        let now = self.get_current_time();
-        check_batch_boundary!(global_config, now, agg_share_req.batch_selector);
+        // Ensure the batch boundaries are valid and that the batch doesn't overlap with previosuly
+        // collected batches.
+        //
+        // TODO(issue #100) For "fixed_size" we need to check if we recognize the batch ID.
+        check_batch(
+            self,
+            task_config,
+            &agg_share_req.task_id,
+            &agg_share_req.batch_selector,
+            &agg_share_req.agg_param,
+            now,
+        )
+        .await?;
 
-        // Check that the batch does not overlap with previous boundaries.
-        if self
-            .is_batch_overlapping(&agg_share_req.task_id, &agg_share_req.batch_selector)
-            .await?
-        {
-            return Err(DapAbort::BatchOverlap);
-        }
         let agg_share = self
             .get_agg_share(&agg_share_req.task_id, &agg_share_req.batch_selector)
             .await?;
@@ -773,3 +728,78 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
         })
     }
 }
+
+fn check_part_batch(
+    task_config: &DapTaskConfig,
+    batch_param: &BatchParameter,
+    agg_param: &[u8],
+) -> Result<(), DapAbort> {
+    match (&task_config.query, batch_param) {
+        (DapQueryConfig::TimeInterval { .. }, BatchParameter::TimeInterval)
+        | (DapQueryConfig::FixedSize { .. }, BatchParameter::FixedSize { .. }) => (),
+        _ => return Err(DapAbort::QueryMismatch),
+    };
+
+    // Check that the aggreation parameter is suitable for the given VDAF.
+    if !task_config.vdaf.is_valid_agg_param(agg_param) {
+        // TODO spec: Define this behavior.
+        return Err(DapAbort::UnrecognizedMessage);
+    }
+
+    Ok(())
+}
+
+async fn check_batch<'a, S>(
+    agg: &impl DapAggregator<'a, S>,
+    task_config: &DapTaskConfig,
+    task_id: &Id,
+    batch_sel: &BatchSelector,
+    agg_param: &[u8],
+    now: Time,
+) -> Result<(), DapAbort> {
+    let global_config = agg.get_global_config();
+    let batch_overlapping = agg.is_batch_overlapping(task_id, batch_sel);
+
+    // Check that the aggreation parameter is suitable for the given VDAF.
+    if !task_config.vdaf.is_valid_agg_param(agg_param) {
+        // TODO spec: Define this behavior.
+        return Err(DapAbort::UnrecognizedMessage);
+    }
+
+    // Check that the batch boundaries are valid.
+    match (&task_config.query, batch_sel) {
+        (DapQueryConfig::TimeInterval { .. }, BatchSelector::TimeInterval { batch_interval }) => {
+            if batch_interval.start % task_config.time_precision != 0
+                || batch_interval.duration % task_config.time_precision != 0
+                || batch_interval.duration < task_config.time_precision
+            {
+                return Err(DapAbort::BatchInvalid);
+            }
+
+            if batch_interval.duration > global_config.max_batch_duration {
+                return Err(DapAbort::BadRequest("batch interval too large".to_string()));
+            }
+
+            if now.abs_diff(batch_interval.start) > global_config.min_batch_interval_start {
+                return Err(DapAbort::BadRequest(
+                    "batch interval too far into past".to_string(),
+                ));
+            }
+
+            if now.abs_diff(batch_interval.end()) > global_config.max_batch_interval_end {
+                return Err(DapAbort::BadRequest(
+                    "batch interval too far into future".to_string(),
+                ));
+            }
+        }
+        (DapQueryConfig::FixedSize { .. }, BatchSelector::FixedSize { .. }) => (),
+        _ => return Err(DapAbort::QueryMismatch),
+    };
+
+    // Check that the batch does not overlap with any previously collected batch.
+    if batch_overlapping.await? {
+        return Err(DapAbort::BatchOverlap);
+    }
+
+    Ok(())
+}


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index cff727f..92c3fc4 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -466,6 +466,38 @@ impl Test {
     }
 }
 
+// Test that the Helper properly handles the batch parameter in the AggregateInitializeReq.
+#[tokio::test]
+async fn http_post_aggregate_invalid_batch_selector() {
+    let mut rng = thread_rng();
+    let t = Test::new();
+    let task_id = &t.time_interval_task_id;
+    let task_config = t.leader.tasks.get(task_id).unwrap();
+
+    // Helper expects "time_interval" query, but Leader indicates "fixed_size".
+    let req = t
+        .leader_authorized_req(
+            task_id,
+            task_config.version,
+            MEDIA_TYPE_AGG_INIT_REQ,
+            AggregateInitializeReq {
+                task_id: task_id.clone(),
+                agg_job_id: Id(rng.gen()),
+                agg_param: Vec::default(),
+                batch_param: BatchParameter::FixedSize {
+                    batch_id: Id(rng.gen()),
+                },
+                report_shares: Vec::default(),
+            },
+            task_config.helper_url.join("aggregate").unwrap(),
+        )
+        .await;
+    assert_matches!(
+        t.helper.http_post_aggregate(&req).await.unwrap_err(),
+        DapAbort::QueryMismatch
+    );
+}
+
 #[tokio::test]
 async fn http_post_aggregate_init_unauthorized_request() {
     let t = Test::new();
@@ -572,6 +604,38 @@ async fn http_post_aggregate_share_unauthorized_request() {
     );
 }
 
+// Test that the Helper handles the batch selector sent from the Leader properly.
+#[tokio::test]
+async fn http_post_aggregate_share_invalid_batch_selector() {
+    let mut rng = thread_rng();
+    let t = Test::new();
+    let task_id = &t.time_interval_task_id;
+    let task_config = t.leader.tasks.get(task_id).unwrap();
+
+    // Helper expects "time_interval" query, but Leader sent "fixed_size".
+    let req = t
+        .leader_authorized_req(
+            task_id,
+            task_config.version,
+            MEDIA_TYPE_AGG_SHARE_REQ,
+            AggregateShareReq {
+                task_id: task_id.clone(),
+                batch_selector: BatchSelector::FixedSize {
+                    batch_id: Id(rng.gen()),
+                },
+                agg_param: Vec::default(),
+                report_count: 0,
+                checksum: [0; 32],
+            },
+            task_config.helper_url.join("aggregate_share").unwrap(),
+        )
+        .await;
+    assert_matches!(
+        t.helper.http_post_aggregate_share(&req).await.unwrap_err(),
+        DapAbort::QueryMismatch
+    );
+}
+
 #[tokio::test]
 async fn http_post_collect_unauthorized_request() {
     let t = Test::new();
@@ -1125,6 +1189,35 @@ async fn http_post_collect_success() {
     );
 }
 
+// Test that the Leader handles queries from the Collector properly.
+#[tokio::test]
+async fn http_post_collect_invalid_query() {
+    let mut rng = thread_rng();
+    let t = Test::new();
+    let task_id = &t.time_interval_task_id;
+    let task_config = t.leader.tasks.get(task_id).unwrap();
+
+    // Leader expects "time_interval" query, but Collector sent "fixed_size".
+    let req = t
+        .collector_authorized_req(
+            task_config.version,
+            MEDIA_TYPE_COLLECT_REQ,
+            CollectReq {
+                task_id: task_id.clone(),
+                query: Query::FixedSize {
+                    batch_id: Id(rng.gen()),
+                },
+                agg_param: Vec::default(),
+            },
+            task_config.leader_url.join("collect").unwrap(),
+        )
+        .await;
+    assert_matches!(
+        t.leader.http_post_collect(&req).await.unwrap_err(),
+        DapAbort::QueryMismatch
+    );
+}
+
 // Test HTTP POST requests with a wrong DAP version.
 #[tokio::test]
 async fn http_post_fail_wrong_dap_version() {
diff --git a/daphne_worker_test/tests/e2e.rs b/daphne_worker_test/tests/e2e.rs
index b5bd731..be8f9d4 100644
--- a/daphne_worker_test/tests/e2e.rs
+++ b/daphne_worker_test/tests/e2e.rs
@@ -582,7 +582,7 @@ async fn e2e_leader_collect_abort_invalid_batch_interval() {
         constants::MEDIA_TYPE_COLLECT_REQ,
         collect_req.get_encoded(),
         400,
-        "invalidBatchInterval",
+        "batchInvalid",
     )
     .await;
 
@@ -604,7 +604,7 @@ async fn e2e_leader_collect_abort_invalid_batch_interval() {
         constants::MEDIA_TYPE_COLLECT_REQ,
         collect_req.get_encoded(),
         400,
-        "invalidBatchInterval",
+        "batchInvalid",
     )
     .await;
 }