diff --git a/collector/Cargo.toml b/collector/Cargo.toml index 0cdd65ba5..b75d8b7b1 100644 --- a/collector/Cargo.toml +++ b/collector/Cargo.toml @@ -26,7 +26,6 @@ janus_core.workspace = true janus_messages.workspace = true mime.workspace = true prio.workspace = true -rand = { workspace = true } reqwest = { workspace = true, features = ["json"] } retry-after = { workspace = true } serde = { workspace = true } @@ -42,3 +41,4 @@ base64.workspace = true janus_collector = { workspace = true, features = ["test-util"] } janus_core = { workspace = true, features = ["test-util"] } mockito = { workspace = true } +rand = { workspace = true } diff --git a/collector/src/lib.rs b/collector/src/lib.rs index 733588381..fb4c32448 100644 --- a/collector/src/lib.rs +++ b/collector/src/lib.rs @@ -12,7 +12,10 @@ //! use std::{fs::File, str::FromStr}; //! //! use janus_collector::{Collector, ConfiguredVdaf, PrivateCollectorCredential}; -//! use janus_messages::{BatchConfig, Duration, Interval, Query, TaskId, Time, TimePrecision, Url}; +//! use janus_messages::{ +//! BatchConfig, Duration, Interval, Query, TaskId, Time, TimePrecision, Url, +//! batch_mode::LeaderSelectedQueryConfig, +//! }; //! //! # async fn run() { //! # const TIME_PRECISION: u64 = 3600; @@ -62,8 +65,12 @@ //! .await //! .unwrap(); //! -//! // Or if this is a leader-selected task, make a leader-selected query. -//! let query = Query::new_leader_selected(todo!("idempotency key")); +//! // Or if this is a leader-selected task, make a leader-selected query. You must choose an +//! // idempotency key to use as the query configuration. The idempotency key allows collection job +//! // creation to be safely retried, so you may wish to persist their values before sending +//! // requests. +//! let idempotency_key = LeaderSelectedQueryConfig::new(vec![0x00, 0x01, 0x02, 0x03]).unwrap(); +//! let query = Query::new_leader_selected(idempotency_key); //! let aggregation_result = collector.collection(query, &()).collect().await.unwrap(); //! # } //! ``` @@ -81,17 +88,19 @@ pub use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder}; use chrono::{DateTime, Duration, Utc}; pub use credential::PrivateCollectorCredential; use educe::Educe; -pub use janus_core::{auth_tokens::AuthenticationToken, hpke::HpkeKeypair, vdaf::ConfiguredVdaf}; use janus_core::{ + UrlExt, hpke::{self, HpkeApplicationInfo}, http::{HttpErrorResponse, ReqwestAuthenticationToken, check_content_type}, retries::{ - ExponentialWithTotalDelayBuilder, http_request_exponential_backoff, retry_http_request, + ExponentialWithTotalDelayBuilder, HeaderMapExt, HeaderMapExtError, + http_request_exponential_backoff, retry_http_request, }, task_config::build_task_configuration, time::TimeExt, url_for_join, }; +pub use janus_core::{auth_tokens::AuthenticationToken, hpke::HpkeKeypair, vdaf::ConfiguredVdaf}; use janus_messages::{ AggregateShareAad, BatchConfig, CollectionJobExtension, CollectionJobId, CollectionJobReq, CollectionJobResp, Interval, MediaType, Query, Role, TaskConfiguration, TaskId, TimePrecision, @@ -102,10 +111,9 @@ use prio::{ codec::{Decode, Encode, ParameterizedDecode}, vdaf, }; -use rand::random; use reqwest::{ StatusCode, - header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderValue, RETRY_AFTER, ToStrError}, + header::{CONTENT_LENGTH, CONTENT_TYPE, RETRY_AFTER}, }; pub use retry_after; use retry_after::{FromHeaderValueError, RetryAfter}; @@ -124,10 +132,8 @@ pub enum Error { Url(#[from] url::ParseError), #[error("missing Location header in See Other response")] MissingLocationHeader, - #[error("invalid bytes in header")] - InvalidHeader(#[from] ToStrError), - #[error("wrong Content-Type header: {0:?}")] - BadContentType(Option), + #[error("invalid header")] + InvalidHeader(#[from] HeaderMapExtError), #[error("invalid Retry-After header value: {0}")] InvalidRetryAfterHeader(#[from] FromHeaderValueError), #[error("codec error: {0}")] @@ -223,23 +229,16 @@ impl CollectionJob { /// A builder for a collection request, returned by [`Collector::collection`]. /// -/// Configure the optional collection job ID ([`Self::with_id`]) and extensions -/// ([`Self::with_extensions`]), then send the request with [`Self::start`] or [`Self::collect`]. +/// Configure the optional extensions ([`Self::with_extensions`]), then send the request with +/// [`Self::start`] or [`Self::collect`]. pub struct CollectionRequestBuilder<'a, V: vdaf::Collector, B: BatchMode> { collector: &'a Collector, query: Query, aggregation_parameter: &'a V::AggregationParam, - collection_job_id: Option, extensions: Vec, } impl<'a, V: vdaf::Collector, B: BatchMode> CollectionRequestBuilder<'a, V, B> { - /// Use a caller-chosen collection job ID instead of a randomly generated one. - pub fn with_id(mut self, collection_job_id: CollectionJobId) -> Self { - self.collection_job_id = Some(collection_job_id); - self - } - /// Set the collection job extensions. pub fn with_extensions(mut self, extensions: Vec) -> Self { self.extensions = extensions; @@ -250,14 +249,8 @@ impl<'a, V: vdaf::Collector, B: BatchMode> CollectionRequestBuilder<'a, V, B> { /// [`CollectionJob`] that must be polled separately using [`Collector::poll_once`] or /// [`Collector::poll_until_complete`]. pub async fn start(self) -> Result, Error> { - let collection_job_id = self.collection_job_id.unwrap_or_else(random); self.collector - .send_collection_request( - collection_job_id, - self.query, - self.aggregation_parameter, - self.extensions, - ) + .send_collection_request(self.query, self.aggregation_parameter, self.extensions) .await } @@ -649,22 +642,27 @@ impl Collector { &self.task_configuration } + /// Construct a URI for collection job creation. + fn collection_job_creation_uri(&self) -> Result { + Ok( + url_for_join(self.task_configuration.leader_aggregator_endpoint())? + .join(&format!("tasks/{}/collection_jobs", self.task_id))?, + ) + } + /// Construct a URI for a collection. fn collection_job_uri(&self, collection_job_id: CollectionJobId) -> Result { // Joined at request time so the endpoint bytes bound into HPKE AADs stay verbatim // (DAP ยง4.1). - Ok( - url_for_join(self.task_configuration.leader_aggregator_endpoint())?.join(&format!( - "tasks/{}/collection_jobs/{collection_job_id}", - self.task_id - ))?, - ) + Ok(self + .collection_job_creation_uri()? + .ensure_trailing_slash() + .join(&collection_job_id.to_string())?) } /// Construct a collection request to send to the leader aggregator. /// - /// Returns a [`CollectionRequestBuilder`]; configure the optional collection job ID - /// ([`CollectionRequestBuilder::with_id`]) and extensions + /// Returns a [`CollectionRequestBuilder`]; configure the optional extensions /// ([`CollectionRequestBuilder::with_extensions`]), then send the request with /// [`CollectionRequestBuilder::start`] (returns a [`CollectionJob`] to poll) or /// [`CollectionRequestBuilder::collect`] (sends and waits for the result). @@ -677,7 +675,6 @@ impl Collector { collector: self, query, aggregation_parameter, - collection_job_id: None, extensions: Vec::new(), } } @@ -687,7 +684,6 @@ impl Collector { #[tracing::instrument(skip(self, aggregation_parameter, extensions), err)] async fn send_collection_request( &self, - collection_job_id: CollectionJobId, query: Query, aggregation_parameter: &V::AggregationParam, extensions: Vec, @@ -696,12 +692,12 @@ impl Collector { CollectionJobReq::new(query.clone(), aggregation_parameter.get_encoded()?) .with_extensions(extensions) .get_encoded()?; - let collection_job_url = self.collection_job_uri(collection_job_id)?; + let collection_job_url = self.collection_job_creation_uri()?; let response_res = retry_http_request(self.http_request_retry_parameters.build(), || async { self.http_client - .put(collection_job_url.clone()) + .post(collection_job_url.clone()) .header(CONTENT_TYPE, CollectionJobReq::::MEDIA_TYPE) .body(collect_request.clone()) .authentication_token(&self.authentication) @@ -710,14 +706,20 @@ impl Collector { }) .await; - match response_res { + let collection_job_id = match response_res { // Successful response. Ok(response) => { let status = response.status(); - if status != StatusCode::CREATED { + if status != StatusCode::OK { // Incorrect success status code. return Err(Error::Http(Box::new(status.into()))); } + + let (_, collection_job_id) = response + .headers() + .ids_from_collection_job_location()? + .ok_or_else(|| Error::MissingLocationHeader)?; + collection_job_id } // HTTP-level error. @@ -997,7 +999,7 @@ mod tests { use rand::random; use reqwest::{ StatusCode, - header::{AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE}, + header::{AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE, LOCATION}, }; use retry_after::RetryAfter; @@ -1029,15 +1031,20 @@ mod tests { .unwrap() } + /// Matches on the relative path for collection job creation. + fn collection_creation_uri_matcher(task_id: &TaskId) -> Matcher { + Matcher::Regex(format!("^/tasks/{task_id}/collection_jobs")) + } + + /// Matches on the relative path for a collection job resource. fn collection_uri_regex_matcher(task_id: &TaskId) -> Matcher { - // Matches on the relative path for a collection job resource. The Base64 URL-safe encoding - // of a collection ID is always 22 characters. + // The Base64 URL-safe encoding of a collection ID is always 22 characters. Matcher::Regex(format!( "^/tasks/{task_id}/collection_jobs/[A-Za-z0-9-_]{{22}}$" )) } - fn build_collect_response_time( + fn build_collect_response_time_interval( transcript: &VdafTranscript, collector: &Collector, aggregation_parameter: &V::AggregationParam, @@ -1075,9 +1082,10 @@ mod tests { } } - fn build_collect_response_fixed( + fn build_collect_response_leader_selected( transcript: &VdafTranscript, collector: &Collector, + idempotency_key: &LeaderSelectedQueryConfig, aggregation_parameter: &V::AggregationParam, ) -> CollectionJobResp where @@ -1088,7 +1096,7 @@ mod tests { collector.task_id, collector.task_configuration().clone(), CollectionJobReq::new( - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(idempotency_key.clone()), aggregation_parameter.get_encoded().unwrap(), ), ); @@ -1280,6 +1288,10 @@ mod tests { .as_str(), endpoint ); + assert_eq!( + collector.collection_job_creation_uri().unwrap().as_str(), + format!("{joined_base}tasks/{}/collection_jobs", collector.task_id) + ); assert_eq!( collector .collection_job_uri(collection_job_id) @@ -1307,7 +1319,12 @@ mod tests { &random(), &true, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1315,11 +1332,11 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_error = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1329,13 +1346,14 @@ mod tests { .create_async() .await; let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) .match_authentication_token(&collector.authentication) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1350,11 +1368,8 @@ mod tests { let job = job.unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_error = server .mock("GET", collection_job_path.as_str()) .with_status(500) @@ -1412,7 +1427,12 @@ mod tests { &random(), &144, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1420,16 +1440,17 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1440,12 +1461,9 @@ mod tests { .await .unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1485,7 +1503,12 @@ mod tests { &random(), &3, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1493,16 +1516,17 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1513,13 +1537,10 @@ mod tests { .await .unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1559,37 +1580,39 @@ mod tests { &random(), &true, ); + let idempotency_key = random(); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); - let collect_resp = build_collect_response_fixed(&transcript, &collector, &()); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collect_resp = + build_collect_response_leader_selected(&transcript, &collector, &idempotency_key, &()); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; let job = collector - .collection( - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), - &(), - ) + .collection(Query::new_leader_selected(idempotency_key), &()) .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1647,6 +1670,11 @@ mod tests { .with_collect_poll_backoff(test_http_request_exponential_backoff()) .build() .unwrap(); + let collection_job_id: CollectionJobId = random(); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1654,17 +1682,18 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) .match_header(AUTHORIZATION.as_str(), "Bearer AAAAAAAAAAAAAAAA") - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1677,11 +1706,8 @@ mod tests { mocked_collect_start_success.assert_async().await; let job = job.unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .match_header(AUTHORIZATION.as_str(), "Bearer AAAAAAAAAAAAAAAA") @@ -1715,10 +1741,10 @@ mod tests { let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_server_error = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1746,7 +1772,7 @@ mod tests { mock_server_error.assert_async().await; let mock_server_error_details = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1772,7 +1798,7 @@ mod tests { mock_server_error_details.assert_async().await; let mock_bad_request = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1809,16 +1835,22 @@ mod tests { initialize_rustls(); let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_collect_start = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1839,6 +1871,7 @@ mod tests { .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); let error = collector.poll_once(&job).await.unwrap_err(); assert_matches!(error, Error::Http(error_response) => { assert_eq!(error_response.status(), StatusCode::INTERNAL_SERVER_ERROR); @@ -1848,10 +1881,6 @@ mod tests { mock_collect_start.assert_async().await; mock_collection_job_server_error.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mock_collection_job_server_error_details = server .mock("GET", collection_job_path.as_str()) .with_status(500) @@ -2057,16 +2086,22 @@ mod tests { initialize_rustls(); let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); + let collection_job_id = random(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_collect_start = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -2080,13 +2115,9 @@ mod tests { .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); mock_collect_start.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); - let mock_collect_poll_no_retry_after = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -2255,7 +2286,7 @@ mod tests { let collection_job_id = random(); let collection_job = CollectionJob::new( collection_job_id, - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(random()), dummy::AggregationParam(1), ); let matcher = collection_uri_regex_matcher(&collector.task_id); @@ -2293,7 +2324,7 @@ mod tests { let collection_job_id = random(); let collection_job = CollectionJob::new( collection_job_id, - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(random()), dummy::AggregationParam(1), ); let matcher = collection_uri_regex_matcher(&collector.task_id); @@ -2338,7 +2369,7 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); let job = CollectionJob { collection_job_id: random(), diff --git a/integration_tests/tests/integration/simulation/run.rs b/integration_tests/tests/integration/simulation/run.rs index e07cb392f..815ae259a 100644 --- a/integration_tests/tests/integration/simulation/run.rs +++ b/integration_tests/tests/integration/simulation/run.rs @@ -442,7 +442,6 @@ impl Simulation { .components .collector .collection(query, &()) - .with_id(*collection_job_id) .start() .await { @@ -462,7 +461,6 @@ impl Simulation { .components .collector .collection(query, &()) - .with_id(*collection_job_id) .start() .await { diff --git a/messages/src/batch_mode.rs b/messages/src/batch_mode.rs index 6832cea72..04edebd6f 100644 --- a/messages/src/batch_mode.rs +++ b/messages/src/batch_mode.rs @@ -215,6 +215,12 @@ impl Decode for LeaderSelectedQueryConfig { } } +impl AsRef<[u8]> for LeaderSelectedQueryConfig { + fn as_ref(&self) -> &[u8] { + self.idempotency_key.as_ref() + } +} + /// DAP protocol message indicating a batch mode. #[derive(Copy, Clone, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)] #[repr(u8)] diff --git a/tools/src/bin/collect.rs b/tools/src/bin/collect.rs index 4c821516a..a88634ae4 100644 --- a/tools/src/bin/collect.rs +++ b/tools/src/bin/collect.rs @@ -19,13 +19,13 @@ use janus_core::{ }; use janus_messages::{ BatchConfig, CollectionJobId, Duration, HpkeConfig, Interval, Query, TaskConfiguration, TaskId, - Time, batch_mode::BatchMode, + Time, + batch_mode::{BatchMode, LeaderSelectedQueryConfig}, }; use prio::{ codec::Decode, vdaf::{self, Vdaf, VdafError}, }; -use rand::random; use tracing_log::LogTracer; use tracing_subscriber::{EnvFilter, Registry, prelude::*}; @@ -132,6 +132,37 @@ fn private_collector_credential_parser( serde_json::from_str(s) } +#[derive(Debug, Args, PartialEq, Eq)] +struct AuthenticationOptions { + /// Authentication token for the DAP-Auth-Token HTTP header + #[clap( + long, + value_parser = StringValueParser::new().try_map( + AuthenticationToken::new_dap_auth_token_from_string, + ), + env, + hide_env_values = true, + help_heading = "Authorization", + display_order = 0, + conflicts_with = "authorization_bearer_token" + )] + dap_auth_token: Option, + + /// Authentication token for the "Authorization: Bearer ..." HTTP header + #[clap( + long, + value_parser = StringValueParser::new().try_map( + AuthenticationToken::new_bearer_token_from_string, + ), + env, + hide_env_values = true, + help_heading = "Authorization", + display_order = 1, + conflicts_with = "dap_auth_token" + )] + authorization_bearer_token: Option, +} + #[derive(Clone)] struct TaskConfigurationValueParser { inner: NonEmptyStringValueParser, @@ -163,35 +194,35 @@ impl TypedValueParser for TaskConfigurationValueParser { } } -#[derive(Debug, Args, PartialEq, Eq)] -struct AuthenticationOptions { - /// Authentication token for the DAP-Auth-Token HTTP header - #[clap( - long, - value_parser = StringValueParser::new().try_map( - AuthenticationToken::new_dap_auth_token_from_string, - ), - env, - hide_env_values = true, - help_heading = "Authorization", - display_order = 0, - conflicts_with = "authorization_bearer_token" - )] - dap_auth_token: Option, +#[derive(Clone)] +struct IdempotencyKeyValueParser { + inner: StringValueParser, +} - /// Authentication token for the "Authorization: Bearer ..." HTTP header - #[clap( - long, - value_parser = StringValueParser::new().try_map( - AuthenticationToken::new_bearer_token_from_string, - ), - env, - hide_env_values = true, - help_heading = "Authorization", - display_order = 1, - conflicts_with = "dap_auth_token" - )] - authorization_bearer_token: Option, +impl IdempotencyKeyValueParser { + fn new() -> Self { + Self { + inner: StringValueParser::new(), + } + } +} + +impl TypedValueParser for IdempotencyKeyValueParser { + type Value = LeaderSelectedQueryConfig; + + fn parse_ref( + &self, + cmd: &clap::Command, + arg: Option<&clap::Arg>, + value: &std::ffi::OsStr, + ) -> Result { + let input = self.inner.parse_ref(cmd, arg, value)?; + let bytes = URL_SAFE_NO_PAD + .decode(input) + .map_err(|err| clap::Error::raw(ErrorKind::ValueValidation, err))?; + LeaderSelectedQueryConfig::new(bytes) + .map_err(|err| clap::Error::raw(ErrorKind::ValueValidation, err)) + } } #[derive(Debug, Args, PartialEq, Eq)] @@ -211,6 +242,13 @@ struct QueryOptions { help_heading = "Collect Request Parameters (Time Interval)" )] batch_interval_duration: Option, + + /// Idempotency key for collection job creation, as base64 encoded bytes. + /// + /// If not set, a random idempotency key is used. + #[clap(long, value_parser = IdempotencyKeyValueParser::new(), + help_heading = "Collect Request Parameters (Leader Selected)")] + idempotency_key: Option, } #[derive(Debug, Args, PartialEq, Eq)] @@ -273,12 +311,7 @@ enum Subcommands { /// Initialize a new collection job /// /// Outputs collection job ID to stdout. - NewJob { - /// Job ID to use for the new collection job. If absent, an ID is randomly generated - /// - /// A valid ID consists of 16 randomly selected bytes, encoded with unpadded base64url. - collection_job_id: Option, - }, + NewJob, /// Poll an existing collection job once /// /// The supplied query options must exactly match the ones used to create the collection job, @@ -438,7 +471,7 @@ async fn main() -> anyhow::Result<()> { /// The batch to collect, resolved against the task's batch configuration. enum ResolvedQuery { TimeInterval(Interval), - LeaderSelected, + LeaderSelected(LeaderSelectedQueryConfig), } /// Resolves the query arguments against the batch configuration in `--task-config`. The batch mode @@ -450,8 +483,9 @@ fn resolve_query(options: &Options) -> Result { batch_config, options.query.batch_interval_start, options.query.batch_interval_duration, + options.query.idempotency_key.clone(), ) { - (BatchConfig::TimeInterval, Some(start), Some(duration)) => { + (BatchConfig::TimeInterval, Some(start), Some(duration), _) => { let time_precision = options.task_config.time_precision(); Ok(ResolvedQuery::TimeInterval( Interval::new( @@ -461,23 +495,25 @@ fn resolve_query(options: &Options) -> Result { .map_err(|err| Error::Anyhow(err.into()))?, )) } - (BatchConfig::TimeInterval, None, None) => Err(clap::Error::raw( + (BatchConfig::TimeInterval, None, None, _) => Err(clap::Error::raw( ErrorKind::MissingRequiredArgument, "the task is a time-interval task, so --batch-interval-start and \ --batch-interval-duration are required\n", ) .into()), - (BatchConfig::TimeInterval, _, _) => { + (BatchConfig::TimeInterval, _, _, _) => { unreachable!("clap requires both batch interval arguments together") } - (BatchConfig::LeaderSelected, None, None) => Ok(ResolvedQuery::LeaderSelected), - (BatchConfig::LeaderSelected, _, _) => Err(clap::Error::raw( + (BatchConfig::LeaderSelected, None, None, Some(idempotency_key)) => { + Ok(ResolvedQuery::LeaderSelected(idempotency_key)) + } + (BatchConfig::LeaderSelected, _, _, _) => Err(clap::Error::raw( ErrorKind::ArgumentConflict, "the task is a leader-selected task, so --batch-interval-start and \ --batch-interval-duration must not be provided\n", ) .into()), - (batch_config, _, _) => Err(clap::Error::raw( + (batch_config, _, _, _) => Err(clap::Error::raw( ErrorKind::ValueValidation, format!( "unsupported batch mode in --task-config: {}\n", @@ -495,10 +531,8 @@ macro_rules! options_query_dispatch { let $query = Query::new_time_interval(batch_interval); $body } - ResolvedQuery::LeaderSelected => { - let $query = Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ); + ResolvedQuery::LeaderSelected(idempotency_key) => { + let $query = Query::new_leader_selected(idempotency_key); $body } } @@ -540,9 +574,8 @@ async fn run(options: Options) -> Result<(), Error> { // Every VDAF this tool dispatches to has a trivial aggregation parameter. let agg_param: &::AggregationParam = &Default::default(); match options.subcommand { - Some(Subcommands::NewJob { collection_job_id }) => { - let collection_job_id = collection_job_id.unwrap_or_else(random); - run_new_job(options, vdaf, http_client, query, agg_param, collection_job_id).await + Some(Subcommands::NewJob) => { + run_new_job(options, vdaf, http_client, query, agg_param).await } Some(Subcommands::PollJob { collection_job_id }) => { run_poll_job(options, vdaf, http_client, query, agg_param, collection_job_id).await @@ -577,14 +610,12 @@ async fn run_new_job( http_client: reqwest::Client, query: Query, agg_param: &V::AggregationParam, - collection_job_id: CollectionJobId, ) -> Result<(), Error> where V::AggregateResult: Debug, { let collection = new_collector(options, vdaf, http_client)? .collection(query, agg_param) - .with_id(collection_job_id) .start() .await .map_err(|err| Error::Anyhow(err.into()))?; @@ -701,7 +732,10 @@ mod tests { initialize_rustls, test_util::install_test_trace_subscriber, }; - use janus_messages::{BatchConfig, TaskConfiguration, TaskId, TimePrecision, VdafConfig}; + use janus_messages::{ + BatchConfig, TaskConfiguration, TaskId, TimePrecision, VdafConfig, + batch_mode::LeaderSelectedQueryConfig, + }; use prio::{codec::Encode, vdaf::prio3::Prio3}; use rand::random; use tempfile::NamedTempFile; @@ -780,6 +814,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); @@ -989,6 +1024,8 @@ mod tests { let encoded_hpke_config = URL_SAFE_NO_PAD.encode(hpke_keypair.config().get_encoded().unwrap()); let encoded_private_key = URL_SAFE_NO_PAD.encode(hpke_keypair.private_key().as_ref()); + let encoded_idempotency_key = + URL_SAFE_NO_PAD.encode(LeaderSelectedQueryConfig::test_default().as_ref()); let auth_token = AuthenticationToken::DapAuth(random()); let task_config_argument = task_config_argument(BatchConfig::LeaderSelected); @@ -1010,6 +1047,7 @@ mod tests { query: QueryOptions { batch_interval_start: None, batch_interval_duration: None, + idempotency_key: Some(LeaderSelectedQueryConfig::test_default()), }, }; let correct_arguments = [ @@ -1019,6 +1057,7 @@ mod tests { &format!("--dap-auth-token={}", auth_token.as_str()), &format!("--hpke-config={encoded_hpke_config}"), &format!("--hpke-private-key={encoded_private_key}"), + &format!("--idempotency-key={encoded_idempotency_key}"), ]; match Options::try_parse_from(correct_arguments) { Ok(got) => assert_eq!(got, expected), @@ -1313,9 +1352,7 @@ mod tests { let task_config_argument = task_config_argument(BatchConfig::TimeInterval); let mut expected = Options { - subcommand: Some(Subcommands::NewJob { - collection_job_id: None, - }), + subcommand: Some(Subcommands::NewJob), task_id, task_config: task_config(BatchConfig::TimeInterval), authentication: AuthenticationOptions { @@ -1331,6 +1368,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); @@ -1352,10 +1390,7 @@ mod tests { Err(e) => panic!("{e}\narguments were {correct_arguments:?}"), } - let collection_job_id = random(); - expected.subcommand = Some(Subcommands::NewJob { - collection_job_id: Some(collection_job_id), - }); + expected.subcommand = Some(Subcommands::NewJob); let correct_arguments = [ "collect", &format!("--task-id={task_id_encoded}"), @@ -1368,8 +1403,6 @@ mod tests { "--batch-interval-duration", "1000", "new-job", - "--", // prevent ID from being interpreted as a flag, in case it starts with a hyphen. - &format!("{collection_job_id}"), ]; match Options::try_parse_from(correct_arguments) { Ok(got) => assert_eq!(got, expected), @@ -1405,6 +1438,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); diff --git a/tools/tests/cmd/collect.trycmd b/tools/tests/cmd/collect.trycmd index 3e5848b83..74cf36b5e 100644 --- a/tools/tests/cmd/collect.trycmd +++ b/tools/tests/cmd/collect.trycmd @@ -69,4 +69,10 @@ Collect Request Parameters (Time Interval): --batch-interval-duration Duration of the collection batch interval, in seconds +Collect Request Parameters (Leader Selected): + --idempotency-key + Idempotency key for collection job creation, as base64 encoded bytes. + + If not set, a random idempotency key is used. + ```