From 74a0937b8c7948e694bb9afcd28ec14f8227ba0b Mon Sep 17 00:00:00 2001
From: Suhaha <jklopsdfw@gmail.com>
Date: Thu, 11 Nov 2021 13:02:42 +0800
Subject: [PATCH] refine: feature flag middleware

---
 pkg/apiserver/conprof/service.go       | 13 ++++---
 pkg/apiserver/utils/mw_feature_flag.go | 48 ++++++++++++++++++++++++++
 pkg/utils/feature_flag.go              | 17 ---------
 3 files changed, 56 insertions(+), 22 deletions(-)
 create mode 100644 pkg/apiserver/utils/mw_feature_flag.go

diff --git a/pkg/apiserver/conprof/service.go b/pkg/apiserver/conprof/service.go
index 80edca9f84..45e3b7b517 100644
--- a/pkg/apiserver/conprof/service.go
+++ b/pkg/apiserver/conprof/service.go
@@ -30,8 +30,9 @@ import (
 	"golang.org/x/sync/singleflight"
 
 	"github.com/pingcap/tidb-dashboard/pkg/apiserver/user"
-	"github.com/pingcap/tidb-dashboard/pkg/apiserver/utils"
+	apiutils "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils"
 	"github.com/pingcap/tidb-dashboard/pkg/config"
+	"github.com/pingcap/tidb-dashboard/pkg/utils"
 	"github.com/pingcap/tidb-dashboard/pkg/utils/topology"
 )
 
@@ -81,7 +82,9 @@ func newService(lc fx.Lifecycle, p ServiceParams) *Service {
 func registerRouter(r *gin.RouterGroup, auth *user.AuthService, s *Service) {
 	endpoint := r.Group("/continuous_profiling")
 
-	endpoint.Use(FeatureFlagConprof.Middleware(s.params.Config.FeatureVersion))
+	endpoint.Use(apiutils.MWForbidByFeatureFlag([]*utils.FeatureFlag{
+		FeatureFlagConprof,
+	}, s.params.Config.FeatureVersion))
 	{
 		endpoint.GET("/config", auth.MWAuthRequired(), s.reverseProxy("/config"), s.conprofConfig)
 		endpoint.POST("/config", auth.MWAuthRequired(), auth.MWRequireWritePriv(), s.reverseProxy("/config"), s.updateConprofConfig)
@@ -107,9 +110,9 @@ func (s *Service) reverseProxy(targetPath string) gin.HandlerFunc {
 		c.Request.URL.Path = targetPath
 		token := c.Query("token")
 		if token != "" {
-			queryStr, err := utils.ParseJWTString("conprof", token)
+			queryStr, err := apiutils.ParseJWTString("conprof", token)
 			if err != nil {
-				utils.MakeInvalidRequestErrorFromError(c, err)
+				apiutils.MakeInvalidRequestErrorFromError(c, err)
 				return
 			}
 			c.Request.URL.RawQuery = queryStr
@@ -295,7 +298,7 @@ func (s *Service) conprofGroupProfileDetail(c *gin.Context) {
 // @Failure 500 {object} utils.APIError
 func (s *Service) genConprofActionToken(c *gin.Context) {
 	q := c.Query("q")
-	token, err := utils.NewJWTString("conprof", q)
+	token, err := apiutils.NewJWTString("conprof", q)
 	if err != nil {
 		_ = c.Error(err)
 		return
diff --git a/pkg/apiserver/utils/mw_feature_flag.go b/pkg/apiserver/utils/mw_feature_flag.go
new file mode 100644
index 0000000000..331d73788a
--- /dev/null
+++ b/pkg/apiserver/utils/mw_feature_flag.go
@@ -0,0 +1,48 @@
+// Copyright 2021 Suhaha
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/pingcap/tidb-dashboard/pkg/utils"
+)
+
+var ErrFeatureNotSupported = ErrNS.NewType("feature_not_supported")
+
+func MWForbidByFeatureFlag(featureFlags []*utils.FeatureFlag, targetVersion string) gin.HandlerFunc {
+	supported := true
+	unsupportedFeatures := make([]string, len(featureFlags))
+	for _, ff := range featureFlags {
+		if !ff.IsSupported(targetVersion) {
+			supported = false
+			unsupportedFeatures = append(unsupportedFeatures, ff.Name)
+			continue
+		}
+	}
+
+	return func(c *gin.Context) {
+		if !supported {
+			_ = c.Error(ErrFeatureNotSupported.New("unsupported features: %v", unsupportedFeatures))
+			c.Status(http.StatusForbidden)
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}
diff --git a/pkg/utils/feature_flag.go b/pkg/utils/feature_flag.go
index 0226060463..7e3b395b80 100644
--- a/pkg/utils/feature_flag.go
+++ b/pkg/utils/feature_flag.go
@@ -15,12 +15,9 @@
 package utils
 
 import (
-	"fmt"
-	"net/http"
 	"strings"
 
 	"github.com/Masterminds/semver"
-	"github.com/gin-gonic/gin"
 
 	"github.com/pingcap/tidb-dashboard/pkg/utils/version"
 )
@@ -60,17 +57,3 @@ func (ff *FeatureFlag) IsSupported(targetVersion string) bool {
 	}
 	return false
 }
-
-func (ff *FeatureFlag) Middleware(targetVersion string) gin.HandlerFunc {
-	isSupported := !ff.IsSupported(targetVersion)
-	return func(c *gin.Context) {
-		if isSupported {
-			_ = c.Error(fmt.Errorf("the feature is not supported"))
-			c.Status(http.StatusForbidden)
-			c.Abort()
-			return
-		}
-
-		c.Next()
-	}
-}