Restrict proto3 strings to valid UTF-8

Fixes google#43
vitalybuka · Jun 10, 2017 · c183b85 · c183b85
1 parent f9f0498
commit c183b85
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 11 deletions.
diff --git a/src/field_instance.h b/src/field_instance.h
@@ -177,6 +177,12 @@ class ConstFieldInstance {
     return descriptor_->message_type();
   }
 
+  bool EnforceUtf8() const {
+    return descriptor_->type() == protobuf::FieldDescriptor::TYPE_STRING &&
+           descriptor()->file()->syntax() ==
+               protobuf::FileDescriptor::SYNTAX_PROTO3;
+  }
+
  protected:
   bool is_repeated() const { return descriptor_->is_repeated(); }
 
@@ -303,12 +309,6 @@ class FieldInstance : public ConstFieldInstance {
     if (value) mutable_message->CopyFrom(*value);
   }
 
-  bool EnforceUtf8() const {
-    return descriptor()->type() == protobuf::FieldDescriptor::TYPE_STRING &&
-           descriptor()->file()->syntax() ==
-               protobuf::FileDescriptor::SYNTAX_PROTO3;
-  }
-
  private:
   template <class T>
   void InsertRepeated(const T& value) const {

diff --git a/src/mutator.cc b/src/mutator.cc
@@ -20,6 +20,7 @@
 #include <string>
 
 #include "src/field_instance.h"
+#include "src/utf8_fix.h"
 #include "src/weighted_reservoir_sampler.h"
 
 namespace protobuf_mutator {
@@ -310,12 +311,14 @@ class DataSourceSampler {
         if (int field_size = reflection->FieldSize(*message, field)) {
           ConstFieldInstance source(message, field,
                                     GetRandomIndex(random_, field_size));
+          if (match_.EnforceUtf8() && !source.EnforceUtf8()) continue;
           if (!IsEqualValueField()(match_, source))
             sampler_.Try(field_size, source);
         }
       } else {
         if (reflection->HasField(*message, field)) {
           ConstFieldInstance source(message, field);
+          if (match_.EnforceUtf8() && !source.EnforceUtf8()) continue;
           if (!IsEqualValueField()(match_, source)) sampler_.Try(1, source);
         }
       }
@@ -333,9 +336,10 @@ class DataSourceSampler {
 class FieldMutator {
  public:
   FieldMutator(size_t size_increase_hint, bool enforce_changes,
-               Mutator* mutator)
+               bool enforce_utf8_strings, Mutator* mutator)
       : size_increase_hint_(size_increase_hint),
         enforce_changes_(enforce_changes),
+        enforce_utf8_strings_(enforce_utf8_strings),
         mutator_(mutator) {}
 
   void Mutate(int32_t* value) const {
@@ -374,8 +378,13 @@ class FieldMutator {
   }
 
   void Mutate(std::string* value) const {
-    RepeatMutate(value, std::bind(&Mutator::MutateString, mutator_, _1,
-                                  size_increase_hint_));
+    if (enforce_utf8_strings_) {
+      RepeatMutate(value, std::bind(&Mutator::MutateUtf8String, mutator_, _1,
+                                    size_increase_hint_));
+    } else {
+      RepeatMutate(value, std::bind(&Mutator::MutateString, mutator_, _1,
+                                    size_increase_hint_));
+    }
   }
 
   void Mutate(std::unique_ptr<Message>* message) const {
@@ -402,6 +411,7 @@ class FieldMutator {
 
   size_t size_increase_hint_;
   size_t enforce_changes_;
+  bool enforce_utf8_strings_;
   Mutator* mutator_;
 };
 
@@ -413,7 +423,8 @@ struct MutateField : public FieldFunction<MutateField> {
                Mutator* mutator) const {
     T value;
     field.Load(&value);
-    FieldMutator(size_increase_hint, true, mutator).Mutate(&value);
+    FieldMutator(size_increase_hint, true, field.EnforceUtf8(), mutator)
+        .Mutate(&value);
     field.Store(value);
   }
 };
@@ -426,7 +437,8 @@ struct CreateField : public FieldFunction<CreateField> {
     T value;
     field.GetDefault(&value);
     FieldMutator field_mutator(size_increase_hint,
-                               false /* defaults could be useful */, mutator);
+                               false /* defaults could be useful */,
+                               field.EnforceUtf8(), mutator);
     field_mutator.Mutate(&value);
     field.Create(value);
   }
@@ -647,4 +659,11 @@ std::string Mutator::MutateString(const std::string& value,
   return result;
 }
 
+std::string Mutator::MutateUtf8String(const std::string& value,
+                                      size_t size_increase_hint) {
+  std::string str = MutateString(value, size_increase_hint);
+  FixUtf8String(&str, random_);
+  return str;
+}
+
 }  // namespace protobuf_mutator
diff --git a/src/mutator.h b/src/mutator.h
@@ -81,6 +81,8 @@ class Mutator {
   void InitializeMessage(protobuf::Message* message, size_t max_depth);
   void CrossOverImpl(const protobuf::Message& message1,
                      protobuf::Message* message2);
+  std::string MutateUtf8String(const std::string& value,
+                               size_t size_increase_hint);
 
   bool keep_initialized_ = true;
   RandomEngine* random_;

diff --git a/src/mutator_test.cc b/src/mutator_test.cc
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "port/gtest.h"
+#include "src/binary_format.h"
 #include "src/mutator_test_proto2.pb.h"
 #include "src/mutator_test_proto3.pb.h"
 #include "src/text_format.h"
@@ -576,6 +577,23 @@ TYPED_TEST(MutatorTypedTest, FailedMutations) {
   EXPECT_LT(crossovers, 100);
 }
 
+TYPED_TEST(MutatorTypedTest, Serialization) {
+  TestMutator mutator(false);
+  for (int i = 0; i < 10000; ++i) {
+    typename TestFixture::Message message;
+    for (int j = 0; j < 5; ++j) {
+      mutator.Mutate(&message, 1000);
+      typename TestFixture::Message parsed;
+
+      EXPECT_TRUE(ParseTextMessage(SaveMessageAsText(message), &parsed));
+      EXPECT_TRUE(MessageDifferencer::Equals(parsed, message));
+
+      EXPECT_TRUE(ParseBinaryMessage(SaveMessageAsBinary(message), &parsed));
+      EXPECT_TRUE(MessageDifferencer::Equals(parsed, message));
+    }
+  }
+}
+
 class MutatorMessagesTest : public MutatorTest {};
 INSTANTIATE_TEST_CASE_P(Proto2, MutatorMessagesTest,
                         ValuesIn(GetMessageTestParams<Msg>({kMessages})));