diff --git a/aggregator/src/aggregator/collection_job_tests.rs b/aggregator/src/aggregator/collection_job_tests.rs index 6c27913c7..4dc0ec9d5 100644 --- a/aggregator/src/aggregator/collection_job_tests.rs +++ b/aggregator/src/aggregator/collection_job_tests.rs @@ -29,7 +29,7 @@ use janus_core::{ }; use janus_messages::{ AggregateShareAad, AggregationJobStep, BatchId, CollectionJobId, CollectionJobReq, - CollectionJobResp, Interval, MediaType, Query, Role, TaskId, Time, TimePrecision, + CollectionJobResp, Duration, Interval, MediaType, Query, Role, TaskId, Time, TimePrecision, batch_mode::{ BatchMode as BatchModeTrait, LeaderSelected, LeaderSelectedQueryConfig, TimeInterval, }, @@ -58,6 +58,79 @@ pub(crate) struct CollectionJobTestCase { } impl CollectionJobTestCase { + pub(crate) async fn setup(role: Role, batch_mode: BatchMode) -> Self { + Self::setup_with_report_expiry_age(role, batch_mode, None).await + } + + pub(crate) async fn setup_with_report_expiry_age( + role: Role, + batch_mode: BatchMode, + report_expiry_age: Option, + ) -> Self { + install_test_trace_subscriber(); + + let task = TaskBuilder::new( + batch_mode, + AggregationMode::Synchronous, + VdafInstance::Fake { rounds: 1 }, + ) + .with_report_expiry_age(report_expiry_age) + .build(); + let role_task = task.view_for_role(role).unwrap(); + let clock = MockClock::default(); + let ephemeral_datastore = ephemeral_datastore().await; + let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await); + + datastore.put_aggregator_task(&role_task).await.unwrap(); + datastore.put_hpke_key().await.unwrap(); + + let router = AggregatorHandlerBuilder::new( + Arc::clone(&datastore), + TestRuntime::default(), + &noop_meter(), + Config { + batch_aggregation_shard_count: BATCH_AGGREGATION_SHARD_COUNT, + ..Default::default() + }, + ) + .await + .unwrap() + .build() + .unwrap(); + + CollectionJobTestCase { + task, + clock, + router, + datastore, + _ephemeral_datastore: ephemeral_datastore, + } + } + + pub(super) async fn setup_leader_selected_current_batch( + &mut self, + ) -> (BatchId, BatchId, Interval) { + // Fill the datastore with the necessary data so that there are two outstanding batches to be + // collected. + let time = self + .clock + .now() + .to_time(self.task.time_precision()); + let batch_id_1 = self + .setup_leader_selected_batch(time, self.task.min_batch_size() + 1) + .await; + let batch_id_2 = self + .setup_leader_selected_batch(time, self.task.min_batch_size() + 1) + .await; + let client_timestamp_interval = Interval::minimal(time).unwrap(); + + (batch_id_1, batch_id_2, client_timestamp_interval) + } + + pub(super) fn clock(&self) -> &MockClock { + &self.clock + } + pub(super) async fn post_collection_job_with_auth_token( &self, request: &CollectionJobReq, @@ -124,6 +197,29 @@ impl CollectionJobTestCase { .await } + pub(super) async fn delete_collection_job( + &self, + collection_job_id: &CollectionJobId, + ) -> http::Response { + self.router + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri( + self.task + .collection_job_uri(&collection_job_id) + .unwrap() + .path(), + ) + .with_authentication_token(self.task.collector_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + } + /// Seed the database with reports and various aggregation artifacts from aggregating them in a /// leader-selected batch. pub(super) async fn setup_leader_selected_batch( @@ -252,148 +348,70 @@ impl CollectionJobTestCase { .await .unwrap() } -} - -/// Parse task ID and collection job ID out of response location header. -/// -/// Panics if header is missing from response or value is malformed. -pub(super) fn ids_from_collection_job_response( - response: &http::Response, -) -> (TaskId, CollectionJobId) { - response - .headers() - .ids_from_collection_job_location() - .unwrap() - .unwrap() -} -pub(crate) async fn setup_collection_job_test_case( - role: Role, - batch_mode: BatchMode, -) -> CollectionJobTestCase { - install_test_trace_subscriber(); + pub(super) async fn simulate_successful_collection_job_time_interval( + &mut self, + report_time: Time, + queried_interval: Interval, + ) -> (CollectionJobReq, CollectionJobId) { + self.setup_time_interval_batch(report_time).await; - let task = TaskBuilder::new( - batch_mode, - AggregationMode::Synchronous, - VdafInstance::Fake { rounds: 1 }, - ) - .build(); - let role_task = task.view_for_role(role).unwrap(); - let clock = MockClock::default(); - let ephemeral_datastore = ephemeral_datastore().await; - let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await); - - datastore.put_aggregator_task(&role_task).await.unwrap(); - datastore.put_hpke_key().await.unwrap(); - - let router = AggregatorHandlerBuilder::new( - Arc::clone(&datastore), - TestRuntime::default(), - &noop_meter(), - Config { - batch_aggregation_shard_count: BATCH_AGGREGATION_SHARD_COUNT, - ..Default::default() - }, - ) - .await - .unwrap() - .build() - .unwrap(); - - CollectionJobTestCase { - task, - clock, - router, - datastore, - _ephemeral_datastore: ephemeral_datastore, - } -} + let aggregation_param = dummy::AggregationParam::default(); + let leader_aggregate_share = dummy::AggregateShare(0); + let helper_aggregate_share = dummy::AggregateShare(1); -async fn setup_leader_selected_current_batch_collection_job_test_case() --> (CollectionJobTestCase, BatchId, BatchId, Interval) { - let test_case = setup_collection_job_test_case( - Role::Leader, - BatchMode::LeaderSelected { - batch_time_window_size: None, - }, - ) - .await; - - // Fill the datastore with the necessary data so that there are two outstanding batches to be - // collected. - let time = test_case - .clock - .now() - .to_time(test_case.task.time_precision()); - let batch_id_1 = test_case - .setup_leader_selected_batch(time, test_case.task.min_batch_size() + 1) - .await; - let batch_id_2 = test_case - .setup_leader_selected_batch(time, test_case.task.min_batch_size() + 1) - .await; - let client_timestamp_interval = Interval::minimal(time).unwrap(); - - (test_case, batch_id_1, batch_id_2, client_timestamp_interval) -} - -#[tokio::test] -async fn collection_job_success_leader_selected() { - // This test drives two current batch collection jobs to completion, verifying that distinct - // collection job IDs are minted and distinct batch IDs are collected each time. Then, we - // attempt to collect another current batch, which must fail as there are no more outstanding - // batches. - let (test_case, batch_id_1, batch_id_2, spanned_interval) = - setup_leader_selected_current_batch_collection_job_test_case().await; - - let first_idempotency_key: LeaderSelectedQueryConfig = random(); - let mut saw_batch_id_1 = false; - let mut saw_batch_id_2 = false; - let mut collection_job_ids = vec![]; - let mut collection_job_resps = vec![]; - let vdaf = dummy::Vdaf::new(1); - let leader_aggregate_share = dummy::AggregateShare(0); - let helper_aggregate_share = dummy::AggregateShare(1); - let aggregation_param = dummy::AggregationParam::default(); - - for idempotency_key in [first_idempotency_key.clone(), random()] { let request = CollectionJobReq::new( - Query::new_leader_selected(idempotency_key), + Query::new_time_interval(queried_interval), aggregation_param.get_encoded().unwrap(), ); - let mut response = test_case.post_collection_job(&request).await; + + let mut response = self.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::OK); let (task_id, collection_job_id) = ids_from_collection_job_response(&response); - collection_job_ids.push(collection_job_id); - assert_eq!(test_case.task.id(), &task_id); + assert_eq!(&task_id, self.task.id()); + + let got_collection_job = self + .datastore + .run_unnamed_tx(|tx| { + let task_id = *self.task.id(); + + Box::pin(async move { + Ok(tx + .get_collection_job(&dummy::Vdaf::new(1), &task_id, &collection_job_id) + .await + .unwrap() + .unwrap()) + }) + }) + .await + .unwrap(); + + let want_collection_job = CollectionJob::<0, TimeInterval, dummy::Vdaf>::new( + *self.task.id(), + collection_job_id, + // This is chosen by the leader, we have to copy it + *got_collection_job.aggregate_share_id(), + Query::new_time_interval(queried_interval), + aggregation_param, + queried_interval, + CollectionJobState::Start, + ); + + assert_eq!(want_collection_job, got_collection_job); assert!(take_response_body(&mut response).await.is_empty()); - let mut response = test_case.get_collection_job(&collection_job_id).await; + let mut response = self.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::OK); assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); assert!(take_response_body(&mut response).await.is_empty()); - // Update the collection job with the aggregate shares. collection job should now be - // complete. - let batch_id = test_case - .datastore + // Update the collection job with the aggregate shares and some aggregation jobs. collection + // job should now be complete. + self.datastore .run_unnamed_tx(|tx| { - let task = test_case.task.clone(); - let vdaf = vdaf.clone(); + let task = self.task.clone(); let helper_aggregate_share_bytes = helper_aggregate_share.get_encoded().unwrap(); - let request = request.clone(); Box::pin(async move { - let collection_job = tx - .get_collection_job::<0, LeaderSelected, dummy::Vdaf>( - &vdaf, - task.id(), - &collection_job_id, - ) - .await - .unwrap() - .unwrap(); - let batch_id = *collection_job.batch_identifier(); - let encrypted_helper_aggregate_share = hpke::seal( task.collector_hpke_keypair().config(), &HpkeApplicationInfo::new( @@ -402,43 +420,50 @@ async fn collection_job_success_leader_selected() { &Role::Collector, ), &helper_aggregate_share_bytes, - &AggregateShareAad::new(*task.id(), task.task_configuration(), request) - .get_encoded() - .unwrap(), + &AggregateShareAad::new( + *task.id(), + task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), + ) + .get_encoded() + .unwrap(), ) .unwrap(); - tx.update_collection_job::<0, LeaderSelected, dummy::Vdaf>( - &collection_job.with_state(CollectionJobState::Finished { - report_count: task.min_batch_size() + 1, - client_timestamp_interval: spanned_interval, + let collection_job = tx + .get_collection_job::<0, TimeInterval, dummy::Vdaf>( + &dummy::Vdaf::new(1), + task.id(), + &collection_job_id, + ) + .await + .unwrap() + .unwrap() + .with_state(CollectionJobState::Finished { + report_count: 12, + client_timestamp_interval: queried_interval, encrypted_helper_aggregate_share, leader_aggregate_share, - }), - ) - .await - .unwrap(); + }); - Ok(batch_id) + tx.update_collection_job::<0, TimeInterval, dummy::Vdaf>(&collection_job) + .await + .unwrap(); + Ok(()) }) }) .await .unwrap(); - if batch_id.eq(&batch_id_1) { - saw_batch_id_1 = true; - } else if batch_id.eq(&batch_id_2) { - saw_batch_id_2 = true; - } else { - panic!("unexpected batch ID"); - } - - let mut response = test_case.get_collection_job(&collection_job_id).await; + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.headers().get("content-type").unwrap(), CollectionJobResp::MEDIA_TYPE, ); - let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; let ( report_count, @@ -447,31 +472,29 @@ async fn collection_job_success_leader_selected() { helper_encrypted_aggregate_share, ) = assert_matches!( collect_resp.clone(), - CollectionJobResp { + CollectionJobResp{ report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share, .. - } => ( - report_count, - interval, - leader_encrypted_agg_share, - helper_encrypted_agg_share - ) + } => (report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share) ); - assert_eq!(report_count, test_case.task.min_batch_size() + 1); - assert_eq!(interval, spanned_interval); + assert_eq!(report_count, 12); + assert_eq!(interval, queried_interval); let decrypted_leader_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), + self.task.collector_hpke_keypair(), &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), &leader_encrypted_aggregate_share, &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - request.clone(), + *self.task.id(), + self.task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), ) .get_encoded() .unwrap(), @@ -483,13 +506,16 @@ async fn collection_job_success_leader_selected() { ); let decrypted_helper_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), + self.task.collector_hpke_keypair(), &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), &helper_encrypted_aggregate_share, &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - request.clone(), + *self.task.id(), + self.task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), ) .get_encoded() .unwrap(), @@ -500,20 +526,231 @@ async fn collection_job_success_leader_selected() { dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() ); - collection_job_resps.push(collect_resp); + (request, collection_job_id) + } + + pub(super) async fn simulate_successful_collection_job_leader_selected( + &mut self, + batch_id_1: BatchId, + batch_id_2: BatchId, + spanned_interval: Interval, + ) -> (CollectionJobReq, CollectionJobId) { + let mut saw_batch_id_1 = false; + let mut saw_batch_id_2 = false; + let mut collection_job_ids = vec![]; + let mut collection_job_resps = vec![]; + let vdaf = dummy::Vdaf::new(1); + let leader_aggregate_share = dummy::AggregateShare(0); + let helper_aggregate_share = dummy::AggregateShare(1); + let aggregation_param = dummy::AggregationParam::default(); + let requests = [ + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + ]; + + for request in requests.clone() { + let mut response = self.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (task_id, collection_job_id) = ids_from_collection_job_response(&response); + collection_job_ids.push(collection_job_id); + assert_eq!(self.task.id(), &task_id); + assert!(take_response_body(&mut response).await.is_empty()); + + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); + assert!(take_response_body(&mut response).await.is_empty()); + + // Update the collection job with the aggregate shares. collection job should now be + // complete. + let batch_id = self + .datastore + .run_unnamed_tx(|tx| { + let task = self.task.clone(); + let vdaf = vdaf.clone(); + let helper_aggregate_share_bytes = + helper_aggregate_share.get_encoded().unwrap(); + let request = request.clone(); + Box::pin(async move { + let collection_job = tx + .get_collection_job::<0, LeaderSelected, dummy::Vdaf>( + &vdaf, + task.id(), + &collection_job_id, + ) + .await + .unwrap() + .unwrap(); + let batch_id = *collection_job.batch_identifier(); + + let encrypted_helper_aggregate_share = hpke::seal( + task.collector_hpke_keypair().config(), + &HpkeApplicationInfo::new( + &Label::AggregateShare, + &Role::Helper, + &Role::Collector, + ), + &helper_aggregate_share_bytes, + &AggregateShareAad::new(*task.id(), task.task_configuration(), request) + .get_encoded() + .unwrap(), + ) + .unwrap(); + + tx.update_collection_job::<0, LeaderSelected, dummy::Vdaf>( + &collection_job.with_state(CollectionJobState::Finished { + report_count: task.min_batch_size() + 1, + client_timestamp_interval: spanned_interval, + encrypted_helper_aggregate_share, + leader_aggregate_share, + }), + ) + .await + .unwrap(); + + Ok(batch_id) + }) + }) + .await + .unwrap(); + + if batch_id.eq(&batch_id_1) { + saw_batch_id_1 = true; + } else if batch_id.eq(&batch_id_2) { + saw_batch_id_2 = true; + } else { + panic!("unexpected batch ID"); + } + + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!( + response.headers().get("content-type").unwrap(), + CollectionJobResp::MEDIA_TYPE, + ); + + let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; + let ( + report_count, + interval, + leader_encrypted_aggregate_share, + helper_encrypted_aggregate_share, + ) = assert_matches!( + collect_resp.clone(), + CollectionJobResp { + report_count, + interval, + leader_encrypted_agg_share, + helper_encrypted_agg_share, + .. + } => ( + report_count, + interval, + leader_encrypted_agg_share, + helper_encrypted_agg_share + ) + ); + + assert_eq!(report_count, self.task.min_batch_size() + 1); + assert_eq!(interval, spanned_interval); + + let decrypted_leader_aggregate_share = hpke::open( + self.task.collector_hpke_keypair(), + &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), + &leader_encrypted_aggregate_share, + &AggregateShareAad::new( + *self.task.id(), + self.task.task_configuration(), + request.clone(), + ) + .get_encoded() + .unwrap(), + ) + .unwrap(); + assert_eq!( + leader_aggregate_share, + dummy::AggregateShare::get_decoded(decrypted_leader_aggregate_share.as_ref()) + .unwrap() + ); + + let decrypted_helper_aggregate_share = hpke::open( + self.task.collector_hpke_keypair(), + &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), + &helper_encrypted_aggregate_share, + &AggregateShareAad::new( + *self.task.id(), + self.task.task_configuration(), + request.clone(), + ) + .get_encoded() + .unwrap(), + ) + .unwrap(); + assert_eq!( + helper_aggregate_share, + dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()) + .unwrap() + ); + + collection_job_resps.push(collect_resp); + } + + assert!(saw_batch_id_1 && saw_batch_id_2); + // Collection job IDs should be distinct + assert_eq!(collection_job_ids.len(), 2); + assert_ne!(collection_job_ids[0], collection_job_ids[1]); + + (requests[0].clone(), collection_job_ids[0]) } +} + +/// Parse task ID and collection job ID out of response location header. +/// +/// Panics if header is missing from response or value is malformed. +pub(super) fn ids_from_collection_job_response( + response: &http::Response, +) -> (TaskId, CollectionJobId) { + response + .headers() + .ids_from_collection_job_location() + .unwrap() + .unwrap() +} - assert!(saw_batch_id_1 && saw_batch_id_2); - // Collection job IDs should be distinct - assert_eq!(collection_job_ids.len(), 2); - assert_ne!(collection_job_ids[0], collection_job_ids[1]); +#[tokio::test] +async fn collection_job_success_leader_selected() { + // This test drives two current batch collection jobs to completion, verifying that distinct + // collection job IDs are minted and distinct batch IDs are collected each time. Then, we + // attempt to collect another current batch, which must fail as there are no more outstanding + // batches. + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; // We have run the two ready batches to completion. Further attempts to collect current batch // ought to fail. let mut response = test_case .post_collection_job(&CollectionJobReq::new( Query::new_leader_selected(random()), - aggregation_param.get_encoded().unwrap(), + dummy::AggregationParam::default().get_encoded().unwrap(), )) .await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); @@ -529,14 +766,10 @@ async fn collection_job_success_leader_selected() { // Re-send the very first collection job creation request. We should get the representation of // the completed collection job without needing a further GET. - let request = CollectionJobReq::new( - Query::new_leader_selected(first_idempotency_key), - aggregation_param.get_encoded().unwrap(), - ); let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::OK); - let (_, collection_job_id) = ids_from_collection_job_response(&response); - assert_eq!(collection_job_id, collection_job_ids[0]); + let (_, collection_job_id_again) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id, collection_job_id_again); assert_eq!( response.headers().get("content-type").unwrap(), CollectionJobResp::MEDIA_TYPE, @@ -546,7 +779,7 @@ async fn collection_job_success_leader_selected() { #[tokio::test] async fn collection_job_post_idempotence_time_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // Seed DB with reports, aggregations, etc. corresponding to the two time intervals used later // in the test test_case @@ -637,7 +870,7 @@ async fn collection_job_post_idempotence_time_interval() { async fn collection_job_post_idempotence_time_interval_mutate_aggregation_param() { // Create a collection job for some time interval, then try to create another, varying the // aggregation parameter. It should fail due to batch overlap. - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_time_precision_units(0)) .await; @@ -670,8 +903,14 @@ async fn collection_job_post_idempotence_time_interval_mutate_aggregation_param( #[tokio::test] async fn collection_job_post_idempotence_leader_selected() { - let (test_case, batch_id_1, batch_id_2, _) = - setup_leader_selected_current_batch_collection_job_test_case().await; + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, _) = test_case.setup_leader_selected_current_batch().await; let mut seen_batch_id = None; let mut collection_job_id = None; @@ -769,9 +1008,16 @@ async fn collection_job_post_idempotence_leader_selected() { #[tokio::test] async fn collection_job_post_idempotence_leader_selected_mutate_aggregation_param() { + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + test_case.setup_leader_selected_current_batch().await; // Create a collection job for some idempotency key, then try to create another, varying the // aggregation parameter. It should fail because the batch is marked as aggregated. - let (test_case, _, _, _) = setup_leader_selected_current_batch_collection_job_test_case().await; let idempotency_key: LeaderSelectedQueryConfig = random(); @@ -803,8 +1049,15 @@ async fn collection_job_post_idempotence_leader_selected_mutate_aggregation_para #[tokio::test] async fn collection_job_post_idempotence_leader_selected_no_extra_reports() { - let (test_case, _batch_id_1, _batch_id_2, _) = - setup_leader_selected_current_batch_collection_job_test_case().await; + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + + test_case.setup_leader_selected_current_batch().await; let request_1 = CollectionJobReq::new( Query::new_leader_selected(random()), @@ -843,7 +1096,7 @@ async fn collection_job_post_idempotence_leader_selected_no_extra_reports() { #[tokio::test] async fn collection_job_batch_mode_misaligned() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs b/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs index 5a2bcbcb3..1811aaff7 100644 --- a/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs +++ b/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs @@ -1,5 +1,6 @@ use assert_matches::assert_matches; use axum::{Router, body::Body}; +use chrono::TimeDelta; use futures::future::try_join_all; use http::{Request, StatusCode}; use janus_aggregator_core::{ @@ -17,7 +18,7 @@ use janus_core::{ auth_tokens::test_util::WithAuthenticationToken, hpke::{self, HpkeApplicationInfo, Label}, report_id::ReportIdChecksumExt, - time::{Clock, DateTimeExt}, + time::{Clock, DateTimeExt, TimeExt}, vdaf::VdafInstance, }; use janus_messages::{ @@ -37,7 +38,9 @@ use tower::ServiceExt; use crate::aggregator::{ error::BatchMismatch, - http_handlers::test_util::{HttpHandlerTest, decode_response_body, take_problem_details}, + http_handlers::test_util::{ + HttpHandlerTest, decode_response_body, take_problem_details, take_response_body, + }, }; pub(crate) async fn put_aggregate_share_request( @@ -997,16 +1000,15 @@ async fn aggregate_share_request_get_poll_after_put() { assert_eq!(response.status(), StatusCode::NO_CONTENT); } -#[tokio::test] -async fn aggregate_share_request_get_poll_after_put_leader_selected() { - let HttpHandlerTest { - clock: _, - ephemeral_datastore: _ephemeral_datastore, - datastore, - router, - .. - } = HttpHandlerTest::new().await; - +async fn setup_leader_selected_task_with_batch_aggregation( + datastore: &Datastore, + batch_id: BatchId, + client_timestamp_interval: Interval, + aggregation_param: dummy::AggregationParam, + report_count: u64, + checksum: ReportIdChecksum, + report_expiry_age: Option, +) -> Task { let task = TaskBuilder::new( BatchMode::LeaderSelected { batch_time_window_size: None, @@ -1014,18 +1016,13 @@ async fn aggregate_share_request_get_poll_after_put_leader_selected() { AggregationMode::Synchronous, VdafInstance::Fake { rounds: 1 }, ) + .with_report_expiry_age(report_expiry_age) .with_helper_aggregator_endpoint("https://helper.example.com/".parse().unwrap()) .build(); let helper_task = task.helper_view().unwrap(); datastore.put_aggregator_task(&helper_task).await.unwrap(); - let batch_id = BatchId::from([9u8; 32]); - let client_timestamp_interval = Interval::minimal(Time::from_time_precision_units(0)).unwrap(); - let aggregation_param = dummy::AggregationParam(0); - let report_count = 5; - let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); - datastore .run_unnamed_tx(|tx| { let helper_task = helper_task.clone(); @@ -1056,6 +1053,36 @@ async fn aggregate_share_request_get_poll_after_put_leader_selected() { .await .unwrap(); + task +} + +#[tokio::test] +async fn aggregate_share_request_get_poll_after_put_leader_selected() { + let HttpHandlerTest { + clock: _, + ephemeral_datastore: _ephemeral_datastore, + datastore, + router, + .. + } = HttpHandlerTest::new().await; + + let batch_id = BatchId::from([9u8; 32]); + let client_timestamp_interval = Interval::minimal(Time::from_time_precision_units(0)).unwrap(); + let aggregation_param = dummy::AggregationParam(0); + let report_count = 5; + let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); + + let task = setup_leader_selected_task_with_batch_aggregation( + &datastore, + batch_id, + client_timestamp_interval, + aggregation_param, + report_count, + checksum, + None, + ) + .await; + let request = AggregateShareReq::new( CollectionJobReq::new( Query::new_leader_selected(random()), @@ -1110,6 +1137,25 @@ async fn setup_time_interval_task_with_batch_aggregation( aggregation_param: dummy::AggregationParam, report_count: u64, checksum: ReportIdChecksum, +) -> Task { + setup_time_interval_task_with_batch_aggregation_and_report_expiry( + datastore, + batch_interval, + aggregation_param, + report_count, + checksum, + None, + ) + .await +} + +async fn setup_time_interval_task_with_batch_aggregation_and_report_expiry( + datastore: &Datastore, + batch_interval: Interval, + aggregation_param: dummy::AggregationParam, + report_count: u64, + checksum: ReportIdChecksum, + report_expiry_age: Option, ) -> Task { let task = TaskBuilder::new( BatchMode::TimeInterval, @@ -1117,6 +1163,7 @@ async fn setup_time_interval_task_with_batch_aggregation( VdafInstance::Fake { rounds: 1 }, ) .with_helper_aggregator_endpoint("https://helper.example.com/".parse().unwrap()) + .with_report_expiry_age(report_expiry_age) .build(); let helper_task = task.helper_view().unwrap(); @@ -1767,3 +1814,381 @@ async fn aggregate_share_request_unsupported_extension() { }) ); } + +#[tokio::test] +async fn recreate_deleted_aggregate_share_time_interval() { + let HttpHandlerTest { + clock: _, + ephemeral_datastore: _ephemeral_datastore, + datastore, + router, + .. + } = HttpHandlerTest::new().await; + + let selected_interval = Interval::minimal(Time::from_time_precision_units(3)).unwrap(); + let queried_interval = Interval::new( + Time::from_time_precision_units(0), + Duration::from_time_precision_units(10), + ) + .unwrap(); + let aggregation_param = dummy::AggregationParam(0); + let report_count = 5; + let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); + + let task = setup_time_interval_task_with_batch_aggregation( + &datastore, + selected_interval, + aggregation_param, + report_count, + checksum, + ) + .await; + + // Drive aggregate share request to successful completion. + let request = AggregateShareReq::new( + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), + BatchSelector::new_time_interval(selected_interval), + report_count, + checksum, + ); + + let aggregate_share_id = AggregateShareId::from([42u8; 16]); + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + let mut response = router + .clone() + .oneshot( + Request::delete( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(take_response_body(&mut response).await.is_empty()); + + // Re-create the deleted aggregate share + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_deleted_aggregate_share_leader_selected() { + let HttpHandlerTest { + clock: _, + ephemeral_datastore: _ephemeral_datastore, + datastore, + router, + .. + } = HttpHandlerTest::new().await; + + let selected_interval = Interval::minimal(Time::from_time_precision_units(3)).unwrap(); + let aggregation_param = dummy::AggregationParam(0); + let report_count = 5; + let batch_id = random(); + let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); + + let task = setup_leader_selected_task_with_batch_aggregation( + &datastore, + batch_id, + selected_interval, + aggregation_param, + report_count, + checksum, + None, + ) + .await; + + // Drive aggregate share request to successful completion. + let request = AggregateShareReq::new( + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + BatchSelector::new_leader_selected(batch_id), + report_count, + checksum, + ); + + let aggregate_share_id = AggregateShareId::from([42u8; 16]); + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + let mut response = router + .clone() + .oneshot( + Request::delete( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(take_response_body(&mut response).await.is_empty()); + + // Re-create the deleted aggregate share + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_aggregate_share_time_interval() { + let HttpHandlerTest { + clock, + ephemeral_datastore: _ephemeral_datastore, + datastore, + router, + .. + } = HttpHandlerTest::new().await; + + // Default TaskBuilder time precision + let time_precision = TimePrecision::from_hours(8); + let now = Time::from_seconds_since_epoch(clock.now().as_seconds_since_epoch(), &time_precision); + + let selected_interval = Interval::minimal( + now.sub_duration(&Duration::from_time_precision_units(3)) + .unwrap(), + ) + .unwrap(); + let queried_interval = Interval::new( + now.sub_duration(&Duration::from_time_precision_units(10)) + .unwrap(), + Duration::from_time_precision_units(10), + ) + .unwrap(); + let aggregation_param = dummy::AggregationParam(0); + let report_count = 5; + let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); + let expiry = Duration::from_time_precision_units(20); + + let task = setup_time_interval_task_with_batch_aggregation_and_report_expiry( + &datastore, + selected_interval, + aggregation_param, + report_count, + checksum, + Some(expiry), + ) + .await; + assert_eq!(&time_precision, task.time_precision()); + + // Drive aggregate share request to successful completion. + let request = AggregateShareReq::new( + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), + BatchSelector::new_time_interval(selected_interval), + report_count, + checksum, + ); + + let aggregate_share_id = AggregateShareId::from([42u8; 16]); + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + let mut response = router + .clone() + .oneshot( + Request::get( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + // Advance clock far enough for the aggregate share to expire. + clock.advance(TimeDelta::seconds( + expiry.as_seconds(task.time_precision()) as i64 + )); + assert_eq!( + router + .clone() + .oneshot( + Request::get( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + .status(), + StatusCode::NOT_FOUND, + ); + + // Re-create the deleted aggregate share + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "detail": "aggregate share request too late", + "type": "https://docs.divviup.org/references/janus-errors#aggregate-share-request-rejected", + "title": "Aggregate share request rejected.", + "taskid": format!("{}", task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_aggregate_share_leader_selected() { + let HttpHandlerTest { + clock, + ephemeral_datastore: _ephemeral_datastore, + datastore, + router, + .. + } = HttpHandlerTest::new().await; + + // Default TaskBuilder time precision + let time_precision = TimePrecision::from_hours(8); + let now = Time::from_seconds_since_epoch(clock.now().as_seconds_since_epoch(), &time_precision); + + let selected_interval = Interval::minimal( + now.sub_duration(&Duration::from_time_precision_units(3)) + .unwrap(), + ) + .unwrap(); + let aggregation_param = dummy::AggregationParam(0); + let batch_id = random(); + let report_count = 5; + let checksum = ReportIdChecksum::get_decoded(&[3; 32]).unwrap(); + let expiry = Duration::from_time_precision_units(20); + + let task = setup_leader_selected_task_with_batch_aggregation( + &datastore, + batch_id, + selected_interval, + aggregation_param, + report_count, + checksum, + Some(expiry), + ) + .await; + assert_eq!(&time_precision, task.time_precision()); + + // Drive aggregate share request to successful completion. + let request = AggregateShareReq::new( + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + BatchSelector::new_leader_selected(batch_id), + report_count, + checksum, + ); + + let aggregate_share_id = AggregateShareId::from([42u8; 16]); + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + let mut response = router + .clone() + .oneshot( + Request::get( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + decode_response_body::(&mut response).await; + + // Advance clock far enough for the aggregate share to expire. + clock.advance(TimeDelta::seconds( + expiry.as_seconds(task.time_precision()) as i64 + )); + assert_eq!( + router + .clone() + .oneshot( + Request::get( + task.aggregate_shares_uri(&aggregate_share_id) + .unwrap() + .path(), + ) + .with_authentication_token(task.aggregator_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + .status(), + StatusCode::NOT_FOUND, + ); + + // Re-create the deleted aggregate share + let mut response = + put_aggregate_share_request(&task, &request, &aggregate_share_id, &router).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "detail": "aggregate share request too late", + "type": "https://docs.divviup.org/references/janus-errors#aggregate-share-request-rejected", + "title": "Aggregate share request rejected.", + "taskid": format!("{}", task.id()), + }) + ); +} diff --git a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs index 9c1f9fde9..48fc5f791 100644 --- a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs +++ b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs @@ -1,39 +1,29 @@ -use assert_matches::assert_matches; use axum::body::Body; +use chrono::TimeDelta; use http::{Request, StatusCode}; -use janus_aggregator_core::{ - batch_mode::AccumulableBatchMode, - datastore::models::{CollectionJob, CollectionJobState}, - task::{AggregationMode, BatchMode, test_util::TaskBuilder}, -}; +use janus_aggregator_core::task::{AggregationMode, BatchMode, test_util::TaskBuilder}; use janus_core::{ auth_tokens::test_util::WithAuthenticationToken, - hpke::{self, HpkeApplicationInfo, Label}, + time::{Clock, DateTimeExt, TimeExt}, vdaf::VdafInstance, }; use janus_messages::{ - AggregateShareAad, CollectionJobExtension, CollectionJobExtensionType, CollectionJobId, - CollectionJobReq, CollectionJobResp, Duration, Interval, MediaType, Query, Role, Time, - batch_mode::TimeInterval, -}; -use prio::{ - codec::{Decode, Encode}, - vdaf::dummy, + CollectionJobExtension, CollectionJobExtensionType, CollectionJobId, CollectionJobReq, + CollectionJobResp, Duration, Interval, MediaType, Query, Role, Time, batch_mode::TimeInterval, }; +use prio::{codec::Encode, vdaf::dummy}; use rand::random; use serde_json::json; use tower::ServiceExt; use crate::aggregator::{ - collection_job_tests::{ids_from_collection_job_response, setup_collection_job_test_case}, - http_handlers::test_util::{ - HttpHandlerTest, decode_response_body, take_problem_details, take_response_body, - }, + collection_job_tests::{CollectionJobTestCase, ids_from_collection_job_response}, + http_handlers::test_util::{HttpHandlerTest, decode_response_body, take_problem_details}, }; #[tokio::test] async fn collection_job_post_request_to_helper() { - let test_case = setup_collection_job_test_case(Role::Helper, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Helper, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -70,7 +60,7 @@ async fn collection_job_post_request_to_helper() { #[tokio::test] async fn collection_job_post_request_invalid_batch_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -109,7 +99,7 @@ async fn collection_job_post_request_invalid_batch_interval() { #[tokio::test] async fn collection_job_post_request_invalid_aggregation_parameter() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -146,7 +136,7 @@ async fn collection_job_post_request_invalid_aggregation_parameter() { #[tokio::test] async fn collection_job_post_request_unsupported_extension() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // No collection job extension types are defined, so any extension is unsupported. let request = CollectionJobReq::new( @@ -181,7 +171,7 @@ async fn collection_job_post_request_unsupported_extension() { #[tokio::test] async fn collection_job_post_request_invalid_extension_order() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // Extension types 0x0002 then 0x0001 are not in strictly increasing order. let request = CollectionJobReq::new( @@ -271,7 +261,7 @@ async fn collection_job_post_request_invalid_batch_size() { #[tokio::test] async fn collection_job_post_request_unauthenticated() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -310,7 +300,7 @@ async fn collection_job_post_request_unauthenticated() { #[tokio::test] async fn collection_job_get_request_unauthenticated_collection_jobs() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -358,188 +348,14 @@ async fn collection_job_get_request_unauthenticated_collection_jobs() { #[tokio::test] async fn collection_job_success_time_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; - test_case - .setup_time_interval_batch(Time::from_seconds_since_epoch( - 0, - test_case.task.time_precision(), - )) - .await; - - let batch_interval = - TimeInterval::to_batch_identifier(&(), &Time::from_time_precision_units(0)).unwrap(); - - let aggregation_param = dummy::AggregationParam::default(); - let leader_aggregate_share = dummy::AggregateShare(0); - let helper_aggregate_share = dummy::AggregateShare(1); - - let request = CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ); - - let mut response = test_case.post_collection_job(&request).await; - assert_eq!(response.status(), StatusCode::OK); - let (task_id, collection_job_id) = ids_from_collection_job_response(&response); - assert_eq!(&task_id, test_case.task.id()); - - let got_collection_job = test_case - .datastore - .run_unnamed_tx(|tx| { - let task_id = *test_case.task.id(); - - Box::pin(async move { - Ok(tx - .get_collection_job(&dummy::Vdaf::new(1), &task_id, &collection_job_id) - .await - .unwrap() - .unwrap()) - }) - }) - .await - .unwrap(); - - let want_collection_job = CollectionJob::<0, TimeInterval, dummy::Vdaf>::new( - *test_case.task.id(), - collection_job_id, - // This is chosen by the leader, we have to copy it - *got_collection_job.aggregate_share_id(), - Query::new_time_interval(batch_interval), - aggregation_param, - batch_interval, - CollectionJobState::Start, - ); - - assert_eq!(want_collection_job, got_collection_job); - assert!(take_response_body(&mut response).await.is_empty()); - - let mut response = test_case.get_collection_job(&collection_job_id).await; - assert_eq!(response.status(), StatusCode::OK); - assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); - assert!(take_response_body(&mut response).await.is_empty()); - - // Update the collection job with the aggregate shares and some aggregation jobs. collection - // job should now be complete. - test_case - .datastore - .run_unnamed_tx(|tx| { - let task = test_case.task.clone(); - let helper_aggregate_share_bytes = helper_aggregate_share.get_encoded().unwrap(); - Box::pin(async move { - let encrypted_helper_aggregate_share = hpke::seal( - task.collector_hpke_keypair().config(), - &HpkeApplicationInfo::new( - &Label::AggregateShare, - &Role::Helper, - &Role::Collector, - ), - &helper_aggregate_share_bytes, - &AggregateShareAad::new( - *task.id(), - task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), - ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - - let collection_job = tx - .get_collection_job::<0, TimeInterval, dummy::Vdaf>( - &dummy::Vdaf::new(1), - task.id(), - &collection_job_id, - ) - .await - .unwrap() - .unwrap() - .with_state(CollectionJobState::Finished { - report_count: 12, - client_timestamp_interval: batch_interval, - encrypted_helper_aggregate_share, - leader_aggregate_share, - }); - - tx.update_collection_job::<0, TimeInterval, dummy::Vdaf>(&collection_job) - .await - .unwrap(); - Ok(()) - }) - }) - .await - .unwrap(); - - let mut response = test_case.get_collection_job(&collection_job_id).await; - - assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - response.headers().get("content-type").unwrap(), - CollectionJobResp::MEDIA_TYPE, - ); - let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; - let ( - report_count, - interval, - leader_encrypted_aggregate_share, - helper_encrypted_aggregate_share, - ) = assert_matches!( - collect_resp.clone(), - CollectionJobResp{ - report_count, - interval, - leader_encrypted_agg_share, - helper_encrypted_agg_share, - .. - } => (report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share) - ); + let mut test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; - assert_eq!(report_count, 12); - assert_eq!(interval, batch_interval); - - let decrypted_leader_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), - &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), - &leader_encrypted_aggregate_share, - &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval( + Time::from_time_precision_units(10), + Interval::minimal(Time::from_time_precision_units(10)).unwrap(), ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - assert_eq!( - leader_aggregate_share, - dummy::AggregateShare::get_decoded(decrypted_leader_aggregate_share.as_ref()).unwrap() - ); - - let decrypted_helper_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), - &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), - &helper_encrypted_aggregate_share, - &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), - ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - assert_eq!( - helper_aggregate_share, - dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() - ); + .await; // Re-send the very first collection job creation request. We should get the representation of // the completed collection job without needing a further GET. @@ -556,7 +372,7 @@ async fn collection_job_success_time_interval() { #[tokio::test] async fn collection_job_get_request_no_such_collection_job() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -574,7 +390,7 @@ async fn collection_job_get_request_no_such_collection_job() { #[tokio::test] async fn collection_job_post_request_batch_queried_multiple_times() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -612,7 +428,7 @@ async fn collection_job_post_request_batch_queried_multiple_times() { #[tokio::test] async fn collection_job_post_request_batch_overlap() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -657,7 +473,7 @@ async fn collection_job_post_request_batch_overlap() { #[tokio::test] async fn delete_collection_job() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let batch_interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -666,19 +482,7 @@ async fn delete_collection_job() { .await; // Try to delete a collection job that doesn't exist - let response = test_case - .router - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri(test_case.task.collection_job_uri(&random()).unwrap().path()) - .with_authentication_token(test_case.task.collector_auth_token()) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); + let response = test_case.delete_collection_job(&random()).await; assert_eq!(response.status(), StatusCode::NOT_FOUND); // Create a collection job @@ -692,28 +496,169 @@ async fn delete_collection_job() { let (_, collection_job_id) = ids_from_collection_job_response(&response); // Cancel the job - let response = test_case - .router - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri( - test_case - .task - .collection_job_uri(&collection_job_id) - .unwrap() - .path(), - ) - .with_authentication_token(test_case.task.collector_auth_token()) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); + let response = test_case.delete_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::NO_CONTENT); // Get the job again let response = test_case.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::NO_CONTENT); } + +#[tokio::test] +async fn recreate_deleted_collection_job_time_interval() { + let mut test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; + + // Run a collection to completion, then delete the job. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval( + Time::from_time_precision_units(10), + Interval::minimal(Time::from_time_precision_units(10)).unwrap(), + ) + .await; + let response = test_case.delete_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job should conflict with the deleted job. + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_deleted_collection_job_leader_selected() { + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + + // Run a collection to completion, then delete the job. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; + + let response = test_case.delete_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job should conflict with the deleted job. + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_collection_job_time_interval() { + let expiry = Duration::from_time_precision_units(20); + let mut test_case = CollectionJobTestCase::setup_with_report_expiry_age( + Role::Leader, + BatchMode::TimeInterval, + Some(expiry), + ) + .await; + + let now = Time::from_seconds_since_epoch( + test_case.clock().now().as_seconds_since_epoch(), + test_case.task.time_precision(), + ); + + let report_time = now + .sub_duration(&Duration::from_time_precision_units(3)) + .unwrap(); + let queried_interval = Interval::new( + now.sub_duration(&Duration::from_time_precision_units(10)) + .unwrap(), + Duration::from_time_precision_units(10), + ) + .unwrap(); + + // Run a collection to completion. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval(report_time, queried_interval) + .await; + + // Advance clock far enough for the collection job to expire. + test_case.clock().advance(TimeDelta::seconds( + expiry.as_seconds(test_case.task.time_precision()) as i64, + )); + let response = test_case.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + // Re-creating the job should fail + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_collection_job_leader_selected() { + let expiry = Duration::from_time_precision_units(20); + let mut test_case = CollectionJobTestCase::setup_with_report_expiry_age( + Role::Leader, + BatchMode::TimeInterval, + Some(expiry), + ) + .await; + + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; + + // Advance clock far enough for the collection job to expire. + test_case.clock().advance(TimeDelta::seconds( + expiry.as_seconds(test_case.task.time_precision()) as i64, + )); + let response = test_case.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + // Re-creating the job should fail + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +}