diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__init__.py b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__init__.py index 05698233..62f7fc95 100644 --- a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__init__.py +++ b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__init__.py @@ -1 +1,2 @@ +from .helpers import RegisterPostProcessor from .helpers import Setup diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/__init__.cpython-313.pyc b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 00000000..edfd02d5 Binary files /dev/null and b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/__init__.cpython-313.pyc differ diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/helpers.cpython-313.pyc b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/helpers.cpython-313.pyc new file mode 100644 index 00000000..1e0afcd7 Binary files /dev/null and b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/helpers.cpython-313.pyc differ diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/proto_fuzz_test.cpython-313.pyc b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/proto_fuzz_test.cpython-313.pyc new file mode 100644 index 00000000..3ca47f9a Binary files /dev/null and b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/__pycache__/proto_fuzz_test.cpython-313.pyc differ diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/helpers.py b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/helpers.py index 8abc08f9..e8dbab72 100644 --- a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/helpers.py +++ b/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) diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/mutator.cc b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/mutator.cc index 36b68c19..58529b3e 100644 --- a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/mutator.cc +++ b/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 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>(); + message->CopyFrom(*updated); + } + }); + }); } } // namespace protobuf_mutator diff --git a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/proto_fuzz_test.py b/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/proto_fuzz_test.py index 1a43a46c..5c69138c 100644 --- a/contrib/libprotobuf_mutator/atheris_libprotobuf_mutator/proto_fuzz_test.py +++ b/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()