Repository: zeromicro/go-zero
Pull request #1557 (2022-02-18 11:42:39)

Feature: Add goctl env
- goctl environment variable view
- goctl environment detection
- goctl dependency installation


Patch:
diff --git a/rest/engine.go b/rest/engine.go
index 8ba13378..46e92f0c 100644
--- a/rest/engine.go
+++ b/rest/engine.go
@@ -14,6 +14,7 @@ import (
 	"github.com/zeromicro/go-zero/rest/handler"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zeromicro/go-zero/rest/internal"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 // use 1000m to represent 100%
@@ -154,6 +155,27 @@ func (ng *engine) getShedder(priority bool) load.Shedder {
 	return ng.shedder
 }
 
+// notFoundHandler returns a middleware that handles 404 not found requests.
+func (ng *engine) notFoundHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		chain := alice.New(
+			handler.TracingHandler(ng.conf.Name, ""),
+			ng.getLogHandler(),
+		)
+
+		var h http.Handler
+		if next != nil {
+			h = chain.Then(next)
+		} else {
+			h = chain.Then(http.NotFoundHandler())
+		}
+
+		cw := response.NewHeaderOnceResponseWriter(w)
+		h.ServeHTTP(cw, r)
+		cw.WriteHeader(http.StatusNotFound)
+	})
+}
+
 func (ng *engine) setTlsConfig(cfg *tls.Config) {
 	ng.tlsConfig = cfg
 }
diff --git a/rest/handler/authhandler.go b/rest/handler/authhandler.go
index 26aacca0..9099e62b 100644
--- a/rest/handler/authhandler.go
+++ b/rest/handler/authhandler.go
@@ -1,15 +1,14 @@
 package handler
 
 import (
-	"bufio"
 	"context"
 	"errors"
-	"net"
 	"net/http"
 	"net/http/httputil"
 
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 	"github.com/zeromicro/go-zero/rest/token"
 )
 
@@ -105,7 +104,7 @@ func detailAuthLog(r *http.Request, reason string) {
 }
 
 func unauthorized(w http.ResponseWriter, r *http.Request, err error, callback UnauthorizedCallback) {
-	writer := newGuardedResponseWriter(w)
+	writer := response.NewHeaderOnceResponseWriter(w)
 
 	if err != nil {
 		detailAuthLog(r, err.Error())
@@ -121,47 +120,3 @@ func unauthorized(w http.ResponseWriter, r *http.Request, err error, callback Un
 	// if user not setting HTTP header, we set header with 401
 	writer.WriteHeader(http.StatusUnauthorized)
 }
-
-type guardedResponseWriter struct {
-	writer      http.ResponseWriter
-	wroteHeader bool
-}
-
-func newGuardedResponseWriter(w http.ResponseWriter) *guardedResponseWriter {
-	return &guardedResponseWriter{
-		writer: w,
-	}
-}
-
-func (grw *guardedResponseWriter) Flush() {
-	if flusher, ok := grw.writer.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
-func (grw *guardedResponseWriter) Header() http.Header {
-	return grw.writer.Header()
-}
-
-// Hijack implements the http.Hijacker interface.
-// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
-func (grw *guardedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := grw.writer.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (grw *guardedResponseWriter) Write(body []byte) (int, error) {
-	return grw.writer.Write(body)
-}
-
-func (grw *guardedResponseWriter) WriteHeader(statusCode int) {
-	if grw.wroteHeader {
-		return
-	}
-
-	grw.wroteHeader = true
-	grw.writer.WriteHeader(statusCode)
-}
diff --git a/rest/handler/breakerhandler.go b/rest/handler/breakerhandler.go
index 516aba76..e9d7243a 100644
--- a/rest/handler/breakerhandler.go
+++ b/rest/handler/breakerhandler.go
@@ -9,7 +9,7 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/rest/httpx"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const breakerSeparator = "://"
@@ -28,7 +28,7 @@ func BreakerHandler(method, path string, metrics *stat.Metrics) func(http.Handle
 				return
 			}
 
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				if cw.Code < http.StatusInternalServerError {
 					promise.Accept()
diff --git a/rest/handler/prometheushandler.go b/rest/handler/prometheushandler.go
index a1a286a0..639d59c0 100644
--- a/rest/handler/prometheushandler.go
+++ b/rest/handler/prometheushandler.go
@@ -8,7 +8,7 @@ import (
 	"github.com/zeromicro/go-zero/core/metric"
 	"github.com/zeromicro/go-zero/core/prometheus"
 	"github.com/zeromicro/go-zero/core/timex"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const serverNamespace = "http_server"
@@ -41,7 +41,7 @@ func PrometheusHandler(path string) func(http.Handler) http.Handler {
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			startTime := timex.Now()
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				metricServerReqDur.Observe(int64(timex.Since(startTime)/time.Millisecond), path)
 				metricServerReqCodeTotal.Inc(path, strconv.Itoa(cw.Code))
diff --git a/rest/handler/sheddinghandler.go b/rest/handler/sheddinghandler.go
index ffbee34f..977824dc 100644
--- a/rest/handler/sheddinghandler.go
+++ b/rest/handler/sheddinghandler.go
@@ -8,7 +8,7 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/rest/httpx"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const serviceType = "api"
@@ -41,7 +41,7 @@ func SheddingHandler(shedder load.Shedder, metrics *stat.Metrics) func(http.Hand
 				return
 			}
 
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				if cw.Code == http.StatusServiceUnavailable {
 					promise.Fail()
diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go
index 0a8cba9e..0b1cdf7e 100644
--- a/rest/handler/tracinghandler.go
+++ b/rest/handler/tracinghandler.go
@@ -18,12 +18,16 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			ctx := propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+			spanName := path
+			if len(spanName) == 0 {
+				spanName = r.URL.Path
+			}
 			spanCtx, span := tracer.Start(
 				ctx,
-				path,
+				spanName,
 				oteltrace.WithSpanKind(oteltrace.SpanKindServer),
 				oteltrace.WithAttributes(semconv.HTTPServerAttributesFromHTTPRequest(
-					serviceName, path, r)...),
+					serviceName, spanName, r)...),
 			)
 			defer span.End()
 
diff --git a/rest/internal/cors/handlers.go b/rest/internal/cors/handlers.go
index c613b27d..7bb3f077 100644
--- a/rest/internal/cors/handlers.go
+++ b/rest/internal/cors/handlers.go
@@ -1,10 +1,9 @@
 package cors
 
 import (
-	"bufio"
-	"errors"
-	"net"
 	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const (
@@ -30,7 +29,7 @@ const (
 // At most one origin can be specified, other origins are ignored if given, default to be *.
 func NotAllowedHandler(fn func(w http.ResponseWriter), origins ...string) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		gw := &guardedResponseWriter{w: w}
+		gw := response.NewHeaderOnceResponseWriter(w)
 		checkAndSetHeaders(gw, r, origins)
 		if fn != nil {
 			fn(gw)
@@ -62,44 +61,6 @@ func Middleware(fn func(w http.Header), origins ...string) func(http.HandlerFunc
 	}
 }
 
-type guardedResponseWriter struct {
-	w           http.ResponseWriter
-	wroteHeader bool
-}
-
-func (w *guardedResponseWriter) Flush() {
-	if flusher, ok := w.w.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
-func (w *guardedResponseWriter) Header() http.Header {
-	return w.w.Header()
-}
-
-// Hijack implements the http.Hijacker interface.
-// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
-func (w *guardedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := w.w.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (w *guardedResponseWriter) Write(bytes []byte) (int, error) {
-	return w.w.Write(bytes)
-}
-
-func (w *guardedResponseWriter) WriteHeader(code int) {
-	if w.wroteHeader {
-		return
-	}
-
-	w.w.WriteHeader(code)
-	w.wroteHeader = true
-}
-
 func checkAndSetHeaders(w http.ResponseWriter, r *http.Request, origins []string) {
 	setVaryHeaders(w, r)
 
diff --git a/rest/internal/response/headeronceresponsewriter.go b/rest/internal/response/headeronceresponsewriter.go
new file mode 100644
index 00000000..62ba4442
--- /dev/null
+++ b/rest/internal/response/headeronceresponsewriter.go
@@ -0,0 +1,57 @@
+package response
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+)
+
+// HeaderOnceResponseWriter is a http.ResponseWriter implementation
+// that only the first WriterHeader takes effect.
+type HeaderOnceResponseWriter struct {
+	w           http.ResponseWriter
+	wroteHeader bool
+}
+
+// NewHeaderOnceResponseWriter returns a HeaderOnceResponseWriter.
+func NewHeaderOnceResponseWriter(w http.ResponseWriter) http.ResponseWriter {
+	return &HeaderOnceResponseWriter{w: w}
+}
+
+// Flush flushes the response writer.
+func (w *HeaderOnceResponseWriter) Flush() {
+	if flusher, ok := w.w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+// Header returns the http header.
+func (w *HeaderOnceResponseWriter) Header() http.Header {
+	return w.w.Header()
+}
+
+// Hijack implements the http.Hijacker interface.
+// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
+func (w *HeaderOnceResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hijacked, ok := w.w.(http.Hijacker); ok {
+		return hijacked.Hijack()
+	}
+
+	return nil, nil, errors.New("server doesn't support hijacking")
+}
+
+// Write writes bytes into w.
+func (w *HeaderOnceResponseWriter) Write(bytes []byte) (int, error) {
+	return w.w.Write(bytes)
+}
+
+// WriteHeader writes code into w, and not sealing the writer.
+func (w *HeaderOnceResponseWriter) WriteHeader(code int) {
+	if w.wroteHeader {
+		return
+	}
+
+	w.w.WriteHeader(code)
+	w.wroteHeader = true
+}
diff --git a/rest/internal/security/withcoderesponsewriter.go b/rest/internal/response/withcoderesponsewriter.go
similarity index 85%
rename from rest/internal/security/withcoderesponsewriter.go
rename to rest/internal/response/withcoderesponsewriter.go
index 2496e301..4fa9631e 100644
--- a/rest/internal/security/withcoderesponsewriter.go
+++ b/rest/internal/response/withcoderesponsewriter.go
@@ -1,7 +1,8 @@
-package security
+package response
 
 import (
 	"bufio"
+	"errors"
 	"net"
 	"net/http"
 )
@@ -27,7 +28,11 @@ func (w *WithCodeResponseWriter) Header() http.Header {
 // Hijack implements the http.Hijacker interface.
 // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
 func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	return w.Writer.(http.Hijacker).Hijack()
+	if hijacked, ok := w.Writer.(http.Hijacker); ok {
+		return hijacked.Hijack()
+	}
+
+	return nil, nil, errors.New("server doesn't support hijacking")
 }
 
 // Write writes bytes into w.
diff --git a/rest/server.go b/rest/server.go
index 391ee75c..0c053ceb 100644
--- a/rest/server.go
+++ b/rest/server.go
@@ -49,6 +49,7 @@ func NewServer(c RestConf, opts ...RunOption) (*Server, error) {
 		router: router.NewRouter(),
 	}
 
+	opts = append([]RunOption{WithNotFoundHandler(nil)}, opts...)
 	for _, opt := range opts {
 		opt(server)
 	}
@@ -163,7 +164,8 @@ func WithMiddleware(middleware Middleware, rs ...Route) []Route {
 // WithNotFoundHandler returns a RunOption with not found handler set to given handler.
 func WithNotFoundHandler(handler http.Handler) RunOption {
 	return func(server *Server) {
-		server.router.SetNotFoundHandler(handler)
+		notFoundHandler := server.ngin.notFoundHandler(handler)
+		server.router.SetNotFoundHandler(notFoundHandler)
 	}
 }
 
diff --git a/tools/goctl/env/check.go b/tools/goctl/env/check.go
new file mode 100644
index 00000000..723e9b96
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/env/env.go b/tools/goctl/env/env.go
new file mode 100644
index 00000000..04b468e2
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/goctl.go b/tools/goctl/goctl.go
index b685f6c4..304387c7 100644
--- a/tools/goctl/goctl.go
+++ b/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)
diff --git a/tools/goctl/internal/errorx/errorx.go b/tools/goctl/internal/errorx/errorx.go
index c513b814..5e667904 100644
--- a/tools/goctl/internal/errorx/errorx.go
+++ b/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
-goctl version: %s
+var errorFormat = `goctl error: %+v
+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, v, detail)
+	return fmt.Sprintf(errorFormat, e.err, env.Print(), detail)
 }
 
 // Wrap wraps an error with goctl version and message.
diff --git a/tools/goctl/pkg/collection/sortedmap.go b/tools/goctl/pkg/collection/sortedmap.go
new file mode 100644
index 00000000..90ed2ab1
--- /dev/null
+++ b/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)
+	}
+}
diff --git a/tools/goctl/pkg/downloader/downloader.go b/tools/goctl/pkg/downloader/downloader.go
new file mode 100644
index 00000000..f17a9a20
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/pkg/env/env.go b/tools/goctl/pkg/env/env.go
new file mode 100644
index 00000000..32dea269
--- /dev/null
+++ b/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)
+}
diff --git a/tools/goctl/pkg/goctl/goctl.go b/tools/goctl/pkg/goctl/goctl.go
new file mode 100644
index 00000000..7729648a
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/pkg/golang/bin.go b/tools/goctl/pkg/golang/bin.go
new file mode 100644
index 00000000..c5278682
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/pkg/golang/install.go b/tools/goctl/pkg/golang/install.go
new file mode 100644
index 00000000..465b2e0f
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/pkg/protoc/protoc.go b/tools/goctl/pkg/protoc/protoc.go
new file mode 100644
index 00000000..0db64de9
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/pkg/protocgengo/protocgengo.go b/tools/goctl/pkg/protocgengo/protocgengo.go
new file mode 100644
index 00000000..23420834
--- /dev/null
+++ b/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
+	}
+}
diff --git a/tools/goctl/pkg/protocgengogrpc/protocgengogrpc.go b/tools/goctl/pkg/protocgengogrpc/protocgengogrpc.go
new file mode 100644
index 00000000..62b0f499
--- /dev/null
+++ b/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
+}
diff --git a/tools/goctl/util/env/env.go b/tools/goctl/util/env/env.go
index 8938dc34..93c95314 100644
--- a/tools/goctl/util/env/env.go
+++ b/tools/goctl/util/env/env.go
@@ -11,10 +11,11 @@ import (
 )
 
 const (
-	bin            = "bin"
-	binGo          = "go"
-	binProtoc      = "protoc"
-	binProtocGenGo = "protoc-gen-go"
+	bin                = "bin"
+	binGo              = "go"
+	binProtoc          = "protoc"
+	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
diff --git a/tools/goctl/util/pathx/file.go b/tools/goctl/util/pathx/file.go
index 0608acab..4fe25283 100644
--- a/tools/goctl/util/pathx/file.go
+++ b/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
+}
diff --git a/tools/goctl/util/stringx/string.go b/tools/goctl/util/stringx/string.go
index ae148c2f..5f3f63cc 100644
--- a/tools/goctl/util/stringx/string.go
+++ b/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...)
+}
diff --git a/tools/goctl/util/zipx/zipx.go b/tools/goctl/util/zipx/zipx.go
new file mode 100644
index 00000000..2b1619dd
--- /dev/null
+++ b/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
+}


Test patch:
diff --git a/rest/engine_test.go b/rest/engine_test.go
index 799309e6..ca540149 100644
--- a/rest/engine_test.go
+++ b/rest/engine_test.go
@@ -1,13 +1,17 @@
 package rest
 
 import (
+	"context"
 	"errors"
 	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/conf"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 func TestNewEngine(t *testing.T) {
@@ -190,6 +194,75 @@ func TestEngine_checkedTimeout(t *testing.T) {
 	}
 }
 
+func TestEngine_notFoundHandler(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	ts := httptest.NewServer(ng.notFoundHandler(nil))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+}
+
+func TestEngine_notFoundHandlerNotNil(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	var called int32
+	ts := httptest.NewServer(ng.notFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&called, 1)
+	})))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&called))
+}
+
+func TestEngine_notFoundHandlerNotNilWriteHeader(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	var called int32
+	ts := httptest.NewServer(ng.notFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&called, 1)
+		w.WriteHeader(http.StatusExpectationFailed)
+	})))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusExpectationFailed, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&called))
+}
+
 type mockedRouter struct{}
 
 func (m mockedRouter) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
diff --git a/rest/handler/authhandler_test.go b/rest/handler/authhandler_test.go
index 1587b35f..e9496e10 100644
--- a/rest/handler/authhandler_test.go
+++ b/rest/handler/authhandler_test.go
@@ -90,26 +90,6 @@ func TestAuthHandler_NilError(t *testing.T) {
 	})
 }
 
-func TestAuthHandler_Flush(t *testing.T) {
-	resp := httptest.NewRecorder()
-	handler := newGuardedResponseWriter(resp)
-	handler.Flush()
-	assert.True(t, resp.Flushed)
-}
-
-func TestAuthHandler_Hijack(t *testing.T) {
-	resp := httptest.NewRecorder()
-	writer := newGuardedResponseWriter(resp)
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = newGuardedResponseWriter(mockedHijackable{resp})
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-}
-
 func buildToken(secretKey string, payloads map[string]interface{}, seconds int64) (string, error) {
 	now := time.Now().Unix()
 	claims := make(jwt.MapClaims)
diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go
index b25edaea..94af55e0 100644
--- a/rest/handler/tracinghandler_test.go
+++ b/rest/handler/tracinghandler_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/justinas/alice"
 	"github.com/stretchr/testify/assert"
 	ztrace "github.com/zeromicro/go-zero/core/trace"
 	"go.opentelemetry.io/otel"
@@ -21,28 +22,31 @@ func TestOtelHandler(t *testing.T) {
 		Sampler:  1.0,
 	})
 
-	ts := httptest.NewServer(
-		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
-			spanCtx := trace.SpanContextFromContext(ctx)
-			assert.Equal(t, true, spanCtx.IsValid())
-		}),
-	)
-	defer ts.Close()
-
-	client := ts.Client()
-	err := func(ctx context.Context) error {
-		ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
-		defer span.End()
-
-		req, _ := http.NewRequest("GET", ts.URL, nil)
-		otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
-
-		res, err := client.Do(req)
-		assert.Equal(t, err, nil)
-		_ = res.Body.Close()
-		return nil
-	}(context.Background())
-
-	assert.Equal(t, err, nil)
+	for _, test := range []string{"", "bar"} {
+		t.Run(test, func(t *testing.T) {
+			h := alice.New(TracingHandler("foo", test)).Then(
+				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+					spanCtx := trace.SpanContextFromContext(ctx)
+					assert.True(t, spanCtx.IsValid())
+				}))
+			ts := httptest.NewServer(h)
+			defer ts.Close()
+
+			client := ts.Client()
+			err := func(ctx context.Context) error {
+				ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
+				defer span.End()
+
+				req, _ := http.NewRequest("GET", ts.URL, nil)
+				otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
+
+				res, err := client.Do(req)
+				assert.Nil(t, err)
+				return res.Body.Close()
+			}(context.Background())
+
+			assert.Nil(t, err)
+		})
+	}
 }
diff --git a/rest/internal/cors/handlers_test.go b/rest/internal/cors/handlers_test.go
index 047fdb98..dea7bb4e 100644
--- a/rest/internal/cors/handlers_test.go
+++ b/rest/internal/cors/handlers_test.go
@@ -1,8 +1,6 @@
 package cors
 
 import (
-	"bufio"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -131,48 +129,3 @@ func TestCorsHandlerWithOrigins(t *testing.T) {
 		}
 	}
 }
-
-func TestGuardedResponseWriter_Flush(t *testing.T) {
-	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
-	handler := NotAllowedHandler(func(w http.ResponseWriter) {
-		w.Header().Set("X-Test", "test")
-		w.WriteHeader(http.StatusServiceUnavailable)
-		_, err := w.Write([]byte("content"))
-		assert.Nil(t, err)
-
-		flusher, ok := w.(http.Flusher)
-		assert.True(t, ok)
-		flusher.Flush()
-	}, "foo.com")
-
-	resp := httptest.NewRecorder()
-	handler.ServeHTTP(resp, req)
-	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
-	assert.Equal(t, "test", resp.Header().Get("X-Test"))
-	assert.Equal(t, "content", resp.Body.String())
-}
-
-func TestGuardedResponseWriter_Hijack(t *testing.T) {
-	resp := httptest.NewRecorder()
-	writer := &guardedResponseWriter{
-		w: resp,
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = &guardedResponseWriter{
-		w: mockedHijackable{resp},
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-}
-
-type mockedHijackable struct {
-	*httptest.ResponseRecorder
-}
-
-func (m mockedHijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	return nil, nil, nil
-}
diff --git a/rest/internal/response/headeronceresponsewriter_test.go b/rest/internal/response/headeronceresponsewriter_test.go
new file mode 100644
index 00000000..cbc8fe1d
--- /dev/null
+++ b/rest/internal/response/headeronceresponsewriter_test.go
@@ -0,0 +1,58 @@
+package response
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHeaderOnceResponseWriter_Flush(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cw := NewHeaderOnceResponseWriter(w)
+		cw.Header().Set("X-Test", "test")
+		cw.WriteHeader(http.StatusServiceUnavailable)
+		cw.WriteHeader(http.StatusExpectationFailed)
+		_, err := cw.Write([]byte("content"))
+		assert.Nil(t, err)
+
+		flusher, ok := cw.(http.Flusher)
+		assert.True(t, ok)
+		flusher.Flush()
+	})
+
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
+	assert.Equal(t, "test", resp.Header().Get("X-Test"))
+	assert.Equal(t, "content", resp.Body.String())
+}
+
+func TestHeaderOnceResponseWriter_Hijack(t *testing.T) {
+	resp := httptest.NewRecorder()
+	writer := &HeaderOnceResponseWriter{
+		w: resp,
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &HeaderOnceResponseWriter{
+		w: mockedHijackable{resp},
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+}
+
+type mockedHijackable struct {
+	*httptest.ResponseRecorder
+}
+
+func (m mockedHijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, nil
+}
diff --git a/rest/internal/security/withcoderesponsewriter_test.go b/rest/internal/response/withcoderesponsewriter_test.go
similarity index 71%
rename from rest/internal/security/withcoderesponsewriter_test.go
rename to rest/internal/response/withcoderesponsewriter_test.go
index 3a627798..a4fcfae6 100644
--- a/rest/internal/security/withcoderesponsewriter_test.go
+++ b/rest/internal/response/withcoderesponsewriter_test.go
@@ -1,4 +1,4 @@
-package security
+package response
 
 import (
 	"net/http"
@@ -31,3 +31,20 @@ func TestWithCodeResponseWriter(t *testing.T) {
 	assert.Equal(t, "test", resp.Header().Get("X-Test"))
 	assert.Equal(t, "content", resp.Body.String())
 }
+
+func TestWithCodeResponseWriter_Hijack(t *testing.T) {
+	resp := httptest.NewRecorder()
+	writer := &WithCodeResponseWriter{
+		Writer: resp,
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &WithCodeResponseWriter{
+		Writer: mockedHijackable{resp},
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+}
diff --git a/tools/goctl/pkg/collection/sortedmap_test.go b/tools/goctl/pkg/collection/sortedmap_test.go
new file mode 100644
index 00000000..788804bd
--- /dev/null
+++ b/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()
+	})
+}