Exactly one aggregator and collector auth token per task

divviup · Sep 21, 2023 · bde5aa8 · bde5aa8
1 parent 9cb721a
commit bde5aa8
Show file tree

Hide file tree

Showing 18 changed files with 406 additions and 479 deletions.
diff --git a/aggregator/src/aggregator/aggregate_init_tests.rs b/aggregator/src/aggregator/aggregate_init_tests.rs
@@ -169,6 +169,7 @@ async fn setup_aggregate_init_test_for_vdaf<
         vdaf_instance,
         aggregation_param,
         measurement,
+        false,
     )
     .await;
 
@@ -203,10 +204,18 @@ async fn setup_aggregate_init_test_without_sending_request<
     vdaf_instance: VdafInstance,
     aggregation_param: V::AggregationParam,
     measurement: V::Measurement,
+    use_dap_auth_token: bool,
 ) -> AggregationJobInitTestCase<VERIFY_KEY_SIZE, V> {
     install_test_trace_subscriber();
 
-    let task = TaskBuilder::new(QueryType::TimeInterval, vdaf_instance, Role::Helper).build();
+    let mut task_builder = TaskBuilder::new(QueryType::TimeInterval, vdaf_instance, Role::Helper);
+
+    if use_dap_auth_token {
+        task_builder =
+            task_builder.with_aggregator_auth_token(Some(AuthenticationToken::DapAuth(random())));
+    }
+
+    let task = task_builder.build();
     let clock = MockClock::default();
     let ephemeral_datastore = ephemeral_datastore().await;
     let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
@@ -258,11 +267,12 @@ pub(crate) async fn put_aggregation_job(
     aggregation_job: &AggregationJobInitializeReq<TimeInterval>,
     handler: &impl Handler,
 ) -> TestConn {
+    let (header, value) = task
+        .aggregator_auth_token()
+        .unwrap()
+        .request_authentication();
     put(task.aggregation_job_uri(aggregation_job_id).unwrap().path())
-        .with_request_header(
-            DAP_AUTH_HEADER,
-            task.primary_aggregator_auth_token().as_ref().to_owned(),
-        )
+        .with_request_header(header, value)
         .with_request_header(
             KnownHeaderName::ContentType,
             AggregationJobInitializeReq::<TimeInterval>::MEDIA_TYPE,
@@ -279,14 +289,13 @@ async fn aggregation_job_init_authorization_dap_auth_token() {
         VdafInstance::Fake,
         dummy_vdaf::AggregationParam(0),
         (),
+        true,
     )
     .await;
-    // Find a DapAuthToken among the task's aggregator auth tokens
+
     let (auth_header, auth_value) = test_case
         .task
-        .aggregator_auth_tokens()
-        .iter()
-        .find(|auth| matches!(auth, AuthenticationToken::DapAuth(_)))
+        .aggregator_auth_token()
         .unwrap()
         .request_authentication();
 
@@ -317,6 +326,7 @@ async fn aggregation_job_init_malformed_authorization_header(#[case] header_valu
         VdafInstance::Fake,
         dummy_vdaf::AggregationParam(0),
         (),
+        false,
     )
     .await;
 
@@ -333,7 +343,8 @@ async fn aggregation_job_init_malformed_authorization_header(#[case] header_valu
         DAP_AUTH_HEADER,
         test_case
             .task
-            .primary_aggregator_auth_token()
+            .aggregator_auth_token()
+            .unwrap()
             .as_ref()
             .to_owned(),
     )
@@ -480,19 +491,18 @@ async fn aggregation_job_init_wrong_query() {
         test_case.prepare_inits,
     );
 
+    let (header, value) = test_case
+        .task
+        .aggregator_auth_token()
+        .unwrap()
+        .request_authentication();
+
     let mut response = put(test_case
         .task
         .aggregation_job_uri(&random())
         .unwrap()
         .path())
-    .with_request_header(
-        DAP_AUTH_HEADER,
-        test_case
-            .task
-            .primary_aggregator_auth_token()
-            .as_ref()
-            .to_owned(),
-    )
+    .with_request_header(header, value)
     .with_request_header(
         KnownHeaderName::ContentType,
         AggregationJobInitializeReq::<TimeInterval>::MEDIA_TYPE,

diff --git a/aggregator/src/aggregator/aggregation_job_continue.rs b/aggregator/src/aggregator/aggregation_job_continue.rs
@@ -302,11 +302,12 @@ pub mod test_util {
         request: &AggregationJobContinueReq,
         handler: &impl Handler,
     ) -> TestConn {
+        let (header, value) = task
+            .aggregator_auth_token()
+            .unwrap()
+            .request_authentication();
         post(task.aggregation_job_uri(aggregation_job_id).unwrap().path())
-            .with_request_header(
-                "DAP-Auth-Token",
-                task.primary_aggregator_auth_token().as_ref().to_owned(),
-            )
+            .with_request_header(header, value)
             .with_request_header(
                 KnownHeaderName::ContentType,
                 AggregationJobContinueReq::MEDIA_TYPE,