Don't overwrite customized field nullness with WithOneofAlwaysSet.

PiperOrigin-RevId: 608603586
google · Feb 20, 2024 · ad3dadf · ad3dadf
1 parent cc15acf
commit ad3dadf
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/domain_tests/arbitrary_domains_protobuf_test.cc b/domain_tests/arbitrary_domains_protobuf_test.cc
@@ -562,19 +562,33 @@ TEST(ProtocolBuffer, WithFieldsAlwaysSetResetsWithMaxRepeatedFieldsSize) {
                   Gt(1))));
 }
 
+bool IsInt64(const FieldDescriptor* field) {
+  return field->type() == FieldDescriptor::TYPE_INT64;
+}
+
 TEST(ProtocolBuffer, ValidationRejectsIncorrectlySetOneofField) {
   Domain<TestProtobuf> domain_a = Arbitrary<TestProtobuf>();
   Domain<TestProtobuf> domain_b = Arbitrary<TestProtobuf>()
+                                      .WithFieldsUnset(IsInt64)
                                       .WithOneofAlwaysSet("oneof_field")
                                       .WithFieldUnset("oneof_u32");
-  TestProtobuf user_value;
-  user_value.set_oneof_u32(1);
-  auto corpus_value = domain_a.FromValue(user_value);
+  TestProtobuf user_value_1;
+  user_value_1.set_oneof_u32(1);
+  auto corpus_value_1 = domain_a.FromValue(user_value_1);
 
   EXPECT_THAT(
-      domain_b.ValidateCorpusValue(*corpus_value),
+      domain_b.ValidateCorpusValue(*corpus_value_1),
       IsInvalid(
           "Invalid value for field oneof_u32 >> Optional value must be null"));
+
+  TestProtobuf user_value_2;
+  user_value_2.set_oneof_i64(1);
+  auto corpus_value_2 = domain_a.FromValue(user_value_2);
+
+  EXPECT_THAT(
+      domain_b.ValidateCorpusValue(*corpus_value_2),
+      IsInvalid(
+          "Invalid value for field oneof_i64 >> Optional value must be null"));
 }
 
 TEST(ProtocolBuffer, ValidationRejectsUnsetOneofsWithOneofAlwaysSet) {

diff --git a/fuzztest/internal/domains/protobuf_domain_impl.h b/fuzztest/internal/domains/protobuf_domain_impl.h
@@ -1136,7 +1136,11 @@ class ProtobufDomainUntypedImpl
         "\") should be called before customizing sub-fields.");
     always_set_oneofs_.insert(oneof->index());
     for (int i = 0; i < oneof->field_count(); ++i) {
-      SetOneofFieldPolicy(oneof->field(i), OptionalPolicy::kWithoutNull);
+      // Don't set previously unset fields.
+      if (GetPolicy().GetOptionalPolicy(oneof->field(i)) !=
+          OptionalPolicy::kAlwaysNull) {
+        SetOneofFieldPolicy(oneof->field(i), OptionalPolicy::kWithoutNull);
+      }
     }
   }
 
@@ -1227,7 +1231,6 @@ class ProtobufDomainUntypedImpl
   }
 
   ProtoPolicy<Message>& GetPolicy() {
-    CheckIfPolicyCanBeUpdated();
     return policy_;
   }