Feature: Add goctl env (#1557)

tools/goctl/env/check.go

@@ -0,0 +1,112 @@
+package env
+
+import (
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protoc"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengo"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengogrpc"
+	"github.com/zeromicro/go-zero/tools/goctl/util/console"
+)
+
+type bin struct {
+	name   string
+	exists bool
+	get    func(cacheDir string) (string, error)
+}
+
+var bins = []bin{
+	{
+		name:   "protoc",
+		exists: protoc.Exists(),
+		get:    protoc.Install,
+	},
+	{
+		name:   "protoc-gen-go",
+		exists: protocgengo.Exists(),
+		get:    protocgengo.Install,
+	},
+	{
+		name:   "protoc-gen-go-grpc",
+		exists: protocgengogrpc.Exists(),
+		get:    protocgengogrpc.Install,
+	},
+}
+
+func Check(ctx *cli.Context) error {
+	install := ctx.Bool("install")
+	force := ctx.Bool("force")
+	return check(install, force)
+}
+
+func check(install, force bool) error {
+	var pending = true
+	console.Info("[goctl-env]: preparing to check env")
+	defer func() {
+		if p := recover(); p != nil {
+			console.Error("%+v", p)
+			return
+		}
+		if pending {
+			console.Success("\n[goctl-env]: congratulations! your goctl environment is ready!")
+		} else {
+			console.Error(`
+[goctl-env]: check env finish, some dependencies is not found in PATH, you can execute
+command 'goctl env check --install' or 'goctl env install' to install it, for details, 
+please see 'goctl env check --help' or 'goctl env install --help'`)
+		}
+	}()
+	for _, e := range bins {
+		time.Sleep(200 * time.Millisecond)
+		console.Info("")
+		console.Info("[goctl-env]: looking up %q", e.name)
+		if e.exists {
+			console.Success("[goctl-env]: %q is installed", e.name)
+			continue
+		}
+		console.Warning("[goctl-env]: %q is not found in PATH", e.name)
+		if install {
+			install := func() {
+				console.Info("[goctl-env]: preparing to install %q", e.name)
+				path, err := e.get(env.Get(env.GoctlCache))
+				if err != nil {
+					console.Error("[goctl-env]: an error interrupted the installation: %+v", err)
+					pending = false
+				} else {
+					console.Success("[goctl-env]: %q is already installed in %q", e.name, path)
+				}
+			}
+			if force {
+				install()
+				continue
+			}
+			console.Info("[goctl-env]: do you want to install %q [y: YES, n: No]", e.name)
+			for {
+				var in string
+				fmt.Scanln(&in)
+				var brk bool
+				switch {
+				case strings.EqualFold(in, "y"):
+					install()
+					brk = true
+				case strings.EqualFold(in, "n"):
+					pending = false
+					console.Info("[goctl-env]: %q installation is ignored", e.name)
+					brk = true
+				default:
+					console.Error("[goctl-env]: invalid input, input 'y' for yes, 'n' for no")
+				}
+				if brk {
+					break
+				}
+			}
+		} else {
+			pending = false
+		}
+	}
+	return nil
+}

tools/goctl/env/env.go

@@ -0,0 +1,17 @@
+package env
+
+import (
+	"fmt"
+
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
+)
+
+func Action(c *cli.Context) error {
+	write := c.StringSlice("write")
+	if len(write) > 0 {
+		return env.WriteEnv(write)
+	}
+	fmt.Println(env.Print())
+	return nil
+}

tools/goctl/goctl.go

@@ -22,6 +22,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/bug"
 	"github.com/zeromicro/go-zero/tools/goctl/completion"
 	"github.com/zeromicro/go-zero/tools/goctl/docker"
+	"github.com/zeromicro/go-zero/tools/goctl/env"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/errorx"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
 	"github.com/zeromicro/go-zero/tools/goctl/kube"
@@ -47,6 +48,33 @@ var commands = []cli.Command{
 		Usage:  "upgrade goctl to latest version",
 		Action: upgrade.Upgrade,
 	},
+	{
+		Name: "env",
+		Flags: []cli.Flag{
+			cli.StringSliceFlag{
+				Name:  "write, w",
+				Usage: "edit goctl env",
+			},
+		},
+		Subcommands: []cli.Command{
+			{
+				Name:  "check",
+				Usage: "detect goctl env and dependency tools",
+				Flags: []cli.Flag{
+					cli.BoolFlag{
+						Name:  "install, i",
+						Usage: "install dependencies if not found",
+					},
+					cli.BoolFlag{
+						Name:  "force, f",
+						Usage: "silent installation of non-existent dependencies",
+					},
+				},
+				Action: env.Check,
+			},
+		},
+		Action: env.Action,
+	},
 	{
 		Name:        "migrate",
 		Usage:       "migrate from tal-tech to zeromicro",
@@ -829,7 +857,7 @@ func main() {
 	app.Version = fmt.Sprintf("%s %s/%s", version.BuildVersion, runtime.GOOS, runtime.GOARCH)
 	app.Commands = commands
 
-	// cli already print error messages
+	// cli already print error messages.
 	if err := app.Run(os.Args); err != nil {
 		fmt.Println(aurora.Red(errorx.Wrap(err).Error()))
 		os.Exit(codeFailure)

tools/goctl/internal/errorx/errorx.go

@@ -2,14 +2,14 @@ package errorx
 
 import (
 	"fmt"
-	"runtime"
 	"strings"
 
-	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
 )
 
-var errorFormat = `goctl: generation error: %+v
+var errorFormat = `goctl error: %+v
-goctl version: %s
+goctl env:
+%s
 %s`
 
 // GoctlError represents a goctl error.
@@ -20,8 +20,7 @@ type GoctlError struct {
 
 func (e *GoctlError) Error() string {
 	detail := wrapMessage(e.message...)
-	v := fmt.Sprintf("%s %s/%s", version.BuildVersion, runtime.GOOS, runtime.GOARCH)
+	return fmt.Sprintf(errorFormat, e.err, env.Print(), detail)
-	return fmt.Sprintf(errorFormat, e.err, v, detail)
 }
 
 // Wrap wraps an error with goctl version and message.

tools/goctl/pkg/collection/sortedmap.go

@@ -0,0 +1,208 @@
+package sortedmap
+
+import (
+	"container/list"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/stringx"
+)
+
+var ErrInvalidKVExpression = errors.New(`invalid key-value expression`)
+var ErrInvalidKVS = errors.New("the length of kv must be a even number")
+
+type KV []interface{}
+
+type SortedMap struct {
+	kv   *list.List
+	keys map[interface{}]*list.Element
+}
+
+func New() *SortedMap {
+	return &SortedMap{
+		kv:   list.New(),
+		keys: make(map[interface{}]*list.Element),
+	}
+}
+
+func (m *SortedMap) SetExpression(expression string) (key interface{}, value interface{}, err error) {
+	idx := strings.Index(expression, "=")
+	if idx == -1 {
+		return "", "", ErrInvalidKVExpression
+	}
+	key = expression[:idx]
+	if len(expression) == idx {
+		value = ""
+	} else {
+		value = expression[idx+1:]
+	}
+	if keys, ok := key.(string); ok && stringx.ContainsWhiteSpace(keys) {
+		return "", "", ErrInvalidKVExpression
+	}
+	if values, ok := value.(string); ok && stringx.ContainsWhiteSpace(values) {
+		return "", "", ErrInvalidKVExpression
+	}
+	if len(key.(string)) == 0 {
+		return "", "", ErrInvalidKVExpression
+	}
+
+	m.SetKV(key, value)
+	return
+}
+
+func (m *SortedMap) SetKV(key, value interface{}) {
+	e, ok := m.keys[key]
+	if !ok {
+		e = m.kv.PushBack(KV{
+			key, value,
+		})
+	} else {
+		e.Value.(KV)[1] = value
+	}
+	m.keys[key] = e
+}
+
+func (m *SortedMap) Set(kv KV) error {
+	if len(kv) == 0 {
+		return nil
+	}
+	if len(kv)%2 != 0 {
+		return ErrInvalidKVS
+	}
+	for idx := 0; idx < len(kv); idx += 2 {
+		m.SetKV(kv[idx], kv[idx+1])
+	}
+	return nil
+}
+
+func (m *SortedMap) Get(key interface{}) (interface{}, bool) {
+	e, ok := m.keys[key]
+	if !ok {
+		return nil, false
+	}
+	return e.Value.(KV)[1], true
+}
+
+func (m *SortedMap) GetOr(key interface{}, dft interface{}) interface{} {
+	e, ok := m.keys[key]
+	if !ok {
+		return dft
+	}
+	return e.Value.(KV)[1]
+}
+
+func (m *SortedMap) GetString(key interface{}) (string, bool) {
+	value, ok := m.Get(key)
+	if !ok {
+		return "", false
+	}
+	vs, ok := value.(string)
+	return vs, ok
+}
+
+func (m *SortedMap) GetStringOr(key interface{}, dft string) string {
+	value, ok := m.GetString(key)
+	if !ok {
+		return dft
+	}
+	return value
+}
+
+func (m *SortedMap) HasKey(key interface{}) bool {
+	_, ok := m.keys[key]
+	return ok
+}
+
+func (m *SortedMap) HasValue(value interface{}) bool {
+	var contains bool
+	m.RangeIf(func(key, v interface{}) bool {
+		if value == v {
+			contains = true
+			return false
+		}
+		return true
+	})
+	return contains
+}
+
+func (m *SortedMap) Keys() []interface{} {
+	keys := make([]interface{}, 0)
+	next := m.kv.Front()
+	for next != nil {
+		keys = append(keys, next.Value.(KV)[0])
+		next = next.Next()
+	}
+	return keys
+}
+
+func (m *SortedMap) Values() []interface{} {
+	keys := m.Keys()
+	values := make([]interface{}, len(keys))
+	for idx, key := range keys {
+		values[idx] = m.keys[key].Value.(KV)[1]
+	}
+	return values
+}
+
+func (m *SortedMap) Range(iterator func(key, value interface{})) {
+	next := m.kv.Front()
+	for next != nil {
+		value := next.Value.(KV)
+		iterator(value[0], value[1])
+		next = next.Next()
+	}
+}
+
+func (m *SortedMap) RangeIf(iterator func(key, value interface{}) bool) {
+	next := m.kv.Front()
+	for next != nil {
+		value := next.Value.(KV)
+		loop := iterator(value[0], value[1])
+		if !loop {
+			return
+		}
+		next = next.Next()
+	}
+}
+
+func (m *SortedMap) Remove(key interface{}) (value interface{}, ok bool) {
+	v, ok := m.keys[key]
+	if !ok {
+		return nil, false
+	}
+	value = v.Value.(KV)[1]
+	ok = true
+	m.kv.Remove(v)
+	delete(m.keys, key)
+	return
+}
+
+func (m *SortedMap) Insert(sm *SortedMap) {
+	sm.Range(func(key, value interface{}) {
+		m.SetKV(key, value)
+	})
+}
+
+func (m *SortedMap) Copy() *SortedMap {
+	sm := New()
+	m.Range(func(key, value interface{}) {
+		sm.SetKV(key, value)
+	})
+	return sm
+}
+
+func (m *SortedMap) Format() []string {
+	var format = make([]string, 0)
+	m.Range(func(key, value interface{}) {
+		format = append(format, fmt.Sprintf("%s=%s", key, value))
+	})
+	return format
+}
+
+func (m *SortedMap) Reset() {
+	m.kv.Init()
+	for key := range m.keys {
+		delete(m.keys, key)
+	}
+}

tools/goctl/pkg/collection/sortedmap_test.go

@@ -0,0 +1,235 @@
+package sortedmap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_SortedMap(t *testing.T) {
+	sm := New()
+	t.Run("SetExpression", func(t *testing.T) {
+		_, _, err := sm.SetExpression("")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo= ")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression(" foo=")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo =")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("=")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("=bar")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		key, value, err := sm.SetExpression("foo=bar")
+		assert.Nil(t, err)
+		assert.Equal(t, "foo", key)
+		assert.Equal(t, "bar", value)
+		key, value, err = sm.SetExpression("foo=")
+		assert.Nil(t, err)
+		assert.Equal(t, value, sm.GetOr(key, ""))
+		sm.Reset()
+	})
+
+	t.Run("SetKV", func(t *testing.T) {
+		sm.SetKV("foo", "bar")
+		assert.Equal(t, "bar", sm.GetOr("foo", ""))
+		sm.SetKV("foo", "bar-changed")
+		assert.Equal(t, "bar-changed", sm.GetOr("foo", ""))
+		sm.Reset()
+	})
+
+	t.Run("Set", func(t *testing.T) {
+		err := sm.Set(KV{})
+		assert.Nil(t, err)
+		err = sm.Set(KV{"foo"})
+		assert.ErrorIs(t, ErrInvalidKVS, err)
+		err = sm.Set(KV{"foo", "bar", "bar", "foo"})
+		assert.Nil(t, err)
+		assert.Equal(t, "bar", sm.GetOr("foo", ""))
+		assert.Equal(t, "foo", sm.GetOr("bar", ""))
+		sm.Reset()
+	})
+
+	t.Run("Get", func(t *testing.T) {
+		_, ok := sm.Get("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.Get("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		sm.Reset()
+	})
+
+	t.Run("GetString", func(t *testing.T) {
+		_, ok := sm.GetString("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.GetString("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		sm.Reset()
+	})
+
+	t.Run("GetStringOr", func(t *testing.T) {
+		value := sm.GetStringOr("foo", "bar")
+		assert.Equal(t, "bar", value)
+		sm.SetKV("foo", "foo")
+		value = sm.GetStringOr("foo", "bar")
+		assert.Equal(t, "foo", value)
+		sm.Reset()
+	})
+
+	t.Run("GetOr", func(t *testing.T) {
+		value := sm.GetOr("foo", "bar")
+		assert.Equal(t, "bar", value)
+		sm.SetKV("foo", "foo")
+		value = sm.GetOr("foo", "bar")
+		assert.Equal(t, "foo", value)
+		sm.Reset()
+	})
+
+	t.Run("HasKey", func(t *testing.T) {
+		ok := sm.HasKey("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "")
+		assert.True(t, sm.HasKey("foo"))
+		sm.Reset()
+	})
+
+	t.Run("HasValue", func(t *testing.T) {
+		assert.False(t, sm.HasValue("bar"))
+		sm.SetKV("foo", "bar")
+		assert.True(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Keys", func(t *testing.T) {
+		keys := sm.Keys()
+		assert.Equal(t, 0, len(keys))
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, "")
+		}
+		keys = sm.Keys()
+		var actual []string
+		for _, key := range keys {
+			actual = append(actual, key.(string))
+		}
+
+		assert.Equal(t, expected, actual)
+		sm.Reset()
+	})
+
+	t.Run("Values", func(t *testing.T) {
+		values := sm.Values()
+		assert.Equal(t, 0, len(values))
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		values = sm.Values()
+		var actual []string
+		for _, value := range values {
+			actual = append(actual, value.(string))
+		}
+
+		assert.Equal(t, expected, actual)
+		sm.Reset()
+	})
+
+	t.Run("Range", func(t *testing.T) {
+		var keys, values []string
+		sm.Range(func(key, value interface{}) {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+		})
+		assert.Len(t, keys, 0)
+		assert.Len(t, values, 0)
+
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		sm.Range(func(key, value interface{}) {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+		})
+		assert.Equal(t, expected, keys)
+		assert.Equal(t, expected, values)
+		sm.Reset()
+	})
+
+	t.Run("RangeIf", func(t *testing.T) {
+		var keys, values []string
+		sm.RangeIf(func(key, value interface{}) bool {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+			return true
+		})
+		assert.Len(t, keys, 0)
+		assert.Len(t, values, 0)
+
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		sm.RangeIf(func(key, value interface{}) bool {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+			if key.(string) == "foo1" {
+				return false
+			}
+			return true
+		})
+		assert.Equal(t, []string{"foo1"}, keys)
+		assert.Equal(t, []string{"foo1"}, values)
+		sm.Reset()
+	})
+
+	t.Run("Remove", func(t *testing.T) {
+		_, ok := sm.Remove("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.Remove("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		assert.False(t, sm.HasKey("foo"))
+		assert.False(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Insert", func(t *testing.T) {
+		data := New()
+		data.SetKV("foo", "bar")
+		sm.SetKV("foo1", "bar1")
+		sm.Insert(data)
+		assert.True(t, sm.HasKey("foo"))
+		assert.True(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Copy", func(t *testing.T) {
+		sm.SetKV("foo", "bar")
+		data := sm.Copy()
+		assert.True(t, data.HasKey("foo"))
+		assert.True(t, data.HasValue("bar"))
+		sm.SetKV("foo", "bar1")
+		assert.True(t, data.HasKey("foo"))
+		assert.True(t, data.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Format", func(t *testing.T) {
+		format := sm.Format()
+		assert.Equal(t, []string{}, format)
+		sm.SetKV("foo1", "bar1")
+		sm.SetKV("foo2", "bar2")
+		sm.SetKV("foo3", "")
+		format = sm.Format()
+		assert.Equal(t, []string{"foo1=bar1", "foo2=bar2", "foo3="}, format)
+		sm.Reset()
+	})
+}

tools/goctl/pkg/downloader/downloader.go

@@ -0,0 +1,23 @@
+package downloader
+
+import (
+	"io"
+	"net/http"
+	"os"
+)
+
+func Download(url string, filename string) error {
+	resp, err := http.Get(url)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	f, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = io.Copy(f, resp.Body)
+	return err
+}

tools/goctl/pkg/env/env.go

@@ -0,0 +1,147 @@
+package env
+
+import (
+	"fmt"
+	"io/ioutil"
+	"log"
+	"path/filepath"
+	"runtime"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
+	sortedmap "github.com/zeromicro/go-zero/tools/goctl/pkg/collection"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protoc"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengo"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengogrpc"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+var goctlEnv *sortedmap.SortedMap
+
+const (
+	GoctlOS                = "GOCTL_OS"
+	GoctlArch              = "GOCTL_ARCH"
+	GoctlHome              = "GOCTL_HOME"
+	GoctlDebug             = "GOCTL_DEBUG"
+	GoctlCache             = "GOCTL_CACHE"
+	GoctlVersion           = "GOCTL_VERSION"
+	ProtocVersion          = "PROTOC_VERSION"
+	ProtocGenGoVersion     = "PROTOC_GEN_GO_VERSION"
+	ProtocGenGoGRPCVersion = "PROTO_GEN_GO_GRPC_VERSION"
+
+	envFileDir = "env"
+)
+
+// init initializes the goctl environment variables, the environment variables of the function are set in order,
+// please do not change the logic order of the code.
+func init() {
+	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
+	if err != nil {
+		log.Fatalln(err)
+	}
+	goctlEnv = sortedmap.New()
+	goctlEnv.SetKV(GoctlOS, runtime.GOOS)
+	goctlEnv.SetKV(GoctlArch, runtime.GOARCH)
+	existsEnv := readEnv(defaultGoctlHome)
+	if existsEnv != nil {
+		goctlHome, ok := existsEnv.GetString(GoctlHome)
+		if ok && len(goctlHome) > 0 {
+			goctlEnv.SetKV(GoctlHome, goctlHome)
+		}
+		if debug := existsEnv.GetOr(GoctlDebug, "").(string); debug != "" {
+			if strings.EqualFold(debug, "true") || strings.EqualFold(debug, "false") {
+				goctlEnv.SetKV(GoctlDebug, debug)
+			}
+		}
+		if value := existsEnv.GetStringOr(GoctlCache, ""); value != "" {
+			goctlEnv.SetKV(GoctlCache, value)
+		}
+	}
+	if !goctlEnv.HasKey(GoctlHome) {
+		goctlEnv.SetKV(GoctlHome, defaultGoctlHome)
+	}
+	if !goctlEnv.HasKey(GoctlDebug) {
+		goctlEnv.SetKV(GoctlDebug, "False")
+	}
+
+	if !goctlEnv.HasKey(GoctlCache) {
+		cacheDir, _ := pathx.GetCacheDir()
+		goctlEnv.SetKV(GoctlCache, cacheDir)
+	}
+
+	goctlEnv.SetKV(GoctlVersion, version.BuildVersion)
+	protocVer, _ := protoc.Version()
+	goctlEnv.SetKV(ProtocVersion, protocVer)
+
+	protocGenGoVer, _ := protocgengo.Version()
+	goctlEnv.SetKV(ProtocGenGoVersion, protocGenGoVer)
+
+	protocGenGoGrpcVer, _ := protocgengogrpc.Version()
+	goctlEnv.SetKV(ProtocGenGoGRPCVersion, protocGenGoGrpcVer)
+}
+
+func Print() string {
+	return strings.Join(goctlEnv.Format(), "\n")
+}
+
+func Get(key string) string {
+	return GetOr(key, "")
+}
+
+func GetOr(key string, def string) string {
+	return goctlEnv.GetStringOr(key, def)
+}
+
+func readEnv(goctlHome string) *sortedmap.SortedMap {
+	envFile := filepath.Join(goctlHome, envFileDir)
+	data, err := ioutil.ReadFile(envFile)
+	if err != nil {
+		return nil
+	}
+	dataStr := string(data)
+	lines := strings.Split(dataStr, "\n")
+	sm := sortedmap.New()
+	for _, line := range lines {
+		_, _, err = sm.SetExpression(line)
+		if err != nil {
+			continue
+		}
+	}
+	return sm
+}
+
+func WriteEnv(kv []string) error {
+	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
+	if err != nil {
+		log.Fatalln(err)
+	}
+	data := sortedmap.New()
+	for _, e := range kv {
+		_, _, err := data.SetExpression(e)
+		if err != nil {
+			return err
+		}
+	}
+	data.RangeIf(func(key, value interface{}) bool {
+		switch key.(string) {
+		case GoctlHome, GoctlCache:
+			path := value.(string)
+			if !pathx.FileExists(path) {
+				err = fmt.Errorf("[writeEnv]: path %q is not exists", path)
+				return false
+			}
+		}
+		if goctlEnv.HasKey(key) {
+			goctlEnv.SetKV(key, value)
+			return true
+		} else {
+			err = fmt.Errorf("[writeEnv]: invalid key: %v", key)
+			return false
+		}
+	})
+	if err != nil {
+		return err
+	}
+	envFile := filepath.Join(defaultGoctlHome, envFileDir)
+	return ioutil.WriteFile(envFile, []byte(strings.Join(goctlEnv.Format(), "\n")), 0777)
+}

tools/goctl/pkg/goctl/goctl.go

@@ -0,0 +1,41 @@
+package goctl
+
+import (
+	"path/filepath"
+	"runtime"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/util/console"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+	"github.com/zeromicro/go-zero/tools/goctl/vars"
+)
+
+func Install(cacheDir, name string, installFn func(dest string) (string, error)) (string, error) {
+	goBin := golang.GoBin()
+	cacheFile := filepath.Join(cacheDir, name)
+	binFile := filepath.Join(goBin, name)
+
+	goos := runtime.GOOS
+	if goos == vars.OsWindows {
+		cacheFile = cacheFile + ".exe"
+		binFile = binFile + ".exe"
+	}
+	// read cache.
+	err := pathx.Copy(cacheFile, binFile)
+	if err == nil {
+		console.Info("%q installed from cache", name)
+		return binFile, nil
+	}
+
+	binFile, err = installFn(binFile)
+	if err != nil {
+		return "", err
+	}
+
+	// write cache.
+	err = pathx.Copy(binFile, cacheFile)
+	if err != nil {
+		console.Warning("write cache error: %+v", err)
+	}
+	return binFile, nil
+}

tools/goctl/pkg/golang/bin.go

@@ -0,0 +1,27 @@
+package golang
+
+import (
+	"go/build"
+	"os"
+	"path/filepath"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+// GoBin returns a path of GOBIN.
+func GoBin() string {
+	def := build.Default
+	goroot := os.Getenv("GOROOT")
+	bin := filepath.Join(goroot, "bin")
+	if !pathx.FileExists(bin) {
+		gopath := os.Getenv("GOPATH")
+		bin = filepath.Join(gopath, "bin")
+	}
+	if !pathx.FileExists(bin) {
+		bin = os.Getenv("GOBIN")
+	}
+	if !pathx.FileExists(bin) {
+		bin = filepath.Join(def.GOPATH, "bin")
+	}
+	return bin
+}

tools/goctl/pkg/golang/install.go

@@ -0,0 +1,17 @@
+package golang
+
+import (
+	"os"
+	"os/exec"
+)
+
+func Install(git string) error {
+	cmd := exec.Command("go", "install", git)
+	env := os.Environ()
+	env = append(env, "GO111MODULE=on", "GOPROXY=https://goproxy.cn,direct")
+	cmd.Env = env
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Run()
+	return err
+}

tools/goctl/pkg/protoc/protoc.go

@@ -0,0 +1,79 @@
+package protoc
+
+import (
+	"archive/zip"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/downloader"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+	"github.com/zeromicro/go-zero/tools/goctl/util/zipx"
+	"github.com/zeromicro/go-zero/tools/goctl/vars"
+)
+
+var url = map[string]string{
+	"linux_32":   "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-linux-x86_32.zip",
+	"linux_64":   "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-linux-x86_64.zip",
+	"darwin":     "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-osx-x86_64.zip",
+	"windows_32": "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-win32.zip",
+	"windows_64": "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-win64.zip",
+}
+
+const (
+	Name        = "protoc"
+	ZipFileName = Name + ".zip"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		goos := runtime.GOOS
+		tempFile := filepath.Join(os.TempDir(), ZipFileName)
+		bit := 32 << (^uint(0) >> 63)
+		var downloadUrl string
+		switch goos {
+		case vars.OsMac:
+			downloadUrl = url[vars.OsMac]
+		case vars.OsWindows:
+			downloadUrl = url[fmt.Sprintf("%s_%d", vars.OsWindows, bit)]
+		case vars.OsLinux:
+			downloadUrl = url[fmt.Sprintf("%s_%d", vars.OsLinux, bit)]
+		default:
+			return "", fmt.Errorf("unsupport OS: %q", goos)
+		}
+
+		err := downloader.Download(downloadUrl, tempFile)
+		if err != nil {
+			return "", err
+		}
+
+		return dest, zipx.Unpacking(tempFile, filepath.Dir(dest), func(f *zip.File) bool {
+			return filepath.Base(f.Name) == filepath.Base(dest)
+		})
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtoc()
+	return err == nil
+}
+
+func Version() (string, error) {
+	path, err := env.LookUpProtoc()
+	if err != nil {
+		return "", err
+	}
+	version, err := execx.Run(path+" --version", "")
+	if err != nil {
+		return "", err
+	}
+	fields := strings.Fields(version)
+	if len(fields) > 1 {
+		return fields[1], nil
+	}
+	return "", nil
+}

tools/goctl/pkg/protocgengo/protocgengo.go

@@ -0,0 +1,53 @@
+package protocgengo
+
+import (
+	"strings"
+	"time"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+)
+
+const (
+	Name = "protoc-gen-go"
+	url  = "google.golang.org/protobuf/cmd/protoc-gen-go@latest"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		err := golang.Install(url)
+		return dest, err
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtocGenGo()
+	return err == nil
+}
+
+// Version is used to get the version of the protoc-gen-go plugin. For older versions, protoc-gen-go does not support
+// version fetching, so if protoc-gen-go --version is executed, it will cause the process to block, so it is controlled
+// by a timer to prevent the older version process from blocking.
+func Version() (string, error) {
+	path, err := env.LookUpProtocGenGo()
+	if err != nil {
+		return "", err
+	}
+	versionC := make(chan string)
+	go func(c chan string) {
+		version, _ := execx.Run(path+" --version", "")
+		fields := strings.Fields(version)
+		if len(fields) > 1 {
+			c <- fields[1]
+		}
+	}(versionC)
+	t := time.NewTimer(time.Second)
+	select {
+	case <-t.C:
+		return "", nil
+	case version := <-versionC:
+		return version, nil
+	}
+}

tools/goctl/pkg/protocgengogrpc/protocgengogrpc.go

@@ -0,0 +1,44 @@
+package protocgengogrpc
+
+import (
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+)
+
+const (
+	Name = "protoc-gen-go-grpc"
+	url  = "google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		err := golang.Install(url)
+		return dest, err
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtocGenGoGrpc()
+	return err == nil
+}
+
+// Version is used to get the version of the protoc-gen-go-grpc plugin.
+func Version() (string, error) {
+	path, err := env.LookUpProtocGenGoGrpc()
+	if err != nil {
+		return "", err
+	}
+	version, err := execx.Run(path+" --version", "")
+	if err != nil {
+		return "", err
+	}
+	fields := strings.Fields(version)
+	if len(fields) > 1 {
+		return fields[1], nil
+	}
+	return "", nil
+}

tools/goctl/util/env/env.go

@@ -11,10 +11,11 @@ import (
 )
 
 const (
-	bin            = "bin"
+	bin                = "bin"
-	binGo          = "go"
+	binGo              = "go"
-	binProtoc      = "protoc"
+	binProtoc          = "protoc"
-	binProtocGenGo = "protoc-gen-go"
+	binProtocGenGo     = "protoc-gen-go"
+	binProtocGenGrpcGo = "protoc-gen-go-grpc"
 )
 
 // LookUpGo searches an executable go in the directories
@@ -46,6 +47,14 @@ func LookUpProtocGenGo() (string, error) {
 	return LookPath(xProtocGenGo)
 }
 
+// LookUpProtocGenGoGrpc searches an executable protoc-gen-go-grpc in the directories
+// named by the PATH environment variable.
+func LookUpProtocGenGoGrpc() (string, error) {
+	suffix := getExeSuffix()
+	xProtocGenGoGrpc := binProtocGenGrpcGo + suffix
+	return LookPath(xProtocGenGoGrpc)
+}
+
 // LookPath searches for an executable named file in the
 // directories named by the PATH environment variable,
 // for the os windows, the named file will be spliced with the

tools/goctl/util/pathx/file.go

@@ -3,6 +3,7 @@ package pathx
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
@@ -13,22 +14,23 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
 )
 
-// NL defines a new line
+// NL defines a new line.
 const (
 	NL              = "\n"
 	goctlDir        = ".goctl"
 	gitDir          = ".git"
 	autoCompleteDir = ".auto_complete"
+	cacheDir        = "cache"
 )
 
 var goctlHome string
 
-// RegisterGoctlHome register goctl home path
+// RegisterGoctlHome register goctl home path.
 func RegisterGoctlHome(home string) {
 	goctlHome = home
 }
 
-// CreateIfNotExist creates a file if it is not exists
+// CreateIfNotExist creates a file if it is not exists.
 func CreateIfNotExist(file string) (*os.File, error) {
 	_, err := os.Stat(file)
 	if !os.IsNotExist(err) {
@@ -38,7 +40,7 @@ func CreateIfNotExist(file string) (*os.File, error) {
 	return os.Create(file)
 }
 
-// RemoveIfExist deletes the specified file if it is exists
+// RemoveIfExist deletes the specified file if it is exists.
 func RemoveIfExist(filename string) error {
 	if !FileExists(filename) {
 		return nil
@@ -47,7 +49,7 @@ func RemoveIfExist(filename string) error {
 	return os.Remove(filename)
 }
 
-// RemoveOrQuit deletes the specified file if read a permit command from stdin
+// RemoveOrQuit deletes the specified file if read a permit command from stdin.
 func RemoveOrQuit(filename string) error {
 	if !FileExists(filename) {
 		return nil
@@ -60,23 +62,29 @@ func RemoveOrQuit(filename string) error {
 	return os.Remove(filename)
 }
 
-// FileExists returns true if the specified file is exists
+// FileExists returns true if the specified file is exists.
 func FileExists(file string) bool {
 	_, err := os.Stat(file)
 	return err == nil
 }
 
-// FileNameWithoutExt returns a file name without suffix
+// FileNameWithoutExt returns a file name without suffix.
 func FileNameWithoutExt(file string) string {
 	return strings.TrimSuffix(file, filepath.Ext(file))
 }
 
-// GetGoctlHome returns the path value of the goctl home where Join $HOME with .goctl
+// GetGoctlHome returns the path value of the goctl, the default path is ~/.goctl, if the path has
+// been set by calling the RegisterGoctlHome method, the user-defined path refers to.
 func GetGoctlHome() (string, error) {
 	if len(goctlHome) != 0 {
 		return goctlHome, nil
 	}
 
+	return GetDefaultGoctlHome()
+}
+
+// GetDefaultGoctlHome returns the path value of the goctl home where Join $HOME with .goctl.
+func GetDefaultGoctlHome() (string, error) {
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", err
@@ -104,7 +112,17 @@ func GetAutoCompleteHome() (string, error) {
 	return filepath.Join(goctlH, autoCompleteDir), nil
 }
 
-// GetTemplateDir returns the category path value in GoctlHome where could get it by GetGoctlHome
+// GetCacheDir returns the cache dit of goctl.
+func GetCacheDir() (string, error) {
+	goctlH, err := GetGoctlHome()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(goctlH, cacheDir), nil
+}
+
+// GetTemplateDir returns the category path value in GoctlHome where could get it by GetGoctlHome.
 func GetTemplateDir(category string) (string, error) {
 	home, err := GetGoctlHome()
 	if err != nil {
@@ -112,7 +130,7 @@ func GetTemplateDir(category string) (string, error) {
 	}
 	if home == goctlHome {
 		// backward compatible, it will be removed in the feature
-		// backward compatible start
+		// backward compatible start.
 		beforeTemplateDir := filepath.Join(home, version.GetGoctlVersion(), category)
 		fs, _ := ioutil.ReadDir(beforeTemplateDir)
 		var hasContent bool
@@ -124,7 +142,7 @@ func GetTemplateDir(category string) (string, error) {
 		if hasContent {
 			return beforeTemplateDir, nil
 		}
-		// backward compatible end
+		// backward compatible end.
 
 		return filepath.Join(home, category), nil
 	}
@@ -132,7 +150,7 @@ func GetTemplateDir(category string) (string, error) {
 	return filepath.Join(home, version.GetGoctlVersion(), category), nil
 }
 
-// InitTemplates creates template files GoctlHome where could get it by GetGoctlHome
+// InitTemplates creates template files GoctlHome where could get it by GetGoctlHome.
 func InitTemplates(category string, templates map[string]string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -152,7 +170,7 @@ func InitTemplates(category string, templates map[string]string) error {
 	return nil
 }
 
-// CreateTemplate writes template into file even it is exists
+// CreateTemplate writes template into file even it is exists.
 func CreateTemplate(category, name, content string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -161,7 +179,7 @@ func CreateTemplate(category, name, content string) error {
 	return createTemplate(filepath.Join(dir, name), content, true)
 }
 
-// Clean deletes all templates and removes the parent directory
+// Clean deletes all templates and removes the parent directory.
 func Clean(category string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -170,7 +188,7 @@ func Clean(category string) error {
 	return os.RemoveAll(dir)
 }
 
-// LoadTemplate gets template content by the specified file
+// LoadTemplate gets template content by the specified file.
 func LoadTemplate(category, file, builtin string) (string, error) {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -223,7 +241,7 @@ func createTemplate(file, content string, force bool) error {
 	return err
 }
 
-// MustTempDir creates a temporary directory
+// MustTempDir creates a temporary directory.
 func MustTempDir() string {
 	dir, err := ioutil.TempDir("", "")
 	if err != nil {
@@ -232,3 +250,25 @@ func MustTempDir() string {
 
 	return dir
 }
+
+func Copy(src, dest string) error {
+	f, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	dir := filepath.Dir(dest)
+	err = MkdirIfNotExist(dir)
+	if err != nil {
+		return err
+	}
+	w, err := os.Create(dest)
+	if err != nil {
+		return err
+	}
+	w.Chmod(os.ModePerm)
+	defer w.Close()
+	_, err = io.Copy(w, f)
+	return err
+}

tools/goctl/util/stringx/string.go

@@ -6,6 +6,8 @@ import (
 	"unicode"
 )
 
+var WhiteSpace = []rune{'\n', '\t', '\f', '\v', ' '}
+
 // String  provides for converting the source text into other spell case,like lower,snake,camel
 type String struct {
 	source string
@@ -114,3 +116,24 @@ func (s String) splitBy(fn func(r rune) bool, remove bool) []string {
 	}
 	return list
 }
+
+func ContainsAny(s string, runes ...rune) bool {
+	if len(runes) == 0 {
+		return true
+	}
+	tmp := make(map[rune]struct{}, len(runes))
+	for _, r := range runes {
+		tmp[r] = struct{}{}
+	}
+
+	for _, r := range s {
+		if _, ok := tmp[r]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+func ContainsWhiteSpace(s string) bool {
+	return ContainsAny(s, WhiteSpace...)
+}

tools/goctl/util/zipx/zipx.go

@@ -0,0 +1,51 @@
+package zipx
+
+import (
+	"archive/zip"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+func Unpacking(name, destPath string, mapper func(f *zip.File) bool) error {
+	r, err := zip.OpenReader(name)
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+
+	for _, file := range r.File {
+		ok := mapper(file)
+		if ok {
+			err = fileCopy(file, destPath)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func fileCopy(file *zip.File, destPath string) error {
+	rc, err := file.Open()
+	if err != nil {
+		return err
+	}
+	defer rc.Close()
+	filename := filepath.Join(destPath, filepath.Base(file.Name))
+	dir := filepath.Dir(filename)
+	err = pathx.MkdirIfNotExist(dir)
+	if err != nil {
+		return err
+	}
+
+	w, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer w.Close()
+	_, err = io.Copy(w, rc)
+	return err
+}