| // Copyright 2018 The Chromium Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| // Plugin for the protobuf compiler (protoc) that ensures proto definitions are |
| // compiled in a way that they can be used with libprotobuf-mutator. Compiles |
| // protobufs to C++ like the normal protoc (using the cpp plugin). |
| |
| #include <assert.h> |
| #include <string> |
| #include <unordered_set> |
| #include <vector> |
| |
| #include "third_party/protobuf/src/google/protobuf/compiler/cpp/cpp_generator.h" |
| #include "third_party/protobuf/src/google/protobuf/compiler/plugin.h" |
| #include "third_party/protobuf/src/google/protobuf/descriptor.h" |
| #include "third_party/protobuf/src/google/protobuf/descriptor.pb.h" |
| |
| using google::protobuf::FileDescriptor; |
| using google::protobuf::compiler::GeneratorContext; |
| using google::protobuf::compiler::cpp::CppGenerator; |
| |
| // Class that generates C++ code that can be used by LPM from proto libraries. |
| class LpmCppCodeGenerator : public CppGenerator { |
| public: |
| // Overrides the GenerateAll method on CppGenerator. This method essentially |
| // does the same thing except it ensures that files are not optimized for |
| // LITE_RUNTIME. |
| virtual bool GenerateAll(const std::vector<const FileDescriptor*>& files, |
| const std::string& parameter, |
| GeneratorContext* generator_context, |
| std::string* error) const { |
| if (files.size() == 0) |
| return true; |
| |
| // Created a DescriptorPool once here so that modified files will use the |
| // modified versions when importing. |
| google::protobuf::DescriptorPool descriptor_pool(files[0]->pool()); |
| |
| // Keep a list of files we have generated already, so that |
| // GenerateFileAndDependencies won't generate the same file twice. |
| std::unordered_set<const FileDescriptor*> prev_generated; |
| |
| // Mostly copied from GenerateAll from |
| // //third_party/protobuf/src/google/protobuf/compiler/code_generator.cc |
| bool succeeded = true; |
| for (size_t idx = 0; idx < files.size(); idx++) { |
| const FileDescriptor* file = files[idx]; |
| succeeded = |
| GenerateFileAndDependencies(file, parameter, generator_context, error, |
| &descriptor_pool, &prev_generated); |
| |
| if (!succeeded && error && error->empty()) { |
| *error = |
| "Code generator returned false but provided no error " |
| "description."; |
| } |
| if (error && !error->empty()) { |
| *error = file->name() + ": " + *error; |
| break; |
| } |
| if (!succeeded) |
| break; |
| } |
| return succeeded; |
| } |
| |
| // Ensures that file and its dependancies are optimized for LPM by making them |
| // optimized for speed (as opposed to LITE_RUNTIME which would make file |
| // usable for LPM) then returns the result of a call to Generate on the |
| // modified file and the other arguments to this method. Needs to modify |
| // dependencies before file because protobuf doesn't allow a file to import |
| // another if file is not optimized_for LITE_RUNTIME but the dependency is. |
| // Returns true if file is in prev_generated. |
| virtual bool GenerateFileAndDependencies( |
| const FileDescriptor* file, |
| const std::string& parameter, |
| GeneratorContext* generator_context, |
| std::string* error, |
| google::protobuf::DescriptorPool* descriptor_pool, |
| std::unordered_set<const FileDescriptor*>* prev_generated) const { |
| if (prev_generated->find(file) != prev_generated->end()) |
| return true; |
| |
| // Make a copy of the file that we can modify. |
| google::protobuf::FileDescriptorProto file_proto; |
| file->CopyTo(&file_proto); |
| |
| // Fix all dependencies before fixing this file (A file must be optimized |
| // for the lite runtime if it imports files that are. |
| for (int idx = 0; idx < file->dependency_count(); idx++) { |
| const FileDescriptor* dependent_file = file->dependency(idx); |
| assert(dependent_file); |
| bool result = GenerateFileAndDependencies( |
| dependent_file, parameter, generator_context, error, descriptor_pool, |
| prev_generated); |
| assert(result); |
| if (!result) |
| return result; |
| } |
| |
| // Base case: |
| // Now make sure we aren't using the LITE_RUNTIME. |
| file_proto.mutable_options()->set_optimize_for( |
| google::protobuf::FileOptions::SPEED); |
| |
| // Convert it back to a FileDescriptor and pass it to the parent Generate |
| // method. |
| const FileDescriptor* modified_file = |
| descriptor_pool->BuildFile(file_proto); |
| assert(modified_file); |
| |
| // Ensure we only generate code once for file. |
| prev_generated->insert(file); |
| return CppGenerator::Generate(modified_file, parameter, generator_context, |
| error); |
| } |
| }; |
| |
| int main(int argc, char** argv) { |
| // Invoke our lightly modified C++ code generator on the inputs. |
| LpmCppCodeGenerator generator; |
| return PluginMain(argc, argv, &generator); |
| } |