github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/pluginmode/thriftgo/convertor.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 thriftgo

import (
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/cloudwego/thriftgo/generator/backend"
	"github.com/cloudwego/thriftgo/generator/golang"
	"github.com/cloudwego/thriftgo/generator/golang/streaming"
	"github.com/cloudwego/thriftgo/parser"
	"github.com/cloudwego/thriftgo/plugin"
	"github.com/cloudwego/thriftgo/semantic"

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

var (
	prelude  = map[string]bool{"client": true, "server": true, "callopt": true, "context": true, "thrift": true, "kitex": true}
	keyWords = []string{"client", "server", "callopt", "context", "thrift", "kitex"}
)

type converter struct {
	Warnings []string
	Utils    *golang.CodeUtils
	Config   generator.Config
	Package  generator.PackageInfo
	Services []*generator.ServiceInfo
	svc2ast  map[*generator.ServiceInfo]*parser.Thrift
}

func (c *converter) init(req *plugin.Request) error {
	if req.Language != "go" {
		return fmt.Errorf("expect language to be 'go'. Encountered '%s'", req.Language)
	}

	// resotre the arguments for kitex
	if err := c.Config.Unpack(req.PluginParameters); err != nil {
		return err
	}

	c.Utils = golang.NewCodeUtils(c.initLogs())
	c.Utils.HandleOptions(req.GeneratorParameters)

	return nil
}

func (c *converter) initLogs() backend.LogFunc {
	lf := backend.LogFunc{
		Info: func(v ...interface{}) {},
		Warn: func(v ...interface{}) {
			c.Warnings = append(c.Warnings, fmt.Sprint(v...))
		},
		MultiWarn: func(warns []string) {
			c.Warnings = append(c.Warnings, warns...)
		},
	}
	if c.Config.Verbose {
		lf.Info = lf.Warn
	}

	internal_log.SetDefaultLogger(internal_log.Logger{
		Println: func(w io.Writer, a ...interface{}) (n int, err error) {
			if w != os.Stdout || c.Config.Verbose {
				c.Warnings = append(c.Warnings, fmt.Sprint(a...))
			}
			return 0, nil
		},
		Printf: func(w io.Writer, format string, a ...interface{}) (n int, err error) {
			if w != os.Stdout || c.Config.Verbose {
				c.Warnings = append(c.Warnings, fmt.Sprintf(format, a...))
			}
			return 0, nil
		},
	})
	return lf
}

func (c *converter) fail(err error) int {
	res := plugin.BuildErrorResponse(err.Error(), c.Warnings...)
	return exit(res)
}

func (c *converter) avoidIncludeConflict(ast *parser.Thrift, ref string) (*parser.Thrift, string) {
	fn := filepath.Base(ast.Filename)
	for _, inc := range ast.Includes {
		if filepath.Base(inc.Path) == fn { // will cause include conflict
			ref = "kitex_faked_idl"
			faked := *ast
			faked.Filename = util.JoinPath(filepath.Dir(faked.Filename), ref+".thrift")
			_, hasNamespace := ast.GetNamespace("go")
			if !hasNamespace {
				faked.Namespaces = append(faked.Namespaces, &parser.Namespace{
					Language: "go",
					Name:     ast.GetNamespaceOrReferenceName("go"),
				})
			}
			return &faked, ref
		}
	}
	return ast, ref
}

// TODO: copy by marshal & unmarshal? to avoid missing fields.
func (c *converter) copyTreeWithRef(ast *parser.Thrift, ref string) *parser.Thrift {
	ast, ref = c.avoidIncludeConflict(ast, ref)

	t := &parser.Thrift{
		Filename: ast.Filename,
		Namespaces: []*parser.Namespace{
			{Language: "*", Name: "fake"},
		},
	}
	t.Includes = append(t.Includes, &parser.Include{Path: ast.Filename, Reference: ast})
	t.Includes = append(t.Includes, ast.Includes...)

	for _, s := range ast.Services {
		ss := &parser.Service{
			Name:    s.Name,
			Extends: s.Extends,
		}
		for _, f := range s.Functions {
			ff := c.copyFunctionWithRef(f, ref)
			ss.Functions = append(ss.Functions, ff)
		}
		t.Services = append(t.Services, ss)
	}
	return t
}

func (c *converter) copyFunctionWithRef(f *parser.Function, ref string) *parser.Function {
	ff := &parser.Function{
		Name:         f.Name,
		Oneway:       f.Oneway,
		Void:         f.Void,
		FunctionType: c.copyTypeWithRef(f.FunctionType, ref),
		Annotations:  c.copyAnnotations(f.Annotations),
	}
	for _, x := range f.Arguments {
		y := *x
		y.Type = c.copyTypeWithRef(x.Type, ref)
		ff.Arguments = append(ff.Arguments, &y)
	}
	for _, x := range f.Throws {
		y := *x
		y.Type = c.copyTypeWithRef(x.Type, ref)
		ff.Throws = append(ff.Throws, &y)
	}
	return ff
}

func (c *converter) copyTypeWithRef(t *parser.Type, ref string) (res *parser.Type) {
	switch t.Name {
	case "void":
		return t
	case "bool", "byte", "i8", "i16", "i32", "i64", "double", "string", "binary":
		return t
	case "map":
		return &parser.Type{
			Name:      t.Name,
			KeyType:   c.copyTypeWithRef(t.KeyType, ref),
			ValueType: c.copyTypeWithRef(t.ValueType, ref),
		}
	case "set", "list":
		return &parser.Type{
			Name:      t.Name,
			ValueType: c.copyTypeWithRef(t.ValueType, ref),
		}
	default:
		if strings.Contains(t.Name, ".") {
			return &parser.Type{
				Name:      t.Name,
				KeyType:   t.KeyType,
				ValueType: t.ValueType,
			}
		}
		return &parser.Type{
			Name: ref + "." + t.Name,
		}
	}
}

func (c *converter) getImports(t *parser.Type) (res []generator.PkgInfo) {
	switch t.Name {
	case "void":
		return nil
	case "bool", "byte", "i8", "i16", "i32", "i64", "double", "string", "binary":
		return nil
	case "map":
		res = append(res, c.getImports(t.KeyType)...)
		fallthrough
	case "set", "list":
		res = append(res, c.getImports(t.ValueType)...)
		return res
	default:
		if ref := t.GetReference(); ref != nil {
			inc := c.Utils.RootScope().Includes().ByIndex(int(ref.GetIndex()))
			res = append(res, generator.PkgInfo{
				PkgRefName: inc.PackageName,
				ImportPath: inc.ImportPath,
			})
		}
		return
	}
}

func (c *converter) fixImportConflicts() {
	for pkg, pth := range c.Package.Dependencies {
		if prelude[pkg] {
			delete(c.Package.Dependencies, pkg)
			c.Package.Dependencies[pkg+"0"] = pth
		}
	}
	var objs []interface{}
	for _, s := range c.Services {
		objs = append(objs, s)
	}

	fix := func(p *generator.PkgInfo) {
		if prelude[p.PkgRefName] {
			p.PkgRefName += "0"
		}
	}
	kw := strings.Join(keyWords, "|")
	re := regexp.MustCompile(`^(\*?)(` + kw + `)\.([^.]+)$`)
	for len(objs) > 0 {
		switch v := objs[0].(type) {
		case *generator.ServiceInfo:
			fix(&v.PkgInfo)
			if v.Base != nil {
				objs = append(objs, v.Base)
			}
			for _, m := range v.Methods {
				objs = append(objs, m)
			}
		case *generator.MethodInfo:
			fix(&v.PkgInfo)
			for _, a := range v.Args {
				objs = append(objs, a)
			}
			if !v.Void {
				objs = append(objs, v.Resp)
			}
			for _, e := range v.Exceptions {
				objs = append(objs, e)
			}
			v.ArgStructName = re.ReplaceAllString(v.ArgStructName, "${1}${2}0.${3}")
			v.ResStructName = re.ReplaceAllString(v.ResStructName, "${1}${2}0.${3}")
		case *generator.Parameter:
			for i := 0; i < len(v.Deps); i++ {
				fix(&v.Deps[i])
			}
			v.Type = re.ReplaceAllString(v.Type, "${1}${2}0.${3}")
		}
		objs = objs[1:]
	}
}

type ast2svc map[string][]*generator.ServiceInfo

func (t ast2svc) findService(ast *parser.Thrift, name string) *generator.ServiceInfo {
	var list []*generator.ServiceInfo
	for filename, l := range t {
		if filename == ast.Filename {
			list = l
			break
		}
	}
	for _, s := range list {
		if s.RawServiceName == name {
			return s
		}
	}
	return nil
}

func (c *converter) convertTypes(req *plugin.Request) error {
	var all ast2svc = make(map[string][]*generator.ServiceInfo)

	c.svc2ast = make(map[*generator.ServiceInfo]*parser.Thrift)
	for ast := range req.AST.DepthFirstSearch() {
		ref, pkg, pth := c.Utils.ParseNamespace(ast)
		// make the current ast as an include to produce correct type references.
		fake := c.copyTreeWithRef(ast, ref)
		fake.Name2Category = nil
		if err := semantic.ResolveSymbols(fake); err != nil {
			return fmt.Errorf("resolve fakse ast '%s': %w", ast.Filename, err)
		}
		used := true
		fake.ForEachInclude(func(v *parser.Include) bool {
			v.Used = &used // mark all includes used to force renaming for conflict IDLs in thriftgo
			return true
		})

		scope, err := golang.BuildScope(c.Utils, fake)
		if err != nil {
			return fmt.Errorf("build scope for fake ast '%s': %w", ast.Filename, err)
		}
		c.Utils.SetRootScope(scope)
		pi := generator.PkgInfo{
			PkgName:    pkg,
			PkgRefName: pkg,
			ImportPath: util.CombineOutputPath(c.Config.PackagePrefix, pth),
		}
		for _, svc := range scope.Services() {
			si, err := c.makeService(pi, svc)
			if err != nil {
				return fmt.Errorf("%s: makeService '%s': %w", ast.Filename, svc.Name, err)
			}
			si.ServiceFilePath = ast.Filename
			all[ast.Filename] = append(all[ast.Filename], si)
			c.svc2ast[si] = ast
		}
		// fill .Base
		for i, svc := range ast.Services {
			if len(svc.Extends) > 0 {
				si := all[ast.Filename][i]
				parts := semantic.SplitType(svc.Extends)
				switch len(parts) {
				case 1:
					si.Base = all.findService(ast, parts[0])
				case 2:
					tmp, found := ast.GetReference(parts[0])
					if !found {
						break
					}
					si.Base = all.findService(tmp, parts[1])
				}
				if len(parts) > 0 && si.Base == nil {
					return fmt.Errorf("base service '%s' %d not found for '%s'", svc.Extends, len(parts), svc.Name)
				}
			}
		}

		c.fixStreamingForExtendedServices(ast, all)

		// combine service
		if ast == req.AST && c.Config.CombineService && len(ast.Services) > 0 {
			var (
				svcs    []*generator.ServiceInfo
				methods []*generator.MethodInfo
			)
			hasStreaming := false
			for _, s := range all[ast.Filename] {
				svcs = append(svcs, s)
				hasStreaming = hasStreaming || s.HasStreaming
				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 {
					return fmt.Errorf("combine service method %s in %s conflicts with %s in %s", m.Name, m.ServiceName, m.Name, mm[m.Name].ServiceName)
				}
				mm[m.Name] = m
			}
			svcName := c.getCombineServiceName("CombineService", all[ast.Filename])
			si := &generator.ServiceInfo{
				PkgInfo:         pi,
				ServiceName:     svcName,
				RawServiceName:  svcName,
				CombineServices: svcs,
				Methods:         methods,
				ServiceFilePath: ast.Filename,
				HasStreaming:    hasStreaming,
			}

			if c.IsHessian2() {
				si.Protocol = transport.HESSIAN2.String()
			}

			si.HandlerReturnKeepResp = c.Config.HandlerReturnKeepResp
			si.UseThriftReflection = c.Utils.Features().WithReflection
			si.ServiceTypeName = func() string { return si.ServiceName }
			all[ast.Filename] = append(all[ast.Filename], si)
			c.svc2ast[si] = ast
		}

		c.Services = append(c.Services, all[ast.Filename]...)
	}
	return nil
}

func (c *converter) fixStreamingForExtendedServices(ast *parser.Thrift, all ast2svc) {
	for i, svc := range ast.Services {
		if svc.Extends == "" {
			continue
		}
		si := all[ast.Filename][i]
		if si.Base != nil {
			si.FixHasStreamingForExtendedService()
		}
	}
}

func (c *converter) makeService(pkg generator.PkgInfo, svc *golang.Service) (*generator.ServiceInfo, error) {
	si := &generator.ServiceInfo{
		PkgInfo:        pkg,
		ServiceName:    svc.GoName().String(),
		RawServiceName: svc.Name,
	}
	si.ServiceTypeName = func() string { return si.PkgRefName + "." + si.ServiceName }

	for _, f := range svc.Functions() {
		if strings.HasPrefix(f.Name, "_") {
			continue
		}
		mi, err := c.makeMethod(si, f)
		if err != nil {
			return nil, err
		}
		si.Methods = append(si.Methods, mi)
	}

	if c.IsHessian2() {
		si.Protocol = transport.HESSIAN2.String()
	}
	si.HandlerReturnKeepResp = c.Config.HandlerReturnKeepResp
	si.UseThriftReflection = c.Utils.Features().WithReflection
	return si, nil
}

func (c *converter) makeMethod(si *generator.ServiceInfo, f *golang.Function) (*generator.MethodInfo, error) {
	st, err := streaming.ParseStreaming(f.Function)
	if err != nil {
		return nil, err
	}
	mi := &generator.MethodInfo{
		PkgInfo:            si.PkgInfo,
		ServiceName:        si.ServiceName,
		Name:               f.GoName().String(),
		RawName:            f.Name,
		Oneway:             f.Oneway,
		Void:               f.Void,
		ArgStructName:      f.ArgType().GoName().String(),
		GenArgResultStruct: false,
		Streaming:          st,
		ClientStreaming:    st.ClientStreaming,
		ServerStreaming:    st.ServerStreaming,
		ArgsLength:         len(f.Arguments()),
	}
	if st.IsStreaming {
		si.HasStreaming = true
	}

	if !f.Oneway {
		mi.ResStructName = f.ResType().GoName().String()
	}
	if !f.Void {
		typeName := f.ResponseGoTypeName().String()
		mi.Resp = &generator.Parameter{
			Deps: c.getImports(f.FunctionType),
			Type: typeName,
		}
		mi.IsResponseNeedRedirect = "*"+typeName == f.ResType().Fields()[0].GoTypeName().String()
	}

	for _, a := range f.Arguments() {
		arg := &generator.Parameter{
			Deps:    c.getImports(a.Type),
			Name:    f.ArgType().Field(a.Name).GoName().String(),
			RawName: a.GoName().String(),
			Type:    a.GoTypeName().String(),
		}
		mi.Args = append(mi.Args, arg)
	}
	for _, t := range f.Throws() {
		ex := &generator.Parameter{
			Deps:    c.getImports(t.Type),
			Name:    f.ResType().Field(t.Name).GoName().String(),
			RawName: t.GoName().String(),
			Type:    t.GoTypeName().String(),
		}
		mi.Exceptions = append(mi.Exceptions, ex)
	}
	return mi, nil
}

func (c *converter) persist(res *plugin.Response) error {
	for _, c := range res.Contents {
		full := c.GetName()
		content := []byte(c.Content)
		if filepath.Ext(full) == ".go" {
			if formatted, err := format.Source([]byte(c.Content)); err != nil {
				internal_log.Warn(fmt.Sprintf("Failed to format %s: %s", full, err.Error()))
			} else {
				content = formatted
			}
		}

		internal_log.Info("Write", full)
		path := filepath.Dir(full)
		if err := os.MkdirAll(path, 0o755); err != nil && !os.IsExist(err) {
			return fmt.Errorf("failed to create path '%s': %w", path, err)
		}
		if err := ioutil.WriteFile(full, content, 0o644); err != nil {
			return fmt.Errorf("failed to write file '%s': %w", full, err)
		}
	}
	return nil
}

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

func (c *converter) IsHessian2() bool {
	return strings.EqualFold(c.Config.Protocol, transport.HESSIAN2.String())
}

func (c *converter) copyAnnotations(annotations parser.Annotations) parser.Annotations {
	copied := make(parser.Annotations, 0, len(annotations))
	for _, annotation := range annotations {
		values := make([]string, 0, len(annotation.Values))
		values = append(values, annotation.Values...)
		copied = append(copied, &parser.Annotation{
			Key:    annotation.Key,
			Values: values,
		})
	}
	return copied
}