From f9aaca05efa2249f26889907aaa0f37cbd67c858 Mon Sep 17 00:00:00 2001 From: DPF Team Date: Fri, 6 Nov 2026 00:36:33 -0800 Subject: [PATCH] Validate untrusted DpfKey fields and fix DPF value type helpers. ProtoValidator::ValidateDpfKey checked the structure of a DpfKey but not its contents: it accepted party values other than 0 and 1, which evaluation inconsistently treats as party 1 for control bits and as party 0 for output negation, and it did not check that value corrections are valid elements of the level's value type. Uint128To only range-checked the low 64 bits, so a 128-bit proto integer with non-zero high bits was silently truncated when converted to a narrower type. BitsNeeded indexed the first num_other tuple elements instead of the non-IntModN ones, so a tuple listing IntModN elements first was sized from the wrong element types. A malformed key, which only affects the submitting client's own evaluation, was evaluated into meaningless output shares instead of being rejected, and the same correction could be encoded in several ways; these are robustness and canonical-encoding fixes, not a privacy fix. The BitsNeeded fix changes PRG output sizing, and therefore generated keys, only for tuples in which IntModN elements precede other elements. No production value type is affected: the Privacy Sandbox ReachTuple types are all-integer or all-IntModN. Reject keys whose party is not 0 or 1 or whose value corrections fail ValidateValue, reject non-zero high 64 bits in Uint128To, and skip IntModN elements when sizing the remaining tuple elements. A stale 34-bit value for a 32-bit type in the validator's test key is corrected. Tests cover invalid parties at validation and evaluation time, oversized and high-bit corrections, high-bit conversion for every integer width, and order-independent BitsNeeded for mixed tuples. PiperOrigin-RevId: 1015333534 --- dpf/distributed_point_function_test.cc | 11 +++++ dpf/internal/proto_validator.cc | 12 ++++++ dpf/internal/proto_validator_test.cc | 47 +++++++++++++++++++++ dpf/internal/proto_validator_test.textproto | 2 +- dpf/internal/value_type_helpers.cc | 17 +++++--- dpf/internal/value_type_helpers.h | 5 ++- dpf/internal/value_type_helpers_test.cc | 33 +++++++++++++++ 7 files changed, 117 insertions(+), 10 deletions(-) diff --git a/dpf/distributed_point_function_test.cc b/dpf/distributed_point_function_test.cc index c651818..e8212e3 100644 --- a/dpf/distributed_point_function_test.cc +++ b/dpf/distributed_point_function_test.cc @@ -1185,5 +1185,16 @@ TYPED_TEST(DpfEvaluationTest, EvaluateAndApplyFailsWithInvalidKey) { StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("key"))); } +TYPED_TEST(DpfEvaluationTest, EvaluateAndApplyFailsWithInvalidKeyParty) { + DpfKey key = this->keys_.first; + key.set_party(2); + + EXPECT_THAT(this->dpf_->template EvaluateAndApply( + absl::MakeConstSpan(&key, 1), {0}, + [](absl::Span) { return true; }), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("key.party must be 0 or 1"))); +} + } // namespace } // namespace distributed_point_functions diff --git a/dpf/internal/proto_validator.cc b/dpf/internal/proto_validator.cc index 678b3de..63e9520 100644 --- a/dpf/internal/proto_validator.cc +++ b/dpf/internal/proto_validator.cc @@ -217,6 +217,14 @@ absl::Status ProtoValidator::ValidateDpfKey(const DpfKey& key) const { "Malformed DpfKey: expected ", tree_levels_needed_ - 1, " correction words, but got ", key.correction_words_size())); } + if (key.party() != 0 && key.party() != 1) { + return absl::InvalidArgumentError("key.party must be 0 or 1"); + } + // Value corrections of untrusted keys must be valid elements of the value + // type of the hierarchy level they belong to. + for (const Value& value : key.last_level_value_correction()) { + DPF_RETURN_IF_ERROR(ValidateValue(value, parameters_.back().value_type())); + } for (int i = 0; i < static_cast(hierarchy_to_tree_.size()); ++i) { if (hierarchy_to_tree_[i] == tree_levels_needed_ - 1) { // The output correction of the last tree level is always stored in @@ -231,6 +239,10 @@ absl::Status ProtoValidator::ValidateDpfKey(const DpfKey& key) const { "Malformed DpfKey: expected correction_words[", hierarchy_to_tree_[i], "] to contain the value correction of hierarchy level ", i)); } + for (const Value& value : + key.correction_words(hierarchy_to_tree_[i]).value_correction()) { + DPF_RETURN_IF_ERROR(ValidateValue(value, parameters_[i].value_type())); + } } return absl::OkStatus(); } diff --git a/dpf/internal/proto_validator_test.cc b/dpf/internal/proto_validator_test.cc index 62e9d70..3a7502d 100644 --- a/dpf/internal/proto_validator_test.cc +++ b/dpf/internal/proto_validator_test.cc @@ -214,6 +214,53 @@ TEST_F(ProtoValidatorTest, ValidateDpfKeyFailsIfOutputCorrectionIsMissing) { StartsWith("Malformed DpfKey: expected correction_words"))); } +TEST_F(ProtoValidatorTest, ValidateDpfKeySucceedsForWellFormedKey) { + DPF_EXPECT_OK(proto_validator_->ValidateDpfKey(dpf_key_)); +} + +TEST_F(ProtoValidatorTest, ValidateDpfKeyFailsIfPartyIsInvalid) { + dpf_key_.set_party(2); + + EXPECT_THAT( + proto_validator_->ValidateDpfKey(dpf_key_), + StatusIs(absl::StatusCode::kInvalidArgument, "key.party must be 0 or 1")); +} + +TEST_F(ProtoValidatorTest, + ValidateDpfKeyFailsIfLastLevelValueCorrectionTooLarge) { + auto value_64 = uint64_t{1} << 32; + dpf_key_.mutable_last_level_value_correction(0) + ->mutable_integer() + ->set_value_uint64(value_64); + + EXPECT_THAT( + proto_validator_->ValidateDpfKey(dpf_key_), + StatusIs(absl::StatusCode::kInvalidArgument, + absl::StrFormat( + "Value (= %d) too large for ValueType with bitsize = %d", + value_64, 32))); +} + +TEST_F(ProtoValidatorTest, + ValidateDpfKeyFailsIfIntermediateValueCorrectionHasHighBitsSet) { + // Setting the upper 64 bits of an otherwise in-range value must not be + // accepted for a 32-bit value type. + int num_modified = 0; + for (CorrectionWord& cw : *(dpf_key_.mutable_correction_words())) { + for (Value& value : *(cw.mutable_value_correction())) { + Block* block = value.mutable_integer()->mutable_value_uint128(); + block->set_high(1); + block->set_low(23); + ++num_modified; + } + } + ASSERT_GT(num_modified, 0); + + EXPECT_THAT(proto_validator_->ValidateDpfKey(dpf_key_), + StatusIs(absl::StatusCode::kInvalidArgument, + testing::HasSubstr("too large for ValueType"))); +} + TEST_F(ProtoValidatorTest, ValidateEvaluationContextFailsIfKeyIsMissing) { ctx_.clear_key(); diff --git a/dpf/internal/proto_validator_test.textproto b/dpf/internal/proto_validator_test.textproto index 5e411d7..8ee6ece 100644 --- a/dpf/internal/proto_validator_test.textproto +++ b/dpf/internal/proto_validator_test.textproto @@ -101,7 +101,7 @@ key { } last_level_value_correction: { integer: { - value_uint64: 8471844854 + value_uint64: 4176877558 } } } diff --git a/dpf/internal/value_type_helpers.cc b/dpf/internal/value_type_helpers.cc index 0704cce..9a1a2e4 100644 --- a/dpf/internal/value_type_helpers.cc +++ b/dpf/internal/value_type_helpers.cc @@ -103,13 +103,16 @@ absl::StatusOr BitsNeeded(const ValueType& value_type, } } if (num_other > 0) { - for (int i = 0; i < num_other; ++i) { - double per_element_security_parameter = - security_parameter + std::log2(static_cast(num_other)); - DPF_ASSIGN_OR_RETURN(int el_bitsize, - BitsNeeded(value_type.tuple().elements(i), - per_element_security_parameter)); - bitsize_other += el_bitsize; + double per_element_security_parameter = + security_parameter + std::log2(static_cast(num_other)); + // Only consider the elements that are not of type IntModN, since those + // were accounted for above. + for (const ValueType& el : value_type.tuple().elements()) { + if (el.type_case() != ValueType::kIntModN) { + DPF_ASSIGN_OR_RETURN(int el_bitsize, + BitsNeeded(el, per_element_security_parameter)); + bitsize_other += el_bitsize; + } } } if (num_ints_mod_n > 0) { diff --git a/dpf/internal/value_type_helpers.h b/dpf/internal/value_type_helpers.h index 1e28348..942f166 100644 --- a/dpf/internal/value_type_helpers.h +++ b/dpf/internal/value_type_helpers.h @@ -169,8 +169,9 @@ template >> absl::StatusOr Uint128To(absl::uint128 in) { // Check whether value is in range if it's smaller than 128 bits. if (!std::is_same::value && - absl::Uint128Low64(in) > - static_cast(std::numeric_limits::max())) { + (absl::Uint128High64(in) != 0 || + absl::Uint128Low64(in) > + static_cast(std::numeric_limits::max()))) { return absl::InvalidArgumentError(absl::StrCat( "Value (= ", absl::Uint128Low64(in), ") too large for the given type T (size ", sizeof(T), ")")); diff --git a/dpf/internal/value_type_helpers_test.cc b/dpf/internal/value_type_helpers_test.cc index 1abec0a..1b0b7f6 100644 --- a/dpf/internal/value_type_helpers_test.cc +++ b/dpf/internal/value_type_helpers_test.cc @@ -126,6 +126,22 @@ TYPED_TEST(ValueTypeIntegerTest, ValueConversionFailsIfValueOutOfRange) { } } +TYPED_TEST(ValueTypeIntegerTest, ValueConversionFailsIfHigh64BitsSet) { + Value value; + *(value.mutable_integer()) = + Uint128ToValueInteger(absl::MakeUint128(/*high=*/1, /*low=*/23)); + + if (sizeof(TypeParam) == sizeof(absl::uint128)) { + DPF_ASSERT_OK_AND_ASSIGN(TypeParam result, + ValueTypeHelper::FromValue(value)); + EXPECT_EQ(result, absl::MakeUint128(1, 23)); + } else { + EXPECT_THAT(ValueTypeHelper::FromValue(value), + StatusIs(absl::StatusCode::kInvalidArgument, + testing::HasSubstr("too large for the given type T"))); + } +} + template class ValueTypeTupleTest : public testing::Test {}; @@ -227,6 +243,23 @@ TEST(ValueTypeTupleTest, TestValueTypesAreNotEqual) { EXPECT_FALSE(equal); } +TEST(ValueTypeTupleTest, BitsNeededForMixedTupleIsIndependentOfElementOrder) { + using MyIntModN = IntModN; + ValueType int_mod_n_first = + ValueTypeHelper>::ToValueType(); + ValueType int_mod_n_last = + ValueTypeHelper>::ToValueType(); + DPF_ASSERT_OK_AND_ASSIGN(int bits_int_mod_n, + BitsNeeded(ValueTypeHelper::ToValueType(), + kDefaultSecurityParameter)); + + // The uint8_t element needs 8 bits, regardless of its position in the tuple. + EXPECT_THAT(BitsNeeded(int_mod_n_first, kDefaultSecurityParameter), + IsOkAndHolds(bits_int_mod_n + 8)); + EXPECT_THAT(BitsNeeded(int_mod_n_last, kDefaultSecurityParameter), + IsOkAndHolds(bits_int_mod_n + 8)); +} + TEST(ValueTypeTupleTest, TestFromBytesWithConcreteExample) { std::string bytes = "A 128 bit string";