Skip to content

Persist CollectionJobReq on the Helper, and allow narrower batch selectors #4766

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jcjones merged 5 commits into from
Aug 5, 2026
Merged

Persist CollectionJobReq on the Helper, and allow narrower batch selectors #4766

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0d9aa93
Persist CollectionJobReq on the Helper, and allow narrower batch sele…
jcjones Jul 30, 2026
fd2004a
allow migrations
jcjones Aug 3, 2026
b49a4d2
Review updates from Tim
jcjones Aug 4, 2026
b73f8d5
Remove unnecessary note
jcjones Aug 4, 2026
afefe0e
Review updates
jcjones Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 35 additions & 29 deletions aggregator/src/aggregator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3636,18 +3636,7 @@ impl VdafOps {
&AggregateShareAad::new(
*task.id(),
task.task_configuration()?,
// On the poll path we no longer have the Collector's original CollectionJobReq
// (there is no incoming AggregateShareReq), so reconstruct it from the stored job.
// This must be byte-identical to the request the Collector sent; see
// `query_for_collection_identifier` (and Issue #4743). When we support
// extensions (Issue #4715) this will need to change.
CollectionJobReq::new(
B::query_for_collection_identifier(aggregate_share_job.batch_identifier()),
aggregate_share_job
.aggregation_parameter()
.get_encoded()
.map_err(Error::MessageEncode)?,
),
aggregate_share_job.collection_job_req().clone(),
)
.get_encoded()
.map_err(Error::MessageEncode)?,
Expand Down Expand Up @@ -3811,14 +3800,11 @@ impl VdafOps {
)?;

// DAP-19 §4.6.4: the Helper verifies the Leader-chosen batch selector against the query in
// the CollectionJobReq, which the AAD binds. The spec defines consistency per batch mode
// and would permit any time-interval selector within the queried interval, but Janus
// requires exact round-tripping, otherwise the poll path rebuilds the AAD's query from
// the stored selector alone, so a narrower selector would decrypt here and fail there.
if B::query_for_collection_identifier(
// the CollectionJobReq, which the AAD binds.
if !B::is_batch_identifier_consistent_with_query(
aggregate_share_req.batch_selector().batch_identifier(),
) != *aggregate_share_req.collection_job_req().query()
{
aggregate_share_req.collection_job_req().query(),
) {
return Err(Error::BatchInvalid(
*task.id(),
format!(
Expand Down Expand Up @@ -3885,20 +3871,18 @@ impl VdafOps {
)
.await?
{
// Duplicate aggregate share job found - verify the aggregate share ID
// matches. Note, when we add collection extensions, those will need
// to be checked here as well. (Issue #4715)
// DAP requires duplicate requests to be identical.
if aggregate_share_job.aggregate_share_id() != &aggregate_share_id
|| aggregate_share_job.collection_job_req()
!= aggregate_share_req.collection_job_req()
Comment on lines +3874 to +3877

@divergentdave divergentdave Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should check the batch selector as well, just in case. (The report count and checksum are okay to skip, as those are a diagnostic tool in the first place)

@jcjones jcjones Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The literal check is a no-op, I think. BatchSelector<B> wraps nothing but B::BatchIdentifier so comparing the selectors is comparing the identifiers.

But, uh aggregate_share_id doesn't have anything ensuring its uniqueness. The dedup lookup is keyed on batch+param, so a leader that PUTs the same aggregate share ID against a different batch identifier doesn't hit our dedup path at all; it creates a second row sharing that ID.

The poll path then does WHERE aggregate_share_id = $2 via query_opt, which would return error for more than one row. Which becomes a 500.

I guess it's always been this way, but this PR makes the exposure a bit wider. Before, the helper derived the query from the stored batch identifier. A batch identifier had to equal its query interval exactly. Now that narrower selectors are consistent, one collection_job_req legitimately maps to many batch identifiers, so duping IDs is a more ... achievable? error condition.

I feel like I should guard against this on the cache miss scenario and do ForbiddenMutation, but I haven't figured out where yet.

@divergentdave divergentdave Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was referring to the batch_selector field of AggregateShareReq. That may be different than what is wrapped in the collection job request's query.

Ah, some of this may be left over from before aggregate shares had their own IDs. Doing a separate check by ID for an existing aggregate share would make sense I think.

I think we should also ensure that if a leader sends two identical aggregate share requests under different IDs, we send the same ciphertext in response to the second one instead of returning an error for overlapping batches. This could be useful for recovering from operational issues on the leader side without unnecessary data loss.

@jcjones jcjones Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps a real fix would be to add a UNIQUE(task_id, aggregate_share_id) constraint to the aggregate_share_jobs table.

@jcjones jcjones Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#4768

@jcjones jcjones Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was referring to the batch_selector field of AggregateShareReq. That may be different than what is wrapped in the collection job request's query.

Reasonably confident here that that would still be doing a comparison against its own clone. The case I think you're guarding against -- same query, different selector -- isn't going to make it here. A different selector is a different key, so it'll go down the cache-miss path.

I'm adding alookup by aggregate share ID there, which rejects binding one ID to two batches (with a 409), and that has turned up three existing tests that are re-using a single ID across different batches, all of which would have caused 500s if they hit the poll path. Which they didn't.

I'm also adding a test aggregate_share_request_same_id_different_batch to explicitly check this. Take a look at afefe0e.

{
// Mismatch here indicates a duplicate request with a different
// aggregate share ID. This violates the DAP protocol requirement
// that duplicate requests must be identical.
return Err(datastore::Error::User(
Error::AggregateShareRequestRejected(
*task.id(),
"aggregate share request is a duplicate but uses a different aggregate share ID"
Error::ForbiddenMutation {
resource_type: "aggregate share job",
identifier: aggregate_share_job
.aggregate_share_id()
.to_string(),
)
}
.into(),
));
}
Expand All @@ -3910,6 +3894,27 @@ impl VdafOps {
return Ok(aggregate_share_job);
}

// No job exists for this batch & aggregation parameter, so any job already
// using this aggregate share ID is for a different batch. Reject it: the poll
// path looks jobs up by ID alone and cannot disambiguate.
if tx
.get_aggregate_share_job_by_id::<SEED_SIZE, B, A>(
vdaf.as_ref(),
task.id(),
aggregate_share_id,
)
.await?
.is_some()
{
return Err(datastore::Error::User(
Error::ForbiddenMutation {
resource_type: "aggregate share job",
identifier: aggregate_share_id.to_string(),
}
.into(),
));
}

// This is a new aggregate share request, compute & validate the response.
debug!(
?aggregate_share_req,
Expand Down Expand Up @@ -3988,6 +3993,7 @@ impl VdafOps {
aggregate_share_id,
aggregate_share.report_count,
aggregate_share.checksum,
aggregate_share_req.collection_job_req().clone(),
);

tx.put_aggregate_share_job(&aggregate_share_job).await?;
Expand Down
14 changes: 11 additions & 3 deletions aggregator/src/aggregator/garbage_collector.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,11 +195,11 @@ mod tests {
vdaf::VdafInstance,
};
use janus_messages::{
AggregationJobStep, Duration, HpkeCiphertext, HpkeConfigId, Interval, Query,
ReportIdChecksum, ReportMetadata, ReportShare, Role, TimePrecision,
AggregationJobStep, CollectionJobReq, Duration, HpkeCiphertext, HpkeConfigId, Interval,
Query, ReportIdChecksum, ReportMetadata, ReportShare, Role, TimePrecision,
batch_mode::{LeaderSelected, TimeInterval},
};
use prio::vdaf::dummy;
use prio::{codec::Encode, vdaf::dummy};
use rand::random;

use crate::aggregator::garbage_collector::GarbageCollector;
Expand Down Expand Up @@ -495,6 +495,10 @@ mod tests {
random(),
0,
ReportIdChecksum::default(),
CollectionJobReq::new(
Query::new_time_interval(batch_identifier),
dummy::AggregationParam(0).get_encoded().unwrap(),
),
),
)
.await
Expand Down Expand Up @@ -889,6 +893,10 @@ mod tests {
random(),
0,
ReportIdChecksum::default(),
CollectionJobReq::new(
Query::new_leader_selected(),
dummy::AggregationParam(0).get_encoded().unwrap(),
),
),
)
.await
Expand Down
23 changes: 21 additions & 2 deletions aggregator/src/aggregator/http_handlers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,13 +186,32 @@ impl Error {
| Error::HttpClient(_)
| Error::Http { .. }
| Error::TaskParameters(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
Error::AggregateShareRequestRejected(_, _) => StatusCode::BAD_REQUEST.into_response(),
Error::AggregateShareRequestRejected(task_id, detail) => ProblemDocument::new(
"https://docs.divviup.org/references/janus-errors#aggregate-share-request-rejected",
"Aggregate share request rejected.",
StatusCode::BAD_REQUEST,
)
.with_task_id(task_id)
.with_detail(detail)
.into_response(),
Error::EmptyAggregation(task_id) => {
ProblemDocument::new_dap(DapProblemType::InvalidMessage)
.with_task_id(task_id)
.into_response()
}
Error::ForbiddenMutation { .. } => StatusCode::CONFLICT.into_response(),
Error::ForbiddenMutation {
resource_type,
identifier,
} => ProblemDocument::new(
"https://docs.divviup.org/references/janus-errors#forbidden-mutation",
"Forbidden mutation of an immutable resource.",
StatusCode::CONFLICT,
)
.with_detail(&format!(
"The {resource_type} {identifier} already exists and cannot be modified. Use a new \
identifier instead of re-sending this one with changed parameters."
))
.into_response(),
Error::BadContentType(_) => StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response(),
Error::BadRequest(detail) => ProblemDocument::new(
"about:blank",
Expand Down
Loading
Loading