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)
+	}
+}

internal/protoc/protoc-base/templates.go

@@ -0,0 +1,106 @@
+package protocbase
+
+import (
+	"bytes"
+	"fmt"
+	"text/template"
+	"time"
+
+	"github.com/Masterminds/sprig"
+	"github.com/golang/protobuf/proto"
+	"github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor"
+	"golang.org/x/tools/imports"
+)
+
+var extensions = map[string]*proto.ExtensionDesc{}
+
+type BaseTemplateData struct {
+	Now  time.Time
+	File *descriptor.File
+
+	registry *descriptor.Registry
+}
+
+var templateFuncs = map[string]interface{}{
+	"option": getOption,
+}
+
+func RegisterTmplFunc(name string, f interface{}) {
+	if _, existing := templateFuncs[name]; existing {
+		panic(fmt.Sprintf("func with name %v is already registered", name))
+	}
+
+	templateFuncs[name] = f
+}
+
+func RegisterExtension(ext *proto.ExtensionDesc) {
+	extensions[ext.Name] = ext
+}
+
+func GetBaseTemplateData(registry *descriptor.Registry, file *descriptor.File) *BaseTemplateData {
+	return &BaseTemplateData{
+		Now:      time.Now().UTC(),
+		File:     file,
+		registry: registry,
+	}
+}
+
+func getOption(opts proto.Message, extName string) interface{} {
+	extDesc := extensions[extName]
+
+	if !proto.HasExtension(opts, extDesc) {
+		return nil
+	}
+
+	ext, err := proto.GetExtension(opts, extDesc)
+	if err != nil {
+		panic(err)
+	}
+
+	return ext
+}
+
+func (data *BaseTemplateData) ResolveMsgType(msgType string) string {
+	msg, err := data.registry.LookupMsg(data.File.GetPackage(), msgType)
+	if err != nil {
+		panic(err)
+	}
+
+	return msg.GoType(data.File.GoPkg.Path)
+}
+
+func (data *BaseTemplateData) ResolveFile(fileName string) *descriptor.File {
+	file, err := data.registry.LookupFile(fileName)
+	if err != nil {
+		panic(err)
+	}
+
+	return file
+}
+
+func LoadTemplate(templateData []byte, err error) *template.Template {
+	if err != nil {
+		panic(err)
+	}
+
+	return template.Must(template.New("").
+		Funcs(sprig.TxtFuncMap()).
+		Funcs(templateFuncs).
+		Parse(string(templateData)))
+}
+
+func GenerateFromTemplate(tmpl *template.Template, data interface{}) (string, error) {
+	var tpl bytes.Buffer
+	err := tmpl.Execute(&tpl, data)
+	if err != nil {
+		return "", err
+	}
+
+	tmplResult := tpl.Bytes()
+	tmplResult, err = imports.Process(".", tmplResult, nil)
+	return string(tmplResult), err
+}
+
+func GenerateFromBaseTemplate(tmpl *template.Template, registry *descriptor.Registry, file *descriptor.File) (string, error) {
+	return GenerateFromTemplate(tmpl, GetBaseTemplateData(registry, file))
+}