Validate untrusted DpfKey fields and fix DPF value type helpers.

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<TypeParam>(
+                  absl::MakeConstSpan(&key, 1), {0},
+                  [](absl::Span<const TypeParam>) { return true; }),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("key.party must be 0 or 1")));
+}
+
 }  // namespace
 }  // namespace distributed_point_functions

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<int>(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();
 }

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();
 

dpf/internal/proto_validator_test.textproto

@@ -101,7 +101,7 @@ key {
   }
   last_level_value_correction: {
     integer: {
-      value_uint64: 8471844854
+      value_uint64: 4176877558
     }
   }
 }

dpf/internal/value_type_helpers.cc

@@ -103,13 +103,16 @@ absl::StatusOr<int> 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<double>(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<double>(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) {

dpf/internal/value_type_helpers.h

@@ -169,8 +169,9 @@ template <typename T, typename = absl::enable_if_t<is_unsigned_integer_v<T>>>
 absl::StatusOr<T> Uint128To(absl::uint128 in) {
   // Check whether value is in range if it's smaller than 128 bits.
   if (!std::is_same<T, absl::uint128>::value &&
-      absl::Uint128Low64(in) >
-          static_cast<uint64_t>(std::numeric_limits<T>::max())) {
+      (absl::Uint128High64(in) != 0 ||
+       absl::Uint128Low64(in) >
+           static_cast<uint64_t>(std::numeric_limits<T>::max()))) {
     return absl::InvalidArgumentError(absl::StrCat(
         "Value (= ", absl::Uint128Low64(in),
         ") too large for the given type T (size ", sizeof(T), ")"));

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<TypeParam>::FromValue(value));
+    EXPECT_EQ(result, absl::MakeUint128(1, 23));
+  } else {
+    EXPECT_THAT(ValueTypeHelper<TypeParam>::FromValue(value),
+                StatusIs(absl::StatusCode::kInvalidArgument,
+                         testing::HasSubstr("too large for the given type T")));
+  }
+}
+
 template <typename T>
 class ValueTypeTupleTest : public testing::Test {};
 
@@ -227,6 +243,23 @@ TEST(ValueTypeTupleTest, TestValueTypesAreNotEqual) {
   EXPECT_FALSE(equal);
 }
 
+TEST(ValueTypeTupleTest, BitsNeededForMixedTupleIsIndependentOfElementOrder) {
+  using MyIntModN = IntModN<uint32_t, 4294967291u>;
+  ValueType int_mod_n_first =
+      ValueTypeHelper<Tuple<MyIntModN, uint8_t>>::ToValueType();
+  ValueType int_mod_n_last =
+      ValueTypeHelper<Tuple<uint8_t, MyIntModN>>::ToValueType();
+  DPF_ASSERT_OK_AND_ASSIGN(int bits_int_mod_n,
+                           BitsNeeded(ValueTypeHelper<MyIntModN>::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";