Merge pull request #1 from WineChord/cherry_picks

cherry pick picks cherries

bindata/compress/tar_test.go

@@ -66,13 +66,13 @@ func TestTar_Install(t *testing.T) {
 	require.Nil(t, err)
 
 	// compare file list
-	srcFileSet := map[string]struct{}{}
+	srcFileSet := make(map[string]struct{})
 	filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
 		path = strings.TrimPrefix(path, src)
 		srcFileSet[path] = struct{}{}
 		return nil
 	})
-	dstFileSet := map[string]struct{}{}
+	dstFileSet := make(map[string]struct{})
 	filepath.Walk(dst, func(path string, info os.FileInfo, err error) error {
 		path = strings.TrimPrefix(path, dst)
 		dstFileSet[path] = struct{}{}

cmd/apidocs/apidocs.go

@@ -50,6 +50,9 @@ The field information of input and output also includes the leading and trailing
 
 	// Flag for alias.
 	apidocsCmd.Flags().Bool("alias", false, "Use alias mode for rpcname")
+	// Preserve original rpcname.
+	apidocsCmd.Flags().BoolP("keep-orig-rpcname", "k", true,
+		"Preserve the original rpcname (if --alias=true), set it to false if you only want alias names.")
 
 	// The rules of documents order.
 	apidocsCmd.Flags().Bool("order-by-pbname", false,
@@ -120,7 +123,10 @@ func loadAPIDocsOptions(flagSet *pflag.FlagSet) (*params.Option, error) {
 		return nil, err
 	}
 	option.Protodirs, _ = flagSet.GetStringArray("protodir")
+	// Always append the current working directory.
+	option.Protodirs = append(option.Protodirs, ".")
 	option.AliasOn, _ = flagSet.GetBool("alias")
+	option.KeepOrigRPCName, _ = flagSet.GetBool("keep-orig-rpcname")
 
 	// Check if the pb file is valid.
 	if len(option.Protofile) == 0 {

cmd/create/cmdflags.go

@@ -30,6 +30,16 @@ func AddCreateFlags(createCmd *cobra.Command) {
 		"Whether to generate stub code for dependencies, only effective when --rpconly=true, defaults to false")
 	createCmd.Flags().Bool("nogomod", false,
 		"Do not generate go.mod file in the stub code, only effective when --rpconly=true, defaults to false")
+	createCmd.Flags().Bool("secvenabled", true,
+		"Enable generation of validate.go file using protoc-gen-secv, defaults to true")
+	createCmd.Flags().String("kvfile", "",
+		"Provide a json file path to unmarshal into key-value pairs (KVs) for usage in template files")
+	createCmd.Flags().String("kvrawjson", "",
+		"Provide raw json content to unmarshal into key-value pairs (KVs) for usage in template files")
+	createCmd.Flags().String("app", "",
+		"Provide custom app name to use in stub code")
+	createCmd.Flags().String("server", "",
+		"Provide custom server name to use in stub code")
 
 	// Add functionality similar to "protoc --go_out=. testdesc.proto --descriptor_set_in=testdesc.pb".
 	createCmd.Flags().StringP("descriptor_set_in", "", "",
@@ -70,6 +80,7 @@ func addOutputFlags(createCmd *cobra.Command) {
 
 	// Enable rpcname aliases.
 	createCmd.Flags().Bool("alias", false, "Use rpcname aliases")
+	createCmd.Flags().Bool("alias-as-client-rpcname", true, "Use alias name as client rpcname in stub code")
 
 	// Customize pb message struct tag.
 	createCmd.Flags().Bool("gotag", true, "Generate custom pb struct tag")

cmd/create/create.go

@@ -51,11 +51,25 @@ func CMD() *cobra.Command {
 // New initializes a create command.
 func New(preRunHook func() error) *Create {
 	return &Create{
-		options:    &params.Option{},
+		options: &params.Option{
+			Envs: getEnvMap(),
+		},
 		preRunHook: preRunHook,
 	}
 }
 
+func getEnvMap() map[string]string {
+	envs := os.Environ()
+	m := make(map[string]string, len(envs))
+	for _, kv := range envs {
+		vals := strings.Split(kv, "=")
+		if len(vals) == 2 {
+			m[vals[0]] = vals[1]
+		}
+	}
+	return m
+}
+
 // PreRunE provides *cobra.Command.PreRunE.
 func (c *Create) PreRunE(cmd *cobra.Command, args []string) error {
 	if err := c.loadOptions(cmd.Flags()); err != nil {
@@ -72,6 +86,9 @@ func (c *Create) PreRunE(cmd *cobra.Command, args []string) error {
 	var err error
 	opts := []parser.Option{
 		parser.WithAliasOn(c.options.AliasOn),
+		parser.WithAPPName(c.options.CustomAPPName),
+		parser.WithServerName(c.options.CustomServerName),
+		parser.WithAliasAsClientRPCName(c.options.AliasAsClientRPCName),
 		parser.WithLanguage(c.options.Language),
 		parser.WithRPCOnly(c.options.RPCOnly),
 		parser.WithMultiVersion(c.options.MultiVersion),

cmd/create/generate.go

@@ -123,10 +123,14 @@ func runProtoc(fd *FD, pbOutDir string, option *params.Option) ([]string, error)
 		return nil, fmt.Errorf("run protoc --$lang_out err: %v", err)
 	}
 
-	// run protoc-gen-secv
-	opts = append(opts, pb.WithSecvEnabled(true))
-	if err := pb.Protoc(option.Protodirs, option.Protofile, option.Language, pbOutDir, opts...); err != nil {
-		return nil, fmt.Errorf("run protoc-gen-secv err: %v", err)
+	// Run protoc-gen-secv.
+	if support, ok := config.SupportValidate[option.Language]; ok && support && option.SecvEnabled {
+		if err := pb.Protoc(
+			option.Protodirs, option.Protofile, option.Language, pbOutDir,
+			append(opts, pb.WithSecvEnabled(true))...,
+		); err != nil {
+			return nil, fmt.Errorf("run protoc-gen-secv for %s err: %w", option.Protofile, err)
+		}
 	}
 	log.Debug("pbOutDir = %s", pbOutDir)
 	// collect the generated files
@@ -330,10 +334,14 @@ func (g *genDependencyRPCStubParam) genDependencyRPCStubPB() error {
 		return fmt.Errorf("GenerateFiles: %v", err)
 	}
 
-	// run protoc-gen-secv
-	opts = append(opts, pb.WithSecvEnabled(true))
-	if err = pb.Protoc(searchPath, g.fname, g.option.Language, g.outputDir, opts...); err != nil {
-		return fmt.Errorf("GenerateFiles: %v", err)
+	// Run protoc-gen-secv.
+	if support, ok := config.SupportValidate[g.option.Language]; ok && support && g.option.SecvEnabled {
+		if err := pb.Protoc(
+			searchPath, g.fname, g.option.Language, g.outputDir,
+			append(opts, pb.WithSecvEnabled(true))...,
+		); err != nil {
+			return fmt.Errorf("generate validation file for %s err: %w", g.fname, err)
+		}
 	}
 	if g.option.DescriptorSetIn != "" {
 		return nil // skip copy if descriptor_set_in is passed.

cmd/create/options.go

@@ -1,7 +1,9 @@
 package create
 
 import (
+	"encoding/json"
 	"fmt"
+	"os"
 	"path/filepath"
 
 	"github.com/pkg/errors"
@@ -140,7 +142,7 @@ func (c *Create) fixProtoDirs() error {
 
 	p, err := paths.Locate(pb.ProtoTRPC)
 	if err != nil {
-		return fmt.Errorf("paths local trpc failed err: %w", err)
+		return fmt.Errorf("paths locate %s failed err: %w", pb.ProtoTRPC, err)
 	}
 
 	c.options.Protodirs = fs.UniqFilePath(append(append(c.options.Protodirs, p),
@@ -251,6 +253,10 @@ func (c *Create) parseInputOptions(flags *pflag.FlagSet) error {
 	if err != nil {
 		return fmt.Errorf("flags parse alias bool err: %w", err)
 	}
+	c.options.AliasAsClientRPCName, err = flags.GetBool("alias-as-client-rpcname")
+	if err != nil {
+		return fmt.Errorf("flags parse alias-as-client-rpcname bool err: %w", err)
+	}
 	c.options.GoMod, err = flags.GetString("mod")
 	if err != nil {
 		return fmt.Errorf("flags parse mod string err: %w", err)
@@ -263,6 +269,14 @@ func (c *Create) parseInputOptions(flags *pflag.FlagSet) error {
 	if err != nil {
 		return fmt.Errorf("flags parse trpcgoversion string err: %w", err)
 	}
+	c.options.CustomAPPName, err = flags.GetString("app")
+	if err != nil {
+		return fmt.Errorf("flags parse app string err: %w", err)
+	}
+	c.options.CustomServerName, err = flags.GetString("server")
+	if err != nil {
+		return fmt.Errorf("flags parse server string err: %w", err)
+	}
 	c.options.DescriptorSetIn, err = flags.GetString("descriptor_set_in")
 	if err != nil {
 		return fmt.Errorf("flags parse descriptor_set_in string err: %w", err)
@@ -291,6 +305,32 @@ func (c *Create) parseOutputOptions(flags *pflag.FlagSet) error {
 		return fmt.Errorf("flags parse nogomod bool err: %w", err)
 	}
 	c.options.KeepOrigRPCName = true // Always true.
+	c.options.SecvEnabled, err = flags.GetBool("secvenabled")
+	if err != nil {
+		return fmt.Errorf("flags parse secvenabled bool err: %w", err)
+	}
+	kvFile, err := flags.GetString("kvfile")
+	if err != nil {
+		return fmt.Errorf("flags parse kvfile string err: %w", err)
+	}
+	if kvFile != "" {
+		bs, err := os.ReadFile(kvFile)
+		if err != nil {
+			return fmt.Errorf("read kv file %s err: %w", kvFile, err)
+		}
+		if err := json.Unmarshal(bs, &c.options.KVs); err != nil {
+			return fmt.Errorf("json unmarshal kv file %s into %T err: %w", kvFile, c.options.KVs, err)
+		}
+	}
+	kvRawJSON, err := flags.GetString("kvrawjson")
+	if err != nil {
+		return fmt.Errorf("flags parse kvrawjson string err: %w", err)
+	}
+	if kvRawJSON != "" {
+		if err := json.Unmarshal([]byte(kvRawJSON), &c.options.KVs); err != nil {
+			return fmt.Errorf("json unmarshal kv raw json %s into %T err: %w", kvRawJSON, c.options.KVs, err)
+		}
+	}
 	c.options.Force, err = flags.GetBool("force")
 	if err != nil {
 		return fmt.Errorf("flags parse force bool err: %w", err)
@@ -359,7 +399,8 @@ func (c *Create) parsePBIDLOptions(flags *pflag.FlagSet) error {
 	if err != nil {
 		return fmt.Errorf("flags get protodir string array failed err: %w", err)
 	}
-	c.options.Protodirs = fs.UniqFilePath(dirs)
+	// Always append the current working directory.
+	c.options.Protodirs = fs.UniqFilePath(append(dirs, "."))
 	c.options.Protofile, err = flags.GetString("protofile")
 	if err != nil {
 		return fmt.Errorf("flags get protofile string failed err: %w", err)
@@ -379,7 +420,8 @@ func (c *Create) parseFBIDLOptions(flags *pflag.FlagSet) error {
 	if err != nil {
 		return fmt.Errorf("flags get fbsdir string array failed err: %w", err)
 	}
-	c.options.Protodirs = fs.UniqFilePath(dirs)
+	// Always append the current working directory.
+	c.options.Protodirs = fs.UniqFilePath(append(dirs, "."))
 	c.options.Protofile, err = flags.GetString("fbs")
 	if err != nil {
 		return fmt.Errorf("flags get fbs string failed err: %w", err)

cmd/create/post.go

@@ -16,32 +16,26 @@ func (c *Create) PostRunE(cmd *cobra.Command, args []string) error {
 	wd, _ := os.Getwd()
 	defer os.Chdir(wd)
 
-	var allp []plugin.Plugin
-	allp = append(allp, plugin.Plugins...)
-	allp = append(allp, plugin.PluginsExt[c.options.Language]...)
-
-	err := os.Chdir(c.options.OutputDir)
-	if err != nil {
+	if err := os.Chdir(c.options.OutputDir); err != nil {
 		return err
 	}
-	for _, p := range allp {
+
+	for _, p := range append(plugin.Plugins, plugin.PluginsExt[c.options.Language]...) {
 		if !p.Check(c.fileDescriptor, c.options) {
 			continue
 		}
 
-		err = p.Run(c.fileDescriptor, c.options)
-		if err != nil {
+		if err := p.Run(c.fileDescriptor, c.options); err != nil {
 			return fmt.Errorf(
 				"running plugin `%s`, err: %w",
 				p.Name(), err)
 		}
 		if c.options.Verbose {
 			log.Info(
-				"running plugin %s`%s`%s, err: %v",
+				"running plugin %s`%s`%s, succeed",
 				log.ColorRed,
 				p.Name(),
-				log.ColorGreen,
-				err)
+				log.ColorGreen)
 		}
 	}
 

cmd/root.go

@@ -3,6 +3,7 @@ package cmd
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 
 	"github.com/spf13/cobra"
@@ -25,6 +26,7 @@ func Execute() {
 	%+v
 Please run "trpc -h" or "trpc create -h" (or "trpc {some-other-subcommand} -h") for help messages.
 `, err)
+		os.Exit(1) // Exist with non-zero errcode to indicate failure.
 	}
 }
 
@@ -75,7 +77,7 @@ func initConfig() error {
 
 	d, err := config.Init()
 	if err != nil {
-		return err
+		return fmt.Errorf("config init err: %w", err)
 	}
 
 	if cfgFile == defaultConfigFile {
@@ -91,7 +93,7 @@ func initConfig() error {
 
 	// If a config file is found, read it in.
 	if err := viper.ReadInConfig(); err != nil {
-		return err
+		return fmt.Errorf("viper read in config: %w", err)
 	}
 
 	return nil

config/config.go

@@ -61,15 +61,15 @@ func Init(opts ...Option) (string, error) {
 	// List of candidate template paths.
 	paths, err := candidatedInstallPath()
 	if err != nil {
-		return "", fmt.Errorf("get Template search path error: %w", err)
+		return "", fmt.Errorf("get template search path error: %w", err)
 	}
 
 	// Check if the candidate template path exists.
 	var reinstall bool
 	installPath, err := templateInstallPath(paths)
 	if err != nil {
 		if err != errTemplateNotFound {
-			return "", fmt.Errorf("get Template installed path error: %w", err)
+			return "", fmt.Errorf("get template instal path from %+v error: %w", paths, err)
 		}
 		reinstall = true
 		installPath = paths[0]
@@ -82,8 +82,9 @@ func Init(opts ...Option) (string, error) {
 	}
 	// Install the template as needed.
 	if o.force || reinstall {
+		log.Debug("reinstall template to %s", installPath)
 		if err := installTemplate(installPath); err != nil {
-			return "", fmt.Errorf("initialize config error: %w", err)
+			return "", fmt.Errorf("install template to %s error: %w", installPath, err)
 		}
 	}
 	return installPath, nil

config/dependencies.go

@@ -23,7 +23,7 @@ func LoadDependencies(languages ...string) ([]*Dependency, error) {
 	}
 
 	var deps []*Dependency
-	depsSet := map[string]*Dependency{}
+	depsSet := make(map[string]*Dependency)
 
 	deps = append(deps, cfg.IDL[IDLTypeProtobuf.String()])
 	deps = append(deps, cfg.IDL[IDLTypeFlatBuffers.String()])

config/validate.go

@@ -0,0 +1,9 @@
+package config
+
+// SupportValidate defines whether a certain language supports validation.
+// Validation code will be generated only when:
+// 1. The language appears as a key in this map.
+// 2. The corresponding value is true.
+var SupportValidate = map[string]bool{
+	"go": true,
+}

descriptor/desc.go

@@ -168,9 +168,11 @@ func (fd *FileDescriptor) ValidateEnabled() bool {
 
 // ServiceDescriptor provides the description information at the service level.
 type ServiceDescriptor struct {
-	Name string           // Service name.
-	RPC  []*RPCDescriptor // RPC interface definition.
-	RPCx []*RPCDescriptor // RPC interface definition, including the RPC alias and original name.
+	Name       string                      // Service name.
+	RPC        []*RPCDescriptor            // RPC interface definition.
+	RPCx       []*RPCDescriptor            // RPC interface definition, including the RPC alias and original name.
+	MethodRPC  map[string]*RPCDescriptor   // MethodRPC maps method name to RPC definition.
+	MethodRPCx map[string][]*RPCDescriptor // MethodRPCx maps method name to RPCx definition which is actually aliases.
 }
 
 // RPCDescriptor provides the description information at the RPC level.

descriptor/fbsdesc.go

@@ -82,7 +82,7 @@ type FbsAttrs struct {
 func NewFbsAttrs(attrs []string) *FbsAttrs {
 	f := &FbsAttrs{
 		Attrs: attrs,
-		KV:    map[string]string{},
+		KV:    make(map[string]string),
 	}
 	for _, kv := range attrs {
 		strs := strings.Split(kv, "=")

descriptor/protodesc.go

@@ -51,6 +51,9 @@ func (p *ProtoFileDescriptor) GetMessageTypes() []MessageDesc {
 	var descs []MessageDesc
 	for _, md := range p.FD.GetMessageTypes() {
 		descs = append(descs, &ProtoMessageDescriptor{MD: md})
+		for _, nmd := range md.GetNestedMessageTypes() {
+			descs = append(descs, &ProtoMessageDescriptor{MD: nmd})
+		}
 	}
 	return descs
 }