tests: Use global HPKE keys in all tests (#3284)

* Use require_global_hpke_keys flag in unit tests

Remove dependencies on Task::current_hpke_key

Add some tests for task specific keys

* Split upload tests out of aggregator.rs

aggregator/src/aggregator/aggregate_init_tests.rs

@@ -3,7 +3,7 @@ use crate::aggregator::{
         aggregator_handler,
         test_util::{decode_response_body, take_problem_details},
     },
-    tests::generate_helper_report_share,
+    test_util::generate_helper_report_share,
     Config,
 };
 use assert_matches::assert_matches;
@@ -63,10 +63,10 @@ where
     pub(super) fn new(
         clock: MockClock,
         task: AggregatorTask,
+        hpke_config: HpkeConfig,
         vdaf: V,
         aggregation_param: V::AggregationParam,
     ) -> Self {
-        let hpke_config = task.current_hpke_key().config().clone();
         Self {
             clock,
             task,
@@ -77,11 +77,6 @@ where
         }
     }
 
-    pub(super) fn with_hpke_config(mut self, config: HpkeConfig) -> Self {
-        self.hpke_config = config;
-        self
-    }
-
     pub(super) fn with_extensions(mut self, extensions: Vec<Extension>) -> Self {
         self.extensions = extensions;
         self
@@ -260,6 +255,7 @@ async fn setup_aggregate_init_test_without_sending_request<
     let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
 
     datastore.put_aggregator_task(&helper_task).await.unwrap();
+    let keypair = datastore.put_global_hpke_key().await.unwrap();
 
     let handler = aggregator_handler(
         Arc::clone(&datastore),
@@ -274,6 +270,7 @@ async fn setup_aggregate_init_test_without_sending_request<
     let prepare_init_generator = PrepareInitGenerator::new(
         clock.clone(),
         helper_task.clone(),
+        keypair.config().clone(),
         vdaf,
         aggregation_param.clone(),
     );

aggregator/src/aggregator/aggregation_job_continue.rs

@@ -471,6 +471,7 @@ mod tests {
         let ephemeral_datastore = ephemeral_datastore().await;
         let meter = noop_meter();
         let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
+        let keypair = datastore.put_global_hpke_key().await.unwrap();
 
         let aggregation_parameter = Poplar1AggregationParam::try_from_prefixes(Vec::from([
             IdpfInput::from_bools(&[false]),
@@ -479,6 +480,7 @@ mod tests {
         let prepare_init_generator = PrepareInitGenerator::new(
             clock.clone(),
             helper_task.clone(),
+            keypair.config().clone(),
             Poplar1::new_turboshake128(1),
             aggregation_parameter.clone(),
         );

aggregator/src/aggregator/collection_job_tests.rs

@@ -139,6 +139,7 @@ pub(crate) async fn setup_collection_job_test_case(
     let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
 
     datastore.put_aggregator_task(&role_task).await.unwrap();
+    datastore.put_global_hpke_key().await.unwrap();
 
     let handler = aggregator_handler(
         Arc::clone(&datastore),

aggregator/src/aggregator/http_handlers.rs

@@ -759,12 +759,14 @@ pub mod test_util {
     use crate::aggregator::test_util::default_aggregator_config;
     use janus_aggregator_core::{
         datastore::{
+            models::HpkeKeyState,
             test_util::{ephemeral_datastore, EphemeralDatastore},
             Datastore,
         },
         test_util::noop_meter,
     };
     use janus_core::{
+        hpke::HpkeKeypair,
         test_util::{install_test_trace_subscriber, runtime::TestRuntime},
         time::MockClock,
     };
@@ -792,9 +794,60 @@ pub mod test_util {
         serde_json::from_slice(&take_response_body(test_conn).await).unwrap()
     }
 
-    /// Returns structures necessary for completing an HTTP handler test. The returned
+    /// Contains structures necessary for completing an HTTP handler test. The contained
     /// [`EphemeralDatastore`] should be given a variable binding to prevent it being prematurely
     /// dropped.
+    pub struct HttpHandlerTest {
+        pub clock: MockClock,
+        pub ephemeral_datastore: EphemeralDatastore,
+        pub datastore: Arc<Datastore<MockClock>>,
+        pub handler: Box<dyn Handler>,
+        pub hpke_keypair: HpkeKeypair,
+    }
+
+    impl HttpHandlerTest {
+        pub async fn new() -> Self {
+            install_test_trace_subscriber();
+            let clock = MockClock::default();
+            let ephemeral_datastore = ephemeral_datastore().await;
+            let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
+
+            let hpke_keypair = HpkeKeypair::test();
+            datastore
+                .run_unnamed_tx(|tx| {
+                    let hpke_keypair = hpke_keypair.clone();
+                    Box::pin(async move {
+                        tx.put_global_hpke_keypair(&hpke_keypair).await?;
+                        tx.set_global_hpke_keypair_state(
+                            hpke_keypair.config().id(),
+                            &HpkeKeyState::Active,
+                        )
+                        .await
+                    })
+                })
+                .await
+                .unwrap();
+
+            let handler = aggregator_handler(
+                datastore.clone(),
+                clock.clone(),
+                TestRuntime::default(),
+                &noop_meter(),
+                default_aggregator_config(),
+            )
+            .await
+            .unwrap();
+
+            Self {
+                clock,
+                ephemeral_datastore,
+                datastore,
+                handler: Box::new(handler),
+                hpke_keypair,
+            }
+        }
+    }
+
     pub async fn setup_http_handler_test() -> (
         MockClock,
         EphemeralDatastore,
@@ -805,6 +858,7 @@ pub mod test_util {
         let clock = MockClock::default();
         let ephemeral_datastore = ephemeral_datastore().await;
         let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
+        datastore.put_global_hpke_key().await.unwrap();
         let handler = aggregator_handler(
             datastore.clone(),
             clock.clone(),

aggregator/src/aggregator/http_handlers/tests/aggregation_job_continue.rs

@@ -3,9 +3,8 @@ use crate::aggregator::{
         post_aggregation_job_and_decode, post_aggregation_job_expecting_error,
     },
     empty_batch_aggregations,
-    http_handlers::test_util::setup_http_handler_test,
-    test_util::BATCH_AGGREGATION_SHARD_COUNT,
-    tests::generate_helper_report_share,
+    http_handlers::test_util::{setup_http_handler_test, HttpHandlerTest},
+    test_util::{generate_helper_report_share, BATCH_AGGREGATION_SHARD_COUNT},
 };
 use futures::future::try_join_all;
 use janus_aggregator_core::{
@@ -43,15 +42,21 @@ use trillium::Status;
 
 #[tokio::test]
 async fn aggregate_continue() {
-    let (clock, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        clock,
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        hpke_keypair: hpke_key,
+        ..
+    } = HttpHandlerTest::new().await;
 
     let aggregation_job_id = random();
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
     let helper_task = task.helper_view().unwrap();
 
     let vdaf = Arc::new(Poplar1::<XofTurboShake128, 16>::new(1));
     let verify_key: VerifyKey<VERIFY_KEY_LENGTH> = task.vdaf_verify_key().unwrap();
-    let hpke_key = helper_task.current_hpke_key();
     let measurement = IdpfInput::from_bools(&[true]);
     let aggregation_param =
         Poplar1AggregationParam::try_from_prefixes(vec![measurement.clone()]).unwrap();
@@ -371,7 +376,13 @@ async fn aggregate_continue() {
 
 #[tokio::test]
 async fn aggregate_continue_accumulate_batch_aggregation() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        hpke_keypair: hpke_key,
+        ..
+    } = HttpHandlerTest::new().await;
 
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
     let helper_task = task.helper_view().unwrap();
@@ -387,7 +398,6 @@ async fn aggregate_continue_accumulate_batch_aggregation() {
 
     let vdaf = Poplar1::new(1);
     let verify_key: VerifyKey<VERIFY_KEY_LENGTH> = task.vdaf_verify_key().unwrap();
-    let hpke_key = helper_task.current_hpke_key();
     let measurement = IdpfInput::from_bools(&[true]);
     let aggregation_param =
         Poplar1AggregationParam::try_from_prefixes(Vec::from([measurement.clone()])).unwrap();
@@ -1157,7 +1167,13 @@ async fn aggregate_continue_leader_sends_non_continue_or_finish_transition() {
 
 #[tokio::test]
 async fn aggregate_continue_prep_step_fails() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        hpke_keypair: hpke_key,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
@@ -1179,7 +1195,7 @@ async fn aggregate_continue_prep_step_fails() {
     let helper_report_share = generate_helper_report_share::<Poplar1<XofTurboShake128, 16>>(
         *task.id(),
         report_metadata.clone(),
-        helper_task.current_hpke_key().config(),
+        hpke_key.config(),
         &transcript.public_share,
         Vec::new(),
         &transcript.helper_input_share,