Expose RegisterPostProcessor in the libprotobuf-mutator wrapper

contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__init__.py

@@ -1 +1,2 @@
+from .helpers import RegisterPostProcessor
 from .helpers import Setup

contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/helpers.py

@@ -74,3 +74,30 @@ def TestOneProtoInputImpl(data: bytes):
       custom_mutator=_CustomMutator,
       custom_crossover=_CustomCrossOver,
       **kwargs)
+
+
+def RegisterPostProcessor(proto: Callable[..., Any],
+                          callback: Callable[[Any, int], Any]):
+  """Register a post-processor that runs after every mutation of `proto`.
+
+  libprotobuf-mutator mutates fields blindly, so it can produce messages that
+  break invariants the mutator knows nothing about (a length field that has to
+  match a repeated field, a checksum, an enum that's really a bitmask, ...).
+  A post-processor gets a chance to patch those up before the message reaches
+  your test function.
+
+  The callback receives the freshly mutated message and the mutation seed. Edit
+  the message in place, or return a new one of the same type. Returning None
+  keeps whatever edits you made in place. Registration is global per message
+  type, so call this once before atheris.Fuzz().
+
+  Args:
+    proto: the protobuf message type to post-process.
+    callback: called as callback(message, seed) after each mutation.
+  """
+
+  def _wrapped(message: Any, seed: int):
+    result = callback(message, seed)
+    return message if result is None else result
+
+  _mutator.RegisterPostProcessor(proto(), _wrapped)

contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/mutator.cc

@@ -95,6 +95,21 @@ PYBIND11_MODULE(_mutator, m) {
           }
           return std::nullopt;
         });
+  m.def("RegisterPostProcessor",
+        [](std::unique_ptr<protobuf::Message> prototype, py::function callback) {
+          const protobuf::Descriptor* descriptor = prototype->GetDescriptor();
+          libfuzzer::RegisterPostProcessor(
+              descriptor,
+              [callback](protobuf::Message* message, unsigned int seed) {
+                py::gil_scoped_acquire gil;
+                py::object result = callback(message, seed);
+                if (!result.is_none()) {
+                  auto updated =
+                      result.cast<std::unique_ptr<protobuf::Message>>();
+                  message->CopyFrom(*updated);
+                }
+              });
+        });
 }
 
 }  // namespace protobuf_mutator

contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/proto_fuzz_test.py

@@ -22,6 +22,28 @@ def testSimpleProtoComparison(self):
         expected_output=b"Solved",
         timeout=60)
 
+  def testPostProcessor(self):
+    # The post-processor rewrites every mutated message to the solving value, so
+    # the comparison fires on the first input rather than waiting for the
+    # mutator to stumble onto "abc" on its own. If the post-processor's edits
+    # didn't make it back into the message the fuzzer feeds the harness, this
+    # would time out instead.
+    def setup_with_post_processor(argv, test_one_input, **kwargs):
+      def force_value(msg, seed):
+        del seed
+        msg.value = "abc"
+
+      atheris_libprotobuf_mutator.RegisterPostProcessor(
+          wrappers_pb2.StringValue, force_value)
+      return atheris_libprotobuf_mutator.Setup(argv, test_one_input, **kwargs)
+
+    fuzz_test_lib.run_fuzztest(
+        simple_proto_comparison,
+        custom_setup=setup_with_post_processor,
+        setup_kwargs={"proto": wrappers_pb2.StringValue},
+        expected_output=b"Solved",
+        timeout=60)
+
 
 if __name__ == "__main__":
   unittest.main()