zitadel/internal/protoc/protoc-base/protoc_helper.go

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