[DAP-18] TaskConf AAD Part 3: make AggregatorTask self-sufficient

aggregator/src/aggregator.rs

@@ -849,7 +849,10 @@ impl<C: Clock> Aggregator<C> {
         task_config: &TaskConfiguration,
         aggregator_auth_token: Option<&AuthenticationToken>,
     ) -> Result<(), Error> {
-        let (peer_aggregator, leader_url, _) = self
+        // The peer is always the Leader here, so `leader_url` is the peer endpoint and `helper_url`
+        // is this Helper's own endpoint. Both come byte-preserving from the wire TaskConfiguration.
+        debug_assert_eq!(peer_role, &Role::Leader);
+        let (peer_aggregator, leader_url, helper_url) = self
             .taskprov_authorize_request(peer_role, task_id, task_config, aggregator_auth_token)
             .await?;
 
@@ -874,6 +877,7 @@ impl<C: Clock> Aggregator<C> {
             AggregatorTask::new(
                 *task_id,
                 leader_url,
+                helper_url,
                 BatchMode::try_from(task_config.batch_config())?,
                 vdaf_instance,
                 vdaf_verify_key,
@@ -894,7 +898,10 @@ impl<C: Clock> Aggregator<C> {
                     )?,
                 },
             )
-            .map_err(|err| Error::InvalidTask(*task_id, OptOutReason::TaskParameters(err)))?,
+            .map_err(|err| Error::InvalidTask(*task_id, OptOutReason::TaskParameters(err)))?
+            // Bind the received TaskConfiguration verbatim into HPKE AADs; it is not
+            // byte-reconstructible from the stored parameters.
+            .with_taskprov_task_config(task_config.clone()),
         );
         self.datastore
             .run_tx("taskprov_put_task", |tx| {

aggregator/src/aggregator/taskprov_tests.rs

@@ -431,7 +431,13 @@ async fn taskprov_aggregate_init() {
                 .eq(&AggregationJobState::AwaitingRequest)
     );
     let got_task = got_task.unwrap();
-    assert_eq!(test.task.taskprov_helper_view().unwrap(), got_task);
+    assert_eq!(
+        test.task
+            .taskprov_helper_view()
+            .unwrap()
+            .with_taskprov_task_config(test.task_config.clone()),
+        got_task
+    );
     assert_eq!(got_task.task_info(), b"foobar".as_slice());
 }
 
@@ -596,7 +602,13 @@ async fn taskprov_aggregate_init_missing_extension() {
                 .state()
                 .eq(&AggregationJobState::Finished)
     );
-    assert_eq!(test.task.taskprov_helper_view().unwrap(), got_task.unwrap());
+    assert_eq!(
+        test.task
+            .taskprov_helper_view()
+            .unwrap()
+            .with_taskprov_task_config(test.task_config.clone()),
+        got_task.unwrap()
+    );
 }
 
 #[tokio::test]
@@ -691,7 +703,13 @@ async fn taskprov_aggregate_init_malformed_extension() {
                 .state()
                 .eq(&AggregationJobState::Finished)
     );
-    assert_eq!(test.task.taskprov_helper_view().unwrap(), got_task.unwrap());
+    assert_eq!(
+        test.task
+            .taskprov_helper_view()
+            .unwrap()
+            .with_taskprov_task_config(test.task_config.clone()),
+        got_task.unwrap()
+    );
 }
 
 /// The helper should not opt out of the task if the current time is past the task end time.

aggregator/src/binaries/aggregator.rs

@@ -15,12 +15,12 @@ use educe::Educe;
 use janus_aggregator_api::{self, aggregator_api_handler};
 use janus_aggregator_core::datastore::Datastore;
 use janus_core::{TokioRuntime, auth_tokens::AuthenticationToken, time::RealClock};
+use janus_messages::Url as DapUrl;
 use opentelemetry::metrics::Meter;
 use sec1::EcPrivateKey;
 use serde::{Deserialize, Deserializer, Serialize, de};
 use tokio::{spawn, sync::watch, time::interval, try_join};
 use tracing::{error, info};
-use url::Url;
 
 use crate::{
     aggregator::{
@@ -250,7 +250,7 @@ pub struct AggregatorApi {
     /// Resource location at which the DAP service managed by this aggregator api can be found
     /// on the public internet. Required.
     #[educe(Debug(method(std::fmt::Display::fmt)))]
-    pub public_dap_url: Url,
+    pub public_dap_url: DapUrl,
 }
 
 fn deserialize_aggregator_api<'de, D>(deserializer: D) -> Result<Option<AggregatorApi>, D::Error>

aggregator/src/binaries/janus_cli.rs

@@ -1619,6 +1619,7 @@ mod tests {
         // or HPKE keys.
         let serialized_task_yaml = r#"
 - peer_aggregator_endpoint: https://helper
+  own_aggregator_endpoint: https://leader
   batch_mode: TimeInterval
   vdaf: !Prio3Sum
     max_measurement: 4096
@@ -1644,6 +1645,7 @@ mod tests {
     hash: MJOoBO_ysLEuG_lv2C37eEOf1Ngetsr-Ers0ZYj4vdQ
   hpke_keys: []
 - peer_aggregator_endpoint: https://leader
+  own_aggregator_endpoint: https://helper
   batch_mode: TimeInterval
   aggregation_mode: Asynchronous
   vdaf: !Prio3Sum

aggregator_api/src/lib.rs

@@ -20,18 +20,17 @@ use janus_aggregator_core::{
     http_server::{HttpMetrics, http_metrics_middleware},
 };
 use janus_core::{auth_tokens::AuthenticationToken, hpke, http::extract_bearer_token, time::Clock};
-use janus_messages::{HpkeConfigId, RoleParseError, TaskId};
+use janus_messages::{HpkeConfigId, RoleParseError, TaskId, Url as DapUrl};
 use opentelemetry::metrics::Meter;
 use routes::*;
 use tower::ServiceBuilder;
 use tracing::error;
-use url::Url;
 
 /// Represents the configuration for an instance of the Aggregator API.
 #[derive(Clone)]
 pub struct Config {
     pub auth_tokens: Vec<AuthenticationToken>,
-    pub public_dap_url: Url,
+    pub public_dap_url: DapUrl,
 }
 
 /// Content type

aggregator_api/src/models.rs

@@ -34,7 +34,7 @@ pub(crate) enum AggregatorRole {
 pub(crate) struct AggregatorApiConfig {
     pub protocol: &'static str,
     #[educe(Debug(method(std::fmt::Display::fmt)))]
-    pub dap_url: Url,
+    pub dap_url: DapUrl,
     pub role: AggregatorRole,
     pub vdafs: Vec<SupportedVdaf>,
     pub batch_modes: Vec<SupportedBatchMode>,

aggregator_api/src/routes.rs

@@ -205,6 +205,7 @@ pub(super) async fn post_task<C: Clock>(
         AggregatorTask::new(
             task_id,
             /* peer_aggregator_endpoint */ req.peer_aggregator_endpoint,
+            /* own_aggregator_endpoint */ state.config.public_dap_url.clone(),
             /* batch_mode */ req.batch_mode,
             /* vdaf */ req.vdaf,
             vdaf_verify_key,

aggregator_api/src/tests.rs

@@ -2387,6 +2387,7 @@ fn task_resp_serialization() {
     let task = AggregatorTask::new(
         TaskId::from([0u8; 32]),
         "https://helper.com/".parse().unwrap(),
+        "https://leader.com/".parse().unwrap(),
         BatchMode::LeaderSelected {
             batch_time_window_size: None,
         },

aggregator_core/src/datastore.rs

@@ -28,7 +28,8 @@ use janus_core::{
 use janus_messages::{
     AggregateShareId, AggregationJobId, BatchId, CollectionJobId, Duration, Extension,
     HpkeCiphertext, HpkeConfig, HpkeConfigId, Interval, Query, ReportId, ReportIdChecksum,
-    ReportMetadata, Role, TaskId, Time, TimePrecision, VerifyContinue, VerifyInit, VerifyResp,
+    ReportMetadata, Role, TaskConfiguration, TaskId, Time, TimePrecision, VerifyContinue,
+    VerifyInit, VerifyResp,
     batch_mode::{BatchMode, LeaderSelected, TimeInterval},
 };
 use leases::{acquired_aggregation_job_from_row, acquired_collection_job_from_row};
@@ -713,15 +714,16 @@ WHERE success = TRUE ORDER BY version DESC LIMIT(1)",
                 "-- put_aggregator_task()
 INSERT INTO tasks (
     task_id, aggregator_role, aggregation_mode, peer_aggregator_endpoint,
-    batch_mode, vdaf, task_start, task_duration, report_expiry_age, min_batch_size,
+    own_aggregator_endpoint, taskprov_task_config, batch_mode, vdaf, task_start,
+    task_duration, report_expiry_age, min_batch_size,
     time_precision, tolerable_clock_skew, collector_hpke_config,
     vdaf_verify_key, task_info, deactivate_at, aggregator_auth_token_type,
     aggregator_auth_token, aggregator_auth_token_hash,
     collector_auth_token_type, collector_auth_token_hash, created_at,
     updated_at, updated_by)
 VALUES (
     $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18,
-    $19, $20, $21, $22, $23, $24
+    $19, $20, $21, $22, $23, $24, $25, $26
 )
 ON CONFLICT DO NOTHING",
             )
@@ -735,6 +737,13 @@ ON CONFLICT DO NOTHING",
                     /* aggregation_mode */ &task.aggregation_mode().copied(),
                     /* peer_aggregator_endpoint */
                     &task.peer_aggregator_endpoint().as_str(),
+                    /* own_aggregator_endpoint */
+                    &task.own_aggregator_endpoint().as_str(),
+                    /* taskprov_task_config */
+                    &task
+                        .taskprov_task_config()
+                        .map(|c| c.get_encoded())
+                        .transpose()?,
                     /* batch_mode */ &Json(task.batch_mode()),
                     /* vdaf */ &Json(task.vdaf()),
                     /* task_start */
@@ -889,7 +898,8 @@ UPDATE tasks SET
             .prepare_cached(
                 "-- get_aggregator_task()
 SELECT
-    aggregator_role, aggregation_mode, peer_aggregator_endpoint, batch_mode,
+    aggregator_role, aggregation_mode, peer_aggregator_endpoint,
+    own_aggregator_endpoint, taskprov_task_config, batch_mode,
     vdaf, task_start, task_duration, report_expiry_age, min_batch_size,
     time_precision, tolerable_clock_skew, collector_hpke_config,
     vdaf_verify_key, task_info, deactivate_at, aggregator_auth_token_type,
@@ -913,6 +923,7 @@ FROM tasks WHERE task_id = $1",
                 "-- get_aggregator_tasks()
 SELECT
     task_id, aggregator_role, aggregation_mode, peer_aggregator_endpoint,
+    own_aggregator_endpoint, taskprov_task_config,
     batch_mode, vdaf, task_start, task_duration, report_expiry_age, min_batch_size,
     time_precision, tolerable_clock_skew, collector_hpke_config,
     vdaf_verify_key, task_info, deactivate_at, aggregator_auth_token_type,
@@ -934,6 +945,11 @@ FROM tasks",
         // Scalar task parameters.
         let aggregator_role: AggregatorRole = row.get("aggregator_role");
         let peer_aggregator_endpoint = row.get::<_, String>("peer_aggregator_endpoint").parse()?;
+        let own_aggregator_endpoint = row.get::<_, String>("own_aggregator_endpoint").parse()?;
+        let taskprov_task_config = row
+            .get::<_, Option<Vec<u8>>>("taskprov_task_config")
+            .map(|bytes| TaskConfiguration::get_decoded(&bytes))
+            .transpose()?;
         let batch_mode = row.try_get::<_, Json<task::BatchMode>>("batch_mode")?.0;
         let aggregation_mode: Option<AggregationMode> = row.get("aggregation_mode");
         let vdaf = row.try_get::<_, Json<VdafInstance>>("vdaf")?.0;
@@ -1048,6 +1064,7 @@ FROM tasks",
         let task = AggregatorTask::new(
             *task_id,
             peer_aggregator_endpoint,
+            own_aggregator_endpoint,
             batch_mode,
             vdaf,
             vdaf_verify_key,
@@ -1060,6 +1077,10 @@ FROM tasks",
             aggregator_parameters,
         )?
         .with_deactivate_at(deactivate_at);
+        let task = match taskprov_task_config {
+            Some(task_config) => task.with_taskprov_task_config(task_config),
+            None => task,
+        };
         Ok(task)
     }
 

aggregator_core/src/datastore/tests.rs

@@ -341,6 +341,88 @@ async fn roundtrip_task(ephemeral_datastore: EphemeralDatastore) {
     assert_eq!(want_tasks, got_tasks);
 }
 
+#[rstest_reuse::apply(schema_versions_template)]
+#[tokio::test]
+async fn roundtrip_task_own_endpoint_verbatim(ephemeral_datastore: EphemeralDatastore) {
+    // A non-canonical own endpoint (mixed-case host, no trailing slash) must survive a DB
+    // round-trip byte-identical: the bytes are bound into the task's TaskConfiguration for HPKE
+    // AADs, and any re-encoding in the datastore write/read path would break decryption (DAP-18
+    // §4.1).
+    install_test_trace_subscriber();
+    let ds = ephemeral_datastore.datastore(MockClock::default()).await;
+
+    let noncanonical = "https://Example.COM/DAP";
+    let base = TaskBuilder::new(
+        task::BatchMode::TimeInterval,
+        AggregationMode::Synchronous,
+        VdafInstance::Prio3Count,
+    )
+    .build()
+    .leader_view()
+    .unwrap();
+    let task = AggregatorTask::new(
+        *base.id(),
+        base.peer_aggregator_endpoint().clone(),
+        noncanonical.try_into().unwrap(),
+        *base.batch_mode(),
+        base.vdaf().clone(),
+        base.opaque_vdaf_verify_key().clone(),
+        base.task_interval().copied(),
+        base.report_expiry_age().copied(),
+        base.min_batch_size(),
+        *base.time_precision(),
+        *base.tolerable_clock_skew(),
+        base.task_info().to_vec(),
+        base.aggregator_parameters().clone(),
+    )
+    .unwrap();
+    ds.put_aggregator_task(&task).await.unwrap();
+
+    let retrieved = ds
+        .run_unnamed_tx(|tx| {
+            let task_id = *task.id();
+            Box::pin(async move { tx.get_aggregator_task(&task_id).await })
+        })
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(retrieved.own_aggregator_endpoint().as_str(), noncanonical);
+}
+
+#[rstest_reuse::apply(schema_versions_template)]
+#[tokio::test]
+async fn roundtrip_taskprov_task_config(ephemeral_datastore: EphemeralDatastore) {
+    // The verbatim taskprov TaskConfiguration must survive a DB round-trip: it is bound into HPKE
+    // AADs and is not byte-reconstructible from the task's other stored columns.
+    install_test_trace_subscriber();
+    let ds = ephemeral_datastore.datastore(MockClock::default()).await;
+
+    let base = TaskBuilder::new(
+        task::BatchMode::TimeInterval,
+        AggregationMode::Synchronous,
+        VdafInstance::Prio3Count,
+    )
+    .build();
+    // Source a TaskConfiguration to store from a non-taskprov view (synthesizing one for a taskprov
+    // task is now rejected); any valid config exercises the column's byte round-trip.
+    let task_config = base.leader_view().unwrap().task_configuration().unwrap();
+    let task = base
+        .taskprov_helper_view()
+        .unwrap()
+        .with_taskprov_task_config(task_config.clone());
+    ds.put_aggregator_task(&task).await.unwrap();
+
+    let retrieved = ds
+        .run_unnamed_tx(|tx| {
+            let task_id = *task.id();
+            Box::pin(async move { tx.get_aggregator_task(&task_id).await })
+        })
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(retrieved.taskprov_task_config(), Some(&task_config));
+}
+
 #[rstest_reuse::apply(schema_versions_template)]
 #[tokio::test]
 async fn set_task_deactivate_at(ephemeral_datastore: EphemeralDatastore) {