From df0f8ed59e6824dfed42cad5eb05ef606fa56f9a Mon Sep 17 00:00:00 2001
From: anqiansong <anqiansong@gmail.com>
Date: Tue, 18 Jan 2022 11:52:30 +0800
Subject: [PATCH] Fix/issue#1289 (#1460)

* fix #1289

* Add unit test case

* fix `jwtTransKey`

* fix `jwtTransKey`

Co-authored-by: anqiansong <anqiansong@bytedance.com>
---
 tools/goctl/api/gogen/gen_test.go  |  1 +
 tools/goctl/api/gogen/genconfig.go | 13 +++++++++++++
 tools/goctl/api/gogen/genroutes.go | 10 ++++++++++
 tools/goctl/api/gogen/util.go      | 11 +++++++++++
 4 files changed, 35 insertions(+)

diff --git a/tools/goctl/api/gogen/gen_test.go b/tools/goctl/api/gogen/gen_test.go
index ba93856a..eeab8730 100644
--- a/tools/goctl/api/gogen/gen_test.go
+++ b/tools/goctl/api/gogen/gen_test.go
@@ -178,6 +178,7 @@ type Response struct {
 
 @server(
 	jwt: Auth
+	jwtTransition: Trans
 	middleware: TokenValidate
 )
 service A-api {
diff --git a/tools/goctl/api/gogen/genconfig.go b/tools/goctl/api/gogen/genconfig.go
index ce84aba3..1be07dcc 100644
--- a/tools/goctl/api/gogen/genconfig.go
+++ b/tools/goctl/api/gogen/genconfig.go
@@ -19,6 +19,7 @@ import {{.authImport}}
 type Config struct {
 	rest.RestConf
 	{{.auth}}
+	{{.jwtTrans}}
 }
 `
 
@@ -26,6 +27,11 @@ type Config struct {
 		AccessSecret string
 		AccessExpire int64
 	}
+`
+	jwtTransTemplate = ` struct {
+		Secret     string
+		PrevSecret string
+	}
 `
 )
 
@@ -40,6 +46,12 @@ func genConfig(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 	for _, item := range authNames {
 		auths = append(auths, fmt.Sprintf("%s %s", item, jwtTemplate))
 	}
+
+	jwtTransNames := getJwtTrans(api)
+	var jwtTransList []string
+	for _, item := range jwtTransNames {
+		jwtTransList = append(jwtTransList, fmt.Sprintf("%s %s", item, jwtTransTemplate))
+	}
 	authImportStr := fmt.Sprintf("\"%s/rest\"", vars.ProjectOpenSourceURL)
 
 	return genFile(fileGenConfig{
@@ -53,6 +65,7 @@ func genConfig(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 		data: map[string]string{
 			"authImport": authImportStr,
 			"auth":       strings.Join(auths, "\n"),
+			"jwtTrans":   strings.Join(jwtTransList, "\n"),
 		},
 	})
 }
diff --git a/tools/goctl/api/gogen/genroutes.go b/tools/goctl/api/gogen/genroutes.go
index 65c34b78..ef12fd8b 100644
--- a/tools/goctl/api/gogen/genroutes.go
+++ b/tools/goctl/api/gogen/genroutes.go
@@ -17,6 +17,7 @@ import (
 )
 
 const (
+	jwtTransKey    = "jwtTransition"
 	routesFilename = "routes"
 	routesTemplate = `// Code generated by goctl. DO NOT EDIT.
 package handler
@@ -58,6 +59,7 @@ type (
 		authName         string
 		middlewares      []string
 		prefix           string
+		jwtTrans         string
 	}
 	route struct {
 		method  string
@@ -96,6 +98,9 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 		if g.jwtEnabled {
 			jwt = fmt.Sprintf("\n rest.WithJwt(serverCtx.Config.%s.AccessSecret),", g.authName)
 		}
+		if len(g.jwtTrans) > 0 {
+			jwt = jwt + fmt.Sprintf("\n rest.WithJwtTransition(serverCtx.Config.%s.PrevSecret,serverCtx.Config.%s.Secret),", g.jwtTrans, g.jwtTrans)
+		}
 		var signature, prefix string
 		if g.signatureEnabled {
 			signature = "\n rest.WithSignature(serverCtx.Config.Signature),"
@@ -205,6 +210,11 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 			groupedRoutes.authName = jwt
 			groupedRoutes.jwtEnabled = true
 		}
+		jwtTrans := g.GetAnnotation(jwtTransKey)
+		if len(jwtTrans) > 0 {
+			groupedRoutes.jwtTrans = jwtTrans
+		}
+
 		signature := g.GetAnnotation("signature")
 		if signature == "true" {
 			groupedRoutes.signatureEnabled = true
diff --git a/tools/goctl/api/gogen/util.go b/tools/goctl/api/gogen/util.go
index 3eef54ba..cd2f7cce 100644
--- a/tools/goctl/api/gogen/util.go
+++ b/tools/goctl/api/gogen/util.go
@@ -111,6 +111,17 @@ func getAuths(api *spec.ApiSpec) []string {
 	return authNames.KeysStr()
 }
 
+func getJwtTrans(api *spec.ApiSpec) []string {
+	jwtTransList := collection.NewSet()
+	for _, g := range api.Service.Groups {
+		jt := g.GetAnnotation(jwtTransKey)
+		if len(jt) > 0 {
+			jwtTransList.Add(jt)
+		}
+	}
+	return jwtTransList.KeysStr()
+}
+
 func getMiddleware(api *spec.ApiSpec) []string {
 	result := collection.NewSet()
 	for _, g := range api.Service.Groups {