github.com/cloudwego/kitex@v0.9.0/pkg/generic/thrift/parse.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 thrift

import (
	"context"
	"errors"
	"fmt"
	"runtime/debug"

	"github.com/cloudwego/thriftgo/parser"
	"github.com/cloudwego/thriftgo/semantic"

	"github.com/cloudwego/kitex/pkg/generic/descriptor"
	"github.com/cloudwego/kitex/pkg/gofunc"
	"github.com/cloudwego/kitex/pkg/klog"
)

const (
	initRecursionDepth = 0
)

// ParseMode .
type ParseMode int

const (
	// LastServiceOnly forces the parser to parse only the last service definition.
	LastServiceOnly ParseMode = iota

	// FirstServiceOnly forces the parser to parse only the first service definition.
	FirstServiceOnly

	// CombineServices forces the parser to combine methods of all service definitions.
	// Note that method names of the service definitions can not be duplicate.
	CombineServices
)

var defaultParseMode = LastServiceOnly

// DefaultParseMode returns the default parse mode.
func DefaultParseMode() ParseMode {
	return defaultParseMode
}

// SetDefaultParseMode sets the default parse mode.
func SetDefaultParseMode(m ParseMode) {
	defaultParseMode = m
}

// Parse descriptor from parser.Thrift
func Parse(tree *parser.Thrift, mode ParseMode) (*descriptor.ServiceDescriptor, error) {
	if len(tree.Services) == 0 {
		return nil, errors.New("empty serverce from idls")
	}
	if err := semantic.ResolveSymbols(tree); err != nil {
		return nil, err
	}

	sDsc := &descriptor.ServiceDescriptor{
		Functions: map[string]*descriptor.FunctionDescriptor{},
		Router:    descriptor.NewRouter(),
	}

	// support one service
	svcs := tree.Services
	switch mode {
	case LastServiceOnly:
		svcs = svcs[len(svcs)-1:]
		sDsc.Name = svcs[len(svcs)-1].Name
	case FirstServiceOnly:
		svcs = svcs[:1]
		sDsc.Name = svcs[0].Name
	case CombineServices:
		sDsc.Name = "CombinedServices"
	}

	visitedSvcs := make(map[*parser.Service]bool, len(tree.Services))
	for _, svc := range svcs {
		for p := range getAllSvcs(svc, tree, visitedSvcs) {
			svc := p.data.(*parser.Service)
			structsCache := map[string]*descriptor.TypeDescriptor{}
			for _, fn := range svc.Functions {
				if err := addFunction(fn, p.tree, sDsc, structsCache); err != nil {
					return nil, err
				}
			}
		}
	}
	return sDsc, nil
}

type pair struct {
	tree *parser.Thrift
	data interface{}
}

func getAllSvcs(svc *parser.Service, tree *parser.Thrift, visitedSvcs map[*parser.Service]bool) chan *pair {
	svcs := make(chan *pair)
	addSvc := func(tree *parser.Thrift, svc *parser.Service) {
		if exist := visitedSvcs[svc]; !exist {
			svcs <- &pair{tree: tree, data: svc}
			visitedSvcs[svc] = true
		}
	}
	gofunc.GoFunc(context.Background(), func() {
		addSvc(tree, svc)
		for base := svc.Extends; base != ""; base = svc.Extends {
			ref := svc.GetReference()
			if ref != nil {
				idx := ref.GetIndex()
				base = ref.GetName()
				tree = tree.Includes[idx].Reference
			}
			svc, _ = tree.GetService(base)
			addSvc(tree, svc)
		}
		close(svcs)
	})
	return svcs
}

func addFunction(fn *parser.Function, tree *parser.Thrift, sDsc *descriptor.ServiceDescriptor, structsCache map[string]*descriptor.TypeDescriptor) (err error) {
	if sDsc.Functions[fn.Name] != nil {
		return fmt.Errorf("duplicate method name: %s", fn.Name)
	}
	if len(fn.Arguments) == 0 {
		return fmt.Errorf("empty arguments in function: %s", fn.Name)
	}
	// only support single argument
	field := fn.Arguments[0]
	req := &descriptor.TypeDescriptor{
		Type: descriptor.STRUCT,
		Struct: &descriptor.StructDescriptor{
			FieldsByID:   map[int32]*descriptor.FieldDescriptor{},
			FieldsByName: map[string]*descriptor.FieldDescriptor{},
		},
	}
	var reqType *descriptor.TypeDescriptor
	reqType, err = parseType(field.Type, tree, structsCache, initRecursionDepth)
	if err != nil {
		return err
	}
	hasRequestBase := false
	if reqType.Type == descriptor.STRUCT {
		for _, f := range reqType.Struct.FieldsByName {
			if f.Type.IsRequestBase {
				hasRequestBase = true
				break
			}
		}
	}
	reqField := &descriptor.FieldDescriptor{
		Name: field.Name,
		ID:   field.ID,
		Type: reqType,
	}
	req.Struct.FieldsByID[field.ID] = reqField
	req.Struct.FieldsByName[field.Name] = reqField
	// parse response filed
	resp := &descriptor.TypeDescriptor{
		Type: descriptor.STRUCT,
		Struct: &descriptor.StructDescriptor{
			FieldsByID:   map[int32]*descriptor.FieldDescriptor{},
			FieldsByName: map[string]*descriptor.FieldDescriptor{},
		},
	}
	var respType *descriptor.TypeDescriptor
	respType, err = parseType(fn.FunctionType, tree, structsCache, initRecursionDepth)
	if err != nil {
		return err
	}
	respField := &descriptor.FieldDescriptor{
		Type: respType,
	}
	// response has no name or id
	resp.Struct.FieldsByID[0] = respField
	resp.Struct.FieldsByName[""] = respField

	if len(fn.Throws) > 0 {
		// only support single exception
		field := fn.Throws[0]
		var exceptionType *descriptor.TypeDescriptor
		exceptionType, err = parseType(field.Type, tree, structsCache, initRecursionDepth)
		if err != nil {
			return err
		}
		exceptionField := &descriptor.FieldDescriptor{
			Name:        field.Name,
			ID:          field.ID,
			IsException: true,
			Type:        exceptionType,
		}
		resp.Struct.FieldsByID[field.ID] = exceptionField
		resp.Struct.FieldsByName[field.Name] = exceptionField
	}
	fnDsc := &descriptor.FunctionDescriptor{
		Name:           fn.Name,
		Oneway:         fn.Oneway,
		Request:        req,
		Response:       resp,
		HasRequestBase: hasRequestBase,
	}
	defer func() {
		if ret := recover(); ret != nil {
			klog.Errorf("KITEX: router handle failed, err=%v\nstack=%s", ret, string(debug.Stack()))
			err = fmt.Errorf("router handle failed, err=%v", ret)
		}
	}()
	for _, ann := range fn.Annotations {
		for _, v := range ann.GetValues() {
			if handle, ok := descriptor.FindAnnotation(ann.GetKey(), v); ok {
				if nr, ok := handle.(descriptor.NewRoute); ok {
					sDsc.Router.Handle(nr(v, fnDsc))
					break
				}
			}
		}
	}
	sDsc.Functions[fn.Name] = fnDsc
	return nil
}

// reuse builtin types
var builtinTypes = map[string]*descriptor.TypeDescriptor{
	"void":   {Name: "void", Type: descriptor.VOID, Struct: new(descriptor.StructDescriptor)},
	"bool":   {Name: "bool", Type: descriptor.BOOL},
	"byte":   {Name: "byte", Type: descriptor.BYTE},
	"i8":     {Name: "i8", Type: descriptor.I08},
	"i16":    {Name: "i16", Type: descriptor.I16},
	"i32":    {Name: "i32", Type: descriptor.I32},
	"i64":    {Name: "i64", Type: descriptor.I64},
	"double": {Name: "double", Type: descriptor.DOUBLE},
	"string": {Name: "string", Type: descriptor.STRING},
	"binary": {Name: "binary", Type: descriptor.STRING},
}

// arg cache:
// only support self reference on the same file
// cross file self reference complicate matters
func parseType(t *parser.Type, tree *parser.Thrift, cache map[string]*descriptor.TypeDescriptor, recursionDepth int) (*descriptor.TypeDescriptor, error) {
	if ty, ok := builtinTypes[t.Name]; ok {
		return ty, nil
	}

	nextRecursionDepth := recursionDepth + 1

	var err error
	switch t.Name {
	case "list":
		ty := &descriptor.TypeDescriptor{Name: t.Name}
		ty.Type = descriptor.LIST
		ty.Elem, err = parseType(t.ValueType, tree, cache, nextRecursionDepth)
		return ty, err
	case "set":
		ty := &descriptor.TypeDescriptor{Name: t.Name}
		ty.Type = descriptor.SET
		ty.Elem, err = parseType(t.ValueType, tree, cache, nextRecursionDepth)
		return ty, err
	case "map":
		ty := &descriptor.TypeDescriptor{Name: t.Name}
		ty.Type = descriptor.MAP
		if ty.Key, err = parseType(t.KeyType, tree, cache, nextRecursionDepth); err != nil {
			return nil, err
		}
		ty.Elem, err = parseType(t.ValueType, tree, cache, nextRecursionDepth)
		return ty, err
	default:
		// check the cache
		if ty, ok := cache[t.Name]; ok {
			return ty, nil
		}
		typePkg, typeName := splitType(t.Name)
		if typePkg != "" {
			ref, ok := tree.GetReference(typePkg)
			if !ok {
				return nil, fmt.Errorf("miss reference: %s", typePkg)
			}
			tree = ref
			// cross file reference need empty cache
			cache = map[string]*descriptor.TypeDescriptor{}
		}
		if typDef, ok := tree.GetTypedef(typeName); ok {
			return parseType(typDef.Type, tree, cache, nextRecursionDepth)
		}
		if _, ok := tree.GetEnum(typeName); ok {
			return builtinTypes["i32"], nil
		}
		var st *parser.StructLike
		var ok bool
		st, ok = tree.GetUnion(typeName)
		if !ok {
			st, ok = tree.GetStruct(typeName)
		}
		if !ok {
			st, ok = tree.GetException(typeName)
		}
		if !ok {
			return nil, fmt.Errorf("missing type: %s", typeName)
		}
		ty := &descriptor.TypeDescriptor{
			Name: t.Name,
			Type: descriptor.STRUCT,
			Struct: &descriptor.StructDescriptor{
				Name:           typeName,
				FieldsByID:     map[int32]*descriptor.FieldDescriptor{},
				FieldsByName:   map[string]*descriptor.FieldDescriptor{},
				RequiredFields: map[int32]*descriptor.FieldDescriptor{},
				DefaultFields:  map[string]*descriptor.FieldDescriptor{},
			},
			// the first depth of base.Base is Request Base
			IsRequestBase: t.Name == "base.Base" && recursionDepth == 1,
		}
		// cannot cache the request base
		if !ty.IsRequestBase {
			cache[t.Name] = ty
		}
		for _, field := range st.Fields {
			_f := &descriptor.FieldDescriptor{
				ID:       field.ID,
				Name:     field.Name,
				Required: field.Requiredness == parser.FieldType_Required,
				Optional: field.Requiredness == parser.FieldType_Optional,
			}

			for _, ann := range field.Annotations {
				for _, v := range ann.GetValues() {
					if handle, ok := descriptor.FindAnnotation(ann.GetKey(), v); ok {
						switch h := handle.(type) {
						case descriptor.NewHTTPMapping:
							_f.HTTPMapping = h(v)
						case descriptor.NewValueMapping:
							_f.ValueMapping = h(v)
						case descriptor.NewFieldMapping:
							// execute at compile time
							h(v).Handle(_f)
						case nil:
							// none annotation
						default:
							// not supported annotation type
							return nil, fmt.Errorf("not supported handle type: %T", handle)
						}
					}
				}
			}
			if _f.HTTPMapping == nil {
				_f.HTTPMapping = descriptor.DefaultNewMapping(_f.FieldName())
			}
			if _f.Type, err = parseType(field.Type, tree, cache, nextRecursionDepth); err != nil {
				return nil, err
			}
			// set default value
			if field.IsSetDefault() {
				_f.DefaultValue, err = parse(tree, field.Name, field.Type, field.Default)
				if err != nil {
					return nil, err
				}
				ty.Struct.DefaultFields[_f.FieldName()] = _f
			}
			ty.Struct.FieldsByID[field.ID] = _f
			ty.Struct.FieldsByName[_f.FieldName()] = _f
			if _f.Required {
				ty.Struct.RequiredFields[field.ID] = _f
			}
		}
		return ty, nil
	}
}

func parse(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (interface{}, error) {
	switch t.Category {
	case parser.Category_Bool:
		return onBool(tree, name, t, v)
	case parser.Category_Byte:
		val, err := onInt(tree, name, t, v)
		if err != nil {
			return nil, err
		}
		return byte(val), nil
	case parser.Category_I16:
		val, err := onInt(tree, name, t, v)
		if err != nil {
			return nil, err
		}
		return int16(val), nil
	case parser.Category_I32:
		val, err := onInt(tree, name, t, v)
		if err != nil {
			return nil, err
		}
		return int32(val), nil
	case parser.Category_I64:
		return onInt(tree, name, t, v)
	case parser.Category_Double:
		return onDouble(tree, name, t, v)
	case parser.Category_String:
		return onStr(tree, name, t, v)
	case parser.Category_Binary:
		str, err := onStr(tree, name, t, v)
		if err != nil {
			return nil, err
		}
		return []byte(str), nil
	case parser.Category_Enum:
		return onEnum(tree, name, t, v)
	case parser.Category_Set, parser.Category_List:
		return onSetOrList(tree, name, t, v)
	case parser.Category_Map:
		return onMap(tree, name, t, v)
	case parser.Category_Struct, parser.Category_Union, parser.Category_Exception:
		return onStructLike(tree, name, t, v)
	}
	return nil, fmt.Errorf("type error: '%s' was declared as type %s but got value[%v] of category[%s]", name, t, v, t.Category)
}

func onBool(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (bool, error) {
	switch v.Type {
	case parser.ConstType_ConstInt:
		val := v.TypedValue.GetInt()
		return val > 0, nil
	case parser.ConstType_ConstDouble:
		val := v.TypedValue.GetDouble()
		return val > 0, nil
	case parser.ConstType_ConstIdentifier:
		s := v.TypedValue.GetIdentifier()
		if s == "true" {
			return true, nil
		}
		if s == "false" {
			return false, nil
		}
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return false, err
		}
		switch rv := val.(type) {
		case bool:
			return rv, nil
		}
		return false, fmt.Errorf("%s[%T] not match bool type", v.Extra.Name, val)
	}
	return false, fmt.Errorf("type error: '%s' was declared as type %s", name, t)
}

func onInt(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (int64, error) {
	switch v.Type {
	case parser.ConstType_ConstInt:
		val := v.TypedValue.GetInt()
		return val, nil
	case parser.ConstType_ConstIdentifier:
		s := v.TypedValue.GetIdentifier()
		if s == "true" {
			return 1, nil
		}
		if s == "false" {
			return 0, nil
		}
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return 0, err
		}
		switch rv := val.(type) {
		case byte:
			return int64(rv), nil
		case int16:
			return int64(rv), nil
		case int32:
			return int64(rv), nil
		case int64:
			return rv, nil
		}
		return 0, fmt.Errorf("%s[%T] not match int64 type", v.Extra.Name, val)
	}
	return 0, fmt.Errorf("type error: '%s' was declared as type %s", name, t)
}

func onDouble(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (float64, error) {
	switch v.Type {
	case parser.ConstType_ConstInt:
		val := v.TypedValue.GetInt()
		return float64(val), nil
	case parser.ConstType_ConstDouble:
		val := v.TypedValue.GetDouble()
		return val, nil
	case parser.ConstType_ConstIdentifier:
		s := v.TypedValue.GetIdentifier()
		if s == "true" {
			return 1.0, nil
		}
		if s == "false" {
			return 0.0, nil
		}
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return 0, err
		}
		switch rv := val.(type) {
		case float64:
			return rv, nil
		case int64:
			return float64(rv), nil
		}
		return 0, fmt.Errorf("%s[%T] not match float64 type", v.Extra.Name, val)
	}
	return 0, fmt.Errorf("type error: '%s' was declared as type %s", name, t)
}

func onStr(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (string, error) {
	switch v.Type {
	case parser.ConstType_ConstLiteral:
		return v.TypedValue.GetLiteral(), nil
	case parser.ConstType_ConstIdentifier:
		s := v.TypedValue.GetIdentifier()
		if s == "true" || s == "false" {
			break
		}
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return "", err
		}
		switch rv := val.(type) {
		case string:
			return rv, nil
		}
		return "", fmt.Errorf("%s[%T] not match string type", v.Extra.Name, val)
	}
	return "", fmt.Errorf("type error: '%s' was declared as type %s", name, t)
}

func onEnum(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (int64, error) {
	switch v.Type {
	case parser.ConstType_ConstInt:
		return v.TypedValue.GetInt(), nil
	case parser.ConstType_ConstIdentifier:
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return 0, err
		}
		switch rv := val.(type) {
		case int64:
			return rv, nil
		}
		return 0, fmt.Errorf("%s[%T] not match int64 type", v.Extra.Name, val)
	}
	return 0, fmt.Errorf("expect const value for %q is a int or enum, got %+v", name, v)
}

func onSetOrList(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (res []interface{}, err error) {
	switch v.Type {
	case parser.ConstType_ConstList:
		elemName := "element of " + name
		for _, elem := range v.TypedValue.GetList() {
			val, err := parse(tree, elemName, t.ValueType, elem)
			if err != nil {
				return nil, err
			}
			res = append(res, val)
		}
		return res, nil
	case parser.ConstType_ConstIdentifier:
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return nil, err
		}
		switch rv := val.(type) {
		case []interface{}:
			return rv, nil
		}
		return nil, fmt.Errorf("%s[%T] not match []interface{} type", v.Extra.Name, val)
	}
	// fault tolerance
	return
}

func onMap(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (res map[interface{}]interface{}, err error) {
	res = make(map[interface{}]interface{})
	switch v.Type {
	case parser.ConstType_ConstMap:
		for _, mcv := range v.TypedValue.Map {
			keyName := "key of " + name
			key, err := parse(tree, keyName, bin2str(t.KeyType), mcv.Key)
			if err != nil {
				return nil, err
			}
			valName := "value of " + name
			val, err := parse(tree, valName, t.ValueType, mcv.Value)
			if err != nil {
				return nil, err
			}
			res[key] = val
		}
		return res, nil

	case parser.ConstType_ConstIdentifier:
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return nil, err
		}
		switch rv := val.(type) {
		case map[interface{}]interface{}:
			return rv, nil
		}
		return nil, fmt.Errorf("%s[%T] not match map[interface{}]interface{} type", v.Extra.Name, val)
	}
	// fault tolerance
	return
}

func onStructLike(tree *parser.Thrift, name string, t *parser.Type, v *parser.ConstValue) (res map[string]interface{}, err error) {
	if v.Type == parser.ConstType_ConstIdentifier {
		val, err := getIDValue(tree, name, t, v.Extra)
		if err != nil {
			return nil, err
		}
		switch rv := val.(type) {
		case map[string]interface{}:
			return rv, nil
		}
		return nil, fmt.Errorf("%s[%T] not match map[string]interface{} type", v.Extra.Name, val)
	}

	if v.Type != parser.ConstType_ConstMap {
		// constant value of a struct-like must be a map literal
		return nil, fmt.Errorf("type error: '%s' was declared as type %s", name, t)
	}

	// get the target struct-like with typedef dereferenced
	file, st, err := getStructLike(tree, t)
	if err != nil {
		return nil, err
	}

	res = make(map[string]interface{})
	for _, mcv := range v.TypedValue.Map {
		if mcv.Key.Type != parser.ConstType_ConstLiteral {
			return nil, fmt.Errorf("expect literals as keys in default value of struct type '%s', got '%s'", name, mcv.Key.Type)
		}
		n := mcv.Key.TypedValue.GetLiteral()

		f, ok := st.GetField(n)
		if !ok {
			return nil, fmt.Errorf("field %q not found in %q (%q): %v",
				n, st.Name, file.Filename, v,
			)
		}
		val, err := parse(file, st.Name+"."+f.Name, f.Type, mcv.Value)
		if err != nil {
			return nil, err
		}
		res[n] = val
	}

	return res, nil
}

func getStructLike(tree *parser.Thrift, t *parser.Type) (ast *parser.Thrift, s *parser.StructLike, err error) {
	var x *parser.Type
	ast, x, err = semantic.Deref(tree, t)
	if err != nil {
		err = fmt.Errorf("expect %q a typedef or struct-like in %q: %w",
			t.Name, tree.Filename, err)
		return nil, nil, err
	}
	for _, y := range ast.GetStructLikes() {
		if x.Name == y.Name {
			s = y
		}
	}
	if s == nil {
		err = fmt.Errorf("expect %q a struct-like in %q: not found: %v",
			x.Name, ast.Filename, x == t)
		return nil, nil, err
	}
	return
}

func getIDValue(tree *parser.Thrift, name string, t *parser.Type, extra *parser.ConstValueExtra) (interface{}, error) {
	if extra.Index == -1 {
		if extra.IsEnum {
			enum, ok := tree.GetEnum(extra.Sel)
			if !ok {
				return nil, fmt.Errorf("no matching enum of %s", extra.Sel)
			}
			for _, v := range enum.Values {
				if v.Name == extra.Name {
					return v.Value, nil
				}
			}
		} else {
			if con, ok := tree.GetConstant(extra.Name); ok {
				return parse(tree, name, con.Type, con.Value)
			}
		}
	} else {
		tree = tree.Includes[extra.Index].Reference
		extra = &parser.ConstValueExtra{
			Index:  -1,
			IsEnum: extra.IsEnum,
			Name:   extra.Name,
			Sel:    extra.Sel,
		}
		return getIDValue(tree, name, t, extra)
	}
	return nil, fmt.Errorf("undefined value %s", extra.Name)
}

func bin2str(t *parser.Type) *parser.Type {
	if t.Category == parser.Category_Binary {
		r := *t
		r.Category = parser.Category_String
		return &r
	}
	return t
}