feat: add protoc pkg

internal/protoc/protoc-base/protoc_helper.go

@@ -0,0 +1,112 @@
+package protocbase
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"strings"
+	"text/template"
+
+	"github.com/golang/glog"
+	"github.com/golang/protobuf/proto"
+	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
+	"github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor"
+)
+
+type GeneratorFunc func(target string, registry *descriptor.Registry, file *descriptor.File) (string, string, error)
+
+type ProtocGenerator interface {
+	Generate(target string, registry *descriptor.Registry, file *descriptor.File) (string, string, error)
+}
+
+func (f GeneratorFunc) Generate(target string, registry *descriptor.Registry, file *descriptor.File) (string, string, error) {
+	return f(target, registry, file)
+}
+
+func parseReq(r io.Reader) (*plugin.CodeGeneratorRequest, error) {
+	glog.V(1).Info("Parsing code generator request")
+
+	input, err := ioutil.ReadAll(r)
+
+	if err != nil {
+		glog.Errorf("Failed to read code generator request: %v", err)
+		return nil, err
+	}
+
+	req := &plugin.CodeGeneratorRequest{}
+
+	if err = proto.Unmarshal(input, req); err != nil {
+		glog.Errorf("Failed to unmarshal code generator request: %v", err)
+		return nil, err
+	}
+
+	glog.V(1).Info("Parsed code generator request")
+
+	return req, nil
+}
+
+func RunWithBaseTemplate(targetFileNameFmt string, tmpl *template.Template) {
+	Run(GeneratorFunc(func(target string, registry *descriptor.Registry, file *descriptor.File) (string, string, error) {
+		fileName := fmt.Sprintf(targetFileNameFmt, strings.Split(target, ".")[0])
+		fContent, err := GenerateFromBaseTemplate(tmpl, registry, file)
+		return fileName, fContent, err
+	}))
+}
+
+func Run(generator ProtocGenerator) {
+	flag.Parse()
+	defer glog.Flush()
+
+	req, err := parseReq(os.Stdin)
+	if err != nil {
+		glog.Fatal(err)
+	}
+
+	registry := descriptor.NewRegistry()
+	if err = registry.Load(req); err != nil {
+		glog.Fatal(err)
+	}
+
+	var result []*plugin.CodeGeneratorResponse_File
+
+	for _, t := range req.FileToGenerate {
+		file, err := registry.LookupFile(t)
+		if err != nil {
+			EmitError(err)
+			return
+		}
+
+		fName, fContent, err := generator.Generate(t, registry, file)
+		if err != nil {
+			EmitError(err)
+			return
+		}
+
+		result = append(result, &plugin.CodeGeneratorResponse_File{
+			Name:    &fName,
+			Content: &fContent,
+		})
+	}
+
+	EmitFiles(result)
+}
+
+func EmitFiles(out []*plugin.CodeGeneratorResponse_File) {
+	EmitResp(&plugin.CodeGeneratorResponse{File: out})
+}
+
+func EmitError(err error) {
+	EmitResp(&plugin.CodeGeneratorResponse{Error: proto.String(err.Error())})
+}
+
+func EmitResp(resp *plugin.CodeGeneratorResponse) {
+	buf, err := proto.Marshal(resp)
+	if err != nil {
+		glog.Fatal(err)
+	}
+	if _, err := os.Stdout.Write(buf); err != nil {
+		glog.Fatal(err)
+	}
+}