api add middleware support (#140)

* rebase upstream

* rebase

* trim no need line

* trim no need line

* trim no need line

* update doc

* remove update

* remove no need

* remove no need

* goctl add jwt support

* goctl add jwt support

* goctl add jwt support

* goctl support import

* goctl support import

* support return ()

* revert

* refactor and rename folder to group

* remove no need

* add anonymous annotation

* optimized

* rename

* rename

* update test

* api add middleware support: usage:

@server(
    middleware: M1, M2
)

* api add middleware support: usage:

@server(
    middleware: M1, M2
)

* simple logic

* should reverse middlewares

* optimized

* optimized

* rename

Co-authored-by: kingxt <dream4kingxt@163.com>

tools/goctl/api/gogen/genroutes.go

@@ -31,9 +31,9 @@ func RegisterHandlers(engine *rest.Server, serverCtx *svc.ServiceContext) {
 }
 `
 	routesAdditionTemplate = `
-	engine.AddRoutes([]rest.Route{
+	engine.AddRoutes(
 		{{.routes}}
-	}{{.jwt}}{{.signature}})
+	{{.jwt}}{{.signature}})
 `
 )
 
@@ -52,6 +52,7 @@ type (
 		jwtEnabled       bool
 		signatureEnabled bool
 		authName         string
+		middleware       []string
 	}
 	route struct {
 		method  string
@@ -87,8 +88,22 @@ func genRoutes(dir string, api *spec.ApiSpec, force bool) error {
 		if g.signatureEnabled {
 			signature = fmt.Sprintf(", rest.WithSignature(serverCtx.Config.%s.Signature)", g.authName)
 		}
+
+		var routes string
+		if len(g.middleware) > 0 {
+			var params = g.middleware
+			for i := range params {
+				params[i] = "serverCtx." + params[i]
+			}
+			var middlewareStr = strings.Join(params, ", ")
+			routes = fmt.Sprintf("rest.WithMultiMiddleware([]rest.Middleware{ %s }, []rest.Route{\n %s \n}),",
+				middlewareStr, strings.TrimSpace(gbuilder.String()))
+		} else {
+			routes = fmt.Sprintf("[]rest.Route{\n %s \n},", strings.TrimSpace(gbuilder.String()))
+		}
+
 		if err := gt.Execute(&builder, map[string]string{
-			"routes":    strings.TrimSpace(gbuilder.String()),
+			"routes":    routes,
 			"jwt":       jwt,
 			"signature": signature,
 		}); err != nil {
@@ -185,6 +200,11 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 			groupedRoutes.authName = value
 			groupedRoutes.jwtEnabled = true
 		}
+		if value, ok := apiutil.GetAnnotationValue(g.Annotations, "server", "middleware"); ok {
+			for _, item := range strings.Split(value, ",") {
+				groupedRoutes.middleware = append(groupedRoutes.middleware, item)
+			}
+		}
 		routes = append(routes, groupedRoutes)
 	}
 

tools/goctl/api/gogen/gensvc.go

@@ -9,16 +9,20 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/api/util"
 	"github.com/tal-tech/go-zero/tools/goctl/templatex"
 	ctlutil "github.com/tal-tech/go-zero/tools/goctl/util"
+	"github.com/tal-tech/go-zero/tools/goctl/vars"
 )
 
 const (
 	contextFilename = "servicecontext.go"
 	contextTemplate = `package svc
 
-import {{.configImport}}
+import (
+	{{.configImport}}
+)
 
 type ServiceContext struct {
 	Config {{.config}}
+	{{.middleware}}
 }
 
 func NewServiceContext(c {{.config}}) *ServiceContext {
@@ -53,12 +57,22 @@ func genServiceContext(dir string, api *spec.ApiSpec) error {
 		return err
 	}
 
+	var middlewareStr string
+	for _, item := range getMiddleware(api) {
+		middlewareStr += fmt.Sprintf("%s rest.Middleware\n", item)
+	}
+
 	var configImport = "\"" + ctlutil.JoinPackages(parentPkg, configDir) + "\""
+	if len(middlewareStr) > 0 {
+		configImport += fmt.Sprintf("\n\"%s/rest\"", vars.ProjectOpenSourceUrl)
+	}
+
 	t := template.Must(template.New("contextTemplate").Parse(text))
 	buffer := new(bytes.Buffer)
 	err = t.Execute(buffer, map[string]string{
 		"configImport": configImport,
 		"config":       "config.Config",
+		"middleware":   middlewareStr,
 	})
 	if err != nil {
 		return nil

tools/goctl/api/gogen/util.go

@@ -66,6 +66,18 @@ func getAuths(api *spec.ApiSpec) []string {
 	return authNames.KeysStr()
 }
 
+func getMiddleware(api *spec.ApiSpec) []string {
+	result := collection.NewSet()
+	for _, g := range api.Service.Groups {
+		if value, ok := util.GetAnnotationValue(g.Annotations, "server", "middleware"); ok {
+			for _, item := range strings.Split(value, ",") {
+				result.Add(strings.TrimSpace(item))
+			}
+		}
+	}
+	return result.KeysStr()
+}
+
 func formatCode(code string) string {
 	ret, err := goformat.Source([]byte(code))
 	if err != nil {

tools/goctl/api/parser/parser_test.go

@@ -119,6 +119,24 @@ service A-api {
 }
 `
 
+const apiHasMiddleware = `
+type Request struct {
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
+}
+
+type Response struct {
+  Message string ` + "`" + `json:"message"` + "`" + `
+}
+
+@server(
+	middleware: TokenValidate
+)
+service A-api {
+  @handler GreetHandler
+  get /greet/from/:name(Request) returns (Response)
+}
+`
+
 func TestParser(t *testing.T) {
 	filename := "greet.api"
 	err := ioutil.WriteFile(filename, []byte(testApiTemplate), os.ModePerm)
@@ -198,3 +216,16 @@ func TestAnonymousAnnotation(t *testing.T) {
 	assert.Equal(t, len(api.Service.Routes), 1)
 	assert.Equal(t, api.Service.Routes[0].Annotations[0].Value, "GreetHandler")
 }
+
+func TestApiHasMiddleware(t *testing.T) {
+	filename := "greet.api"
+	err := ioutil.WriteFile(filename, []byte(apiHasMiddleware), os.ModePerm)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	parser, err := NewParser(filename)
+	assert.Nil(t, err)
+
+	_, err = parser.Parse()
+	assert.Nil(t, err)
+}