Store mutation weights in the corpus_type for protobuf domains.
PiperOrigin-RevId: 712897464
diff --git a/fuzztest/internal/domains/protobuf_domain_impl.h b/fuzztest/internal/domains/protobuf_domain_impl.h
index 279bf2d..c00e2fb 100644
--- a/fuzztest/internal/domains/protobuf_domain_impl.h
+++ b/fuzztest/internal/domains/protobuf_domain_impl.h
@@ -311,6 +311,10 @@
return max;
}
+ void IncrementDepth() { ++depth_; }
+
+ int depth() const { return depth_; }
+
private:
template <typename T>
struct FilterToValue {
@@ -345,6 +349,7 @@
std::vector<FilterToValue<OptionalPolicy>> optional_policies_;
std::vector<FilterToValue<int64_t>> min_repeated_fields_sizes_;
std::vector<FilterToValue<int64_t>> max_repeated_fields_sizes_;
+ int depth_ = 0;
#define FUZZTEST_INTERNAL_POLICY_MEMBERS(Camel, cpp) \
private: \
@@ -453,6 +458,13 @@
unset_oneof_fields_ = other.unset_oneof_fields_;
}
+ bool ShouldMutate(const FieldDescriptor* field) const {
+ // When all fields in a oneof ae unset, there could be a high chance that
+ // the mutation will lead to noop. So, we exclude them from mutation.
+ return !field->containing_oneof() ||
+ GetOneofFieldPolicy(field) != OptionalPolicy::kAlwaysNull;
+ }
+
corpus_type Init(absl::BitGenRef prng) {
if (auto seed = this->MaybeGetRandomSeed(prng)) return *seed;
FUZZTEST_INTERNAL_CHECK(
@@ -463,7 +475,10 @@
absl::flat_hash_map<int, int> oneof_to_field;
// TODO(b/241124202): Use a valid proto with minimum size.
+ int64_t total_mutation_weight = 0;
for (const FieldDescriptor* field : GetProtobufFields(descriptor)) {
+ int64_t weight = EstimatedMutationWeight(field);
+ if (ShouldMutate(field)) total_mutation_weight += weight;
if (auto* oneof = field->containing_oneof()) {
if (!oneof_to_field.contains(oneof->index())) {
oneof_to_field[oneof->index()] = SelectAFieldIndexInOneof(
@@ -472,17 +487,33 @@
}
if (oneof_to_field[oneof->index()] != field->index()) continue;
} else if (!IsRequired(field) && customized_fields_.empty() &&
- IsFieldRecursive(field)) {
- // We avoid initializing non-required recursive fields by default (if
- // they are not explicitly customized). Otherwise, the initialization
- // may never terminate. If a proto has only non-required recursive
- // fields, the initialization will be deterministic, which violates the
- // assumption on domain Init. However, such cases should be extremely
- // rare and breaking the assumption would not have severe consequences.
- continue;
+ field->message_type() != nullptr) {
+ // We avoid initializing non-required fields after some depth if they
+ // are not explicitly customized). Otherwise, the initialization may
+ // never terminate. We initialize to at most depth 5. This gives a good
+ // balanced between initialization speed and coverage.
+ constexpr int kMaxInitDepth = 5;
+ int acceptable_depth = absl::Uniform(absl::IntervalClosedClosed, prng,
+ int64_t{1}, kMaxInitDepth);
+ if (policy_.depth() >= acceptable_depth) continue;
}
VisitProtobufField(field, InitializeVisitor{prng, *this, val});
+ if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE &&
+ ShouldMutate(field)) {
+ total_mutation_weight -= weight;
+ auto val_it = val.find(field->number());
+ if (field->is_repeated()) {
+ total_mutation_weight +=
+ GetSubDomain<ProtoMessageTag, true>(field).CountNumberOfFields(
+ val_it->second);
+ } else {
+ total_mutation_weight +=
+ GetSubDomain<ProtoMessageTag, false>(field).CountNumberOfFields(
+ val_it->second);
+ }
+ }
}
+ SetMutationWeight(val, total_mutation_weight);
return val;
}
@@ -503,6 +534,7 @@
value_type out(prototype_.Get()->New());
for (auto& [number, data] : value) {
+ if (IsMetadataEntry(number)) continue;
auto* field = GetField(number);
VisitProtobufField(field, GetValueVisitor{*out, *this, data});
}
@@ -585,32 +617,7 @@
}
uint64_t CountNumberOfFields(const corpus_type& val) {
- uint64_t total_weight = 0;
- auto descriptor = prototype_.Get()->GetDescriptor();
- if (GetFieldCount(descriptor) == 0) return total_weight;
-
- for (const FieldDescriptor* field : GetProtobufFields(descriptor)) {
- if (field->containing_oneof() &&
- GetOneofFieldPolicy(field) == OptionalPolicy::kAlwaysNull) {
- continue;
- }
- ++total_weight;
-
- if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
- auto val_it = val.find(field->number());
- if (val_it == val.end()) continue;
- if (field->is_repeated()) {
- total_weight +=
- GetSubDomain<ProtoMessageTag, true>(field).CountNumberOfFields(
- val_it->second);
- } else {
- total_weight +=
- GetSubDomain<ProtoMessageTag, false>(field).CountNumberOfFields(
- val_it->second);
- }
- }
- }
- return total_weight;
+ return GetMutationWeight(val);
}
uint64_t MutateSelectedField(
@@ -620,12 +627,11 @@
uint64_t field_counter = 0;
auto descriptor = prototype_.Get()->GetDescriptor();
if (GetFieldCount(descriptor) == 0) return field_counter;
+ int64_t fields_count = CountNumberOfFields(val);
+ if (fields_count < selected_field_index) return fields_count;
for (const FieldDescriptor* field : GetProtobufFields(descriptor)) {
- if (field->containing_oneof() &&
- GetOneofFieldPolicy(field) == OptionalPolicy::kAlwaysNull) {
- continue;
- }
+ if (!ShouldMutate(field)) continue;
++field_counter;
if (field_counter == selected_field_index) {
VisitProtobufField(
@@ -634,8 +640,15 @@
}
if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
+ --field_counter;
auto val_it = val.find(field->number());
- if (val_it == val.end()) continue;
+ if (val_it == val.end()) {
+ field_counter += EstimatedMutationWeight(field);
+ if (field_counter >= selected_field_index) continue;
+ VisitProtobufField(
+ field, MutateVisitor{prng, metadata, only_shrink, *this, val});
+ return field_counter;
+ }
if (field->is_repeated()) {
field_counter +=
GetSubDomain<ProtoMessageTag, true>(field).MutateSelectedField(
@@ -964,13 +977,18 @@
if (is_present) {
// Mutate the element
+ AddMutationWeight(val, -domain.CountNumberOfFields(it->second));
domain.Mutate(it->second, prng, metadata, only_shrink);
+ AddMutationWeight(val, domain.CountNumberOfFields(it->second));
return;
}
// Add the element
if (!only_shrink) {
+ AddMutationWeight(val, -EstimatedMutationWeight(field));
InitializeFieldValue<T>(prng, self, field, val);
+ AddMutationWeight(val,
+ domain.CountNumberOfFields(val[field->number()]));
}
}
@@ -982,7 +1000,9 @@
if (!is_present) {
if (!only_shrink) {
- val[field->number()] = domain.Init(prng);
+ AddMutationWeight(val, -EstimatedMutationWeight(field));
+ auto& field_val = val[field->number()] = domain.Init(prng);
+ AddMutationWeight(val, domain.CountNumberOfFields(field_val));
}
} else if (field->is_map()) {
// field of synthetic messages of the form:
@@ -1000,6 +1020,7 @@
// TODO(b/231212420): Improve mutation to not add duplicate keys on the
// first place. The current hack is very inefficient.
// Switch the inner domain for maps to use flat_hash_map instead.
+ AddMutationWeight(val, -domain.CountNumberOfFields(it->second));
corpus_type corpus_copy;
auto& copy = corpus_copy[field->number()] = it->second;
domain.Mutate(copy, prng, metadata, only_shrink);
@@ -1012,8 +1033,11 @@
// The number of entries is the same, so accept the change.
it->second = std::move(copy);
}
+ AddMutationWeight(val, domain.CountNumberOfFields(it->second));
} else {
+ AddMutationWeight(val, -domain.CountNumberOfFields(it->second));
domain.Mutate(it->second, prng, metadata, only_shrink);
+ AddMutationWeight(val, domain.CountNumberOfFields(it->second));
}
}
};
@@ -1325,13 +1349,16 @@
return descriptor->field_count() + extensions.size();
}
- static auto GetProtobufFields(const Descriptor* descriptor) {
+ static auto GetProtobufFields(const Descriptor* descriptor,
+ bool include_extensions = true) {
std::vector<const FieldDescriptor*> fields;
fields.reserve(descriptor->field_count());
for (int i = 0; i < descriptor->field_count(); ++i) {
fields.push_back(descriptor->field(i));
}
- descriptor->file()->pool()->FindAllExtensions(descriptor, &fields);
+ if (include_extensions) {
+ descriptor->file()->pool()->FindAllExtensions(descriptor, &fields);
+ }
return fields;
}
@@ -1599,6 +1626,7 @@
field->message_type()),
use_lazy_initialization_);
result.SetPolicy(policy_);
+ result.GetPolicy().IncrementDepth();
return Domain<std::unique_ptr<Message>>(result);
} else {
return Domain<T>(ArbitraryImpl<T>());
@@ -1673,6 +1701,70 @@
/*consider_non_terminating_recursions=*/false);
}
+ // Returns an estimate for the number of sub-fields of a given field to be
+ // used as a heuristic. For simplicity, the recursive sub-fields are ignored.
+ // And for efficiency, proto extensions as well as protos with large number of
+ // fields are ignored.
+ static int64_t EstimatedMutationWeight(const FieldDescriptor* field) {
+ auto descriptor = field->message_type();
+ if (!descriptor) return 1; // non-message fields
+ static absl::flat_hash_map<std::string, int> kFieldCountCache;
+ if (auto it = kFieldCountCache.find(descriptor->full_name());
+ it != kFieldCountCache.end()) {
+ return it->second;
+ }
+ absl::flat_hash_set<decltype(descriptor)> parents;
+ int result = EstimatedMutationWeight(descriptor, parents);
+ kFieldCountCache.insert({descriptor->full_name(), result});
+ return result;
+ }
+
+ template <typename Descriptor>
+ static int EstimatedMutationWeight(
+ const Descriptor* descriptor,
+ absl::flat_hash_set<const Descriptor*>& parents) {
+ int field_count = 0;
+ if (parents.contains(descriptor)) return field_count;
+ parents.insert(descriptor);
+ for (const FieldDescriptor* field :
+ GetProtobufFields(descriptor, /*include_extensions=*/false)) {
+ const auto* child = field->message_type();
+ field_count += child ? EstimatedMutationWeight(child, parents) : 1;
+ constexpr int kMaxFieldCount = 1000;
+ // Avoid expensive estimation operations.
+ if (field_count > kMaxFieldCount) {
+ field_count = kMaxFieldCount;
+ break;
+ }
+ }
+ parents.erase(descriptor);
+ return field_count;
+ }
+
+ static bool IsMetadataEntry(int index) {
+ // Used for storing the mutation weight.
+ return index == -1;
+ }
+
+ // Mutation weight for a proto is defined as the number of fields in the
+ // proto. All fields in a oneof are counted. For sub-messages, if the
+ // sub-message is set, the weight is the mutation weight of the sub-message.
+ // Otherwise, the weight is the predicted number of fields.
+ static void SetMutationWeight(corpus_type& val, int64_t field_count) {
+ val[-1] = GenericDomainCorpusType(std::in_place_type<int64_t>, field_count);
+ }
+
+ static int64_t GetMutationWeight(const corpus_type& val) {
+ auto it = val.find(-1);
+ FUZZTEST_INTERNAL_CHECK_PRECONDITION(it != val.end(),
+ "Invalid corpus value!");
+ return it->second.template GetAs<int64_t>();
+ }
+
+ static void AddMutationWeight(corpus_type& val, int64_t to_add) {
+ SetMutationWeight(val, GetMutationWeight(val) + to_add);
+ }
+
bool IsOneofRecursive(const OneofDescriptor* oneof,
absl::flat_hash_set<const Descriptor*>& parents,
const ProtoPolicy<Message>& policy,
