1 files changed, 48 insertions, 33 deletions

diff --git a/src/google/protobuf/compiler/plugin.cc b/src/google/protobuf/compiler/plugin.cc
index 2bebf1f3..2ff50f61 100644
--- a/src/google/protobuf/compiler/plugin.cc
+++ b/src/google/protobuf/compiler/plugin.cc
@@ -93,31 +93,15 @@ class GeneratorResponseContext : public GeneratorContext {
   const vector<const FileDescriptor*>& parsed_files_;
 };
 
-int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
-
-  if (argc > 1) {
-    std::cerr << argv[0] << ": Unknown option: " << argv[1] << std::endl;
-    return 1;
-  }
-
-#ifdef _WIN32
-  _setmode(STDIN_FILENO, _O_BINARY);
-  _setmode(STDOUT_FILENO, _O_BINARY);
-#endif
-
-  CodeGeneratorRequest request;
-  if (!request.ParseFromFileDescriptor(STDIN_FILENO)) {
-    std::cerr << argv[0] << ": protoc sent unparseable request to plugin."
-              << std::endl;
-    return 1;
-  }
-
+bool GenerateCode(const CodeGeneratorRequest& request,
+    const CodeGenerator& generator, CodeGeneratorResponse* response,
+    string* error_msg) {
   DescriptorPool pool;
   for (int i = 0; i < request.proto_file_size(); i++) {
     const FileDescriptor* file = pool.BuildFile(request.proto_file(i));
     if (file == NULL) {
       // BuildFile() already wrote an error message.
-      return 1;
+      return false;
     }
   }
 
@@ -125,19 +109,18 @@ int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
   for (int i = 0; i < request.file_to_generate_size(); i++) {
     parsed_files.push_back(pool.FindFileByName(request.file_to_generate(i)));
     if (parsed_files.back() == NULL) {
-      std::cerr << argv[0] << ": protoc asked plugin to generate a file but "
-                              "did not provide a descriptor for the file: "
-                << request.file_to_generate(i) << std::endl;
-      return 1;
+      *error_msg = "protoc asked plugin to generate a file but "
+                   "did not provide a descriptor for the file: " +
+                   request.file_to_generate(i);
+      return false;
     }
   }
 
-  CodeGeneratorResponse response;
-  GeneratorResponseContext context(&response, parsed_files);
+  GeneratorResponseContext context(response, parsed_files);
 
-  if (generator->HasGenerateAll()) {
+  if (generator.HasGenerateAll()) {
     string error;
-    bool succeeded = generator->GenerateAll(
+    bool succeeded = generator.GenerateAll(
         parsed_files, request.parameter(), &context, &error);
 
     if (!succeeded && error.empty()) {
@@ -145,14 +128,14 @@ int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
               "description.";
     }
     if (!error.empty()) {
-      response.set_error(error);
+      response->set_error(error);
     }
   } else {
     for (int i = 0; i < parsed_files.size(); i++) {
       const FileDescriptor* file = parsed_files[i];
 
       string error;
-      bool succeeded = generator->Generate(
+      bool succeeded = generator.Generate(
           file, request.parameter(), &context, &error);
 
       if (!succeeded && error.empty()) {
@@ -160,14 +143,46 @@ int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
                 "description.";
       }
       if (!error.empty()) {
-        response.set_error(file->name() + ": " + error);
+        response->set_error(file->name() + ": " + error);
         break;
       }
     }
   }
 
-  if (!response.SerializeToFileDescriptor(STDOUT_FILENO)) {
-    std::cerr << argv[0] << ": Error writing to stdout." << std::endl;
+  return true;
+}
+
+int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
+
+  if (argc > 1) {
+    std::cerr << argv[0] << ": Unknown option: " << argv[1] << std::endl;
+    return 1;
+  }
+
+#ifdef _WIN32
+  _setmode(STDIN_FILENO, _O_BINARY);
+  _setmode(STDOUT_FILENO, _O_BINARY);
+#endif
+
+  CodeGeneratorRequest request;
+  if (!request.ParseFromFileDescriptor(STDIN_FILENO)) {
+    std::cerr << argv[0] << ": protoc sent unparseable request to plugin."
+              << std::endl;
+    return 1;
+  }
+
+  string error_msg;
+  CodeGeneratorResponse response;
+
+  if (GenerateCode(request, *generator, &response, &error_msg)) {
+    if (!response.SerializeToFileDescriptor(STDOUT_FILENO)) {
+      std::cerr << argv[0] << ": Error writing to stdout." << std::endl;
+      return 1;
+    }
+  } else {
+    if (!error_msg.empty()) {
+      std::cerr << argv[0] << ": " << error_msg << std::endl;
+    }
     return 1;
   }