api: add HTTP audit middleware

pkg/audit/audit.go

@@ -0,0 +1,59 @@
+// 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 audit
+
+import (
+	"net/http"
+)
+
+// BackendLabels is used to store some audit backend labels.
+type BackendLabels struct {
+	Labels []string
+}
+
+// LabelMatcher is used to help backend implement audit.Backend
+type LabelMatcher struct {
+	backendLabel string
+}
+
+// Match is used to check whether backendLabel is in the labels
+func (m *LabelMatcher) Match(labels *BackendLabels) bool {
+	for _, item := range labels.Labels {
+		if m.backendLabel == item {
+			return true
+		}
+	}
+	return false
+}
+
+// Sequence is used to help backend implement audit.Backend
+type Sequence struct {
+	before bool
+}
+
+// ProcessBeforeHandler is used to help backend implement audit.Backend
+func (s *Sequence) ProcessBeforeHandler() bool {
+	return s.before
+}
+
+// Backend defines what function audit backend should hold
+type Backend interface {
+	// ProcessHTTPRequest is used to perform HTTP audit process
+	ProcessHTTPRequest(req *http.Request) bool
+	// Match is used to determine if the backend matches
+	Match(*BackendLabels) bool
+	// ProcessBeforeHandler is used to identify whether this backend should execute before handler
+	ProcessBeforeHandler() bool
+}

pkg/audit/audit_test.go

@@ -0,0 +1,39 @@
+// 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 audit
+
+import (
+	"testing"
+
+	. "github.com/pingcap/check"
+)
+
+func Test(t *testing.T) {
+	TestingT(t)
+}
+
+var _ = Suite(&testAuditSuite{})
+
+type testAuditSuite struct {
+}
+
+func (s *testAuditSuite) TestLabelMatcher(c *C) {
+	matcher := &LabelMatcher{"testSuccess"}
+	labels1 := &BackendLabels{Labels: []string{"testFail", "testSuccess"}}
+	c.Assert(matcher.Match(labels1), Equals, true)
+
+	labels2 := &BackendLabels{Labels: []string{"testFail"}}
+	c.Assert(matcher.Match(labels2), Equals, false)
+}

pkg/requestutil/context.go

@@ -22,8 +22,10 @@ import (
 type key int
 
 const (
-	// requestInfoKey is the context key for the request compoenent.
+	// requestInfoKey is the context key for the request info.
 	requestInfoKey key = iota
+	// endTimeKey is the context key for the end time.
+	endTimeKey
 )
 
 // WithRequestInfo returns a copy of parent in which the request info value is set
@@ -36,3 +38,14 @@ func RequestInfoFrom(ctx context.Context) (RequestInfo, bool) {
 	requestInfo, ok := ctx.Value(requestInfoKey).(RequestInfo)
 	return requestInfo, ok
 }
+
+// WithEndTime returns a copy of parent in which the end time value is set
+func WithEndTime(parent context.Context, endTime int64) context.Context {
+	return context.WithValue(parent, endTimeKey, endTime)
+}
+
+// EndTimeFrom returns the value of the excution info key on the ctx
+func EndTimeFrom(ctx context.Context) (int64, bool) {
+	info, ok := ctx.Value(endTimeKey).(int64)
+	return info, ok
+}

pkg/requestutil/context_test.go

@@ -17,6 +17,7 @@ package requestutil
 import (
 	"context"
 	"testing"
+	"time"
 
 	. "github.com/pingcap/check"
 )
@@ -34,15 +35,16 @@ func (s *testRequestContextSuite) TestRequestInfo(c *C) {
 	ctx := context.Background()
 	_, ok := RequestInfoFrom(ctx)
 	c.Assert(ok, Equals, false)
+	timeNow := time.Now().Unix()
 	ctx = WithRequestInfo(ctx,
 		RequestInfo{
-			ServiceLabel: "test label",
-			Method:       "POST",
-			Component:    "pdctl",
-			IP:           "localhost",
-			URLParam:     "{\"id\"=1}",
-			BodyParam:    "{\"state\"=\"Up\"}",
-			TimeStamp:    "2022",
+			ServiceLabel:   "test label",
+			Method:         "POST",
+			Component:      "pdctl",
+			IP:             "localhost",
+			URLParam:       "{\"id\"=1}",
+			BodyParam:      "{\"state\"=\"Up\"}",
+			StartTimeStamp: timeNow,
 		})
 	result, ok := RequestInfoFrom(ctx)
 	c.Assert(result, NotNil)
@@ -53,5 +55,17 @@ func (s *testRequestContextSuite) TestRequestInfo(c *C) {
 	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")
+	c.Assert(result.StartTimeStamp, Equals, timeNow)
+}
+
+func (s *testRequestContextSuite) TestEndTime(c *C) {
+	ctx := context.Background()
+	_, ok := EndTimeFrom(ctx)
+	c.Assert(ok, Equals, false)
+	timeNow := time.Now().Unix()
+	ctx = WithEndTime(ctx, timeNow)
+	result, ok := EndTimeFrom(ctx)
+	c.Assert(result, NotNil)
+	c.Assert(ok, Equals, true)
+	c.Assert(result, Equals, timeNow)
 }

pkg/requestutil/request_info.go

@@ -27,25 +27,25 @@ import (
 
 // RequestInfo holds service information from http.Request
 type RequestInfo struct {
-	ServiceLabel string
-	Method       string
-	Component    string
-	IP           string
-	TimeStamp    string
-	URLParam     string
-	BodyParam    string
+	ServiceLabel   string
+	Method         string
+	Component      string
+	IP             string
+	URLParam       string
+	BodyParam      string
+	StartTimeStamp int64
 }
 
 // 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),
+		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),
+		URLParam:       getURLParam(r),
+		BodyParam:      getBodyParam(r),
+		StartTimeStamp: time.Now().Unix(),
 	}
 }
 

server/api/admin.go

@@ -140,19 +140,19 @@ func (h *adminHandler) UpdateWaitAsyncTime(w http.ResponseWriter, r *http.Reques
 }
 
 // @Tags admin
-// @Summary switch Service Middlewares including ServiceInfo, Audit and rate limit
+// @Summary switch audit middleware
 // @Param enable query string true "enable" Enums(true, false)
 // @Produce json
-// @Success 200 {string} string "Switching Service middleware is successful."
+// @Success 200 {string} string "Switching audit 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) {
+// @Router /admin/audit-middleware [POST]
+func (h *adminHandler) HanldeAuditMiddlewareSwitch(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.")
+	h.svr.SetAuditMiddleware(enable)
+	h.rd.JSON(w, http.StatusOK, "Switching audit middleware is successful.")
 }

server/api/admin_test.go

@@ -193,25 +193,26 @@ 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 (s *testServiceSuite) TestSwitchAuditMiddleware(c *C) {
+	urlPrefix := fmt.Sprintf("%s%s/api/v1/admin/audit-middleware", s.svr.GetAddr(), apiPrefix)
+
+	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(string(res), Equals, "\"Switching audit middleware is successful.\"\n")
 			c.Assert(code, Equals, http.StatusOK)
 		})
 
 	c.Assert(err, IsNil)
-	c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, false)
+	c.Assert(s.svr.IsAuditMiddlewareEnabled(), Equals, true)
 
-	enableURL := fmt.Sprintf("%s?enable=true", urlPrefix)
-	err = postJSON(testDialClient, enableURL, nil,
+	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(string(res), Equals, "\"Switching audit middleware is successful.\"\n")
 			c.Assert(code, Equals, http.StatusOK)
 		})
 
 	c.Assert(err, IsNil)
-	c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, true)
+	c.Assert(s.svr.IsAuditMiddlewareEnabled(), Equals, false)
 }

server/api/middleware.go

@@ -17,8 +17,12 @@ package api
 import (
 	"context"
 	"net/http"
+	"strings"
+	"time"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/audit"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/requestutil"
 	"github.com/tikv/pd/server"
@@ -37,7 +41,7 @@ func newRequestInfoMiddleware(s *server.Server) negroni.Handler {
 }
 
 func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
-	if !rm.svr.IsServiceMiddlewareEnabled() {
+	if !rm.svr.IsAuditMiddlewareEnabled() {
 		next(w, r)
 		return
 	}
@@ -86,3 +90,58 @@ type clusterCtxKey struct{}
 func getCluster(r *http.Request) *cluster.RaftCluster {
 	return r.Context().Value(clusterCtxKey{}).(*cluster.RaftCluster)
 }
+
+type auditMiddleware struct {
+	svr *server.Server
+}
+
+func newAuditMiddleware(s *server.Server) negroni.Handler {
+	return &auditMiddleware{svr: s}
+}
+
+// ServeHTTP is used to implememt negroni.Handler for auditMiddleware
+func (s *auditMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
+	if !s.svr.IsAuditMiddlewareEnabled() {
+		next(w, r)
+		return
+	}
+
+	requestInfo, ok := requestutil.RequestInfoFrom(r.Context())
+	if !ok {
+		log.Error("failed to get request info when auditing")
+		next(w, r)
+	}
+
+	labels := s.svr.GetServiceAuditBackendLabels(requestInfo.ServiceLabel)
+	if labels == nil {
+		next(w, r)
+		return
+	}
+
+	failpoint.Inject("addAuditMiddleware", func() {
+		w.Header().Add("audit-label", strings.Join(labels.Labels, ","))
+	})
+
+	beforeNextBackends := make([]audit.Backend, 0)
+	afterNextBackends := make([]audit.Backend, 0)
+	for _, backend := range s.svr.GetAuditBackend() {
+		if backend.Match(labels) {
+			if backend.ProcessBeforeHandler() {
+				beforeNextBackends = append(beforeNextBackends, backend)
+			} else {
+				afterNextBackends = append(afterNextBackends, backend)
+			}
+		}
+	}
+	for _, backend := range beforeNextBackends {
+		backend.ProcessHTTPRequest(r)
+	}
+
+	next(w, r)
+
+	endTime := time.Now().Unix()
+	r = r.WithContext(requestutil.WithEndTime(r.Context(), endTime))
+	for _, backend := range afterNextBackends {
+		backend.ProcessHTTPRequest(r)
+	}
+}