Address comments #2

Author: jhoyla

crates/daphne-server/tests/e2e/e2e.rs

@@ -557,58 +557,6 @@ async fn leader_upload_taskprov_wrong_version(version: DapVersion) {
 
 async_test_versions!(leader_upload_taskprov_wrong_version);
 
-#[tokio::test]
-async fn leader_upload_taskprov_public() {
-    let version = DapVersion::Latest;
-    let t = TestRunner::default_with_version(version).await;
-    let client = t.http_client();
-    let hpke_config_list = t.get_hpke_configs(version, client).await.unwrap();
-
-    let (task_config, task_id, taskprov_advertisement) = DapTaskParameters {
-        version,
-        min_batch_size: 10,
-        query: DapBatchMode::TimeInterval,
-        leader_url: t.task_config.leader_url.clone(),
-        helper_url: t.task_config.helper_url.clone(),
-        ..Default::default()
-    }
-    .to_config_with_taskprov(
-        b"cool task".to_vec(),
-        t.now,
-        daphne::roles::aggregator::TaskprovConfig {
-            hpke_collector_config: &t.taskprov_collector_hpke_receiver.config,
-            vdaf_verify_key_init: &t.taskprov_vdaf_verify_key_init,
-        },
-    )
-    .unwrap();
-
-    let mut report = task_config
-        .vdaf
-        .produce_report(
-            &hpke_config_list,
-            t.now + 1,
-            &task_id,
-            DapMeasurement::U32Vec(vec![1; 10]),
-            version,
-        )
-        .unwrap();
-    report.report_metadata.public_extensions = Some(vec![Extension::Taskprov]);
-    t.leader_request_expect_ok(
-        client,
-        &format!("tasks/{}/reports", task_id.to_base64url()),
-        &http::Method::POST,
-        DapMediaType::Report,
-        Some(
-            &taskprov_advertisement
-                .serialize_to_header_value(version)
-                .unwrap(),
-        ),
-        report.get_encoded_with_param(&version).unwrap(),
-    )
-    .await
-    .unwrap();
-}
-
 #[tokio::test]
 async fn leader_upload_taksprov_public_errors() {
     let version = DapVersion::Latest;
@@ -635,17 +583,18 @@ async fn leader_upload_taksprov_public_errors() {
     .unwrap();
 
     // Repeated public extension
-    let mut report = task_config
+    let report = task_config
         .vdaf
-        .produce_report(
+        .produce_report_with_extensions(
             &hpke_config_list,
             t.now + 1,
             &task_id,
             DapMeasurement::U32Vec(vec![1; 10]),
+            Some(vec![Extension::Taskprov, Extension::Taskprov]),
+            vec![],
             version,
         )
         .unwrap();
-    report.report_metadata.public_extensions = Some(vec![Extension::Taskprov, Extension::Taskprov]);
     t.leader_request_expect_abort(
         client,
         None,
@@ -665,23 +614,24 @@ async fn leader_upload_taksprov_public_errors() {
     .unwrap();
 
     // Unsupported public extension
-    let mut report = task_config
+    let report = task_config
         .vdaf
-        .produce_report(
+        .produce_report_with_extensions(
             &hpke_config_list,
             t.now + 1,
             &task_id,
             DapMeasurement::U32Vec(vec![1; 10]),
+            Some(vec![
+                Extension::Taskprov,
+                Extension::NotImplemented {
+                    typ: 3,
+                    payload: b"ignore".to_vec(),
+                },
+            ]),
+            vec![],
             version,
         )
         .unwrap();
-    report.report_metadata.public_extensions = Some(vec![
-        Extension::Taskprov,
-        Extension::NotImplemented {
-            typ: 3,
-            payload: b"ignore".to_vec(),
-        },
-    ]);
     t.leader_request_expect_abort(
         client,
         None,
@@ -1514,116 +1464,6 @@ async fn leader_selected() {
     .unwrap();
 }
 
-#[tokio::test]
-async fn leader_collect_taskprov_repeated_abort() {
-    let version = DapVersion::Latest;
-    const DAP_TASKPROV_COLLECTOR_TOKEN: &str = "I-am-the-collector";
-    let t = TestRunner::default_with_version(version).await;
-    let batch_interval = t.batch_interval();
-
-    let client = t.http_client();
-    let hpke_config_list = t.get_hpke_configs(version, client).await.unwrap();
-
-    let (task_config, task_id, taskprov_advertisement) = DapTaskParameters {
-        version,
-        min_batch_size: 10,
-        query: DapBatchMode::TimeInterval,
-        leader_url: t.task_config.leader_url.clone(),
-        helper_url: t.task_config.helper_url.clone(),
-        ..Default::default()
-    }
-    .to_config_with_taskprov(
-        b"cool task".to_vec(),
-        t.now,
-        daphne::roles::aggregator::TaskprovConfig {
-            hpke_collector_config: &t.taskprov_collector_hpke_receiver.config,
-            vdaf_verify_key_init: &t.taskprov_vdaf_verify_key_init,
-        },
-    )
-    .unwrap();
-
-    let path = TestRunner::upload_path_for_task(&task_id);
-    let method = &Method::POST;
-    // The reports are uploaded in the background.
-    let mut rng = thread_rng();
-    for _ in 0..t.task_config.min_batch_size {
-        let extensions = vec![Extension::Taskprov];
-        let now = rng.gen_range(TestRunner::report_interval(&batch_interval));
-        t.leader_request_expect_ok(
-            client,
-            &path,
-            method,
-            DapMediaType::Report,
-            Some(
-                &taskprov_advertisement
-                    .serialize_to_header_value(version)
-                    .unwrap(),
-            ),
-            {
-                let report = task_config
-                    .vdaf
-                    .produce_report_with_extensions(
-                        &hpke_config_list,
-                        now,
-                        &task_id,
-                        DapMeasurement::U32Vec(vec![1; 10]),
-                        Some(vec![Extension::Taskprov]),
-                        extensions,
-                        version,
-                    )
-                    .unwrap();
-                report.get_encoded_with_param(&version).unwrap()
-            },
-        )
-        .await
-        .unwrap();
-    }
-
-    let agg_param = DapAggregationParam::Empty;
-
-    // Get the collect URI.
-    let collect_req = CollectionReq {
-        query: Query::TimeInterval { batch_interval },
-        agg_param: agg_param.get_encoded().unwrap(),
-    };
-    let collect_uri = t
-        .leader_post_collect_using_token(
-            client,
-            DAP_TASKPROV_COLLECTOR_TOKEN,
-            Some(&taskprov_advertisement),
-            Some(&task_id),
-            collect_req.get_encoded_with_param(&t.version).unwrap(),
-        )
-        .await
-        .unwrap();
-    println!("collect_uri: {collect_uri}");
-
-    // Poll the collect URI before the CollectResp is ready.
-    let resp = t
-        .poll_collection_url_using_token(client, &collect_uri, DAP_TASKPROV_COLLECTOR_TOKEN)
-        .await
-        .unwrap();
-    #[expect(clippy::format_in_format_args)]
-    {
-        assert_eq!(
-            resp.status(),
-            202,
-            "response: {} {}",
-            format!("{resp:?}"),
-            resp.text().await.unwrap()
-        );
-    }
-
-    // The reports are aggregated in the background.
-    let agg_telem = t.internal_process(client).await.unwrap();
-    assert_eq!(
-        agg_telem.reports_processed, task_config.min_batch_size,
-        "reports processed"
-    );
-    assert_eq!(agg_telem.reports_aggregated, 0, "reports aggregated");
-    assert_eq!(agg_telem.reports_collected, 0, "reports collected");
-}
-
 async fn leader_collect_taskprov_ok(version: DapVersion) {
     const DAP_TASKPROV_COLLECTOR_TOKEN: &str = "I-am-the-collector";
     let t = TestRunner::default_with_version(version).await;

crates/daphne/src/messages/mod.rs

@@ -256,7 +256,11 @@ impl ParameterizedEncode<DapVersion> for ReportMetadata {
             (DapVersion::Latest, Some(extensions)) => {
                 encode_u16_items(bytes, version, extensions.as_slice())?;
             }
-            _ => return Err(CodecError::UnexpectedValue),
+            _ => {
+                return Err(CodecError::Other(
+                    "encountered incorrectly set public extensions".into(),
+                ))
+            }
         }
 
         Ok(())
@@ -1638,7 +1642,7 @@ mod test {
         };
         assert!(matches!(
             bad_rm.get_encoded_with_param(&version).unwrap_err(),
-            CodecError::UnexpectedValue
+            CodecError::Other(..)
         ));
         let bytes = good_rm.get_encoded_with_param(&version).unwrap();
         assert_eq!(

crates/daphne/src/protocol/client.rs

@@ -3,6 +3,7 @@
 
 use crate::{
     constants::DapAggregatorRole,
+    fatal_error,
     hpke::{info_and_aad, HpkeConfig},
     messages::{Extension, PlaintextInputShare, Report, ReportId, ReportMetadata, TaskId, Time},
     DapError, DapMeasurement, DapVersion, VdafConfig,
@@ -72,13 +73,12 @@ impl VdafConfig {
         private_extensions: Vec<Extension>,
         version: DapVersion,
     ) -> Result<Report, DapError> {
-        match (&public_extensions, version) {
-            (Some(_), DapVersion::Draft09) | (None, DapVersion::Latest) => {
-                return Err(DapError::ReportError(
-                    crate::messages::ReportError::InvalidMessage,
-                ))
-            }
-            _ => (),
+        if let (Some(_), DapVersion::Draft09) | (None, DapVersion::Latest) =
+            (&public_extensions, version)
+        {
+            return Err(fatal_error!(
+                err = format!("public extensions not set correctly for {version:?}")
+            ));
         }
 
         let mut plaintext_input_share = PlaintextInputShare {

crates/daphne/src/roles/leader/mod.rs

@@ -227,8 +227,8 @@ pub async fn handle_upload_req<A: DapLeader>(
     }
 
     if let Some(public_extensions) = &report.report_metadata.public_extensions {
-        // We can be sure at this point that the ReportMetadata is well formed (as
-        // the decoding / error checking happens in the extractor).
+        // We can be sure at this point that the ReportMetadata is well formed
+        // because the decoding / error checking happens in the extractor.
         assert_eq!(DapVersion::Latest, task_config.version);
         let mut unknown_extensions = Vec::<u16>::new();
         if protocol::check_no_duplicates(public_extensions.iter()).is_err() {

crates/daphne/src/roles/mod.rs

@@ -779,11 +779,36 @@ mod test {
         let t = Test::new(version);
         let task_id = &t.time_interval_task_id;
         let task_config = t.leader.unchecked_get_task_config(task_id).await;
-        let mut report = t.gen_test_report(task_id).await;
-        report.report_metadata.public_extensions = Some(vec![Extension::NotImplemented {
-            typ: 0x01,
-            payload: vec![0x01],
-        }]);
+
+        // Construct HPKE config list.
+        let hpke_config_list = [
+            t.leader
+                .get_hpke_config_for(task_config.version, Some(task_id))
+                .await
+                .unwrap()
+                .clone(),
+            t.helper
+                .get_hpke_config_for(task_config.version, Some(task_id))
+                .await
+                .unwrap()
+                .clone(),
+        ];
+
+        let report = task_config
+            .vdaf
+            .produce_report_with_extensions(
+                &hpke_config_list,
+                t.now,
+                task_id,
+                DapMeasurement::U32Vec(vec![1; 10]),
+                Some(vec![Extension::NotImplemented {
+                    typ: 0x01,
+                    payload: vec![0x01],
+                }]),
+                vec![],
+                task_config.version,
+            )
+            .unwrap();
 
         let req = DapRequest {
             meta: DapRequestMeta {
@@ -812,6 +837,8 @@ mod test {
         let task_id = &t.time_interval_task_id;
         let task_config = t.leader.unchecked_get_task_config(task_id).await;
         let mut report = t.gen_test_report(task_id).await;
+        // This change breaks the HPKE decryption, but triggers a failure
+        // before the HPKE data is checked.
         report.report_metadata.public_extensions = Some(vec![]);
 
         let req = DapRequest {