close #4373 : add framework

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>

pkg/apiutil/apiutil.go

@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"strconv"
 
+	"github.com/gorilla/mux"
 	"github.com/pingcap/errcode"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
@@ -127,3 +128,14 @@ func ErrorResp(rd *render.Render, w http.ResponseWriter, err error) {
 		rd.JSON(w, http.StatusInternalServerError, err.Error())
 	}
 }
+
+// GetHTTPRouteName return mux route name registered for ServiceName
+func GetHTTPRouteName(req *http.Request) (string, bool) {
+	route := mux.CurrentRoute(req)
+	if route != nil {
+		if route.GetName() != "" {
+			return route.GetName(), true
+		}
+	}
+	return "", false
+}

pkg/apiutil/serverapi/middleware.go

@@ -79,6 +79,24 @@ func IsServiceAllowed(s *server.Server, group server.ServiceGroup) bool {
 	return false
 }
 
+type selfProtector struct {
+	s *server.Server
+}
+
+// NewSelfProtector handle self-protection
+func NewSelfProtector(s *server.Server) negroni.Handler {
+	return &selfProtector{s: s}
+}
+
+func (protector *selfProtector) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
+	handler := protector.s.GetSelfProtectionHandler()
+	if handler == nil || handler.HandleHTTPSelfProtection(r) {
+		next(w, r)
+	} else {
+		http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+	}
+}
+
 type redirector struct {
 	s *server.Server
 }

server/api/server.go

@@ -37,6 +37,7 @@ func NewHandler(ctx context.Context, svr *server.Server) (http.Handler, server.S
 	router.PathPrefix(apiPrefix).Handler(negroni.New(
 		serverapi.NewRuntimeServiceValidator(svr, group),
 		serverapi.NewRedirector(svr),
+		serverapi.NewSelfProtector(svr),
 		negroni.Wrap(r)),
 	)
 

server/self_protection.go

@@ -0,0 +1,80 @@
+// Copyright 2021 TiKV Project Authors.
+//
+// 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 server
+
+import (
+	"net/http"
+
+	"github.com/tikv/pd/pkg/apiutil"
+)
+
+// SelfProtectionHandler is a framework to handle self protection mechanism
+// Self-protection granularity is a logical service
+type SelfProtectionHandler struct {
+	// ServiceHandlers a
+	ServiceHandlers map[string]*serviceSelfProtectionHandler
+}
+
+// NewSelfProtectionHandler returns a new SelfProtectionHandler with config
+func NewSelfProtectionHandler(server *Server) *SelfProtectionHandler {
+	handler := &SelfProtectionHandler{
+		ServiceHandlers: make(map[string]*serviceSelfProtectionHandler),
+	}
+	return handler
+}
+
+// HandleHTTPSelfProtection is used to handle http api self protection
+func (h *SelfProtectionHandler) HandleHTTPSelfProtection(req *http.Request) bool {
+	serviceName, findName := apiutil.GetHTTPRouteName(req)
+	// if path is not registered in router, go on process
+	if !findName {
+		return true
+	}
+
+	serviceHandler, ok := h.ServiceHandlers[serviceName]
+	// if there is no service handler, go on process
+	if !ok {
+		return true
+	}
+
+	httpHandler := &HTTPServiceSelfProtectionHandler{
+		req:     req,
+		handler: serviceHandler,
+	}
+	return httpHandler.Handle()
+}
+
+// ServiceSelfProtectionHandler is a interface for define self-protection handler by service granularity
+type ServiceSelfProtectionHandler interface {
+	Handle() bool
+}
+
+// HTTPServiceSelfProtectionHandler implement ServiceSelfProtectionHandler to handle http
+type HTTPServiceSelfProtectionHandler struct {
+	req     *http.Request
+	handler *serviceSelfProtectionHandler
+}
+
+// Handle implement ServiceSelfProtectionHandler defined function
+func (h *HTTPServiceSelfProtectionHandler) Handle() bool {
+	// to be implemented
+	return true
+}
+
+// serviceSelfProtectionHandler is a handler which is independent communication mode
+type serviceSelfProtectionHandler struct {
+	// todo APIRateLimiter
+	// todo AuditLogger
+}

server/server.go

@@ -149,6 +149,8 @@ type Server struct {
 	// tsoDispatcher is used to dispatch different TSO requests to
 	// the corresponding forwarding TSO channel.
 	tsoDispatcher sync.Map /* Store as map[string]chan *tsoRequest */
+	// SelfProtectionHandler is used for PD self-pretection
+	selfProtectionHandler *SelfProtectionHandler
 }
 
 // HandlerBuilder builds a server HTTP handler.
@@ -238,6 +240,7 @@ func CreateServer(ctx context.Context, cfg *config.Config, serviceBuilders ...Ha
 	}
 
 	s.handler = newHandler(s)
+	s.selfProtectionHandler = NewSelfProtectionHandler(s)
 
 	// Adjust etcd config.
 	etcdCfg, err := s.cfg.GenEmbedEtcdConfig()
@@ -668,6 +671,11 @@ func (s *Server) stopRaftCluster() {
 	s.cluster.Stop()
 }
 
+// GetAddr returns the server urls for clients.
+func (s *Server) GetSelfProtectionHandler() *SelfProtectionHandler {
+	return s.selfProtectionHandler
+}
+
 // GetAddr returns the server urls for clients.
 func (s *Server) GetAddr() string {
 	return s.cfg.AdvertiseClientUrls

server/server_test.go

@@ -21,11 +21,14 @@ import (
 	"net/http"
 	"testing"
 
+	"github.com/gorilla/mux"
 	. "github.com/pingcap/check"
+	"github.com/tikv/pd/pkg/apiutil"
 	"github.com/tikv/pd/pkg/assertutil"
 	"github.com/tikv/pd/pkg/etcdutil"
 	"github.com/tikv/pd/pkg/testutil"
 	"github.com/tikv/pd/server/config"
+	"github.com/urfave/negroni"
 	"go.etcd.io/etcd/embed"
 	"go.etcd.io/etcd/pkg/types"
 	"go.uber.org/goleak"
@@ -235,3 +238,42 @@ func (s *testServerHandlerSuite) TestRegisterServerHandler(c *C) {
 	bodyString := string(bodyBytes)
 	c.Assert(bodyString, Equals, "Hello World\n")
 }
+
+func (s *testServerHandlerSuite) TestMuxRouterName(c *C) {
+	handler := func(ctx context.Context, s *Server) (http.Handler, ServiceGroup, error) {
+		r := mux.NewRouter()
+		r.HandleFunc("/pd/apis/mok/v1/router", func(w http.ResponseWriter, r *http.Request) {
+			RouterName, _ := apiutil.GetHTTPRouteName(r)
+			fmt.Fprintln(w, RouterName)
+		}).Name("Mux Router")
+		info := ServiceGroup{
+			Name:    "mok",
+			Version: "v1",
+		}
+		router := mux.NewRouter()
+		router.PathPrefix("/pd").Handler(negroni.New(
+			negroni.Wrap(r)),
+		)
+		return router, info, nil
+	}
+	cfg := NewTestSingleConfig(checkerWithNilAssert(c))
+	ctx, cancel := context.WithCancel(context.Background())
+	svr, err := CreateServer(ctx, cfg, handler)
+	c.Assert(err, IsNil)
+	defer func() {
+		cancel()
+		svr.Close()
+		testutil.CleanServer(svr.cfg.DataDir)
+	}()
+	err = svr.Run()
+	c.Assert(err, IsNil)
+	resp, err := http.Get(fmt.Sprintf("%s/pd/apis/mok/v1/router", svr.GetAddr()))
+	c.Assert(err, IsNil)
+	c.Assert(resp.StatusCode, Equals, http.StatusOK)
+	c.Assert(err, IsNil)
+	bodyBytes, err := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	c.Assert(err, IsNil)
+	bodyString := string(bodyBytes)
+	c.Assert(bodyString, Equals, "Mux Router\n")
+}