api: Add Request Info Middleware

pkg/apiutil/apiutil.go

@@ -23,6 +23,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/gorilla/mux"
 	"github.com/pingcap/errcode"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
@@ -190,3 +191,12 @@ func (rt *ComponentSignatureRoundTripper) RoundTrip(req *http.Request) (resp *ht
 	resp, err = rt.proxied.RoundTrip(req)
 	return
 }
+
+// GetRouteName return mux route name registered
+func GetRouteName(req *http.Request) string {
+	route := mux.CurrentRoute(req)
+	if route != nil {
+		return route.GetName()
+	}
+	return ""
+}

pkg/requestutil/context.go

@@ -0,0 +1,38 @@
+// Copyright 2022 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 requestutil
+
+import (
+	"context"
+)
+
+// The key type is unexported to prevent collisions
+type key int
+
+const (
+	// requestInfoKey is the context key for the request compoenent.
+	requestInfoKey key = iota
+)
+
+// WithRequestInfo returns a copy of parent in which the request info value is set
+func WithRequestInfo(parent context.Context, requestInfo RequestInfo) context.Context {
+	return context.WithValue(parent, requestInfoKey, requestInfo)
+}
+
+// RequestInfoFrom returns the value of the request info key on the ctx
+func RequestInfoFrom(ctx context.Context) (RequestInfo, bool) {
+	requestInfo, ok := ctx.Value(requestInfoKey).(RequestInfo)
+	return requestInfo, ok
+}

pkg/requestutil/context_test.go

@@ -0,0 +1,57 @@
+// Copyright 2022 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 requestutil
+
+import (
+	"context"
+	"testing"
+
+	. "github.com/pingcap/check"
+)
+
+func Test(t *testing.T) {
+	TestingT(t)
+}
+
+var _ = Suite(&testRequestContextSuite{})
+
+type testRequestContextSuite struct {
+}
+
+func (s *testRequestContextSuite) TestRequestInfo(c *C) {
+	ctx := context.Background()
+	_, ok := RequestInfoFrom(ctx)
+	c.Assert(ok, Equals, false)
+	ctx = WithRequestInfo(ctx,
+		RequestInfo{
+			ServiceLabel: "test label",
+			Method:       "POST",
+			Component:    "pdctl",
+			IP:           "localhost",
+			URLParam:     "{\"id\"=1}",
+			BodyParam:    "{\"state\"=\"Up\"}",
+			TimeStamp:    "2022",
+		})
+	result, ok := RequestInfoFrom(ctx)
+	c.Assert(result, NotNil)
+	c.Assert(ok, Equals, true)
+	c.Assert(result.ServiceLabel, Equals, "test label")
+	c.Assert(result.Method, Equals, "POST")
+	c.Assert(result.Component, Equals, "pdctl")
+	c.Assert(result.IP, Equals, "localhost")
+	c.Assert(result.URLParam, Equals, "{\"id\"=1}")
+	c.Assert(result.BodyParam, Equals, "{\"state\"=\"Up\"}")
+	c.Assert(result.TimeStamp, Equals, "2022")
+}

pkg/requestutil/request_info.go

@@ -0,0 +1,70 @@
+// Copyright 2022 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 requestutil
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/tikv/pd/pkg/apiutil"
+)
+
+// RequestInfo holds service information from http.Request
+type RequestInfo struct {
+	ServiceLabel string
+	Method       string
+	Component    string
+	IP           string
+	TimeStamp    string
+	URLParam     string
+	BodyParam    string
+}
+
+// GetRequestInfo returns request info needed from http.Request
+func GetRequestInfo(r *http.Request) RequestInfo {
+	return RequestInfo{
+		ServiceLabel: apiutil.GetRouteName(r),
+		Method:       fmt.Sprintf("%s/%s:%s", r.Proto, r.Method, r.URL.Path),
+		Component:    apiutil.GetComponentNameOnHTTP(r),
+		IP:           apiutil.GetIPAddrFromHTTPRequest(r),
+		TimeStamp:    time.Now().Local().String(),
+		URLParam:     getURLParam(r),
+		BodyParam:    getBodyParam(r),
+	}
+}
+
+func getURLParam(r *http.Request) string {
+	buf, err := json.Marshal(r.URL.Query())
+	if err != nil {
+		return ""
+	}
+	return string(buf)
+}
+
+func getBodyParam(r *http.Request) string {
+	if r.Body == nil {
+		return ""
+	}
+	// http request body is a io.Reader between bytes.Reader and strings.Reader, it only has EOF error
+	buf, _ := io.ReadAll(r.Body)
+	r.Body.Close()
+	bodyParam := string(buf)
+	r.Body = io.NopCloser(bytes.NewBuffer(buf))
+	return bodyParam
+}

server/api/admin.go

@@ -138,3 +138,21 @@ func (h *adminHandler) UpdateWaitAsyncTime(w http.ResponseWriter, r *http.Reques
 	cluster.GetReplicationMode().UpdateMemberWaitAsyncTime(memberID)
 	h.rd.JSON(w, http.StatusOK, nil)
 }
+
+// @Tags admin
+// @Summary switch Service Middlewares including ServiceInfo, Audit and rate limit
+// @Param enable query string true "enable" Enums(true, false)
+// @Produce json
+// @Success 200 {string} string "Switching Service middleware is successful."
+// @Failure 400 {string} string "The input is invalid."
+// @Router /admin/service-middleware [POST]
+func (h *adminHandler) HanldeServiceMiddlewareSwitch(w http.ResponseWriter, r *http.Request) {
+	enableStr := r.URL.Query().Get("enable")
+	enable, err := strconv.ParseBool(enableStr)
+	if err != nil {
+		h.rd.JSON(w, http.StatusBadRequest, "The input is invalid.")
+		return
+	}
+	h.svr.SetServiceMiddleware(enable)
+	h.rd.JSON(w, http.StatusOK, "Switching Service middleware is successful.")
+}

server/api/admin_test.go

@@ -173,3 +173,45 @@ func (s *testTSOSuite) TestResetTS(c *C) {
 	c.Assert(err, NotNil)
 	c.Assert(err.Error(), Equals, "\"invalid tso value\"\n")
 }
+
+var _ = Suite(&testServiceSuite{})
+
+type testServiceSuite struct {
+	svr     *server.Server
+	cleanup cleanUpFunc
+}
+
+func (s *testServiceSuite) SetUpSuite(c *C) {
+	s.svr, s.cleanup = mustNewServer(c)
+	mustWaitLeader(c, []*server.Server{s.svr})
+
+	mustBootstrapCluster(c, s.svr)
+	mustPutStore(c, s.svr, 1, metapb.StoreState_Up, nil)
+}
+
+func (s *testServiceSuite) TearDownSuite(c *C) {
+	s.cleanup()
+}
+
+func (s *testServiceSuite) TestSwitchServiceMiddleware(c *C) {
+	urlPrefix := fmt.Sprintf("%s%s/api/v1/admin/service-middleware", s.svr.GetAddr(), apiPrefix)
+	disableURL := fmt.Sprintf("%s?enable=false", urlPrefix)
+	err := postJSON(testDialClient, disableURL, nil,
+		func(res []byte, code int) {
+			c.Assert(string(res), Equals, "\"Switching Service middleware is successful.\"\n")
+			c.Assert(code, Equals, http.StatusOK)
+		})
+
+	c.Assert(err, IsNil)
+	c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, false)
+
+	enableURL := fmt.Sprintf("%s?enable=true", urlPrefix)
+	err = postJSON(testDialClient, enableURL, nil,
+		func(res []byte, code int) {
+			c.Assert(string(res), Equals, "\"Switching Service middleware is successful.\"\n")
+			c.Assert(code, Equals, http.StatusOK)
+		})
+
+	c.Assert(err, IsNil)
+	c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, true)
+}

server/api/middleware.go

@@ -18,12 +18,45 @@ import (
 	"context"
 	"net/http"
 
+	"github.com/pingcap/failpoint"
 	"github.com/tikv/pd/pkg/errs"
+	"github.com/tikv/pd/pkg/requestutil"
 	"github.com/tikv/pd/server"
 	"github.com/tikv/pd/server/cluster"
 	"github.com/unrolled/render"
+	"github.com/urfave/negroni"
 )
 
+// requestInfoMiddleware is used to gather info from requsetInfo
+type requestInfoMiddleware struct {
+	svr *server.Server
+}
+
+func newRequestInfoMiddleware(s *server.Server) negroni.Handler {
+	return &requestInfoMiddleware{svr: s}
+}
+
+func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
+	if !rm.svr.IsServiceMiddlewareEnabled() {
+		next(w, r)
+		return
+	}
+
+	requestInfo := requestutil.GetRequestInfo(r)
+	r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
+
+	failpoint.Inject("addRequestInfoMiddleware", func() {
+		w.Header().Add("service-label", requestInfo.ServiceLabel)
+		w.Header().Add("body-param", requestInfo.BodyParam)
+		w.Header().Add("url-param", requestInfo.URLParam)
+		w.Header().Add("method", requestInfo.Method)
+		w.Header().Add("component", requestInfo.Component)
+		w.Header().Add("ip", requestInfo.IP)
+	})
+
+	next(w, r)
+}
+
 type clusterMiddleware struct {
 	s  *server.Server
 	rd *render.Render