Add option to ignore unrecognized DpMechanism (#2728)

aggregator/src/aggregator.rs

@@ -56,7 +56,7 @@ use janus_core::{
 };
 use janus_messages::{
     query_type::{FixedSize, TimeInterval},
-    taskprov::TaskConfig,
+    taskprov::{DpMechanism, TaskConfig},
     AggregateShare, AggregateShareAad, AggregateShareReq, AggregationJobContinueReq,
     AggregationJobId, AggregationJobInitializeReq, AggregationJobResp, AggregationJobStep,
     BatchSelector, Collection, CollectionJobId, CollectionReq, Duration, ExtensionType, HpkeConfig,
@@ -722,6 +722,21 @@ impl<C: Clock> Aggregator<C> {
         // TODO(#1647): Check whether task config parameters are acceptable for privacy and
         // availability of the system.
 
+        if let DpMechanism::Unrecognized { .. } =
+            task_config.vdaf_config().dp_config().dp_mechanism()
+        {
+            if !self
+                .cfg
+                .taskprov_config
+                .ignore_unknown_differential_privacy_mechanism
+            {
+                return Err(Error::InvalidTask(
+                    *task_id,
+                    OptOutReason::InvalidParameter("unrecognized DP mechanism".into()),
+                ));
+            }
+        }
+
         let vdaf_instance =
             task_config
                 .vdaf_config()

aggregator/src/aggregator/http_handlers.rs

@@ -1074,7 +1074,10 @@ mod tests {
             .unwrap();
 
         let cfg = Config {
-            taskprov_config: TaskprovConfig { enabled: true },
+            taskprov_config: TaskprovConfig {
+                enabled: true,
+                ignore_unknown_differential_privacy_mechanism: false,
+            },
             ..Default::default()
         };
 

aggregator/src/aggregator/taskprov_tests.rs

@@ -122,7 +122,10 @@ impl TaskprovTestCase {
             TestRuntime::default(),
             &noop_meter(),
             Config {
-                taskprov_config: TaskprovConfig { enabled: true },
+                taskprov_config: TaskprovConfig {
+                    enabled: true,
+                    ignore_unknown_differential_privacy_mechanism: false,
+                },
                 ..Default::default()
             },
         )

aggregator/src/binaries/aggregator.rs

@@ -619,7 +619,10 @@ mod tests {
             )
             .unwrap()
             .taskprov_config,
-            TaskprovConfig { enabled: true },
+            TaskprovConfig {
+                enabled: true,
+                ignore_unknown_differential_privacy_mechanism: false
+            },
         );
     }
 

aggregator/src/config.rs

@@ -117,6 +117,16 @@ pub struct TaskprovConfig {
     ///
     /// [spec]: https://datatracker.ietf.org/doc/draft-wang-ppm-dap-taskprov/
     pub enabled: bool,
+
+    /// If true, will silently ignore unknown differential privacy mechanisms, and continue without
+    /// differential privacy noise.
+    ///
+    /// This should only be used for testing purposes. This option will be removed once full
+    /// differential privacy support is implemented.
+    ///
+    /// Defaults to false, i.e. opt out of tasks with unrecognized differential privacy mechanisms.
+    #[serde(default)]
+    pub ignore_unknown_differential_privacy_mechanism: bool,
 }
 
 /// Non-secret configuration options for Janus Job Driver jobs.

aggregator/tests/integration/graceful_shutdown.rs

@@ -253,7 +253,7 @@ async fn aggregator_shutdown() {
             metrics_config: MetricsConfiguration::default(),
             health_check_listen_address: "127.0.0.1:9001".parse().unwrap(),
         },
-        taskprov_config: TaskprovConfig { enabled: false },
+        taskprov_config: TaskprovConfig::default(),
         garbage_collection: None,
         listen_address: aggregator_listen_address,
         aggregator_api: Some(AggregatorApi {
@@ -324,7 +324,7 @@ async fn aggregation_job_driver_shutdown() {
             retry_max_interval_millis: 30_000,
             retry_max_elapsed_time_millis: 300_000,
         },
-        taskprov_config: TaskprovConfig { enabled: false },
+        taskprov_config: TaskprovConfig::default(),
         batch_aggregation_shard_count: 32,
     };
 

messages/src/taskprov.rs

@@ -282,8 +282,8 @@ impl VdafConfig {
         })
     }
 
-    pub fn dp_config(&self) -> DpConfig {
-        self.dp_config
+    pub fn dp_config(&self) -> &DpConfig {
+        &self.dp_config
     }
 
     pub fn vdaf_type(&self) -> &VdafType {
@@ -475,7 +475,7 @@ impl Decode for VdafType {
 /// See [draft-irtf-cfrg-vdaf/#94][1] for discussion.
 ///
 /// [1]: https://github.com/cfrg/draft-irtf-cfrg-vdaf/issues/94
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct DpConfig {
     dp_mechanism: DpMechanism,
 }
@@ -508,12 +508,13 @@ impl Decode for DpConfig {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 #[repr(u8)]
 #[non_exhaustive]
 pub enum DpMechanism {
     Reserved,
     None,
+    Unrecognized { codepoint: u8, payload: Vec<u8> },
 }
 
 impl DpMechanism {
@@ -526,12 +527,22 @@ impl Encode for DpMechanism {
         match self {
             Self::Reserved => Self::RESERVED.encode(bytes)?,
             Self::None => Self::NONE.encode(bytes)?,
+            Self::Unrecognized { codepoint, payload } => {
+                codepoint.encode(bytes)?;
+                bytes.extend_from_slice(payload);
+            }
         };
         Ok(())
     }
 
     fn encoded_len(&self) -> Option<usize> {
-        Some(1)
+        match self {
+            DpMechanism::Reserved | DpMechanism::None => Some(1),
+            DpMechanism::Unrecognized {
+                codepoint: _,
+                payload,
+            } => Some(1 + payload.len()),
+        }
     }
 }
 
@@ -541,10 +552,16 @@ impl Decode for DpMechanism {
         let dp_mechanism = match dp_mechanism_code {
             Self::RESERVED => Self::Reserved,
             Self::NONE => Self::None,
-            val => {
-                return Err(CodecError::Other(
-                    anyhow!("unexpected DpMechanism value {}", val).into(),
-                ))
+            codepoint => {
+                let position = usize::try_from(bytes.position())
+                    .map_err(|_| CodecError::Other(anyhow!("cursor position overflow").into()))?;
+                let inner = bytes.get_ref();
+                let payload = inner[position..].to_vec();
+                bytes
+                    .set_position(u64::try_from(inner.len()).map_err(|_| {
+                        CodecError::Other(anyhow!("cursor length overflow").into())
+                    })?);
+                Self::Unrecognized { codepoint, payload }
             }
         };
 
@@ -573,6 +590,16 @@ mod tests {
                     "01",   // dp_mechanism
                 ),
             ),
+            (
+                DpConfig::new(DpMechanism::Unrecognized {
+                    codepoint: 0xff,
+                    payload: Vec::from([0xde, 0xad, 0xbe, 0xef]),
+                }),
+                concat!(
+                    "FF",       // dp_mechanism
+                    "DEADBEEF"  // uninterpreted DpConfig contents
+                ),
+            ),
         ])
     }
 
@@ -686,6 +713,28 @@ mod tests {
                     ),
                 ),
             ),
+            (
+                VdafConfig::new(
+                    DpConfig::new(DpMechanism::Unrecognized {
+                        codepoint: 0xFF,
+                        payload: Vec::from([0xDE, 0xAD, 0xBE, 0xEF]),
+                    }),
+                    VdafType::Prio3Count,
+                )
+                .unwrap(),
+                concat!(
+                    // dp_config
+                    concat!(
+                        "0005",     // dp_config length
+                        "FF",       // dp_mechanism
+                        "DEADBEEF"  // rest of unrecognized DpConfig
+                    ),
+                    // vdaf_type
+                    concat!(
+                        "00000000" // vdaf_type_code
+                    )
+                ),
+            ),
             (
                 VdafConfig::new(
                     DpConfig::new(DpMechanism::None),