github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/pluginmode/protoc/plugin.go

// Copyright 2021 CloudWeGo Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package protoc

import (
	"bytes"
	"errors"
	"fmt"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	genfastpb "github.com/cloudwego/fastpb/protoc-gen-fastpb/generator"
	"github.com/cloudwego/thriftgo/generator/golang/streaming"
	gengo "google.golang.org/protobuf/cmd/protoc-gen-go/internal_gengo"
	"google.golang.org/protobuf/compiler/protogen"

	"github.com/cloudwego/kitex/tool/internal_pkg/generator"
	"github.com/cloudwego/kitex/tool/internal_pkg/log"
	"github.com/cloudwego/kitex/tool/internal_pkg/util"
)

type protocPlugin struct {
	generator.Config
	generator.PackageInfo
	Services    []*generator.ServiceInfo
	kg          generator.Generator
	err         error
	importPaths map[string]string // file -> import path
}

// Name implements the protobuf_generator.Plugin interface.
func (pp *protocPlugin) Name() string {
	return "kitex-internal"
}

// Init implements the protobuf_generator.Plugin interface.
func (pp *protocPlugin) init() {
	pp.Dependencies = map[string]string{
		"proto": "google.golang.org/protobuf/proto",
	}
}

// parse the 'M*' option
// See https://developers.google.com/protocol-buffers/docs/reference/go-generated for more information.
func (pp *protocPlugin) parseM() {
	pp.importPaths = make(map[string]string)
	for _, po := range pp.Config.ProtobufOptions {
		if po == "" || po[0] != 'M' {
			continue
		}
		idx := strings.Index(po, "=")
		if idx < 0 {
			continue
		}
		key := po[1:idx]
		val := po[idx+1:]
		if val == "" {
			continue
		}
		idx = strings.Index(val, ";")
		if idx >= 0 {
			val = val[:idx]
		}
		pp.importPaths[key] = val
	}
}

var interfaceTemplate = `

// Code generated by Kitex {{.Version}}. DO NOT EDIT.

{{range .Interfaces}}
{{$serviceName := .Name}}
type {{.Name}} interface {
{{- range .Methods}}
{{- if or .ClientStreaming .ServerStreaming}}
	{{.Name}}({{if and .ServerStreaming (not .ClientStreaming)}}req {{.ReqType}}, {{end}}stream {{$serviceName}}_{{.Name}}Server) (err error)
{{- else}}
	{{.Name}}(ctx context.Context, req {{.ReqType}}) (res {{.ResType}}, err error)
{{- end}}
{{- end}}
}

{{range .Methods}}
{{- if or .ClientStreaming .ServerStreaming}}
type {{$serviceName}}_{{.Name}}Server interface {
	streaming.Stream
	{{- if .ClientStreaming}}
	Recv() ({{.ReqType}}, error)
	{{- end}}
	{{- if .ServerStreaming}}
	Send({{.ResType}}) error
	{{- end}}
	{{- if and .ClientStreaming (not .ServerStreaming)}}
	SendAndClose({{.ResType}}) error
	{{- end}}
}
{{- end}}
{{end}}

{{end}}
`

// Generate implements the protobuf_generator.Plugin interface.
func (pp *protocPlugin) GenerateFile(gen *protogen.Plugin, file *protogen.File) {
	if pp.err != nil {
		return
	}
	gopkg := file.Proto.GetOptions().GetGoPackage()
	if !strings.HasPrefix(gopkg, pp.PackagePrefix) {
		log.Warnf("[WARN] %q is skipped because its import path %q is not located in ./kitex_gen. Change the go_package option or use '--protobuf M%s=A-Import-Path-In-kitex_gen' to override it if you want this file to be generated under kitex_gen.\n",
			file.Proto.GetName(), gopkg, file.Proto.GetName())
		return
	}
	log.Infof("[INFO] Generate %q at %q\n", file.Proto.GetName(), gopkg)

	if parts := strings.Split(gopkg, ";"); len(parts) > 1 {
		gopkg = parts[0] // remove package alias from file path
	}
	pp.Namespace = strings.TrimPrefix(gopkg, pp.PackagePrefix)
	pp.IDLName = util.IDLName(pp.Config.IDL)

	ss := pp.convertTypes(file)
	pp.Services = append(pp.Services, ss...)

	if pp.Config.Use != "" {
		return
	}

	hasStreaming := false
	// generate service package
	for _, si := range ss {
		pp.ServiceInfo = si
		fs, err := pp.kg.GenerateService(&pp.PackageInfo)
		if err != nil {
			pp.err = err
			return
		}
		if !hasStreaming && si.HasStreaming {
			hasStreaming = true
		}
		for _, f := range fs {
			gen.NewGeneratedFile(pp.adjustPath(f.Name), "").P(f.Content)
		}
	}
	// generate service interface
	if pp.err == nil {
		fixed := *file
		fixed.GeneratedFilenamePrefix = strings.TrimPrefix(fixed.GeneratedFilenamePrefix, pp.PackagePrefix)
		f := gengo.GenerateFile(gen, &fixed)
		f.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "context"})
		if hasStreaming {
			f.QualifiedGoIdent(protogen.GoIdent{
				GoImportPath: "github.com/cloudwego/kitex/pkg/streaming",
			})
		}
		f.P("var _ context.Context")

		if len(file.Services) != 0 {
			tpl := template.New("interface")
			tpl = template.Must(tpl.Parse(interfaceTemplate))
			var buf bytes.Buffer
			pp.err = tpl.ExecuteTemplate(&buf, tpl.Name(), pp.makeInterfaces(f, file))

			f.P(buf.String())
		}
	}

	// generate fast api
	if !pp.Config.NoFastAPI && pp.err == nil {
		fixed := *file
		fixed.GeneratedFilenamePrefix = strings.TrimPrefix(fixed.GeneratedFilenamePrefix, pp.PackagePrefix)
		genfastpb.GenerateFile(gen, &fixed)
	}
}

func (pp *protocPlugin) process(gen *protogen.Plugin) {
	defer func() {
		if e := recover(); e != nil {
			if err, ok := e.(error); ok {
				gen.Error(err)
			} else {
				gen.Error(fmt.Errorf("%+v", e))
			}
		}
	}()
	if len(gen.Files) == 0 {
		gen.Error(errors.New("no proto file"))
		return
	}
	pp.kg = generator.NewGenerator(&pp.Config, nil)
	// iterate over all proto files
	idl := gen.Request.FileToGenerate[0]
	for _, f := range gen.Files {
		if pp.Config.Use != "" && f.Proto.GetName() != idl {
			continue
		}
		pp.GenerateFile(gen, f)
	}

	if pp.Config.GenerateMain {
		if len(pp.Services) == 0 {
			gen.Error(errors.New("no service defined"))
			return
		}
		pp.ServiceInfo = pp.Services[len(pp.Services)-1]
		fs, err := pp.kg.GenerateMainPackage(&pp.PackageInfo)
		if err != nil {
			pp.err = err
		}
		for _, f := range fs {
			gen.NewGeneratedFile(pp.adjustPath(f.Name), "").P(f.Content)
		}
	}

	if pp.Config.TemplateDir != "" {
		if len(pp.Services) == 0 {
			gen.Error(errors.New("no service defined"))
			return
		}
		pp.ServiceInfo = pp.Services[len(pp.Services)-1]
		fs, err := pp.kg.GenerateCustomPackage(&pp.PackageInfo)
		if err != nil {
			pp.err = err
		}
		for _, f := range fs {
			gen.NewGeneratedFile(pp.adjustPath(f.Name), "").P(f.Content)
		}
	}

	if pp.err != nil {
		gen.Error(pp.err)
	}
	return
}

func (pp *protocPlugin) convertTypes(file *protogen.File) (ss []*generator.ServiceInfo) {
	pth := pp.fixImport(string(file.GoImportPath))
	if pth == "" {
		panic(fmt.Errorf("missing %q option in %q", "go_package", file.Desc.Name()))
	}
	pi := generator.PkgInfo{
		PkgName:    file.Proto.GetPackage(),
		PkgRefName: goSanitized(path.Base(pth)),
		ImportPath: pth,
	}
	for _, service := range file.Services {
		si := &generator.ServiceInfo{
			PkgInfo:        pi,
			ServiceName:    service.GoName,
			RawServiceName: string(service.Desc.Name()),
		}
		si.ServiceTypeName = func() string { return si.PkgRefName + "." + si.ServiceName }
		for _, m := range service.Methods {
			req := pp.convertParameter(m.Input, "Req")
			res := pp.convertParameter(m.Output, "Resp")

			methodName := m.GoName
			mi := &generator.MethodInfo{
				PkgInfo:            pi,
				ServiceName:        si.ServiceName,
				RawName:            string(m.Desc.Name()),
				Name:               methodName,
				Args:               []*generator.Parameter{req},
				Resp:               res,
				ArgStructName:      methodName + "Args",
				ResStructName:      methodName + "Result",
				GenArgResultStruct: true,
				ClientStreaming:    m.Desc.IsStreamingClient(),
				ServerStreaming:    m.Desc.IsStreamingServer(),
			}
			si.Methods = append(si.Methods, mi)
			if !si.HasStreaming && (mi.ClientStreaming || mi.ServerStreaming) {
				si.HasStreaming = true
			}
		}
		for _, m := range si.Methods {
			BuildStreaming(m, si.HasStreaming)
		}
		ss = append(ss, si)
	}
	// combine service
	if pp.Config.CombineService && len(file.Services) > 0 {
		var svcs []*generator.ServiceInfo
		var methods []*generator.MethodInfo
		for _, s := range ss {
			svcs = append(svcs, s)
			methods = append(methods, s.AllMethods()...)
		}
		// check method name conflict
		mm := make(map[string]*generator.MethodInfo)
		for _, m := range methods {
			if _, ok := mm[m.Name]; ok {
				log.Warnf("[WARN] combine service method %s in %s conflicts with %s in %s\n",
					m.Name, m.ServiceName, m.Name, mm[m.Name].ServiceName)
				return
			}
			mm[m.Name] = m
		}
		var hasStreaming bool
		for _, m := range methods {
			if m.ClientStreaming || m.ServerStreaming {
				hasStreaming = true
			}
		}
		svcName := pp.getCombineServiceName("CombineService", ss)
		si := &generator.ServiceInfo{
			PkgInfo:         pi,
			ServiceName:     svcName,
			RawServiceName:  svcName,
			CombineServices: svcs,
			Methods:         methods,
			HasStreaming:    hasStreaming,
		}
		si.ServiceTypeName = func() string { return si.ServiceName }
		ss = append(ss, si)
	}
	return
}

// BuildStreaming builds protobuf MethodInfo.Streaming as for Thrift, to simplify codegen
func BuildStreaming(mi *generator.MethodInfo, serviceHasStreaming bool) {
	s := &streaming.Streaming{
		// pb: if one method is streaming, then the service is streaming, making all methods streaming
		IsStreaming: serviceHasStreaming,
	}
	if mi.ClientStreaming && mi.ServerStreaming {
		s.Mode = streaming.StreamingBidirectional
		s.BidirectionalStreaming = true
		s.ClientStreaming = true
		s.ServerStreaming = true
	} else if mi.ClientStreaming && !mi.ServerStreaming {
		s.Mode = streaming.StreamingClientSide
		s.ClientStreaming = true
	} else if !mi.ClientStreaming && mi.ServerStreaming {
		s.Mode = streaming.StreamingServerSide
		s.ServerStreaming = true
	} else if serviceHasStreaming {
		s.Mode = streaming.StreamingUnary // Unary APIs over HTTP2
	}
	mi.Streaming = s
}

func (pp *protocPlugin) getCombineServiceName(name string, svcs []*generator.ServiceInfo) string {
	for _, svc := range svcs {
		if svc.ServiceName == name {
			return pp.getCombineServiceName(name+"_", svcs)
		}
	}
	return name
}

func (pp *protocPlugin) convertParameter(msg *protogen.Message, paramName string) *generator.Parameter {
	importPath := pp.fixImport(msg.GoIdent.GoImportPath.String())
	pkgRefName := goSanitized(path.Base(importPath))
	res := &generator.Parameter{
		Deps: []generator.PkgInfo{
			{
				PkgRefName: pkgRefName,
				ImportPath: importPath,
			},
		},
		Name:    paramName,
		RawName: paramName,
		Type:    "*" + pkgRefName + "." + msg.GoIdent.GoName,
	}
	return res
}

func (pp *protocPlugin) makeInterfaces(gf *protogen.GeneratedFile, file *protogen.File) interface{} {
	var is []interface{}
	for _, service := range file.Services {
		i := struct {
			Name    string
			Methods []interface{}
		}{
			Name: service.GoName,
		}
		for _, m := range service.Methods {
			i.Methods = append(i.Methods, struct {
				Name            string
				ReqType         string
				ResType         string
				ClientStreaming bool
				ServerStreaming bool
			}{
				m.GoName,
				"*" + gf.QualifiedGoIdent(m.Input.GoIdent),
				"*" + gf.QualifiedGoIdent(m.Output.GoIdent),
				m.Desc.IsStreamingClient(),
				m.Desc.IsStreamingServer(),
			})
		}
		is = append(is, i)
	}
	return struct {
		Version    string
		Interfaces []interface{}
	}{pp.Config.Version, is}
}

func (pp *protocPlugin) adjustPath(path string) (ret string) {
	cur, _ := filepath.Abs(".")
	if pp.Config.Use == "" {
		cur = util.JoinPath(cur, generator.KitexGenPath)
	}
	if filepath.IsAbs(path) {
		path, _ = filepath.Rel(cur, path)
		return path
	}
	if pp.ModuleName == "" {
		gopath := util.GetGOPATH()
		path = util.JoinPath(gopath, "src", path)
		path, _ = filepath.Rel(cur, path)
	} else {
		path, _ = filepath.Rel(pp.ModuleName, path)
	}
	return path
}

func (pp *protocPlugin) fixImport(path string) string {
	path = strings.Trim(path, "\"")
	return path
}